Add instance creation to composd.

At startup we validate the files needed for the current instance exist
and are valid; if anything goes wrong we re-create them.

Bug: 186126194
Test: Manual, via composd_cmd.
Change-Id: Iaf19a74df14a91436c4cb4435f0286fa09307db8
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 9e95ed0..e68deb8 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -55,6 +55,13 @@
 }
 
 impl VmInstance {
+    /// Return a new connection to the Virtualization Service binder interface. This will start the
+    /// service if necessary.
+    pub fn connect_to_virtualization_service() -> Result<Strong<dyn IVirtualizationService>> {
+        wait_for_interface::<dyn IVirtualizationService>("android.system.virtualizationservice")
+            .context("Failed to find VirtualizationService")
+    }
+
     /// Start a new CompOS VM instance using the specified instance image file.
     pub fn start(instance_image: &Path) -> Result<VmInstance> {
         let instance_image =
@@ -84,10 +91,7 @@
             ..Default::default()
         });
 
-        let service = wait_for_interface::<dyn IVirtualizationService>(
-            "android.system.virtualizationservice",
-        )
-        .context("Failed to find VirtualizationService")?;
+        let service = Self::connect_to_virtualization_service()?;
 
         let vm = service.createVm(&config, Some(&log_fd)).context("Failed to create VM")?;
         let vm_state = Arc::new(VmStateMonitor::default());
diff --git a/compos/composd/Android.bp b/compos/composd/Android.bp
index 5c968b8..30d07ab 100644
--- a/compos/composd/Android.bp
+++ b/compos/composd/Android.bp
@@ -9,6 +9,7 @@
     prefer_rlib: true,
     rustlibs: [
         "android.system.composd-rust",
+        "android.system.virtualizationservice-rust",
         "compos_aidl_interface-rust",
         "libandroid_logger",
         "libanyhow",
diff --git a/compos/composd/src/composd_main.rs b/compos/composd/src/composd_main.rs
index 71e8125..3bd4121 100644
--- a/compos/composd/src/composd_main.rs
+++ b/compos/composd/src/composd_main.rs
@@ -22,8 +22,10 @@
 mod odrefresh;
 mod service;
 
+use crate::instance_manager::InstanceManager;
 use android_system_composd::binder::{register_lazy_service, ProcessState};
 use anyhow::{Context, Result};
+use compos_common::compos_client::VmInstance;
 use log::{error, info};
 
 fn try_main() -> Result<()> {
@@ -33,8 +35,10 @@
 
     ProcessState::start_thread_pool();
 
-    let service = service::new_binder();
-    register_lazy_service("android.system.composd", service.as_binder())
+    let virtualization_service = VmInstance::connect_to_virtualization_service()?;
+    let instance_manager = InstanceManager::new(virtualization_service);
+    let composd_service = service::new_binder(instance_manager);
+    register_lazy_service("android.system.composd", composd_service.as_binder())
         .context("Registering service")?;
 
     info!("Registered service, joining threadpool");
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 5d2a7e8..ad9c28a 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -17,13 +17,19 @@
 //! Starts and manages instances of the CompOS VM. At most one instance should be running at
 //! a time.
 
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    IVirtualizationService::IVirtualizationService, PartitionType::PartitionType,
+};
 use anyhow::{bail, Context, Result};
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
-use compos_aidl_interface::binder::Strong;
+use compos_aidl_interface::binder::{ParcelFileDescriptor, Strong};
 use compos_common::compos_client::VmInstance;
-use compos_common::{COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE, PRIVATE_KEY_BLOB_FILE};
+use compos_common::{
+    COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE, PRIVATE_KEY_BLOB_FILE, PUBLIC_KEY_FILE,
+};
+use log::{info, warn};
 use std::fs;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
 
 pub struct CompOsInstance {
@@ -32,25 +38,31 @@
     service: Strong<dyn ICompOsService>,
 }
 
-#[derive(Default)]
-pub struct InstanceManager(Mutex<State>);
+pub struct InstanceManager {
+    service: Strong<dyn IVirtualizationService>,
+    state: Mutex<State>,
+}
 
 impl InstanceManager {
+    pub fn new(service: Strong<dyn IVirtualizationService>) -> Self {
+        Self { service, state: Default::default() }
+    }
+
     pub fn get_running_service(&self) -> Result<Strong<dyn ICompOsService>> {
-        let mut state = self.0.lock().unwrap();
+        let mut state = self.state.lock().unwrap();
         let instance = state.get_running_instance().context("No running instance")?;
         Ok(instance.service.clone())
     }
 
     pub fn start_current_instance(&self) -> Result<Arc<CompOsInstance>> {
-        let mut state = self.0.lock().unwrap();
+        let mut state = self.state.lock().unwrap();
         state.mark_starting()?;
         // Don't hold the lock while we start the instance to avoid blocking other callers.
         drop(state);
 
         let instance = self.try_start_current_instance();
 
-        let mut state = self.0.lock().unwrap();
+        let mut state = self.state.lock().unwrap();
         if let Ok(ref instance) = instance {
             state.mark_started(instance)?;
         } else {
@@ -60,20 +72,134 @@
     }
 
     fn try_start_current_instance(&self) -> Result<Arc<CompOsInstance>> {
-        // TODO: Create instance_image & keys if needed
-        // TODO: Hold on to an IVirtualizationService
-        let instance_image: PathBuf =
-            [COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE].iter().collect();
+        let instance_files = InstanceFiles::new(CURRENT_DIR);
 
-        let vm_instance = VmInstance::start(&instance_image).context("Starting VM")?;
+        let compos_instance = instance_files.create_or_start_instance(&*self.service)?;
+
+        Ok(Arc::new(compos_instance))
+    }
+}
+
+struct InstanceFiles {
+    instance_name: String,
+    instance_root: PathBuf,
+    instance_image: PathBuf,
+    key_blob: PathBuf,
+    public_key: PathBuf,
+}
+
+impl InstanceFiles {
+    fn new(instance_name: &str) -> Self {
+        let instance_root = Path::new(COMPOS_DATA_ROOT).join(instance_name);
+        let instant_root_path = instance_root.as_path();
+        let instance_image = instant_root_path.join(INSTANCE_IMAGE_FILE);
+        let key_blob = instant_root_path.join(PRIVATE_KEY_BLOB_FILE);
+        let public_key = instant_root_path.join(PUBLIC_KEY_FILE);
+        Self {
+            instance_name: instance_name.to_owned(),
+            instance_root,
+            instance_image,
+            key_blob,
+            public_key,
+        }
+    }
+
+    fn create_or_start_instance(
+        &self,
+        service: &dyn IVirtualizationService,
+    ) -> Result<CompOsInstance> {
+        let compos_instance = self.start_instance();
+        match compos_instance {
+            Ok(_) => return compos_instance,
+            Err(e) => warn!("Failed to start {}: {}", self.instance_name, e),
+        }
+
+        self.start_new_instance(service)
+    }
+
+    fn start_instance(&self) -> Result<CompOsInstance> {
+        // No point even trying if the files we need aren't there.
+        self.check_files_exist()?;
+
+        let key_blob = fs::read(&self.key_blob).context("Reading private key blob")?;
+        let public_key = fs::read(&self.public_key).context("Reading public key")?;
+
+        let vm_instance = VmInstance::start(&self.instance_image).context("Starting VM")?;
         let service = vm_instance.get_service().context("Connecting to CompOS")?;
 
-        let key_blob: PathBuf =
-            [COMPOS_DATA_ROOT, CURRENT_DIR, PRIVATE_KEY_BLOB_FILE].iter().collect();
-        let key_blob = fs::read(key_blob).context("Reading private key")?;
-        service.initializeSigningKey(&key_blob).context("Loading key")?;
+        if !service.verifySigningKey(&key_blob, &public_key).context("Verifying key pair")? {
+            bail!("Key pair invalid");
+        }
 
-        Ok(Arc::new(CompOsInstance { vm_instance, service }))
+        // If we get this far then the instance image is valid in the current context (e.g. the
+        // current set of APEXes) and the key blob can be successfully decrypted by the VM. So the
+        // files have not been tampered with and we're good to go.
+
+        service.initializeSigningKey(&key_blob).context("Loading signing key")?;
+
+        Ok(CompOsInstance { vm_instance, service })
+    }
+
+    fn start_new_instance(
+        &self,
+        virtualization_service: &dyn IVirtualizationService,
+    ) -> Result<CompOsInstance> {
+        info!("Creating {} CompOs instance", self.instance_name);
+
+        // Ignore failure here - the directory may already exist.
+        let _ = fs::create_dir(&self.instance_root);
+
+        self.create_instance_image(virtualization_service)?;
+
+        let vm_instance = VmInstance::start(&self.instance_image).context("Starting VM")?;
+        let service = vm_instance.get_service().context("Connecting to CompOS")?;
+
+        let key_data = service.generateSigningKey().context("Generating signing key")?;
+        fs::write(&self.key_blob, &key_data.keyBlob).context("Writing key blob")?;
+        // TODO: Extract public key from cert
+        fs::write(&self.public_key, &key_data.certificate).context("Writing public key")?;
+
+        // We don't need to verify the key, since we just generated it and have it in memory.
+
+        service.initializeSigningKey(&key_data.keyBlob).context("Loading signing key")?;
+
+        Ok(CompOsInstance { vm_instance, service })
+    }
+
+    fn create_instance_image(
+        &self,
+        virtualization_service: &dyn IVirtualizationService,
+    ) -> Result<()> {
+        let instance_image = fs::OpenOptions::new()
+            .create(true)
+            .read(true)
+            .write(true)
+            .open(&self.instance_image)
+            .context("Creating instance image file")?;
+        let instance_image = ParcelFileDescriptor::new(instance_image);
+        // TODO: Where does this number come from?
+        let size = 10 * 1024 * 1024;
+        virtualization_service
+            .initializeWritablePartition(&instance_image, size, PartitionType::ANDROID_VM_INSTANCE)
+            .context("Writing instance image file")?;
+        Ok(())
+    }
+
+    fn check_files_exist(&self) -> Result<()> {
+        if !self.instance_root.is_dir() {
+            bail!("Directory {} not found", self.instance_root.display())
+        };
+        Self::check_file_exists(&self.instance_image)?;
+        Self::check_file_exists(&self.key_blob)?;
+        Self::check_file_exists(&self.public_key)?;
+        Ok(())
+    }
+
+    fn check_file_exists(file: &Path) -> Result<()> {
+        if !file.is_file() {
+            bail!("File {} not found", file.display())
+        };
+        Ok(())
     }
 }
 
@@ -82,6 +208,8 @@
 // Starting: is_starting is true, running_instance is None.
 // Started: is_starting is false, running_instance is Some(x) and there is a strong ref to x.
 // Stopped: is_starting is false and running_instance is None or a weak ref to a dropped instance.
+// The panic calls here should never happen, unless the code above in InstanceManager is buggy.
+// In particular nothing the client does should be able to trigger them.
 #[derive(Default)]
 struct State {
     running_instance: Option<Weak<CompOsInstance>>,
diff --git a/compos/composd/src/service.rs b/compos/composd/src/service.rs
index e3a1be0..fadca6c 100644
--- a/compos/composd/src/service.rs
+++ b/compos/composd/src/service.rs
@@ -30,13 +30,12 @@
 use log::{error, info};
 use std::ffi::CString;
 
-#[derive(Default)]
 pub struct IsolatedCompilationService {
     instance_manager: InstanceManager,
 }
 
-pub fn new_binder() -> Strong<dyn IIsolatedCompilationService> {
-    let service = IsolatedCompilationService::default();
+pub fn new_binder(instance_manager: InstanceManager) -> Strong<dyn IIsolatedCompilationService> {
+    let service = IsolatedCompilationService { instance_manager };
     BnIsolatedCompilationService::new_binder(service, BinderFeatures::default())
 }