Merge "Get odrefresh in composd working."
diff --git a/compos/Android.bp b/compos/Android.bp
index 2597c1c..2cc5131 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -6,6 +6,7 @@
     name: "pvm_exec",
     srcs: ["src/pvm_exec.rs"],
     rustlibs: [
+        "android.system.composd-rust",
         "compos_aidl_interface-rust",
         "libandroid_logger",
         "libanyhow",
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index c525f40..9e95ed0 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -50,6 +50,7 @@
     service: Strong<dyn IVirtualizationService>,
     #[allow(dead_code)] // Keeps the VM alive even if we don`t touch it
     vm: Strong<dyn IVirtualMachine>,
+    #[allow(dead_code)] // TODO: Do we need this?
     cid: i32,
 }
 
diff --git a/compos/common/lib.rs b/compos/common/lib.rs
index 6bea62c..104b8e5 100644
--- a/compos/common/lib.rs
+++ b/compos/common/lib.rs
@@ -18,6 +18,9 @@
 
 pub mod compos_client;
 
+/// Special CID indicating "any".
+pub const VMADDR_CID_ANY: u32 = -1i32 as u32;
+
 /// VSock port that the CompOS server listens on for RPC binder connections. This should be out of
 /// future port range (if happens) that microdroid may reserve for system components.
 pub const COMPOS_VSOCK_PORT: u32 = 6432;
diff --git a/compos/composd/aidl/Android.bp b/compos/composd/aidl/Android.bp
index 90c0de0..0352001 100644
--- a/compos/composd/aidl/Android.bp
+++ b/compos/composd/aidl/Android.bp
@@ -5,6 +5,7 @@
 aidl_interface {
     name: "android.system.composd",
     srcs: ["android/system/composd/*.aidl"],
+    imports: ["compos_aidl_interface"],
     // TODO: Make this stable when the APEX becomes updatable.
     unstable: true,
     backend: {
diff --git a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
index 9240bc6..5ff72fe 100644
--- a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
+++ b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
@@ -15,7 +15,19 @@
  */
 package android.system.composd;
 
+import com.android.compos.CompilationResult;
+import com.android.compos.FdAnnotation;
+
 interface IIsolatedCompilationService {
-    /// Run "odrefresh --force-compile" in CompOS
+    /** Run "odrefresh --force-compile" in CompOS. */
     void runForcedCompile();
+
+    /**
+     * Run dex2oat in the currently running instance of the CompOS VM. This is a simple proxy
+     * to ICompOsService#compile.
+     *
+     * This method can only be called from odrefresh. If there is no currently running instance
+     * an error is returned.
+     */
+    CompilationResult compile(in String[] args, in FdAnnotation fd_annotation);
 }
diff --git a/compos/composd/src/compos_instance.rs b/compos/composd/src/compos_instance.rs
deleted file mode 100644
index e30a8b3..0000000
--- a/compos/composd/src/compos_instance.rs
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-//! Starts and manages instances of the CompOS VM.
-
-use anyhow::{Context, Result};
-use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
-use compos_aidl_interface::binder::Strong;
-use compos_common::compos_client::VmInstance;
-use compos_common::{COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE, PRIVATE_KEY_BLOB_FILE};
-use std::fs;
-use std::path::PathBuf;
-
-#[allow(dead_code)]
-pub struct CompOsInstance {
-    instance: VmInstance,
-    service: Strong<dyn ICompOsService>,
-}
-
-impl CompOsInstance {
-    pub fn start_current_instance() -> Result<CompOsInstance> {
-        let instance_image: PathBuf =
-            [COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE].iter().collect();
-
-        let instance = VmInstance::start(&instance_image).context("Starting VM")?;
-        let service = 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")?;
-
-        Ok(CompOsInstance { instance, service })
-    }
-
-    pub fn cid(&self) -> i32 {
-        self.instance.cid()
-    }
-}
diff --git a/compos/composd/src/composd_main.rs b/compos/composd/src/composd_main.rs
index 33da889..71e8125 100644
--- a/compos/composd/src/composd_main.rs
+++ b/compos/composd/src/composd_main.rs
@@ -18,7 +18,7 @@
 //! responsible for managing the lifecycle of the CompOS VM instances, providing key management for
 //! them, and orchestrating trusted compilation.
 
-mod compos_instance;
+mod instance_manager;
 mod odrefresh;
 mod service;
 
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
new file mode 100644
index 0000000..5d2a7e8
--- /dev/null
+++ b/compos/composd/src/instance_manager.rs
@@ -0,0 +1,140 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! Starts and manages instances of the CompOS VM. At most one instance should be running at
+//! a time.
+
+use anyhow::{bail, Context, Result};
+use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
+use compos_aidl_interface::binder::Strong;
+use compos_common::compos_client::VmInstance;
+use compos_common::{COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE, PRIVATE_KEY_BLOB_FILE};
+use std::fs;
+use std::path::PathBuf;
+use std::sync::{Arc, Mutex, Weak};
+
+pub struct CompOsInstance {
+    #[allow(dead_code)] // Keeps VirtualizationService & the VM alive
+    vm_instance: VmInstance,
+    service: Strong<dyn ICompOsService>,
+}
+
+#[derive(Default)]
+pub struct InstanceManager(Mutex<State>);
+
+impl InstanceManager {
+    pub fn get_running_service(&self) -> Result<Strong<dyn ICompOsService>> {
+        let mut state = self.0.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();
+        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();
+        if let Ok(ref instance) = instance {
+            state.mark_started(instance)?;
+        } else {
+            state.mark_stopped();
+        }
+        instance
+    }
+
+    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 vm_instance = VmInstance::start(&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")?;
+
+        Ok(Arc::new(CompOsInstance { vm_instance, service }))
+    }
+}
+
+// Ensures we only run one instance at a time.
+// Valid states:
+// 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.
+#[derive(Default)]
+struct State {
+    running_instance: Option<Weak<CompOsInstance>>,
+    is_starting: bool,
+}
+
+impl State {
+    // Move to Starting iff we are Stopped.
+    fn mark_starting(&mut self) -> Result<()> {
+        if self.is_starting {
+            bail!("An instance is already starting");
+        }
+        if let Some(weak) = &self.running_instance {
+            if weak.strong_count() != 0 {
+                bail!("An instance is already running");
+            }
+        }
+        self.running_instance = None;
+        self.is_starting = true;
+        Ok(())
+    }
+
+    // Move from Starting to Stopped.
+    fn mark_stopped(&mut self) {
+        if !self.is_starting || self.running_instance.is_some() {
+            panic!("Tried to mark stopped when not starting");
+        }
+        self.is_starting = false;
+    }
+
+    // Move from Starting to Started.
+    fn mark_started(&mut self, instance: &Arc<CompOsInstance>) -> Result<()> {
+        if !self.is_starting {
+            panic!("Tried to mark started when not starting")
+        }
+        if self.running_instance.is_some() {
+            panic!("Attempted to mark started when already started");
+        }
+        self.is_starting = false;
+        self.running_instance = Some(Arc::downgrade(instance));
+        Ok(())
+    }
+
+    // Return the running instance if we are in the Started state.
+    fn get_running_instance(&mut self) -> Option<Arc<CompOsInstance>> {
+        if self.is_starting {
+            return None;
+        }
+        let instance = self.running_instance.as_ref()?.upgrade();
+        if instance.is_none() {
+            // No point keeping an orphaned weak reference
+            self.running_instance = None;
+        }
+        instance
+    }
+}
diff --git a/compos/composd/src/odrefresh.rs b/compos/composd/src/odrefresh.rs
index c0042f0..54da231 100644
--- a/compos/composd/src/odrefresh.rs
+++ b/compos/composd/src/odrefresh.rs
@@ -17,6 +17,7 @@
 //! Handle the details of executing odrefresh to generate compiled artifacts.
 
 use anyhow::{bail, Context, Result};
+use compos_common::VMADDR_CID_ANY;
 use num_derive::FromPrimitive;
 use num_traits::FromPrimitive;
 use std::process::Command;
@@ -36,10 +37,10 @@
     CleanupFailed = EX_MAX + 4,
 }
 
-pub fn run_forced_compile(cid: i32) -> Result<ExitCode> {
+pub fn run_forced_compile() -> Result<ExitCode> {
     // We don`t need to capture stdout/stderr - odrefresh writes to the log
     let mut odrefresh = Command::new(ODREFRESH_BIN)
-        .arg(format!("--use-compilation-os={}", cid))
+        .arg(format!("--use-compilation-os={}", VMADDR_CID_ANY))
         .arg("--force-compile")
         .spawn()
         .context("Running odrefresh")?;
diff --git a/compos/composd/src/service.rs b/compos/composd/src/service.rs
index 7fc9ab0..e3a1be0 100644
--- a/compos/composd/src/service.rs
+++ b/compos/composd/src/service.rs
@@ -17,20 +17,26 @@
 //! Implementation of IIsolatedCompilationService, called from system server when compilation is
 //! desired.
 
-use crate::compos_instance::CompOsInstance;
+use crate::instance_manager::InstanceManager;
 use crate::odrefresh;
 use android_system_composd::aidl::android::system::composd::IIsolatedCompilationService::{
     BnIsolatedCompilationService, IIsolatedCompilationService,
 };
 use android_system_composd::binder::{self, BinderFeatures, Interface, Status, Strong};
 use anyhow::{bail, Context, Result};
+use compos_aidl_interface::aidl::com::android::compos::{
+    CompilationResult::CompilationResult, FdAnnotation::FdAnnotation,
+};
 use log::{error, info};
 use std::ffi::CString;
 
-pub struct IsolatedCompilationService {}
+#[derive(Default)]
+pub struct IsolatedCompilationService {
+    instance_manager: InstanceManager,
+}
 
 pub fn new_binder() -> Strong<dyn IIsolatedCompilationService> {
-    let service = IsolatedCompilationService {};
+    let service = IsolatedCompilationService::default();
     BnIsolatedCompilationService::new_binder(service, BinderFeatures::default())
 }
 
@@ -38,8 +44,18 @@
 
 impl IIsolatedCompilationService for IsolatedCompilationService {
     fn runForcedCompile(&self) -> binder::Result<()> {
+        // TODO - check caller is system or shell/root?
         to_binder_result(self.do_run_forced_compile())
     }
+
+    fn compile(
+        &self,
+        args: &[String],
+        fd_annotation: &FdAnnotation,
+    ) -> binder::Result<CompilationResult> {
+        // TODO - check caller is odrefresh
+        to_binder_result(self.do_compile(args, fd_annotation))
+    }
 }
 
 fn to_binder_result<T>(result: Result<T>) -> binder::Result<T> {
@@ -53,16 +69,26 @@
     fn do_run_forced_compile(&self) -> Result<()> {
         info!("runForcedCompile");
 
-        // TODO: Create instance if need be, handle instance failure, prevent
-        // multiple instances running
-        let comp_os = CompOsInstance::start_current_instance().context("Starting CompOS")?;
+        let comp_os = self.instance_manager.start_current_instance().context("Starting CompOS")?;
 
-        let exit_code = odrefresh::run_forced_compile(comp_os.cid())?;
+        let exit_code = odrefresh::run_forced_compile()?;
 
         if exit_code != odrefresh::ExitCode::CompilationSuccess {
             bail!("Unexpected odrefresh result: {:?}", exit_code);
         }
 
+        // The instance is needed until odrefresh is finished
+        drop(comp_os);
+
         Ok(())
     }
+
+    fn do_compile(
+        &self,
+        args: &[String],
+        fd_annotation: &FdAnnotation,
+    ) -> Result<CompilationResult> {
+        let compos = self.instance_manager.get_running_service()?;
+        compos.compile(args, fd_annotation).context("Compiling")
+    }
 }
diff --git a/compos/src/pvm_exec.rs b/compos/src/pvm_exec.rs
index b6fc729..0b2dbb8 100644
--- a/compos/src/pvm_exec.rs
+++ b/compos/src/pvm_exec.rs
@@ -37,14 +37,23 @@
 use std::path::Path;
 use std::process::exit;
 
+use android_system_composd::{
+    aidl::android::system::composd::IIsolatedCompilationService::IIsolatedCompilationService,
+    binder::wait_for_interface,
+};
 use compos_aidl_interface::aidl::com::android::compos::{
     FdAnnotation::FdAnnotation, ICompOsService::ICompOsService,
 };
 use compos_aidl_interface::binder::Strong;
-use compos_common::COMPOS_VSOCK_PORT;
+use compos_common::{COMPOS_VSOCK_PORT, VMADDR_CID_ANY};
 
 const FD_SERVER_BIN: &str = "/apex/com.android.virt/bin/fd_server";
 
+fn get_composd() -> Result<Strong<dyn IIsolatedCompilationService>> {
+    wait_for_interface::<dyn IIsolatedCompilationService>("android.system.composd")
+        .context("Failed to find IIsolatedCompilationService")
+}
+
 fn get_rpc_binder(cid: u32) -> Result<Strong<dyn ICompOsService>> {
     // SAFETY: AIBinder returned by RpcClient has correct reference count, and the ownership can be
     // safely taken by new_spibinder.
@@ -144,13 +153,7 @@
     Ok(Config { args, fd_annotation: FdAnnotation { input_fds, output_fds }, cid, debuggable })
 }
 
-fn main() -> Result<()> {
-    let debuggable = env!("TARGET_BUILD_VARIANT") != "user";
-    let log_level = if debuggable { log::Level::Trace } else { log::Level::Info };
-    android_logger::init_once(
-        android_logger::Config::default().with_tag("pvm_exec").with_min_level(log_level),
-    );
-
+fn try_main() -> Result<()> {
     // 1. Parse the command line arguments for collect execution data.
     let Config { args, fd_annotation, cid, debuggable } = parse_args()?;
 
@@ -165,8 +168,16 @@
     });
 
     // 3. Send the command line args to the remote to execute.
-    let service = get_rpc_binder(cid)?;
-    let result = service.compile(&args, &fd_annotation).context("Binder call failed")?;
+    let result = if cid == VMADDR_CID_ANY {
+        // Sentinel value that indicates we should use composd
+        let composd = get_composd()?;
+        composd.compile(&args, &fd_annotation)
+    } else {
+        // Call directly into the VM
+        let compos_vm = get_rpc_binder(cid)?;
+        compos_vm.compile(&args, &fd_annotation)
+    };
+    let result = result.context("Binder call failed")?;
 
     // TODO: store/use the signature
     debug!(
@@ -185,3 +196,18 @@
     }
     Ok(())
 }
+
+fn main() {
+    let debuggable = env!("TARGET_BUILD_VARIANT") != "user";
+    let log_level = if debuggable { log::Level::Trace } else { log::Level::Info };
+    android_logger::init_once(
+        android_logger::Config::default().with_tag("pvm_exec").with_min_level(log_level),
+    );
+
+    // Make sure we log and indicate failure if we were unable to run the command and get its exit
+    // code.
+    if let Err(e) = try_main() {
+        error!("{}", e);
+        std::process::exit(-1)
+    }
+}