Merge "Use keystore2_microdroid instead."
diff --git a/authfs/fd_server/src/main.rs b/authfs/fd_server/src/main.rs
index 12f013c..395e2e9 100644
--- a/authfs/fd_server/src/main.rs
+++ b/authfs/fd_server/src/main.rs
@@ -25,7 +25,7 @@
 mod fsverity;
 
 use anyhow::{bail, Result};
-use binder::unstable_api::AsNative;
+use binder_common::rpc_server::run_rpc_server;
 use log::{debug, error};
 use std::cmp::min;
 use std::collections::BTreeMap;
@@ -292,7 +292,12 @@
     Ok((fd, FdConfig::ReadWrite(file)))
 }
 
-fn parse_args() -> Result<BTreeMap<i32, FdConfig>> {
+struct Args {
+    fd_pool: BTreeMap<i32, FdConfig>,
+    ready_fd: Option<File>,
+}
+
+fn parse_args() -> Result<Args> {
     #[rustfmt::skip]
     let matches = clap::App::new("fd_server")
         .arg(clap::Arg::with_name("ro-fds")
@@ -303,6 +308,9 @@
              .long("rw-fds")
              .multiple(true)
              .number_of_values(1))
+        .arg(clap::Arg::with_name("ready-fd")
+            .long("ready-fd")
+            .takes_value(true))
         .get_matches();
 
     let mut fd_pool = BTreeMap::new();
@@ -318,8 +326,13 @@
             fd_pool.insert(fd, config);
         }
     }
-
-    Ok(fd_pool)
+    let ready_fd = if let Some(arg) = matches.value_of("ready-fd") {
+        let fd = arg.parse::<i32>()?;
+        Some(fd_to_file(fd)?)
+    } else {
+        None
+    };
+    Ok(Args { fd_pool, ready_fd })
 }
 
 fn main() -> Result<()> {
@@ -327,18 +340,18 @@
         android_logger::Config::default().with_tag("fd_server").with_min_level(log::Level::Debug),
     );
 
-    let fd_pool = parse_args()?;
+    let args = parse_args()?;
+    let service = FdService::new_binder(args.fd_pool).as_binder();
 
-    let mut service = FdService::new_binder(fd_pool).as_binder();
     debug!("fd_server is starting as a rpc service.");
-    // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
-    // Plus the binder objects are threadsafe.
-    let retval = unsafe {
-        binder_rpc_unstable_bindgen::RunRpcServer(
-            service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder,
-            RPC_SERVICE_PORT,
-        )
-    };
+
+    let mut ready_fd = args.ready_fd;
+    let retval = run_rpc_server(service, RPC_SERVICE_PORT, || {
+        debug!("fd_server is ready");
+        // Close the ready-fd if we were given one to signal our readiness.
+        drop(ready_fd.take());
+    });
+
     if retval {
         debug!("RPC server has shut down gracefully");
         Ok(())
diff --git a/authfs/service/src/authfs.rs b/authfs/service/src/authfs.rs
index 5601738..6d87243 100644
--- a/authfs/service/src/authfs.rs
+++ b/authfs/service/src/authfs.rs
@@ -82,8 +82,11 @@
             &config.outputFdAnnotations,
             debuggable,
         )?;
-        wait_until_authfs_ready(&mountpoint).map_err(|e| {
-            debug!("Wait for authfs: {:?}", child.wait());
+        wait_until_authfs_ready(&child, &mountpoint).map_err(|e| {
+            match child.wait() {
+                Ok(status) => debug!("Wait for authfs: {}", status),
+                Err(e) => warn!("Failed to wait for child: {}", e),
+            }
             e
         })?;
 
@@ -144,13 +147,18 @@
     SharedChild::spawn(&mut command).context("Spawn authfs")
 }
 
-fn wait_until_authfs_ready(mountpoint: &OsStr) -> Result<()> {
+fn wait_until_authfs_ready(child: &SharedChild, mountpoint: &OsStr) -> Result<()> {
     let start_time = Instant::now();
     loop {
         if is_fuse(mountpoint)? {
             break;
         }
+        if let Some(exit_status) = child.try_wait()? {
+            // If the child has exited, we will never become ready.
+            bail!("Child has exited: {}", exit_status);
+        }
         if start_time.elapsed() > AUTHFS_SETUP_TIMEOUT_SEC {
+            let _ = child.kill();
             bail!("Time out mounting authfs");
         }
         sleep(AUTHFS_SETUP_POLL_INTERVAL_MS);
diff --git a/authfs/service/src/main.rs b/authfs/service/src/main.rs
index af8c7f9..890e108 100644
--- a/authfs/service/src/main.rs
+++ b/authfs/service/src/main.rs
@@ -59,7 +59,7 @@
         create_dir(&mountpoint).map_err(|e| {
             new_binder_exception(
                 ExceptionCode::SERVICE_SPECIFIC,
-                format!("Cannot create mount directory {:?}: {}", &mountpoint, e),
+                format!("Cannot create mount directory {:?}: {:?}", &mountpoint, e),
             )
         })?;
 
@@ -109,7 +109,7 @@
     Ok(())
 }
 
-fn main() -> Result<()> {
+fn try_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(
@@ -128,3 +128,10 @@
     ProcessState::join_thread_pool();
     bail!("Unexpected exit after join_thread_pool")
 }
+
+fn main() {
+    if let Err(e) = try_main() {
+        error!("failed with {:?}", e);
+        std::process::exit(1);
+    }
+}
diff --git a/authfs/src/main.rs b/authfs/src/main.rs
index c85d801..ecb0e68 100644
--- a/authfs/src/main.rs
+++ b/authfs/src/main.rs
@@ -28,6 +28,7 @@
 //! e.g. /mountpoint/42.
 
 use anyhow::{bail, Context, Result};
+use log::error;
 use std::collections::BTreeMap;
 use std::convert::TryInto;
 use std::fs::File;
@@ -325,7 +326,7 @@
     Ok(file_pool)
 }
 
-fn main() -> Result<()> {
+fn try_main() -> Result<()> {
     let args = Args::from_args();
 
     let log_level = if args.debug { log::Level::Debug } else { log::Level::Info };
@@ -337,3 +338,10 @@
     fusefs::loop_forever(file_pool, &args.mount_point, &args.extra_options)?;
     bail!("Unexpected exit after the handler loop")
 }
+
+fn main() {
+    if let Err(e) = try_main() {
+        error!("failed with {:?}", e);
+        std::process::exit(1);
+    }
+}
diff --git a/binder_common/Android.bp b/binder_common/Android.bp
index 789a891..209955d 100644
--- a/binder_common/Android.bp
+++ b/binder_common/Android.bp
@@ -9,6 +9,8 @@
     edition: "2018",
     rustlibs: [
         "libbinder_rs",
+        "libbinder_rpc_unstable_bindgen",
+        "liblazy_static",
     ],
     apex_available: [
         "com.android.compos",
diff --git a/binder_common/lazy_service.rs b/binder_common/lazy_service.rs
new file mode 100644
index 0000000..a2b85db
--- /dev/null
+++ b/binder_common/lazy_service.rs
@@ -0,0 +1,77 @@
+/*
+ * 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.
+ */
+
+//! Rust API for lazy (aka dynamic) AIDL services.
+//! See https://source.android.com/devices/architecture/aidl/dynamic-aidl.
+
+use binder::public_api::force_lazy_services_persist;
+use lazy_static::lazy_static;
+use std::sync::Mutex;
+
+// TODO(b/200924402): Move this class to libbinder_rs once the infrastructure needed exists.
+
+/// An RAII object to ensure a server of lazy services is not killed. During the lifetime of any of
+/// these objects the service manager will not not kill the current process even if none of its
+/// lazy services are in use.
+#[must_use]
+#[derive(Debug)]
+pub struct LazyServiceGuard {
+    // Prevent construction outside this module.
+    _private: (),
+}
+
+lazy_static! {
+    // Count of how many LazyServiceGuard objects are in existence.
+    static ref GUARD_COUNT: Mutex<u64> = Mutex::new(0);
+}
+
+impl LazyServiceGuard {
+    /// Create a new LazyServiceGuard to prevent the service manager prematurely killing this
+    /// process.
+    pub fn new() -> Self {
+        let mut count = GUARD_COUNT.lock().unwrap();
+        *count += 1;
+        if *count == 1 {
+            // It's important that we make this call with the mutex held, to make sure
+            // that multiple calls (e.g. if the count goes 1 -> 0 -> 1) are correctly
+            // sequenced. (That also means we can't just use an AtomicU64.)
+            force_lazy_services_persist(true);
+        }
+        Self { _private: () }
+    }
+}
+
+impl Drop for LazyServiceGuard {
+    fn drop(&mut self) {
+        let mut count = GUARD_COUNT.lock().unwrap();
+        *count -= 1;
+        if *count == 0 {
+            force_lazy_services_persist(false);
+        }
+    }
+}
+
+impl Clone for LazyServiceGuard {
+    fn clone(&self) -> Self {
+        Self::new()
+    }
+}
+
+impl Default for LazyServiceGuard {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/binder_common/lib.rs b/binder_common/lib.rs
index 54cb80e..f2391e3 100644
--- a/binder_common/lib.rs
+++ b/binder_common/lib.rs
@@ -16,6 +16,9 @@
 
 //! Common items useful for binder clients and/or servers.
 
+pub mod lazy_service;
+pub mod rpc_server;
+
 use binder::public_api::{ExceptionCode, Status};
 use std::ffi::CString;
 
diff --git a/binder_common/rpc_server.rs b/binder_common/rpc_server.rs
new file mode 100644
index 0000000..36075cf
--- /dev/null
+++ b/binder_common/rpc_server.rs
@@ -0,0 +1,81 @@
+/*
+ * 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.
+ */
+
+//! Helpers for implementing an RPC Binder server.
+
+use binder::public_api::SpIBinder;
+use binder::unstable_api::AsNative;
+use std::os::raw;
+
+/// Run a binder RPC server, serving the supplied binder service implementation on the given vsock
+/// port.
+/// If and when the server is ready for connections (it is listening on the port) on_ready
+/// is called to allow appropriate action to be taken - e.g. to notify clients they
+/// may now attempt to connect.
+/// The current thread is joined to the binder thread pool to handle incoming messages.
+/// Returns true if the server has shutdown normally, false if it failed in some way.
+pub fn run_rpc_server<F>(service: SpIBinder, port: u32, on_ready: F) -> bool
+where
+    F: FnOnce(),
+{
+    let mut ready_notifier = ReadyNotifier(Some(on_ready));
+    ready_notifier.run_server(service, port)
+}
+
+struct ReadyNotifier<F>(Option<F>)
+where
+    F: FnOnce();
+
+impl<F> ReadyNotifier<F>
+where
+    F: FnOnce(),
+{
+    fn run_server(&mut self, mut service: SpIBinder, port: u32) -> bool {
+        let service = service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder;
+        let param = self.as_void_ptr();
+
+        // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
+        // Plus the binder objects are threadsafe.
+        // RunRpcServerCallback does not retain a reference to ready_callback, and only ever
+        // calls it with the param we provide during the lifetime of self.
+        unsafe {
+            binder_rpc_unstable_bindgen::RunRpcServerCallback(
+                service,
+                port,
+                Some(Self::ready_callback),
+                param,
+            )
+        }
+    }
+
+    fn as_void_ptr(&mut self) -> *mut raw::c_void {
+        self as *mut _ as *mut raw::c_void
+    }
+
+    unsafe extern "C" fn ready_callback(param: *mut raw::c_void) {
+        // SAFETY: This is only ever called by RunRpcServerCallback, within the lifetime of the
+        // ReadyNotifier, with param taking the value returned by as_void_ptr (so a properly aligned
+        // non-null pointer to an initialized instance).
+        let ready_notifier = param as *mut Self;
+        ready_notifier.as_mut().unwrap().notify()
+    }
+
+    fn notify(&mut self) {
+        if let Some(on_ready) = self.0.take() {
+            on_ready();
+        }
+    }
+}
diff --git a/compos/aidl/com/android/compos/ICompOsService.aidl b/compos/aidl/com/android/compos/ICompOsService.aidl
index 7904130..29c453b 100644
--- a/compos/aidl/com/android/compos/ICompOsService.aidl
+++ b/compos/aidl/com/android/compos/ICompOsService.aidl
@@ -43,7 +43,22 @@
      * @param fd_annotation Additional file descriptor information of the execution
      * @return a CompilationResult
      */
-    CompilationResult compile(in String[] args, in FdAnnotation fd_annotation);
+    CompilationResult compile_cmd(in String[] args, in FdAnnotation fd_annotation);
+
+    /**
+     * Runs dexopt compilation encoded in the marshaled dexopt arguments.
+     *
+     * To keep ART indepdendantly updatable, the compilation arguments are not stabilized. As a
+     * result, the arguments are marshaled into byte array.  Upon received, the service asks ART to
+     * return relevant information (since ART is able to unmarshal its own encoding), in order to
+     * set up the execution context (mainly file descriptors for compiler input and output) then
+     * invokes the compiler.
+     *
+     * @param marshaledArguments The marshaled dexopt arguments.
+     * @param fd_annotation Additional file descriptor information of the execution.
+     * @return exit code
+     */
+    byte compile(in byte[] marshaledArguments, in FdAnnotation fd_annotation);
 
     /**
      * Generate a new public/private key pair suitable for signing CompOs output files.
diff --git a/compos/apex/Android.bp b/compos/apex/Android.bp
index 547fd44..3a9d5f5 100644
--- a/compos/apex/Android.bp
+++ b/compos/apex/Android.bp
@@ -38,12 +38,19 @@
     platform_apis: true,
 
     binaries: [
+        // Used in Android
         "compos_key_cmd",
         "compos_verify_key",
         "composd",
         "composd_cmd",
+        "pvm_exec", // to be superseded by libcompos_client
+
+        // Used in VM
         "compsvc",
-        "pvm_exec",
+    ],
+
+    native_shared_libs: [
+        "libcompos_client",
     ],
 
     apps: [
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index e68deb8..a69538e 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -46,11 +46,8 @@
 
 /// This owns an instance of the CompOS VM.
 pub struct VmInstance {
-    #[allow(dead_code)] // Prevent service manager from killing the dynamic service
-    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,
 }
 
@@ -63,9 +60,7 @@
     }
 
     /// Start a new CompOS VM instance using the specified instance image file.
-    pub fn start(instance_image: &Path) -> Result<VmInstance> {
-        let instance_image =
-            File::open(instance_image).context("Failed to open instance image file")?;
+    pub fn start(instance_image: File) -> Result<VmInstance> {
         let instance_fd = ParcelFileDescriptor::new(instance_image);
 
         let apex_dir = Path::new(COMPOS_APEX_ROOT);
@@ -113,7 +108,7 @@
 
         let cid = vm_state.wait_until_ready()?;
 
-        Ok(VmInstance { service, vm, cid })
+        Ok(VmInstance { vm, cid })
     }
 
     /// Create and return an RPC Binder connection to the Comp OS service in the VM.
@@ -129,6 +124,7 @@
 
     /// Return the CID of the VM.
     pub fn cid(&self) -> i32 {
+        // TODO: Do we actually need/use this?
         self.cid
     }
 }
@@ -227,9 +223,12 @@
     }
 
     fn wait_until_ready(&self) -> Result<i32> {
+        // 10s is long enough on real hardware, but it can take 90s when using nested
+        // virtualization.
+        // TODO(b/200924405): Reduce timeout/detect nested virtualization
         let (state, result) = self
             .state_ready
-            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(20), |state| {
+            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(120), |state| {
                 state.cid.is_none() && !state.has_died
             })
             .unwrap();
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index e168648..f495816 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -156,7 +156,10 @@
 
     bool waitUntilReady() {
         std::unique_lock lock(mMutex);
-        return mCv.wait_for(lock, std::chrono::seconds(20), [this] { return mReady || mDied; }) &&
+        // 10s is long enough on real hardware, but it can take 90s when using nested
+        // virtualization.
+        // TODO(b/200924405): Reduce timeout/detect nested virtualization
+        return mCv.wait_for(lock, std::chrono::seconds(120), [this] { return mReady || mDied; }) &&
                 !mDied;
     }
 
diff --git a/compos/composd/aidl/Android.bp b/compos/composd/aidl/Android.bp
index 0352001..8116632 100644
--- a/compos/composd/aidl/Android.bp
+++ b/compos/composd/aidl/Android.bp
@@ -18,5 +18,11 @@
                 "com.android.compos",
             ],
         },
+        ndk: {
+            enabled: true,
+            apex_available: [
+                "com.android.compos",
+            ],
+        },
     },
 }
diff --git a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
index 5ff72fe..a1bb92c 100644
--- a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
+++ b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
@@ -24,10 +24,19 @@
 
     /**
      * Run dex2oat in the currently running instance of the CompOS VM. This is a simple proxy
-     * to ICompOsService#compile.
+     * to ICompOsService#compile_cmd.
      *
      * 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);
+    CompilationResult compile_cmd(in String[] args, in FdAnnotation fd_annotation);
+
+    /**
+     * 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 libcompos_client. If there is no currently running
+     * instance an error is returned.
+     */
+    byte compile(in byte[] marshaledArguments, in FdAnnotation fd_annotation);
 }
diff --git a/compos/composd/src/instance_starter.rs b/compos/composd/src/instance_starter.rs
index ec95ff8..1751d35 100644
--- a/compos/composd/src/instance_starter.rs
+++ b/compos/composd/src/instance_starter.rs
@@ -89,8 +89,8 @@
         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 compos_instance = self.start_vm()?;
+        let service = &compos_instance.service;
 
         if !service.verifySigningKey(&key_blob, &public_key).context("Verifying key pair")? {
             bail!("Key pair invalid");
@@ -102,7 +102,7 @@
 
         service.initializeSigningKey(&key_blob).context("Loading signing key")?;
 
-        Ok(CompOsInstance { vm_instance, service })
+        Ok(compos_instance)
     }
 
     fn start_new_instance(
@@ -116,8 +116,8 @@
 
         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 compos_instance = self.start_vm()?;
+        let service = &compos_instance.service;
 
         let key_data = service.generateSigningKey().context("Generating signing key")?;
         fs::write(&self.key_blob, &key_data.keyBlob).context("Writing key blob")?;
@@ -133,6 +133,17 @@
 
         service.initializeSigningKey(&key_data.keyBlob).context("Loading signing key")?;
 
+        Ok(compos_instance)
+    }
+
+    fn start_vm(&self) -> Result<CompOsInstance> {
+        let instance_image = fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(&self.instance_image)
+            .context("Failed to open instance image")?;
+        let vm_instance = VmInstance::start(instance_image).context("Starting VM")?;
+        let service = vm_instance.get_service().context("Connecting to CompOS")?;
         Ok(CompOsInstance { vm_instance, service })
     }
 
diff --git a/compos/composd/src/odrefresh.rs b/compos/composd/src/odrefresh.rs
index 54da231..2d880e2 100644
--- a/compos/composd/src/odrefresh.rs
+++ b/compos/composd/src/odrefresh.rs
@@ -40,7 +40,7 @@
 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={}", VMADDR_CID_ANY))
+        .arg(format!("--use-compilation-os={}", VMADDR_CID_ANY as i32))
         .arg("--force-compile")
         .spawn()
         .context("Running odrefresh")?;
diff --git a/compos/composd/src/service.rs b/compos/composd/src/service.rs
index 2a67a27..be9c30c 100644
--- a/compos/composd/src/service.rs
+++ b/compos/composd/src/service.rs
@@ -47,7 +47,7 @@
         to_binder_result(self.do_run_forced_compile())
     }
 
-    fn compile(
+    fn compile_cmd(
         &self,
         args: &[String],
         fd_annotation: &FdAnnotation,
@@ -55,6 +55,10 @@
         // TODO - check caller is odrefresh
         to_binder_result(self.do_compile(args, fd_annotation))
     }
+
+    fn compile(&self, _marshaled: &[u8], _fd_annotation: &FdAnnotation) -> binder::Result<i8> {
+        Err(new_binder_service_specific_error(-1, "Not yet implemented"))
+    }
 }
 
 fn to_binder_result<T>(result: Result<T>) -> binder::Result<T> {
@@ -89,6 +93,6 @@
         fd_annotation: &FdAnnotation,
     ) -> Result<CompilationResult> {
         let compos = self.instance_manager.get_running_service()?;
-        compos.compile(args, fd_annotation).context("Compiling")
+        compos.compile_cmd(args, fd_annotation).context("Compiling")
     }
 }
diff --git a/compos/libcompos_client/Android.bp b/compos/libcompos_client/Android.bp
new file mode 100644
index 0000000..b6a4ef6
--- /dev/null
+++ b/compos/libcompos_client/Android.bp
@@ -0,0 +1,27 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+cc_library {
+    name: "libcompos_client",
+    srcs: ["libcompos_client.cc"],
+    min_sdk_version: "apex_inherit",
+    shared_libs: [
+        "android.system.composd-ndk",
+        "compos_aidl_interface-ndk",
+        "libbase",
+        "libbinder_ndk",
+        "libbinder_rpc_unstable",
+    ],
+    export_include_dirs: ["include"],
+    stubs: {
+        symbol_file: "libcompos_client.map.txt",
+    },
+    apex_available: [
+        "com.android.compos",
+    ],
+    visibility: [
+        "//packages/modules/Virtualization/compos:__subpackages__",
+        "//art/odrefresh:__subpackages__",
+    ],
+}
diff --git a/compos/libcompos_client/include/libcompos_client.h b/compos/libcompos_client/include/libcompos_client.h
new file mode 100644
index 0000000..171854d
--- /dev/null
+++ b/compos/libcompos_client/include/libcompos_client.h
@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <stdint.h>
+#include <sys/cdefs.h>
+
+__BEGIN_DECLS
+
+/**
+ * Sends request encoded in a marshaled byte buffer to the Compilation OS service, which will
+ * execute the compiler with context encoded in the marshaled byte buffer.
+ *
+ * @param cid the VM's cid to send the request to.
+ * @param marshaled pointer to a marshaled byte buffer.
+ * @param size size of the marshaled byte buffer pointed by `marshaled`.
+ * @param ro_fds pointer to a int array of read-only file descriptor numbers.
+ * @param ro_fds_num size of the array pointed by `ro_fds`.
+ * @param rw_fds pointer to a int array of read-writable file descriptor numbers.
+ * @param rw_fds_num size of the array pointed by `rw_fds`.
+ * @return the exit code of the compiler.
+ *
+ * Available since API level 33.
+ */
+int AComposClient_Request(int cid, const uint8_t* marshaled, size_t size, const int* ro_fds,
+                          size_t ro_fds_num, const int* rw_fds, size_t rw_fds_num)
+        __INTRODUCED_IN(33);
+
+__END_DECLS
diff --git a/compos/libcompos_client/libcompos_client.cc b/compos/libcompos_client/libcompos_client.cc
new file mode 100644
index 0000000..147fcd0
--- /dev/null
+++ b/compos/libcompos_client/libcompos_client.cc
@@ -0,0 +1,209 @@
+/*
+ * 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.
+ */
+
+#include "libcompos_client.h"
+
+#include <android-base/logging.h>
+#include <android-base/strings.h>
+#include <android-base/unique_fd.h>
+#include <android/binder_auto_utils.h>
+#include <android/binder_manager.h>
+#include <binder/IInterface.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <binder_rpc_unstable.hpp>
+#include <memory>
+
+#include "aidl/android/system/composd/IIsolatedCompilationService.h"
+#include "aidl/com/android/compos/FdAnnotation.h"
+#include "aidl/com/android/compos/ICompOsService.h"
+
+using aidl::android::system::composd::IIsolatedCompilationService;
+using aidl::com::android::compos::FdAnnotation;
+using aidl::com::android::compos::ICompOsService;
+using android::base::Join;
+using android::base::Pipe;
+using android::base::unique_fd;
+
+namespace {
+
+constexpr unsigned int kCompsvcRpcPort = 6432;
+constexpr const char* kComposdServiceName = "android.system.composd";
+
+void ExecFdServer(const int* ro_fds, size_t ro_fds_num, const int* rw_fds, size_t rw_fds_num,
+                  unique_fd ready_fd) {
+    // Holder of C Strings, with enough memory reserved to avoid reallocation. Otherwise,
+    // `holder.rbegin()->c_str()` may become invalid.
+    std::vector<std::string> holder;
+    holder.reserve(ro_fds_num + rw_fds_num + 1 /* for --ready-fd */);
+
+    std::vector<char const*> args = {"/apex/com.android.virt/bin/fd_server"};
+    for (int i = 0; i < ro_fds_num; ++i) {
+        args.emplace_back("--ro-fds");
+        holder.emplace_back(std::to_string(*(ro_fds + i)));
+        args.emplace_back(holder.rbegin()->c_str());
+    }
+    for (int i = 0; i < rw_fds_num; ++i) {
+        args.emplace_back("--rw-fds");
+        holder.emplace_back(std::to_string(*(rw_fds + i)));
+        args.emplace_back(holder.rbegin()->c_str());
+    }
+    args.emplace_back("--ready-fd");
+    holder.emplace_back(std::to_string(ready_fd.get()));
+    args.emplace_back(holder.rbegin()->c_str());
+
+    LOG(DEBUG) << "Starting fd_server, args: " << Join(args, ' ');
+    args.emplace_back(nullptr);
+    if (execv(args[0], const_cast<char* const*>(args.data())) < 0) {
+        PLOG(ERROR) << "execv failed";
+    }
+}
+
+class FileSharingSession final {
+public:
+    static std::unique_ptr<FileSharingSession> Create(const int* ro_fds, size_t ro_fds_num,
+                                                      const int* rw_fds, size_t rw_fds_num) {
+        // Create pipe for receiving a ready ping from fd_server.
+        unique_fd pipe_read, pipe_write;
+        if (!Pipe(&pipe_read, &pipe_write, /* flags= */ 0)) {
+            PLOG(ERROR) << "Cannot create pipe";
+            return nullptr;
+        }
+
+        pid_t pid = fork();
+        if (pid < 0) {
+            PLOG(ERROR) << "fork error";
+            return nullptr;
+        } else if (pid > 0) {
+            pipe_write.reset();
+
+            // When fd_server is ready it closes its end of the pipe. And if it exits, the pipe is
+            // also closed. Either way this read will return 0 bytes at that point, and there's no
+            // point waiting any longer.
+            char c;
+            read(pipe_read.get(), &c, sizeof(c));
+
+            std::unique_ptr<FileSharingSession> session(new FileSharingSession(pid));
+            return session;
+        } else if (pid == 0) {
+            pipe_read.reset();
+            ExecFdServer(ro_fds, ro_fds_num, rw_fds, rw_fds_num, std::move(pipe_write));
+            exit(EXIT_FAILURE);
+        }
+        return nullptr;
+    }
+
+    ~FileSharingSession() {
+        if (kill(fd_server_pid_, SIGTERM) < 0) {
+            PLOG(ERROR) << "Cannot kill fd_server (pid " << std::to_string(fd_server_pid_)
+                        << ") with SIGTERM. Retry with SIGKILL.";
+            if (kill(fd_server_pid_, SIGKILL) < 0) {
+                PLOG(ERROR) << "Still cannot terminate with SIGKILL. Give up.";
+                // TODO: it may be the safest if we turn fd_server into a library to run in a
+                // thread.
+            }
+        }
+    }
+
+private:
+    explicit FileSharingSession(pid_t pid) : fd_server_pid_(pid) {}
+
+    pid_t fd_server_pid_;
+};
+
+int MakeRequestToVM(int cid, const uint8_t* marshaled, size_t size, const int* ro_fds,
+                    size_t ro_fds_num, const int* rw_fds, size_t rw_fds_num) {
+    ndk::SpAIBinder binder(RpcClient(cid, kCompsvcRpcPort));
+    std::shared_ptr<ICompOsService> service = ICompOsService::fromBinder(binder);
+    if (!service) {
+        LOG(ERROR) << "Cannot connect to the service";
+        return -1;
+    }
+
+    std::unique_ptr<FileSharingSession> session_raii =
+            FileSharingSession::Create(ro_fds, ro_fds_num, rw_fds, rw_fds_num);
+    if (!session_raii) {
+        LOG(ERROR) << "Cannot start to share FDs";
+        return -1;
+    }
+
+    // Since the input from the C API are raw pointers, we need to duplicate them into vectors in
+    // order to pass to the binder API.
+    std::vector<uint8_t> duplicated_buffer(marshaled, marshaled + size);
+    FdAnnotation fd_annotation = {
+            .input_fds = std::vector<int>(ro_fds, ro_fds + ro_fds_num),
+            .output_fds = std::vector<int>(rw_fds, rw_fds + rw_fds_num),
+    };
+    int8_t exit_code;
+    ndk::ScopedAStatus status = service->compile(duplicated_buffer, fd_annotation, &exit_code);
+    if (!status.isOk()) {
+        LOG(ERROR) << "Compilation failed (exit " << std::to_string(exit_code)
+                   << "): " << status.getDescription();
+        return -1;
+    }
+    return 0;
+}
+
+int MakeRequestToComposd(const uint8_t* marshaled, size_t size, const int* ro_fds,
+                         size_t ro_fds_num, const int* rw_fds, size_t rw_fds_num) {
+    ndk::SpAIBinder binder(AServiceManager_getService(kComposdServiceName));
+    std::shared_ptr<IIsolatedCompilationService> service =
+            IIsolatedCompilationService::fromBinder(binder);
+    if (!service) {
+        LOG(ERROR) << "Cannot connect to the service";
+        return -1;
+    }
+
+    auto session_raii = std::unique_ptr<FileSharingSession>(
+            FileSharingSession::Create(ro_fds, ro_fds_num, rw_fds, rw_fds_num));
+    if (!session_raii) {
+        LOG(ERROR) << "Cannot start to share FDs";
+        return -1;
+    }
+
+    // Since the input from the C API are raw pointers, we need to duplicate them into vectors in
+    // order to pass to the binder API.
+    std::vector<uint8_t> duplicated_buffer(marshaled, marshaled + size);
+    FdAnnotation fd_annotation = {
+            .input_fds = std::vector<int>(ro_fds, ro_fds + ro_fds_num),
+            .output_fds = std::vector<int>(rw_fds, rw_fds + rw_fds_num),
+    };
+    int8_t exit_code;
+    ndk::ScopedAStatus status = service->compile(duplicated_buffer, fd_annotation, &exit_code);
+    if (!status.isOk()) {
+        LOG(ERROR) << "Compilation failed (exit " << std::to_string(exit_code)
+                   << "): " << status.getDescription();
+        return -1;
+    }
+    return 0;
+}
+
+} // namespace
+
+__BEGIN_DECLS
+
+int AComposClient_Request(int cid, const uint8_t* marshaled, size_t size, const int* ro_fds,
+                          size_t ro_fds_num, const int* rw_fds, size_t rw_fds_num) {
+    if (cid == -1 /* VMADDR_CID_ANY */) {
+        return MakeRequestToComposd(marshaled, size, ro_fds, ro_fds_num, rw_fds, rw_fds_num);
+    } else {
+        return MakeRequestToVM(cid, marshaled, size, ro_fds, ro_fds_num, rw_fds, rw_fds_num);
+    }
+}
+
+__END_DECLS
diff --git a/compos/libcompos_client/libcompos_client.map.txt b/compos/libcompos_client/libcompos_client.map.txt
new file mode 100644
index 0000000..9d47c53
--- /dev/null
+++ b/compos/libcompos_client/libcompos_client.map.txt
@@ -0,0 +1,6 @@
+LIBCOMPOS_CLIENT {
+    global:
+        AComposClient_Request; # apex
+    local:
+        *;
+};
diff --git a/compos/src/compilation.rs b/compos/src/compilation.rs
index fec82a6..1499d4b 100644
--- a/compos/src/compilation.rs
+++ b/compos/src/compilation.rs
@@ -52,7 +52,7 @@
 
 /// Runs the compiler with given flags with file descriptors described in `fd_annotation` retrieved
 /// via `authfs_service`. Returns exit code of the compiler process.
-pub fn compile(
+pub fn compile_cmd(
     compiler_path: &Path,
     compiler_args: &[String],
     authfs_service: Strong<dyn IAuthFsService>,
diff --git a/compos/src/compsvc.rs b/compos/src/compsvc.rs
index 954adf5..08f3521 100644
--- a/compos/src/compsvc.rs
+++ b/compos/src/compsvc.rs
@@ -25,7 +25,7 @@
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 
-use crate::compilation::{compile, CompilerOutput};
+use crate::compilation::{compile_cmd, CompilerOutput};
 use crate::compos_key_service::CompOsKeyService;
 use crate::fsverity;
 use authfs_aidl_interface::aidl::com::android::virt::fs::IAuthFsService::IAuthFsService;
@@ -85,14 +85,14 @@
         }
     }
 
-    fn compile(
+    fn compile_cmd(
         &self,
         args: &[String],
         fd_annotation: &FdAnnotation,
     ) -> BinderResult<CompilationResult> {
         let authfs_service = get_authfs_service()?;
         let output =
-            compile(&self.dex2oat_path, args, authfs_service, fd_annotation).map_err(|e| {
+            compile_cmd(&self.dex2oat_path, args, authfs_service, fd_annotation).map_err(|e| {
                 new_binder_exception(
                     ExceptionCode::SERVICE_SPECIFIC,
                     format!("Compilation failed: {}", e),
@@ -124,6 +124,10 @@
         }
     }
 
+    fn compile(&self, _marshaled: &[u8], _fd_annotation: &FdAnnotation) -> BinderResult<i8> {
+        Err(new_binder_exception(ExceptionCode::UNSUPPORTED_OPERATION, "Not yet implemented"))
+    }
+
     fn generateSigningKey(&self) -> BinderResult<CompOsKeyData> {
         self.key_service
             .do_generate()
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 388e79b..6887947 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -30,20 +30,27 @@
 };
 use anyhow::{anyhow, bail, Context, Result};
 use binder::{
-    unstable_api::{new_spibinder, AIBinder, AsNative},
+    unstable_api::{new_spibinder, AIBinder},
     FromIBinder,
 };
+use binder_common::rpc_server::run_rpc_server;
 use compos_common::COMPOS_VSOCK_PORT;
 use log::{debug, error};
 use nix::ioctl_read_bad;
 use std::fs::OpenOptions;
-use std::os::raw;
 use std::os::unix::io::AsRawFd;
 
 /// The CID representing the host VM
 const VMADDR_CID_HOST: u32 = 2;
 
-fn main() -> Result<()> {
+fn main() {
+    if let Err(e) = try_main() {
+        error!("failed with {:?}", e);
+        std::process::exit(1);
+    }
+}
+
+fn try_main() -> Result<()> {
     let args = clap::App::new("compsvc")
         .arg(clap::Arg::with_name("log_to_stderr").long("log_to_stderr"))
         .get_matches();
@@ -55,23 +62,17 @@
         );
     }
 
-    let mut service = compsvc::new_binder()?.as_binder();
+    let service = compsvc::new_binder()?.as_binder();
+    let vm_service = get_vm_service()?;
+    let local_cid = get_local_cid()?;
+
     debug!("compsvc is starting as a rpc service.");
 
-    let mut ready_notifier = ReadyNotifier::new()?;
-
-    // SAFETY: Service ownership is transferring to the server and won't be valid afterward.
-    // Plus the binder objects are threadsafe.
-    // RunRpcServerCallback does not retain a reference to ready_callback, and only ever
-    // calls it with the param we provide during the lifetime of ready_notifier.
-    let retval = unsafe {
-        binder_rpc_unstable_bindgen::RunRpcServerCallback(
-            service.as_native_mut() as *mut binder_rpc_unstable_bindgen::AIBinder,
-            COMPOS_VSOCK_PORT,
-            Some(ReadyNotifier::ready_callback),
-            ready_notifier.as_void_ptr(),
-        )
-    };
+    let retval = run_rpc_server(service, COMPOS_VSOCK_PORT, || {
+        if let Err(e) = vm_service.notifyPayloadReady(local_cid as i32) {
+            error!("Unable to notify ready: {}", e);
+        }
+    });
     if retval {
         debug!("RPC server has shut down gracefully");
         Ok(())
@@ -80,61 +81,32 @@
     }
 }
 
-struct ReadyNotifier {
-    vm_service: Strong<dyn IVirtualMachineService>,
-    local_cid: u32,
+fn get_vm_service() -> Result<Strong<dyn IVirtualMachineService>> {
+    // SAFETY: AIBinder returned by RpcClient has correct reference count, and the ownership
+    // can be safely taken by new_spibinder.
+    let ibinder = unsafe {
+        new_spibinder(binder_rpc_unstable_bindgen::RpcClient(
+            VMADDR_CID_HOST,
+            VM_BINDER_SERVICE_PORT as u32,
+        ) as *mut AIBinder)
+    }
+    .ok_or_else(|| anyhow!("Failed to connect to IVirtualMachineService"))?;
+
+    FromIBinder::try_from(ibinder).context("Connecting to IVirtualMachineService")
 }
 
-impl ReadyNotifier {
-    fn new() -> Result<Self> {
-        Ok(Self { vm_service: Self::get_vm_service()?, local_cid: Self::get_local_cid()? })
-    }
-
-    fn notify(&self) {
-        if let Err(e) = self.vm_service.notifyPayloadReady(self.local_cid as i32) {
-            error!("Unable to notify ready: {}", e);
-        }
-    }
-
-    fn as_void_ptr(&mut self) -> *mut raw::c_void {
-        self as *mut _ as *mut raw::c_void
-    }
-
-    unsafe extern "C" fn ready_callback(param: *mut raw::c_void) {
-        // SAFETY: This is only ever called by RunRpcServerCallback, within the lifetime of the
-        // ReadyNotifier, with param taking the value returned by as_void_ptr (so a properly aligned
-        // non-null pointer to an initialized instance).
-        let ready_notifier = param as *mut Self;
-        ready_notifier.as_ref().unwrap().notify()
-    }
-
-    fn get_vm_service() -> Result<Strong<dyn IVirtualMachineService>> {
-        // SAFETY: AIBinder returned by RpcClient has correct reference count, and the ownership
-        // can be safely taken by new_spibinder.
-        let ibinder = unsafe {
-            new_spibinder(binder_rpc_unstable_bindgen::RpcClient(
-                VMADDR_CID_HOST,
-                VM_BINDER_SERVICE_PORT as u32,
-            ) as *mut AIBinder)
-        }
-        .ok_or_else(|| anyhow!("Failed to connect to IVirtualMachineService"))?;
-
-        FromIBinder::try_from(ibinder).context("Connecting to IVirtualMachineService")
-    }
-
-    // TODO(b/199259751): remove this after VS can check the peer addresses of binder clients
-    fn get_local_cid() -> Result<u32> {
-        let f = OpenOptions::new()
-            .read(true)
-            .write(false)
-            .open("/dev/vsock")
-            .context("Failed to open /dev/vsock")?;
-        let mut cid = 0;
-        // SAFETY: the kernel only modifies the given u32 integer.
-        unsafe { vm_sockets_get_local_cid(f.as_raw_fd(), &mut cid) }
-            .context("Failed to get local CID")?;
-        Ok(cid)
-    }
+// TODO(b/199259751): remove this after VS can check the peer addresses of binder clients
+fn get_local_cid() -> Result<u32> {
+    let f = OpenOptions::new()
+        .read(true)
+        .write(false)
+        .open("/dev/vsock")
+        .context("Failed to open /dev/vsock")?;
+    let mut cid = 0;
+    // SAFETY: the kernel only modifies the given u32 integer.
+    unsafe { vm_sockets_get_local_cid(f.as_raw_fd(), &mut cid) }
+        .context("Failed to get local CID")?;
+    Ok(cid)
 }
 
 // TODO(b/199259751): remove this after VS can check the peer addresses of binder clients
diff --git a/compos/src/pvm_exec.rs b/compos/src/pvm_exec.rs
index fdd9c57..cae0702 100644
--- a/compos/src/pvm_exec.rs
+++ b/compos/src/pvm_exec.rs
@@ -32,8 +32,11 @@
 use clap::{value_t, App, Arg};
 use log::{debug, error, warn};
 use minijail::Minijail;
-use nix::fcntl::{fcntl, FcntlArg::F_GETFD};
-use std::os::unix::io::RawFd;
+use nix::fcntl::{fcntl, FcntlArg::F_GETFD, OFlag};
+use nix::unistd::pipe2;
+use std::fs::File;
+use std::io::Read;
+use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::path::Path;
 use std::process::exit;
 
@@ -69,7 +72,11 @@
     }
 }
 
-fn spawn_fd_server(fd_annotation: &FdAnnotation, debuggable: bool) -> Result<Minijail> {
+fn spawn_fd_server(
+    fd_annotation: &FdAnnotation,
+    ready_file: File,
+    debuggable: bool,
+) -> Result<Minijail> {
     let mut inheritable_fds = if debuggable {
         vec![1, 2] // inherit/redirect stdout/stderr for debugging
     } else {
@@ -87,6 +94,10 @@
         args.push(fd.to_string());
         inheritable_fds.push(*fd);
     }
+    let ready_fd = ready_file.as_raw_fd();
+    args.push("--ready-fd".to_string());
+    args.push(ready_fd.to_string());
+    inheritable_fds.push(ready_fd);
 
     let jail = Minijail::new()?;
     let _pid = jail.run(Path::new(FD_SERVER_BIN), &inheritable_fds, &args)?;
@@ -147,18 +158,37 @@
     let output_fds = results?;
 
     let args: Vec<_> = matches.values_of("args").unwrap().map(|s| s.to_string()).collect();
-    let cid = value_t!(matches, "cid", u32)?;
+    let cid = value_t!(matches, "cid", i32)? as u32;
     let debuggable = matches.is_present("debug");
 
     Ok(Config { args, fd_annotation: FdAnnotation { input_fds, output_fds }, cid, debuggable })
 }
 
+fn create_pipe() -> Result<(File, File)> {
+    let (raw_read, raw_write) = pipe2(OFlag::O_CLOEXEC)?;
+    // SAFETY: We are the sole owners of these fds as they were just created.
+    let read_fd = unsafe { File::from_raw_fd(raw_read) };
+    let write_fd = unsafe { File::from_raw_fd(raw_write) };
+    Ok((read_fd, write_fd))
+}
+
+fn wait_for_fd_server_ready(mut ready_fd: File) -> Result<()> {
+    let mut buffer = [0];
+    // When fd_server is ready it closes its end of the pipe. And if it exits, the pipe is also
+    // closed. Either way this read will return 0 bytes at that point, and there's no point waiting
+    // any longer.
+    let _ = ready_fd.read(&mut buffer).context("Waiting for fd_server to be ready")?;
+    debug!("fd_server is ready");
+    Ok(())
+}
+
 fn try_main() -> Result<()> {
     // 1. Parse the command line arguments for collect execution data.
     let Config { args, fd_annotation, cid, debuggable } = parse_args()?;
 
     // 2. Spawn and configure a fd_server to serve remote read/write requests.
-    let fd_server_jail = spawn_fd_server(&fd_annotation, debuggable)?;
+    let (ready_read_fd, ready_write_fd) = create_pipe()?;
+    let fd_server_jail = spawn_fd_server(&fd_annotation, ready_write_fd, debuggable)?;
     let fd_server_lifetime = scopeguard::guard(fd_server_jail, |fd_server_jail| {
         if let Err(e) = fd_server_jail.kill() {
             if !matches!(e, minijail::Error::Killed(_)) {
@@ -171,11 +201,13 @@
     let result = if cid == VMADDR_CID_ANY {
         // Sentinel value that indicates we should use composd
         let composd = get_composd()?;
-        composd.compile(&args, &fd_annotation)
+        wait_for_fd_server_ready(ready_read_fd)?;
+        composd.compile_cmd(&args, &fd_annotation)
     } else {
         // Call directly into the VM
         let compos_vm = get_rpc_binder(cid)?;
-        compos_vm.compile(&args, &fd_annotation)
+        wait_for_fd_server_ready(ready_read_fd)?;
+        compos_vm.compile_cmd(&args, &fd_annotation)
     };
     let result = result.context("Binder call failed")?;
 
@@ -190,7 +222,7 @@
     // Be explicit about the lifetime, which should last at least until the task is finished.
     drop(fd_server_lifetime);
 
-    if result.exitCode > 0 {
+    if result.exitCode != 0 {
         error!("remote execution failed with exit code {}", result.exitCode);
         exit(result.exitCode as i32);
     }
diff --git a/compos/tests/AndroidTest.xml b/compos/tests/AndroidTest.xml
index 61b6d47..940531b 100644
--- a/compos/tests/AndroidTest.xml
+++ b/compos/tests/AndroidTest.xml
@@ -18,10 +18,6 @@
         <option name="force-root" value="true" />
     </target_preparer>
 
-    <!-- virtualizationservice doesn't have access to shell_data_file. Instead of giving it
-      a test-only permission, run it without selinux -->
-    <target_preparer class="com.android.tradefed.targetprep.DisableSELinuxTargetPreparer"/>
-
     <test class="com.android.compatibility.common.tradefed.testtype.JarHostTest" >
         <option name="jar" value="ComposHostTestCases.jar" />
     </test>
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 40f95c3..01d8ff1 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -22,8 +22,6 @@
 import android.virt.test.CommandRunner;
 import android.virt.test.VirtualizationTestCaseBase;
 
-import com.android.compatibility.common.util.PollingCheck;
-import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandResult;
@@ -37,11 +35,10 @@
 @RunWith(DeviceJUnit4ClassRunner.class)
 public final class ComposTestCase extends VirtualizationTestCaseBase {
 
-    /** Path to odrefresh on Microdroid */
+    // Binaries used in test. (These paths are valid both in host and Microdroid.)
     private static final String ODREFRESH_BIN = "/apex/com.android.art/bin/odrefresh";
-
-    /** Path to compos_key_cmd on Microdroid */
     private static final String COMPOS_KEY_CMD_BIN = "/apex/com.android.compos/bin/compos_key_cmd";
+    private static final String COMPOSD_CMD_BIN = "/apex/com.android.compos/bin/composd_cmd";
 
     /** Output directory of odrefresh */
     private static final String ODREFRESH_OUTPUT_DIR =
@@ -50,38 +47,29 @@
     /** Timeout of odrefresh to finish */
     private static final int ODREFRESH_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes
 
-    /** Wait time for compsvc to be ready on boot */
-    private static final int COMPSVC_READY_LATENCY_MS = 10 * 1000; // 10 seconds
-
     // ExitCode expanded from art/odrefresh/include/odrefresh/odrefresh.h.
     private static final int OKAY = 0;
     private static final int COMPILATION_SUCCESS = 80;
 
-    private String mCid;
+    // Files that define the "current" instance of CompOS
+    private static final String COMPOS_CURRENT_ROOT =
+            "/data/misc/apexdata/com.android.compos/current/";
+    private static final String INSTANCE_IMAGE = COMPOS_CURRENT_ROOT + "instance.img";
+    private static final String PUBLIC_KEY = COMPOS_CURRENT_ROOT + "key.pubkey";
+    private static final String PRIVATE_KEY_BLOB = COMPOS_CURRENT_ROOT + "key.blob";
 
     @Before
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
-
-        prepareVirtualizationTestSetup(getDevice());
-
-        startComposVm();
     }
 
     @After
     public void tearDown() throws Exception {
-        if (mCid != null) {
-            shutdownMicrodroid(getDevice(), mCid);
-            mCid = null;
-        }
-
-        cleanUpVirtualizationTestSetup(getDevice());
+        killVmAndReconnectAdb();
     }
 
     @Test
     public void testOdrefresh() throws Exception {
-        waitForServiceRunning();
-
         CommandRunner android = new CommandRunner(getDevice());
 
         // Prepare the groundtruth. The compilation on Android should finish successfully.
@@ -103,28 +91,18 @@
                 android.runForResultWithTimeout(ODREFRESH_TIMEOUT_MS, ODREFRESH_BIN, "--check");
         assertThat(result.getExitCode()).isEqualTo(OKAY);
 
-        // Initialize the service with the generated key. Should succeed.
-        android.run(
-                COMPOS_KEY_CMD_BIN,
-                "--cid " + mCid,
-                "generate",
-                TEST_ROOT + "test_key.blob",
-                TEST_ROOT + "test_key.pubkey");
-        android.run(COMPOS_KEY_CMD_BIN, "--cid " + mCid, "init-key", TEST_ROOT + "test_key.blob");
+        // Make sure we generate a fresh instance
+        android.tryRun("rm", "-rf", COMPOS_CURRENT_ROOT);
 
         // Expect the compilation in Compilation OS to finish successfully.
         {
             long start = System.currentTimeMillis();
-            result =
-                    android.runForResultWithTimeout(
-                            ODREFRESH_TIMEOUT_MS,
-                            ODREFRESH_BIN,
-                            "--use-compilation-os=" + mCid,
-                            "--force-compile");
+            result = android.runForResultWithTimeout(ODREFRESH_TIMEOUT_MS, COMPOSD_CMD_BIN);
             long elapsed = System.currentTimeMillis() - start;
-            assertThat(result.getExitCode()).isEqualTo(COMPILATION_SUCCESS);
+            assertThat(result.getExitCode()).isEqualTo(0);
             CLog.i("Comp OS compilation took " + elapsed + "ms");
         }
+        killVmAndReconnectAdb();
 
         // Save the actual checksum for the output directory.
         String actualChecksumSnapshot = checksumDirectoryContent(android, ODREFRESH_OUTPUT_DIR);
@@ -140,35 +118,27 @@
         assertThat(actualChecksumSnapshot).isEqualTo(expectedChecksumSnapshot);
     }
 
-    private void startComposVm() throws DeviceNotAvailableException {
-        final String apkName = "CompOSPayloadApp.apk";
-        final String packageName = "com.android.compos.payload";
-        mCid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        apkName,
-                        packageName,
-                        "assets/vm_test_config.json",
-                        /* debug */ false,
-                        /* Use default memory */ 0);
-        adbConnectToMicrodroid(getDevice(), mCid);
+    private void killVmAndReconnectAdb() throws Exception {
+        CommandRunner android = new CommandRunner(getDevice());
+
+        // When a VM exits, we tend to see adb disconnecting. So we attempt to reconnect
+        // when we kill it to avoid problems. Of course VirtualizationService may exit anyway
+        // (it's an on-demand service and all its clients have gone), taking the VM with it,
+        // which makes this a bit unpredictable.
+        reconnectHostAdb(getDevice());
+        android.tryRun("killall", "crosvm");
+        reconnectHostAdb(getDevice());
+        android.tryRun("stop", "virtualizationservice");
+        reconnectHostAdb(getDevice());
+
+        // Delete stale data
+        android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
     }
 
-    private void waitForServiceRunning() {
-        try {
-            PollingCheck.waitFor(COMPSVC_READY_LATENCY_MS, this::isServiceRunning);
-        } catch (Exception e) {
-            throw new RuntimeException("Service unavailable", e);
-        }
-    }
-
-    private boolean isServiceRunning() {
-        return tryRunOnMicrodroid("pidof compsvc") != null;
-    }
-
-    private String checksumDirectoryContent(CommandRunner runner, String path)
-            throws DeviceNotAvailableException {
-        return runner.run("find " + path + " -type f -exec sha256sum {} \\; | sort");
+    private String checksumDirectoryContent(CommandRunner runner, String path) throws Exception {
+        // Sort by filename (second column) to make comparison easier.
+        // TODO(b/192690283): Figure out how to make this work for files odex/oat/art files.
+        return runner.run(
+                "find " + path + " -type f -exec sha256sum {} \\; | grep vdex | sort -k2");
     }
 }
diff --git a/compos/verify_key/verify_key.rs b/compos/verify_key/verify_key.rs
index 8439b97..0cc6473 100644
--- a/compos/verify_key/verify_key.rs
+++ b/compos/verify_key/verify_key.rs
@@ -87,8 +87,9 @@
 
     let blob = read_small_file(blob).context("Failed to read key blob")?;
     let public_key = read_small_file(public_key).context("Failed to read public key")?;
+    let instance_image = File::open(instance_image).context("Failed to open instance image")?;
 
-    let vm_instance = VmInstance::start(&instance_image)?;
+    let vm_instance = VmInstance::start(instance_image)?;
     let service = vm_instance.get_service()?;
 
     let result = service.verifySigningKey(&blob, &public_key).context("Verifying signing key")?;
diff --git a/microdroid/payload/config/src/lib.rs b/microdroid/payload/config/src/lib.rs
index 6dc127b..2547f3d 100644
--- a/microdroid/payload/config/src/lib.rs
+++ b/microdroid/payload/config/src/lib.rs
@@ -30,6 +30,10 @@
     /// APEXes to activate in a VM
     #[serde(default)]
     pub apexes: Vec<ApexConfig>,
+
+    /// Tells VirtualizationService to use staged APEXes if possible
+    #[serde(default)]
+    pub prefer_staged: bool,
 }
 
 /// OS config
diff --git a/pvmfw/pvmfw.img b/pvmfw/pvmfw.img
new file mode 100644
index 0000000..317821f
--- /dev/null
+++ b/pvmfw/pvmfw.img
Binary files differ
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index 24a955b..d62892c 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -73,18 +73,26 @@
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
 
-        // Make sure we're connected to the host adb again (b/194219111)
+        reconnectHostAdb(androidDevice);
+
+        // kill stale VMs and directories
+        android.tryRun("killall", "crosvm");
+        android.tryRun("stop", "virtualizationservice");
+        android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
+    }
+
+    public static void reconnectHostAdb(ITestDevice androidDevice)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(androidDevice);
+
+        // Make sure we're connected to the host adb; this connection seems to get dropped when a VM
+        // exits.
         for (int retry = 0; retry < 3; ++retry) {
             if (android.tryRun("true") != null) {
                 break;
             }
             androidDevice.waitForDeviceOnline(1000);
         }
-
-        // kill stale VMs and directories
-        android.tryRun("killall", "crosvm");
-        android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
-        android.tryRun("stop", "virtualizationservice");
     }
 
     public static void testIfDeviceIsCapable(ITestDevice androidDevice)
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 443436d..18d8ade 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -43,6 +43,8 @@
         "libvmconfig",
         "libzip",
         "libvsock",
+        // TODO(b/202115393) stabilize the interface
+        "packagemanager_aidl-rust",
     ],
     shared_libs: [
         "libbinder_rpc_unstable",
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 76c3a16..dbcc5ce 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -19,13 +19,11 @@
 use crate::payload::add_microdroid_images;
 use crate::{Cid, FIRST_GUEST_CID, SYSPROP_LAST_CID};
 
-use binder_common::new_binder_exception;
+use ::binder::unstable_api::AsNative;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachine::{
-    BnVirtualMachine, IVirtualMachine,
-};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     DiskImage::DiskImage,
+    IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
     IVirtualMachineCallback::IVirtualMachineCallback,
     IVirtualizationService::IVirtualizationService,
     PartitionType::PartitionType,
@@ -36,20 +34,24 @@
     VirtualMachineState::VirtualMachineState,
 };
 use android_system_virtualizationservice::binder::{
-    self, force_lazy_services_persist, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong, ThreadState,
+    self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong,
+    ThreadState,
 };
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
-    VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT, BnVirtualMachineService, IVirtualMachineService,
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
+    IVirtualMachineService::{
+        BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
+        VM_STREAM_SERVICE_PORT,
+    },
 };
 use anyhow::{anyhow, bail, Context, Result};
-use ::binder::unstable_api::AsNative;
+use binder_common::{lazy_service::LazyServiceGuard, new_binder_exception};
 use disk::QcowFile;
-use idsig::{V4Signature, HashAlgorithm};
-use log::{debug, error, warn, info};
+use idsig::{HashAlgorithm, V4Signature};
+use log::{debug, error, info, warn};
 use microdroid_payload_config::VmPayloadConfig;
 use rustutils::system_properties;
 use std::convert::TryInto;
-use std::fs::{File, OpenOptions, create_dir};
+use std::fs::{create_dir, File, OpenOptions};
 use std::io::{Error, ErrorKind, Write};
 use std::num::NonZeroU32;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
@@ -462,14 +464,13 @@
 
     // Microdroid requires an additional payload disk image and the bootconfig partition.
     if os_name == "microdroid" {
-        let apexes = vm_payload_config.apexes.clone();
         add_microdroid_images(
             config,
             temporary_directory,
             apk_file,
             idsig_file,
             instance_file,
-            apexes,
+            &vm_payload_config,
             &mut vm_config,
         )?;
     }
@@ -557,11 +558,13 @@
 #[derive(Debug)]
 struct VirtualMachine {
     instance: Arc<VmInstance>,
+    /// Keeps our service process running as long as this VM instance exists.
+    lazy_service_guard: LazyServiceGuard,
 }
 
 impl VirtualMachine {
     fn create(instance: Arc<VmInstance>) -> Strong<dyn IVirtualMachine> {
-        let binder = VirtualMachine { instance };
+        let binder = VirtualMachine { instance, lazy_service_guard: Default::default() };
         BnVirtualMachine::new_binder(binder, BinderFeatures::default())
     }
 }
@@ -714,19 +717,12 @@
     /// Store a strong VM reference.
     fn debug_hold_vm(&mut self, vm: Strong<dyn IVirtualMachine>) {
         self.debug_held_vms.push(vm);
-        // Make sure our process is not shut down while we hold the VM reference
-        // on behalf of the caller.
-        force_lazy_services_persist(true);
     }
 
     /// Retrieve and remove a strong VM reference.
     fn debug_drop_vm(&mut self, cid: i32) -> Option<Strong<dyn IVirtualMachine>> {
         let pos = self.debug_held_vms.iter().position(|vm| vm.getCid() == Ok(cid))?;
         let vm = self.debug_held_vms.swap_remove(pos);
-        if self.debug_held_vms.is_empty() {
-            // Once we no longer hold any VM references it is ok for our process to be shut down.
-            force_lazy_services_persist(false);
-        }
         Some(vm)
     }
 }
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 9662fa3..3520d9f 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -20,11 +20,15 @@
 };
 use android_system_virtualizationservice::binder::ParcelFileDescriptor;
 use anyhow::{anyhow, Context, Result};
+use binder::{wait_for_interface, Strong};
+use log::{error, info};
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata};
-use microdroid_payload_config::ApexConfig;
+use microdroid_payload_config::{ApexConfig, VmPayloadConfig};
 use once_cell::sync::OnceCell;
+use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
 use serde::Deserialize;
 use serde_xml_rs::from_reader;
+use std::env;
 use std::fs::{File, OpenOptions};
 use std::path::{Path, PathBuf};
 use vmconfig::open_parcel_file;
@@ -35,6 +39,8 @@
 
 const APEX_INFO_LIST_PATH: &str = "/apex/apex-info-list.xml";
 
+const PACKAGE_MANAGER_NATIVE_SERVICE: &str = "package_native";
+
 /// Represents the list of APEXes
 #[derive(Debug, Deserialize)]
 struct ApexInfoList {
@@ -74,19 +80,45 @@
     }
 }
 
+struct PackageManager {
+    service: Strong<dyn IPackageManagerNative>,
+    // TODO(b/199146189) use IPackageManagerNative
+    apex_info_list: &'static ApexInfoList,
+}
+
+impl PackageManager {
+    fn new() -> Result<Self> {
+        let service = wait_for_interface(PACKAGE_MANAGER_NATIVE_SERVICE)
+            .context("Failed to find PackageManager")?;
+        let apex_info_list = ApexInfoList::load()?;
+        Ok(Self { service, apex_info_list })
+    }
+
+    fn get_apex_path(&self, name: &str, prefer_staged: bool) -> Result<PathBuf> {
+        if prefer_staged {
+            let apex_info = self.service.getStagedApexInfo(name)?;
+            if let Some(apex_info) = apex_info {
+                info!("prefer_staged: use {} for {}", apex_info.diskImagePath, name);
+                return Ok(PathBuf::from(apex_info.diskImagePath));
+            }
+        }
+        self.apex_info_list.get_path_for(name)
+    }
+}
+
 fn make_metadata_file(
     config_path: &str,
-    apexes: &[ApexConfig],
+    apex_names: &[String],
     temporary_directory: &Path,
 ) -> Result<ParcelFileDescriptor> {
     let metadata_path = temporary_directory.join("metadata");
     let metadata = Metadata {
         version: 1,
-        apexes: apexes
+        apexes: apex_names
             .iter()
             .enumerate()
-            .map(|(i, apex)| ApexPayload {
-                name: apex.name.clone(),
+            .map(|(i, apex_name)| ApexPayload {
+                name: apex_name.clone(),
                 partition_name: format!("microdroid-apex-{}", i),
                 ..Default::default()
             })
@@ -126,7 +158,8 @@
     apk_file: File,
     idsig_file: File,
     config_path: &str,
-    apexes: &[ApexConfig],
+    apexes: &[String],
+    prefer_staged: bool,
     temporary_directory: &Path,
 ) -> Result<DiskImage> {
     let metadata_file = make_metadata_file(config_path, apexes, temporary_directory)?;
@@ -137,9 +170,9 @@
         writable: false,
     }];
 
-    let apex_info_list = ApexInfoList::load()?;
+    let pm = PackageManager::new()?;
     for (i, apex) in apexes.iter().enumerate() {
-        let apex_path = apex_info_list.get_path_for(&apex.name)?;
+        let apex_path = pm.get_apex_path(apex, prefer_staged)?;
         let apex_file = open_parcel_file(&apex_path, false)?;
         partitions.push(Partition {
             label: format!("microdroid-apex-{}", i),
@@ -161,25 +194,65 @@
     Ok(DiskImage { image: None, partitions, writable: false })
 }
 
+fn find_apex_names_in_classpath_env(classpath_env_var: &str) -> Vec<String> {
+    let val = env::var(classpath_env_var).unwrap_or_else(|e| {
+        error!("Reading {} failed: {}", classpath_env_var, e);
+        String::from("")
+    });
+    val.split(':')
+        .filter_map(|path| {
+            Path::new(path)
+                .strip_prefix("/apex/")
+                .map(|stripped| {
+                    let first = stripped.iter().next().unwrap();
+                    first.to_str().unwrap().to_string()
+                })
+                .ok()
+        })
+        .collect()
+}
+
+// Collect APEX names from config
+fn collect_apex_names(apexes: &[ApexConfig]) -> Vec<String> {
+    // Process pseudo names like "{BOOTCLASSPATH}".
+    // For now we have following pseudo APEX names:
+    // - {BOOTCLASSPATH}: represents APEXes contributing "BOOTCLASSPATH" environment variable
+    // - {DEX2OATBOOTCLASSPATH}: represents APEXes contributing "DEX2OATBOOTCLASSPATH" environment variable
+    // - {SYSTEMSERVERCLASSPATH}: represents APEXes contributing "SYSTEMSERVERCLASSPATH" environment variable
+    let mut apex_names: Vec<String> = apexes
+        .iter()
+        .flat_map(|apex| match apex.name.as_str() {
+            "{BOOTCLASSPATH}" => find_apex_names_in_classpath_env("BOOTCLASSPATH"),
+            "{DEX2OATBOOTCLASSPATH}" => find_apex_names_in_classpath_env("DEX2OATBOOTCLASSPATH"),
+            "{SYSTEMSERVERCLASSPATH}" => find_apex_names_in_classpath_env("SYSTEMSERVERCLASSPATH"),
+            _ => vec![apex.name.clone()],
+        })
+        .collect();
+    // Add required APEXes
+    apex_names.extend(MICRODROID_REQUIRED_APEXES.iter().map(|name| name.to_string()));
+    apex_names.sort();
+    apex_names.dedup();
+    apex_names
+}
+
 pub fn add_microdroid_images(
     config: &VirtualMachineAppConfig,
     temporary_directory: &Path,
     apk_file: File,
     idsig_file: File,
     instance_file: File,
-    mut apexes: Vec<ApexConfig>,
+    vm_payload_config: &VmPayloadConfig,
     vm_config: &mut VirtualMachineRawConfig,
 ) -> Result<()> {
-    apexes.extend(
-        MICRODROID_REQUIRED_APEXES.iter().map(|name| ApexConfig { name: name.to_string() }),
-    );
-    apexes.dedup_by(|a, b| a.name == b.name);
-
+    // collect APEX names from config
+    let apexes = collect_apex_names(&vm_payload_config.apexes);
+    info!("Microdroid payload APEXes: {:?}", apexes);
     vm_config.disks.push(make_payload_disk(
         apk_file,
         idsig_file,
         &config.configPath,
         &apexes,
+        vm_payload_config.prefer_staged,
         temporary_directory,
     )?);
 
@@ -203,3 +276,18 @@
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_find_apex_names_in_classpath_env() {
+        let key = "TEST_BOOTCLASSPATH";
+        let classpath = "/apex/com.android.foo/javalib/foo.jar:/system/framework/framework.jar:/apex/com.android.bar/javalib/bar.jar";
+        env::set_var(key, classpath);
+        assert_eq!(
+            find_apex_names_in_classpath_env(key),
+            vec!["com.android.foo".to_owned(), "com.android.bar".to_owned()]
+        );
+    }
+}