Merge "Remove the i18n APEX from microdroid"
diff --git a/apex/Android.bp b/apex/Android.bp
index c568ae2..b1c5190 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -81,4 +81,5 @@
     name: "com.android.virt.init.rc",
     src: "virtualizationservice.rc",
     filename: "init.rc",
+    installable: false,
 }
diff --git a/apkdmverity/src/main.rs b/apkdmverity/src/main.rs
index cabeb35..b240c85 100644
--- a/apkdmverity/src/main.rs
+++ b/apkdmverity/src/main.rs
@@ -240,8 +240,7 @@
         }
 
         run_test(modified_apk.as_slice(), idsig.as_ref(), "incorrect_apk", |ctx| {
-            let ret = fs::read(&ctx.result.mapper_device).map_err(|e| e.kind());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            fs::read(&ctx.result.mapper_device).expect_err("Should fail");
         });
     }
 
@@ -261,8 +260,7 @@
         }
 
         run_test(apk.as_ref(), modified_idsig.as_slice(), "incorrect_merkle_tree", |ctx| {
-            let ret = fs::read(&ctx.result.mapper_device).map_err(|e| e.kind());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            fs::read(&ctx.result.mapper_device).expect_err("Should fail");
         });
     }
 
@@ -284,9 +282,7 @@
             // Read around the modified location causes an error
             let f = File::open(&ctx.result.mapper_device).unwrap();
             let mut buf = vec![0; 10]; // just read 10 bytes
-            let ret = f.read_at(&mut buf, MODIFIED_OFFSET).map_err(|e| e.kind());
-            assert!(ret.is_err());
-            assert_eq!(ret, Err(std::io::ErrorKind::Other));
+            f.read_at(&mut buf, MODIFIED_OFFSET).expect_err("Should fail");
         });
     }
 
@@ -359,11 +355,17 @@
         let apk = include_bytes!("../testdata/test.apk");
         let idsig = include_bytes!("../testdata/test.apk.idsig");
         let roothash = V4Signature::from(Cursor::new(&idsig)).unwrap().hashing_info.raw_root_hash;
-        run_test_with_hash(apk.as_ref(), idsig.as_ref(), "correct", Some(&roothash), |ctx| {
-            let verity = fs::read(&ctx.result.mapper_device).unwrap();
-            let original = fs::read(&ctx.result.data_device).unwrap();
-            assert_eq!(verity.len(), original.len()); // fail fast
-            assert_eq!(verity.as_slice(), original.as_slice());
-        });
+        run_test_with_hash(
+            apk.as_ref(),
+            idsig.as_ref(),
+            "correct_custom_roothash",
+            Some(&roothash),
+            |ctx| {
+                let verity = fs::read(&ctx.result.mapper_device).unwrap();
+                let original = fs::read(&ctx.result.data_device).unwrap();
+                assert_eq!(verity.len(), original.len()); // fail fast
+                assert_eq!(verity.as_slice(), original.as_slice());
+            },
+        );
     }
 }
diff --git a/apkverify/src/lib.rs b/apkverify/src/lib.rs
index f75913c..71ea857 100644
--- a/apkverify/src/lib.rs
+++ b/apkverify/src/lib.rs
@@ -26,8 +26,14 @@
 use anyhow::Result;
 use std::path::Path;
 
-/// Verifies APK/APEX signing with v2/v3 scheme
-pub fn verify<P: AsRef<Path>>(path: P) -> Result<()> {
+/// Verifies APK/APEX signing with v2/v3 scheme. On success, the public key (in DER format) is
+/// returned.
+pub fn verify<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>> {
     // TODO(jooyung) fallback to v2 when v3 not found
     v3::verify(path)
 }
+
+/// Gets the public key (in DER format) that was used to sign the given APK/APEX file
+pub fn get_public_key_der<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>> {
+    v3::get_public_key_der(path)
+}
diff --git a/apkverify/src/v3.rs b/apkverify/src/v3.rs
index 2f9ce94..797911b 100644
--- a/apkverify/src/v3.rs
+++ b/apkverify/src/v3.rs
@@ -86,40 +86,50 @@
 type X509Certificate = Bytes;
 type AdditionalAttributes = Bytes;
 
-/// Verifies APK Signature Scheme v3 signatures of the provided APK and returns the certificates
-/// associated with each signer.
-pub fn verify<P: AsRef<Path>>(path: P) -> Result<()> {
+/// Verifies APK Signature Scheme v3 signatures of the provided APK and returns the public key
+/// associated with the signer.
+pub fn verify<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>> {
     let f = File::open(path.as_ref())?;
     let mut sections = ApkSections::new(f)?;
-    verify_signature(&mut sections)?;
-    Ok(())
+    find_signer_and_then(&mut sections, |(signer, sections)| signer.verify(sections))
 }
 
-/// Verifies the contents of the provided APK file against the provided APK Signature Scheme v3
-/// Block.
-fn verify_signature<R: Read + Seek>(sections: &mut ApkSections<R>) -> Result<()> {
+/// Finds the supported signer and execute a function on it.
+fn find_signer_and_then<R, U, F>(sections: &mut ApkSections<R>, f: F) -> Result<U>
+where
+    R: Read + Seek,
+    F: FnOnce((&Signer, &mut ApkSections<R>)) -> Result<U>,
+{
     let mut block = sections.find_signature(APK_SIGNATURE_SCHEME_V3_BLOCK_ID)?;
-
     // parse v3 scheme block
     let signers = block.read::<Signers>()?;
 
     // find supported by platform
-    let mut supported =
-        signers.iter().filter(|s| s.sdk_range().contains(&SDK_INT)).collect::<Vec<_>>();
+    let supported = signers.iter().filter(|s| s.sdk_range().contains(&SDK_INT)).collect::<Vec<_>>();
 
     // there should be exactly one
     if supported.len() != 1 {
-        bail!("APK Signature Scheme V3 only supports one signer: {} signers found.", signers.len())
+        bail!(
+            "APK Signature Scheme V3 only supports one signer: {} signers found.",
+            supported.len()
+        )
     }
 
-    // and it should be verified
-    supported.pop().unwrap().verify(sections)?;
+    // Call the supplied function
+    f((supported[0], sections))
+}
 
-    Ok(())
+/// Gets the public key (in DER format) that was used to sign the given APK/APEX file
+pub fn get_public_key_der<P: AsRef<Path>>(path: P) -> Result<Box<[u8]>> {
+    let f = File::open(path.as_ref())?;
+    let mut sections = ApkSections::new(f)?;
+    find_signer_and_then(&mut sections, |(signer, _)| {
+        Ok(signer.public_key.to_vec().into_boxed_slice())
+    })
 }
 
 impl Signer {
-    fn verify<R: Read + Seek>(&self, sections: &mut ApkSections<R>) -> Result<()> {
+    fn verify<R: Read + Seek>(&self, sections: &mut ApkSections<R>) -> Result<Box<[u8]>> {
         // 1. Choose the strongest supported signature algorithm ID from signatures. The strength
         //    ordering is up to each implementation/platform version.
         let strongest: &Signature = self
@@ -181,7 +191,7 @@
         }
 
         // TODO(jooyung) 8. If the proof-of-rotation attribute exists for the signer verify that the struct is valid and this signer is the last certificate in the list.
-        Ok(())
+        Ok(self.public_key.to_vec().into_boxed_slice())
     }
 }
 
diff --git a/compos/Android.bp b/compos/Android.bp
index 2af19c8..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",
@@ -34,6 +35,7 @@
     rustlibs: [
         "android.hardware.security.keymint-V1-rust",
         "android.system.keystore2-V1-rust",
+        "android.system.virtualmachineservice-rust",
         "authfs_aidl_interface-rust",
         "compos_aidl_interface-rust",
         "libandroid_logger",
@@ -42,9 +44,11 @@
         "libbinder_rs",
         "libclap",
         "libcompos_common",
+        "libenv_logger",
         "liblibc",
         "liblog_rust",
         "libminijail_rust",
+        "libnix",
         "libring",
         "libscopeguard",
     ],
diff --git a/compos/apex/Android.bp b/compos/apex/Android.bp
index 853b9f4..547fd44 100644
--- a/compos/apex/Android.bp
+++ b/compos/apex/Android.bp
@@ -40,6 +40,8 @@
     binaries: [
         "compos_key_cmd",
         "compos_verify_key",
+        "composd",
+        "composd_cmd",
         "compsvc",
         "pvm_exec",
     ],
@@ -50,5 +52,13 @@
 
     prebuilts: [
         "CompOSPayloadApp.apk.idsig",
+        "com.android.compos.init.rc",
     ],
 }
+
+prebuilt_etc {
+    name: "com.android.compos.init.rc",
+    src: "composd.rc",
+    filename: "init.rc",
+    installable: false,
+}
diff --git a/compos/apex/composd.rc b/compos/apex/composd.rc
new file mode 100644
index 0000000..3e2efb1
--- /dev/null
+++ b/compos/apex/composd.rc
@@ -0,0 +1,21 @@
+# 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.
+
+service composd /apex/com.android.compos/bin/composd
+    class main
+    user root
+    group system
+    interface aidl android.system.composd
+    disabled
+    oneshot
diff --git a/compos/apk/assets/vm_config.json b/compos/apk/assets/vm_config.json
index 9be60d0..9c71942 100644
--- a/compos/apk/assets/vm_config.json
+++ b/compos/apk/assets/vm_config.json
@@ -5,7 +5,10 @@
   },
   "task": {
     "type": "executable",
-    "command": "/apex/com.android.compos/bin/compsvc"
+    "command": "/apex/com.android.compos/bin/compsvc",
+    "args": [
+      "--log_to_stderr"
+    ]
   },
   "apexes": [
     {
diff --git a/compos/apk/assets/vm_test_config.json b/compos/apk/assets/vm_test_config.json
new file mode 100644
index 0000000..54a0aac
--- /dev/null
+++ b/compos/apk/assets/vm_test_config.json
@@ -0,0 +1,18 @@
+{
+    "version": 1,
+    "os": {
+        "name": "microdroid"
+    },
+    "task": {
+        "type": "executable",
+        "command": "/apex/com.android.compos/bin/compsvc"
+    },
+    "apexes": [
+        {
+            "name": "com.android.art"
+        },
+        {
+            "name": "com.android.compos"
+        }
+    ]
+}
\ No newline at end of file
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 22304f1..9e95ed0 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -34,7 +34,11 @@
     FromIBinder,
 };
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
+use log::{info, warn};
 use std::fs::File;
+use std::io::{BufRead, BufReader};
+use std::os::raw;
+use std::os::unix::io::IntoRawFd;
 use std::path::Path;
 use std::sync::{Arc, Condvar, Mutex};
 use std::thread;
@@ -42,8 +46,11 @@
 
 /// This owns an instance of the CompOS VM.
 pub struct VmInstance {
-    #[allow(dead_code)] // Keeps the vm alive even if we don`t touch it
+    #[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,
 }
 
@@ -100,28 +107,76 @@
 
         vm.start()?;
 
-        let cid = vm_state.wait_for_start()?;
+        let cid = vm_state.wait_until_ready()?;
 
-        // TODO: Use onPayloadReady to avoid this
-        thread::sleep(Duration::from_secs(3));
-
-        Ok(VmInstance { vm, cid })
+        Ok(VmInstance { service, vm, cid })
     }
 
     /// Create and return an RPC Binder connection to the Comp OS service in the VM.
     pub fn get_service(&self) -> Result<Strong<dyn ICompOsService>> {
-        let cid = self.cid as u32;
-        // 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(cid, COMPOS_VSOCK_PORT) as *mut AIBinder
-            )
-        }
-        .ok_or_else(|| anyhow!("Failed to connect to CompOS service"))?;
+        let mut vsock_factory = VsockFactory::new(&*self.vm);
+
+        let ibinder = vsock_factory
+            .connect_rpc_client()
+            .ok_or_else(|| anyhow!("Failed to connect to CompOS service"))?;
 
         FromIBinder::try_from(ibinder).context("Connecting to CompOS service")
     }
+
+    /// Return the CID of the VM.
+    pub fn cid(&self) -> i32 {
+        self.cid
+    }
+}
+
+struct VsockFactory<'a> {
+    vm: &'a dyn IVirtualMachine,
+}
+
+impl<'a> VsockFactory<'a> {
+    fn new(vm: &'a dyn IVirtualMachine) -> Self {
+        Self { vm }
+    }
+
+    fn connect_rpc_client(&mut self) -> Option<binder::SpIBinder> {
+        let param = self.as_void_ptr();
+
+        unsafe {
+            // SAFETY: AIBinder returned by RpcPreconnectedClient has correct reference count, and
+            // the ownership can be safely taken by new_spibinder.
+            // RpcPreconnectedClient does not take ownership of param, only passing it to
+            // request_fd.
+            let binder =
+                binder_rpc_unstable_bindgen::RpcPreconnectedClient(Some(Self::request_fd), param)
+                    as *mut AIBinder;
+            new_spibinder(binder)
+        }
+    }
+
+    fn as_void_ptr(&mut self) -> *mut raw::c_void {
+        self as *mut _ as *mut raw::c_void
+    }
+
+    fn try_new_vsock_fd(&self) -> Result<i32> {
+        let vsock = self.vm.connectVsock(COMPOS_VSOCK_PORT as i32)?;
+        // Ownership of the fd is transferred to binder
+        Ok(vsock.into_raw_fd())
+    }
+
+    fn new_vsock_fd(&self) -> i32 {
+        self.try_new_vsock_fd().unwrap_or_else(|e| {
+            warn!("Connecting vsock failed: {}", e);
+            -1_i32
+        })
+    }
+
+    unsafe extern "C" fn request_fd(param: *mut raw::c_void) -> raw::c_int {
+        // SAFETY: This is only ever called by RpcPreconnectedClient, within the lifetime of the
+        // VsockFactory, with param taking the value returned by as_void_ptr (so a properly aligned
+        // non-null pointer to an initialized instance).
+        let vsock_factory = param as *mut Self;
+        vsock_factory.as_ref().unwrap().new_vsock_fd()
+    }
 }
 
 #[derive(Debug)]
@@ -157,7 +212,7 @@
         self.state_ready.notify_all();
     }
 
-    fn set_started(&self, cid: i32) {
+    fn set_ready(&self, cid: i32) {
         let mut state = self.mutex.lock().unwrap();
         if state.has_died {
             return;
@@ -167,10 +222,10 @@
         self.state_ready.notify_all();
     }
 
-    fn wait_for_start(&self) -> Result<i32> {
+    fn wait_until_ready(&self) -> Result<i32> {
         let (state, result) = self
             .state_ready
-            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(10), |state| {
+            .wait_timeout_while(self.mutex.lock().unwrap(), Duration::from_secs(20), |state| {
                 state.cid.is_none() && !state.has_died
             })
             .unwrap();
@@ -196,16 +251,19 @@
     fn onPayloadStarted(
         &self,
         cid: i32,
-        _stream: Option<&binder::parcel::ParcelFileDescriptor>,
+        stream: Option<&ParcelFileDescriptor>,
     ) -> BinderResult<()> {
-        self.0.set_started(cid);
-        // TODO: Use the stream?
+        if let Some(pfd) = stream {
+            if let Err(e) = start_logging(pfd) {
+                warn!("Can't log vm output: {}", e);
+            };
+        }
         log::info!("VM payload started, cid = {}", cid);
         Ok(())
     }
 
     fn onPayloadReady(&self, cid: i32) -> BinderResult<()> {
-        // TODO: Use this to trigger vsock connection
+        self.0.set_ready(cid);
         log::info!("VM payload ready, cid = {}", cid);
         Ok(())
     }
@@ -218,3 +276,19 @@
         Ok(())
     }
 }
+
+fn start_logging(pfd: &ParcelFileDescriptor) -> Result<()> {
+    let reader = BufReader::new(pfd.as_ref().try_clone().context("Cloning fd failed")?);
+    thread::spawn(move || {
+        for line in reader.lines() {
+            match line {
+                Ok(line) => info!("VM: {}", line),
+                Err(e) => {
+                    warn!("Reading VM output failed: {}", e);
+                    break;
+                }
+            }
+        }
+    });
+    Ok(())
+}
diff --git a/compos/common/lib.rs b/compos/common/lib.rs
index 081f086..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;
@@ -27,3 +30,20 @@
 
 /// The root of the  data directory available for private use by the CompOS APEX.
 pub const COMPOS_DATA_ROOT: &str = "/data/misc/apexdata/com.android.compos";
+
+/// The sub-directory where we store information relating to the pending instance
+/// of CompOS (based on staged APEXes).
+pub const PENDING_DIR: &str = "pending";
+
+/// The sub-directory where we store information relating to the current instance
+/// of CompOS (based on active APEXes).
+pub const CURRENT_DIR: &str = "current";
+
+/// The file that holds the encrypted private key for a CompOS instance.
+pub const PRIVATE_KEY_BLOB_FILE: &str = "key.blob";
+
+/// The file that holds the public key for a CompOS instance.
+pub const PUBLIC_KEY_FILE: &str = "key.pubkey";
+
+/// The file that holds the instance image for a CompOS instance.
+pub const INSTANCE_IMAGE_FILE: &str = "instance.img";
diff --git a/compos/compos_key_cmd/Android.bp b/compos/compos_key_cmd/Android.bp
index e0584f4..36c1b5c 100644
--- a/compos/compos_key_cmd/Android.bp
+++ b/compos/compos_key_cmd/Android.bp
@@ -15,8 +15,8 @@
         "android.system.virtualizationservice-ndk",
         "compos_aidl_interface-ndk",
         "libbase",
-        "libbinder_rpc_unstable",
         "libbinder_ndk",
+        "libbinder_rpc_unstable",
         "libcrypto",
         "libfsverity",
         "libprotobuf-cpp-lite",
diff --git a/compos/compos_key_cmd/compos_key_cmd.cpp b/compos/compos_key_cmd/compos_key_cmd.cpp
index eb11d92..e168648 100644
--- a/compos/compos_key_cmd/compos_key_cmd.cpp
+++ b/compos/compos_key_cmd/compos_key_cmd.cpp
@@ -31,8 +31,10 @@
 #include <openssl/mem.h>
 #include <openssl/sha.h>
 #include <openssl/x509.h>
+#include <stdio.h>
 #include <unistd.h>
 
+#include <binder_rpc_unstable.hpp>
 #include <chrono>
 #include <condition_variable>
 #include <filesystem>
@@ -40,14 +42,10 @@
 #include <mutex>
 #include <string>
 #include <string_view>
+#include <thread>
 
 #include "compos_signature.pb.h"
 
-// From frameworks/native/libs/binder/rust/src/binder_rpc_unstable.hpp
-extern "C" {
-AIBinder* RpcClient(unsigned int cid, unsigned int port);
-}
-
 using namespace std::literals;
 
 using aidl::android::system::virtualizationservice::BnVirtualMachineCallback;
@@ -60,6 +58,7 @@
 using aidl::com::android::compos::ICompOsService;
 using android::base::ErrnoError;
 using android::base::Error;
+using android::base::Fdopen;
 using android::base::Result;
 using android::base::unique_fd;
 using compos::proto::Signature;
@@ -98,22 +97,45 @@
 }
 
 namespace {
+
+void copyToLog(unique_fd&& fd) {
+    FILE* source = Fdopen(std::move(fd), "r");
+    size_t size = 0;
+    char* line = nullptr;
+
+    LOG(INFO) << "Started logging VM output";
+
+    for (;;) {
+        ssize_t len = getline(&line, &size, source);
+        if (len < 0) {
+            LOG(INFO) << "VM logging ended: " << ErrnoError().str();
+            break;
+        }
+        LOG(DEBUG) << "VM: " << std::string_view(line, len);
+    }
+    free(line);
+}
+
 class Callback : public BnVirtualMachineCallback {
 public:
-    ::ndk::ScopedAStatus onPayloadStarted(
-            int32_t in_cid, const ::ndk::ScopedFileDescriptor& /*in_stream*/) override {
-        // TODO: Consider copying stdout somewhere useful?
+    ::ndk::ScopedAStatus onPayloadStarted(int32_t in_cid,
+                                          const ::ndk::ScopedFileDescriptor& stream) override {
         LOG(INFO) << "Payload started! cid = " << in_cid;
-        {
-            std::unique_lock lock(mMutex);
-            mStarted = true;
-        }
-        mCv.notify_all();
+
+        unique_fd stream_fd(dup(stream.get()));
+        std::thread logger([fd = std::move(stream_fd)]() mutable { copyToLog(std::move(fd)); });
+        logger.detach();
+
         return ScopedAStatus::ok();
     }
 
     ::ndk::ScopedAStatus onPayloadReady(int32_t in_cid) override {
         LOG(INFO) << "Payload is ready! cid = " << in_cid;
+        {
+            std::unique_lock lock(mMutex);
+            mReady = true;
+        }
+        mCv.notify_all();
         return ScopedAStatus::ok();
     }
 
@@ -132,16 +154,16 @@
         return ScopedAStatus::ok();
     }
 
-    bool waitForStarted() {
+    bool waitUntilReady() {
         std::unique_lock lock(mMutex);
-        return mCv.wait_for(lock, std::chrono::seconds(10), [this] { return mStarted || mDied; }) &&
+        return mCv.wait_for(lock, std::chrono::seconds(20), [this] { return mReady || mDied; }) &&
                 !mDied;
     }
 
 private:
     std::mutex mMutex;
     std::condition_variable mCv;
-    bool mStarted;
+    bool mReady;
     bool mDied;
 };
 
@@ -241,14 +263,10 @@
         }
         LOG(INFO) << "Started VM";
 
-        if (!mCallback->waitForStarted()) {
+        if (!mCallback->waitUntilReady()) {
             return Error() << "VM Payload failed to start";
         }
 
-        // TODO(b/194677789): Implement a polling loop or find a more reliable
-        // way to detect when the service is listening.
-        sleep(3);
-
         return cid;
     }
 
@@ -259,6 +277,7 @@
     std::shared_ptr<Callback> mCallback;
     std::shared_ptr<IVirtualMachine> mVm;
 };
+
 } // namespace
 
 static Result<std::vector<uint8_t>> extractRsaPublicKey(
diff --git a/compos/composd/Android.bp b/compos/composd/Android.bp
new file mode 100644
index 0000000..5c968b8
--- /dev/null
+++ b/compos/composd/Android.bp
@@ -0,0 +1,24 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_binary {
+    name: "composd",
+    srcs: ["src/composd_main.rs"],
+    edition: "2018",
+    prefer_rlib: true,
+    rustlibs: [
+        "android.system.composd-rust",
+        "compos_aidl_interface-rust",
+        "libandroid_logger",
+        "libanyhow",
+        "libbinder_rs",
+        "libcompos_common",
+        "libnum_traits",
+        "liblog_rust",
+    ],
+    proc_macros: ["libnum_derive"],
+    apex_available: [
+        "com.android.compos",
+    ],
+}
diff --git a/compos/composd/aidl/Android.bp b/compos/composd/aidl/Android.bp
new file mode 100644
index 0000000..0352001
--- /dev/null
+++ b/compos/composd/aidl/Android.bp
@@ -0,0 +1,22 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+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: {
+        java: {
+            apex_available: ["//apex_available:platform"],
+        },
+        rust: {
+            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
new file mode 100644
index 0000000..5ff72fe
--- /dev/null
+++ b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
@@ -0,0 +1,33 @@
+/*
+ * Copyright 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.
+ */
+package android.system.composd;
+
+import com.android.compos.CompilationResult;
+import com.android.compos.FdAnnotation;
+
+interface IIsolatedCompilationService {
+    /** 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/composd_main.rs b/compos/composd/src/composd_main.rs
new file mode 100644
index 0000000..71e8125
--- /dev/null
+++ b/compos/composd/src/composd_main.rs
@@ -0,0 +1,52 @@
+/*
+ * 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.
+ */
+
+//! Exposes an on-demand binder service to perform system compilation tasks using CompOS. It is
+//! responsible for managing the lifecycle of the CompOS VM instances, providing key management for
+//! them, and orchestrating trusted compilation.
+
+mod instance_manager;
+mod odrefresh;
+mod service;
+
+use android_system_composd::binder::{register_lazy_service, ProcessState};
+use anyhow::{Context, Result};
+use log::{error, info};
+
+fn try_main() -> Result<()> {
+    android_logger::init_once(
+        android_logger::Config::default().with_tag("composd").with_min_level(log::Level::Info),
+    );
+
+    ProcessState::start_thread_pool();
+
+    let service = service::new_binder();
+    register_lazy_service("android.system.composd", service.as_binder())
+        .context("Registering service")?;
+
+    info!("Registered service, joining threadpool");
+    ProcessState::join_thread_pool();
+
+    info!("Exiting");
+    Ok(())
+}
+
+fn main() {
+    if let Err(e) = try_main() {
+        error!("{}", e);
+        std::process::exit(1)
+    }
+}
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
new file mode 100644
index 0000000..54da231
--- /dev/null
+++ b/compos/composd/src/odrefresh.rs
@@ -0,0 +1,56 @@
+/*
+ * 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.
+ */
+
+//! 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;
+
+// TODO: What if this changes?
+const EX_MAX: i32 = 78;
+const ODREFRESH_BIN: &str = "/apex/com.android.art/bin/odrefresh";
+
+#[derive(Debug, PartialEq, Eq, FromPrimitive)]
+#[repr(i32)]
+pub enum ExitCode {
+    // Copied from art/odrefresh/include/odrefresh/odrefresh.h
+    Okay = 0i32,
+    CompilationRequired = EX_MAX + 1,
+    CompilationSuccess = EX_MAX + 2,
+    CompilationFailed = EX_MAX + 3,
+    CleanupFailed = EX_MAX + 4,
+}
+
+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("--force-compile")
+        .spawn()
+        .context("Running odrefresh")?;
+
+    // TODO: timeout?
+    let status = odrefresh.wait()?;
+
+    if let Some(exit_code) = status.code().and_then(FromPrimitive::from_i32) {
+        Ok(exit_code)
+    } else {
+        bail!("odrefresh exited with {}", status)
+    }
+}
diff --git a/compos/composd/src/service.rs b/compos/composd/src/service.rs
new file mode 100644
index 0000000..e3a1be0
--- /dev/null
+++ b/compos/composd/src/service.rs
@@ -0,0 +1,94 @@
+/*
+ * 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.
+ */
+
+//! Implementation of IIsolatedCompilationService, called from system server when compilation is
+//! desired.
+
+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;
+
+#[derive(Default)]
+pub struct IsolatedCompilationService {
+    instance_manager: InstanceManager,
+}
+
+pub fn new_binder() -> Strong<dyn IIsolatedCompilationService> {
+    let service = IsolatedCompilationService::default();
+    BnIsolatedCompilationService::new_binder(service, BinderFeatures::default())
+}
+
+impl Interface for IsolatedCompilationService {}
+
+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> {
+    result.map_err(|e| {
+        error!("Returning binder error: {:#}", e);
+        Status::new_service_specific_error(-1, CString::new(format!("{:#}", e)).ok().as_deref())
+    })
+}
+
+impl IsolatedCompilationService {
+    fn do_run_forced_compile(&self) -> Result<()> {
+        info!("runForcedCompile");
+
+        let comp_os = self.instance_manager.start_current_instance().context("Starting CompOS")?;
+
+        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/composd_cmd/Android.bp b/compos/composd_cmd/Android.bp
new file mode 100644
index 0000000..5fef86f
--- /dev/null
+++ b/compos/composd_cmd/Android.bp
@@ -0,0 +1,18 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_binary {
+    name: "composd_cmd",
+    srcs: ["composd_cmd.rs"],
+    edition: "2018",
+    rustlibs: [
+        "android.system.composd-rust",
+        "libanyhow",
+        "libbinder_rs",
+    ],
+    prefer_rlib: true,
+    apex_available: [
+        "com.android.compos",
+    ],
+}
diff --git a/compos/composd_cmd/composd_cmd.rs b/compos/composd_cmd/composd_cmd.rs
new file mode 100644
index 0000000..e4884e3
--- /dev/null
+++ b/compos/composd_cmd/composd_cmd.rs
@@ -0,0 +1,36 @@
+/*
+ * 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.
+ */
+
+//! Simple command-line tool to drive composd for testing and debugging.
+
+use android_system_composd::{
+    aidl::android::system::composd::IIsolatedCompilationService::IIsolatedCompilationService,
+    binder::{wait_for_interface, ProcessState},
+};
+use anyhow::{Context, Result};
+
+fn main() -> Result<()> {
+    ProcessState::start_thread_pool();
+
+    let service = wait_for_interface::<dyn IIsolatedCompilationService>("android.system.composd")
+        .context("Failed to connect to composd service")?;
+
+    service.runForcedCompile().context("Compilation failed")?;
+
+    println!("All Ok!");
+
+    Ok(())
+}
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 9855b53..388e79b 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -22,24 +22,54 @@
 mod fsverity;
 mod signer;
 
-use anyhow::{bail, Result};
-use binder::unstable_api::AsNative;
+use android_system_virtualmachineservice::{
+    aidl::android::system::virtualmachineservice::IVirtualMachineService::{
+        IVirtualMachineService, VM_BINDER_SERVICE_PORT,
+    },
+    binder::Strong,
+};
+use anyhow::{anyhow, bail, Context, Result};
+use binder::{
+    unstable_api::{new_spibinder, AIBinder, AsNative},
+    FromIBinder,
+};
 use compos_common::COMPOS_VSOCK_PORT;
-use log::debug;
+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<()> {
-    android_logger::init_once(
-        android_logger::Config::default().with_tag("compsvc").with_min_level(log::Level::Debug),
-    );
+    let args = clap::App::new("compsvc")
+        .arg(clap::Arg::with_name("log_to_stderr").long("log_to_stderr"))
+        .get_matches();
+    if args.is_present("log_to_stderr") {
+        env_logger::builder().filter_level(log::LevelFilter::Debug).init();
+    } else {
+        android_logger::init_once(
+            android_logger::Config::default().with_tag("compsvc").with_min_level(log::Level::Debug),
+        );
+    }
 
     let mut service = compsvc::new_binder()?.as_binder();
     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::RunRpcServer(
+        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(),
         )
     };
     if retval {
@@ -49,3 +79,69 @@
         bail!("Premature termination of RPC server");
     }
 }
+
+struct ReadyNotifier {
+    vm_service: Strong<dyn IVirtualMachineService>,
+    local_cid: u32,
+}
+
+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
+const IOCTL_VM_SOCKETS_GET_LOCAL_CID: usize = 0x7b9;
+ioctl_read_bad!(
+    /// Gets local cid from /dev/vsock
+    vm_sockets_get_local_cid,
+    IOCTL_VM_SOCKETS_GET_LOCAL_CID,
+    u32
+);
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)
+    }
+}
diff --git a/compos/tests/java/android/compos/test/ComposKeyTestCase.java b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
index 6a3f755..00b002f 100644
--- a/compos/tests/java/android/compos/test/ComposKeyTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposKeyTestCase.java
@@ -158,8 +158,8 @@
                         getBuild(),
                         apkName,
                         packageName,
-                        "assets/vm_config.json",
-                        /* debug */ true);
+                        "assets/vm_test_config.json",
+                        /* debug */ false);
         adbConnectToMicrodroid(getDevice(), mCid);
     }
 
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 8409f44..42a323f 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -149,7 +149,7 @@
                         getBuild(),
                         apkName,
                         packageName,
-                        "assets/vm_config.json",
+                        "assets/vm_test_config.json",
                         /* debug */ false);
         adbConnectToMicrodroid(getDevice(), mCid);
     }
diff --git a/compos/verify_key/verify_key.rs b/compos/verify_key/verify_key.rs
index 2e3d206..8439b97 100644
--- a/compos/verify_key/verify_key.rs
+++ b/compos/verify_key/verify_key.rs
@@ -20,17 +20,14 @@
 use anyhow::{bail, Context, Result};
 use compos_aidl_interface::binder::ProcessState;
 use compos_common::compos_client::VmInstance;
-use compos_common::COMPOS_DATA_ROOT;
+use compos_common::{
+    COMPOS_DATA_ROOT, CURRENT_DIR, INSTANCE_IMAGE_FILE, PENDING_DIR, PRIVATE_KEY_BLOB_FILE,
+    PUBLIC_KEY_FILE,
+};
 use std::fs::{self, File};
 use std::io::Read;
 use std::path::{Path, PathBuf};
 
-const CURRENT_DIR: &str = "current";
-const PENDING_DIR: &str = "pending";
-const PRIVATE_KEY_BLOB_FILE: &str = "key.blob";
-const PUBLIC_KEY_FILE: &str = "key.pubkey";
-const INSTANCE_IMAGE_FILE: &str = "instance.img";
-
 const MAX_FILE_SIZE_BYTES: u64 = 8 * 1024;
 
 fn main() -> Result<()> {
diff --git a/javalib/jni/Android.bp b/javalib/jni/Android.bp
index e141297..2939db5 100644
--- a/javalib/jni/Android.bp
+++ b/javalib/jni/Android.bp
@@ -1,3 +1,7 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
 cc_library_shared {
     name: "libvirtualmachine_jni",
     srcs: ["android_system_virtualmachine_VirtualMachine.cpp"],
diff --git a/libs/libavb_rs/Android.bp b/libs/libavb_rs/Android.bp
new file mode 100644
index 0000000..1035498
--- /dev/null
+++ b/libs/libavb_rs/Android.bp
@@ -0,0 +1,31 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_bindgen {
+    name: "libavb_bindgen",
+    wrapper_src: "bindgen/avb.h",
+    crate_name: "avb_bindgen",
+    source_stem: "bindings",
+    bindgen_flags: [
+        "--size_t-is-usize",
+        "--allowlist-function=.*",
+    ],
+    static_libs: [
+        "libavb",
+    ],
+    shared_libs: [
+        "libcrypto",
+    ],
+    cflags: ["-DBORINGSSL_NO_CXX"],
+}
+
+rust_test {
+    name: "libavb_bindgen_test",
+    srcs: [":libavb_bindgen"],
+    crate_name: "avb_bindgen_test",
+    test_suites: ["general-tests"],
+    auto_gen_config: true,
+    clippy_lints: "none",
+    lints: "none",
+}
diff --git a/libs/libavb_rs/bindgen/avb.h b/libs/libavb_rs/bindgen/avb.h
new file mode 100644
index 0000000..b3d5385
--- /dev/null
+++ b/libs/libavb_rs/bindgen/avb.h
@@ -0,0 +1,19 @@
+/*
+ * 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 <libavb/libavb.h>
diff --git a/microdroid/payload/metadata.proto b/microdroid/payload/metadata.proto
index 4c32dde..5ae2158 100644
--- a/microdroid/payload/metadata.proto
+++ b/microdroid/payload/metadata.proto
@@ -31,14 +31,13 @@
 
 message ApexPayload {
   // Required.
-  // The apex name.
   string name = 1;
-
   string partition_name = 2;
 
   // Optional.
-  // When specified, the public key used to sign the apex should match with it.
+  // When specified, apex payload should be verified with the public key and root digest.
   bytes public_key = 3;
+  bytes root_digest = 4;
 }
 
 message ApkPayload {
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 9957689..721f9fa 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -13,6 +13,7 @@
         "android.system.virtualmachineservice-rust",
         "libanyhow",
         "libapkverify",
+        "libavb_bindgen",
         "libbinder_rpc_unstable_bindgen",
         "libbinder_rs",
         "libbyteorder",
@@ -63,4 +64,5 @@
             enabled: false,
         },
     },
+    data: ["tests/data/*"],
 }
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index 73983a7..47230e3 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -320,7 +320,7 @@
 #[derive(Debug, Serialize, Deserialize, PartialEq)]
 pub struct ApkData {
     pub root_hash: Box<RootHash>,
-    // TODO(b/199143508) add cert
+    pub pubkey: Box<[u8]>,
 }
 
 pub type RootHash = [u8];
@@ -328,5 +328,6 @@
 #[derive(Debug, Serialize, Deserialize, PartialEq)]
 pub struct ApexData {
     pub name: String,
-    pub pubkey: Vec<u8>,
+    pub public_key: Vec<u8>,
+    pub root_digest: Vec<u8>,
 }
diff --git a/microdroid_manager/src/ioutil.rs b/microdroid_manager/src/ioutil.rs
index ab82e05..8ab2413 100644
--- a/microdroid_manager/src/ioutil.rs
+++ b/microdroid_manager/src/ioutil.rs
@@ -15,6 +15,8 @@
 //! IO utilities
 
 use anyhow::{anyhow, Result};
+use log::debug;
+use std::fmt::Debug;
 use std::fs::File;
 use std::io;
 use std::path::Path;
@@ -24,7 +26,8 @@
 const SLEEP_DURATION: Duration = Duration::from_millis(5);
 
 /// waits for a file with a timeout and returns it
-pub fn wait_for_file<P: AsRef<Path>>(path: P, timeout: Duration) -> Result<File> {
+pub fn wait_for_file<P: AsRef<Path> + Debug>(path: P, timeout: Duration) -> Result<File> {
+    debug!("waiting for {:?}...", path);
     let begin = Instant::now();
     loop {
         match File::open(&path) {
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 204feab..319ca2b 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -20,7 +20,7 @@
 
 use crate::instance::{ApkData, InstanceDisk, MicrodroidData, RootHash};
 use anyhow::{anyhow, bail, ensure, Context, Result};
-use apkverify::verify;
+use apkverify::{get_public_key_der, verify};
 use binder::unstable_api::{new_spibinder, AIBinder};
 use binder::{FromIBinder, Strong};
 use idsig::V4Signature;
@@ -88,13 +88,20 @@
     Ok(ret)
 }
 
-fn main() -> Result<()> {
+fn main() {
+    if let Err(e) = try_main() {
+        error!("failed with {}", e);
+        std::process::exit(1);
+    }
+}
+
+fn try_main() -> Result<()> {
     kernlog::init()?;
     info!("started.");
 
-    let metadata = load_metadata()?;
+    let metadata = load_metadata().context("Failed to load payload metadata")?;
 
-    let mut instance = InstanceDisk::new()?;
+    let mut instance = InstanceDisk::new().context("Failed to load instance.img")?;
     let saved_data = instance.read_microdroid_data().context("Failed to read identity data")?;
 
     // Verify the payload before using it.
@@ -162,14 +169,15 @@
     // Start apkdmverity and wait for the dm-verify block
     system_properties::write("ctl.start", "apkdmverity")?;
 
-    // While waiting for apkdmverity to mount APK, gathers APEX pubkeys from payload.
+    // While waiting for apkdmverity to mount APK, gathers public keys and root digests from
+    // APEX payload.
     let apex_data_from_payload = get_apex_data_from_payload(metadata)?;
     if let Some(saved_data) = saved_data.map(|d| &d.apex_data) {
-        // For APEX payload, we don't support updating their pubkeys
+        // We don't support APEX updates. (assuming that update will change root digest)
         ensure!(saved_data == &apex_data_from_payload, "APEX payloads has changed.");
         let apex_metadata = to_metadata(&apex_data_from_payload);
-        // Pass metadata(with pubkeys) to apexd so that it uses the passed metadata
-        // instead of the default one (/dev/block/by-name/payload-metadata)
+        // Pass metadata(with public keys and root digests) to apexd so that it uses the passed
+        // metadata instead of the default one (/dev/block/by-name/payload-metadata)
         OpenOptions::new()
             .create_new(true)
             .write(true)
@@ -187,16 +195,18 @@
     // taken only when the root_hash is un-trustful which can be either when this is the first boot
     // of the VM or APK was updated in the host.
     // TODO(jooyung): consider multithreading to make this faster
-    if !root_hash_trustful {
-        verify(DM_MOUNTED_APK_PATH).context(format!("failed to verify {}", DM_MOUNTED_APK_PATH))?;
-    }
+    let apk_pubkey = if !root_hash_trustful {
+        verify(DM_MOUNTED_APK_PATH).context(format!("failed to verify {}", DM_MOUNTED_APK_PATH))?
+    } else {
+        get_public_key_der(DM_MOUNTED_APK_PATH)?
+    };
 
     info!("payload verification successful. took {:#?}", start_time.elapsed().unwrap());
 
     // At this point, we can ensure that the root_hash from the idsig file is trusted, either by
     // fully verifying the APK or by comparing it with the saved root_hash.
     Ok(MicrodroidData {
-        apk_data: ApkData { root_hash: root_hash_from_idsig },
+        apk_data: ApkData { root_hash: root_hash_from_idsig, pubkey: apk_pubkey },
         apex_data: apex_data_from_payload,
     })
 }
@@ -236,7 +246,8 @@
     info!("notifying payload started");
     service.notifyPayloadStarted(local_cid as i32)?;
 
-    if let Some(code) = child.wait()?.code() {
+    let exit_status = child.wait()?;
+    if let Some(code) = exit_status.code() {
         info!("notifying payload finished");
         service.notifyPayloadFinished(local_cid as i32, code)?;
 
@@ -246,7 +257,7 @@
             error!("task exited with exit code: {}", code);
         }
     } else {
-        error!("task terminated by signal");
+        error!("task terminated: {}", exit_status);
     }
     Ok(())
 }
diff --git a/microdroid_manager/src/payload.rs b/microdroid_manager/src/payload.rs
index bf9d9f9..8ec6f74 100644
--- a/microdroid_manager/src/payload.rs
+++ b/microdroid_manager/src/payload.rs
@@ -17,14 +17,11 @@
 use crate::instance::ApexData;
 use crate::ioutil::wait_for_file;
 use anyhow::Result;
+use apex::verify;
 use log::info;
 use microdroid_metadata::{read_metadata, ApexPayload, Metadata};
-use std::fs::File;
-use std::io::Read;
 use std::time::Duration;
-use zip::ZipArchive;
 
-const APEX_PUBKEY_ENTRY: &str = "apex_pubkey";
 const PAYLOAD_METADATA_PATH: &str = "/dev/block/by-name/payload-metadata";
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
 
@@ -35,29 +32,20 @@
     read_metadata(file)
 }
 
-/// Loads (name, pubkey) from payload APEXes
+/// Loads (name, public_key, root_digest) from payload APEXes
 pub fn get_apex_data_from_payload(metadata: &Metadata) -> Result<Vec<ApexData>> {
     metadata
         .apexes
         .iter()
         .map(|apex| {
             let name = apex.name.clone();
-            let partition = format!("/dev/block/by-name/{}", apex.partition_name);
-            let pubkey = get_pubkey_from_apex(&partition)?;
-            Ok(ApexData { name, pubkey })
+            let apex_path = format!("/dev/block/by-name/{}", apex.partition_name);
+            let result = verify(&apex_path)?;
+            Ok(ApexData { name, public_key: result.public_key, root_digest: result.root_digest })
         })
         .collect()
 }
 
-fn get_pubkey_from_apex(path: &str) -> Result<Vec<u8>> {
-    let f = File::open(path)?;
-    let mut z = ZipArchive::new(f)?;
-    let mut pubkey_file = z.by_name(APEX_PUBKEY_ENTRY)?;
-    let mut pubkey = Vec::new();
-    pubkey_file.read_to_end(&mut pubkey)?;
-    Ok(pubkey)
-}
-
 /// Convert vector of ApexData into Metadata
 pub fn to_metadata(apex_data: &[ApexData]) -> Metadata {
     Metadata {
@@ -65,10 +53,13 @@
             .iter()
             .map(|data| ApexPayload {
                 name: data.name.clone(),
-                public_key: data.pubkey.clone(),
+                public_key: data.public_key.clone(),
+                root_digest: data.root_digest.clone(),
                 ..Default::default()
             })
             .collect(),
         ..Default::default()
     }
 }
+
+mod apex;
diff --git a/microdroid_manager/src/payload/apex.rs b/microdroid_manager/src/payload/apex.rs
new file mode 100644
index 0000000..24c4f05
--- /dev/null
+++ b/microdroid_manager/src/payload/apex.rs
@@ -0,0 +1,225 @@
+// Copyright 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.
+
+//! Routines for handling APEX payload
+
+use anyhow::{anyhow, ensure, Result};
+use avb_bindgen::*;
+use std::ffi::{c_void, CStr};
+use std::fs::File;
+use std::io::{Read, Seek, SeekFrom};
+use std::mem::{size_of, zeroed};
+use std::ops::Deref;
+use std::ptr::null_mut;
+use std::slice::{from_raw_parts, from_raw_parts_mut};
+use zip::ZipArchive;
+
+const APEX_PUBKEY_ENTRY: &str = "apex_pubkey";
+const APEX_PAYLOAD_ENTRY: &str = "apex_payload.img";
+
+/// Verification result holds public key and root digest of apex_payload.img
+pub struct ApexVerificationResult {
+    pub public_key: Vec<u8>,
+    pub root_digest: Vec<u8>,
+}
+
+/// Verify APEX payload by AVB verification and return public key and root digest
+pub fn verify(path: &str) -> Result<ApexVerificationResult> {
+    let apex_file = File::open(path)?;
+    let (public_key, image_offset, image_size) = get_public_key_and_image_info(&apex_file)?;
+    let root_digest = verify_vbmeta(apex_file, image_offset, image_size, &public_key)?;
+    Ok(ApexVerificationResult { public_key, root_digest })
+}
+
+fn get_public_key_and_image_info(apex_file: &File) -> Result<(Vec<u8>, u64, u64)> {
+    let mut z = ZipArchive::new(apex_file)?;
+
+    let mut public_key = Vec::new();
+    z.by_name(APEX_PUBKEY_ENTRY)?.read_to_end(&mut public_key)?;
+
+    let (image_offset, image_size) =
+        z.by_name(APEX_PAYLOAD_ENTRY).map(|f| (f.data_start(), f.size()))?;
+
+    Ok((public_key, image_offset, image_size))
+}
+
+// Manual addition of a missing enum
+#[allow(non_camel_case_types, dead_code)]
+#[repr(u8)]
+enum AvbDescriptorTag {
+    AVB_DESCRIPTOR_TAG_PROPERTY = 0,
+    AVB_DESCRIPTOR_TAG_HASHTREE,
+    AVB_DESCRIPTOR_TAG_HASH,
+    AVB_DESCRIPTOR_TAG_KERNEL_CMDLINE,
+    AVB_DESCRIPTOR_TAG_CHAIN_PARTITION,
+}
+
+const FOOTER_SIZE: usize = size_of::<AvbFooter>();
+const HASHTREE_DESCRIPTOR_SIZE: usize = size_of::<AvbHashtreeDescriptor>();
+
+/// Verify VBmeta image and return root digest
+fn verify_vbmeta<R: Read + Seek>(
+    image: R,
+    offset: u64,
+    size: u64,
+    public_key: &[u8],
+) -> Result<Vec<u8>> {
+    let vbmeta = VbMeta::from(image, offset, size)?;
+    vbmeta.verify(public_key)?;
+    for &descriptor in vbmeta.descriptors()?.iter() {
+        if let Ok(hashtree_descriptor) = HashtreeDescriptor::from(descriptor) {
+            return hashtree_descriptor.root_digest();
+        }
+    }
+    Err(anyhow!("HashtreeDescriptor is not found."))
+}
+
+struct VbMeta {
+    data: Vec<u8>,
+}
+
+impl VbMeta {
+    // Read a VbMeta data from a given image
+    fn from<R: Read + Seek>(mut image: R, offset: u64, size: u64) -> Result<VbMeta> {
+        // Get AvbFooter first
+        image.seek(SeekFrom::Start(offset + size - FOOTER_SIZE as u64))?;
+        // SAFETY: AvbDescriptor is a "repr(C,packed)" struct from bindgen
+        let mut footer: AvbFooter = unsafe { zeroed() };
+        // SAFETY: safe to read because of seek(-FOOTER_SIZE) above
+        unsafe {
+            let footer_slice = from_raw_parts_mut(&mut footer as *mut _ as *mut u8, FOOTER_SIZE);
+            image.read_exact(footer_slice)?;
+            ensure!(avb_footer_validate_and_byteswap(&footer, &mut footer));
+        }
+        // Get VbMeta block
+        image.seek(SeekFrom::Start(offset + footer.vbmeta_offset))?;
+        let vbmeta_size = footer.vbmeta_size as usize;
+        let mut data = vec![0u8; vbmeta_size];
+        image.read_exact(&mut data)?;
+        Ok(VbMeta { data })
+    }
+    // Verify VbMeta image. Its enclosed public key should match with a given public key.
+    fn verify(&self, outer_public_key: &[u8]) -> Result<()> {
+        // SAFETY: self.data points to a valid VBMeta data and avb_vbmeta_image_verify should work fine
+        // with it
+        let public_key = unsafe {
+            let mut pk_ptr: *const u8 = null_mut();
+            let mut pk_len: usize = 0;
+            let res = avb_vbmeta_image_verify(
+                self.data.as_ptr(),
+                self.data.len(),
+                &mut pk_ptr,
+                &mut pk_len,
+            );
+            ensure!(
+                res == AvbVBMetaVerifyResult_AVB_VBMETA_VERIFY_RESULT_OK,
+                CStr::from_ptr(avb_vbmeta_verify_result_to_string(res))
+                    .to_string_lossy()
+                    .into_owned()
+            );
+            from_raw_parts(pk_ptr, pk_len)
+        };
+
+        ensure!(public_key == outer_public_key, "Public key mismatch with a given one.");
+        Ok(())
+    }
+    // Return a slice of AvbDescriptor pointers
+    fn descriptors(&self) -> Result<Descriptors> {
+        let mut num: usize = 0;
+        // SAFETY: ptr will be freed by Descriptor.
+        Ok(unsafe {
+            let ptr = avb_descriptor_get_all(self.data.as_ptr(), self.data.len(), &mut num);
+            ensure!(!ptr.is_null(), "VbMeta has no descriptors.");
+            let all = from_raw_parts(ptr, num);
+            Descriptors { ptr, all }
+        })
+    }
+}
+
+struct HashtreeDescriptor {
+    ptr: *const u8,
+    inner: AvbHashtreeDescriptor,
+}
+
+impl HashtreeDescriptor {
+    fn from(descriptor: *const AvbDescriptor) -> Result<HashtreeDescriptor> {
+        // SAFETY: AvbDescriptor is a "repr(C,packed)" struct from bindgen
+        let mut desc: AvbDescriptor = unsafe { zeroed() };
+        // SAFETY: both points to valid AvbDescriptor pointers
+        unsafe {
+            ensure!(avb_descriptor_validate_and_byteswap(descriptor, &mut desc));
+        }
+        ensure!({ desc.tag } == AvbDescriptorTag::AVB_DESCRIPTOR_TAG_HASHTREE as u64);
+        // SAFETY: AvbHashtreeDescriptor is a "repr(C, packed)" struct from bindgen
+        let mut hashtree_descriptor: AvbHashtreeDescriptor = unsafe { zeroed() };
+        // SAFETY: With tag == AVB_DESCRIPTOR_TAG_HASHTREE, descriptor should point to
+        // a AvbHashtreeDescriptor.
+        unsafe {
+            ensure!(avb_hashtree_descriptor_validate_and_byteswap(
+                descriptor as *const AvbHashtreeDescriptor,
+                &mut hashtree_descriptor,
+            ));
+        }
+        Ok(Self { ptr: descriptor as *const u8, inner: hashtree_descriptor })
+    }
+    fn root_digest(&self) -> Result<Vec<u8>> {
+        // SAFETY: digest_ptr should point to a valid buffer of root_digest_len
+        let root_digest = unsafe {
+            let digest_ptr = self.ptr.offset(
+                HASHTREE_DESCRIPTOR_SIZE as isize
+                    + self.inner.partition_name_len as isize
+                    + self.inner.salt_len as isize,
+            );
+            from_raw_parts(digest_ptr, self.inner.root_digest_len as usize)
+        };
+        Ok(root_digest.to_owned())
+    }
+}
+
+// Wraps pointer to a heap-allocated array of AvbDescriptor pointers
+struct Descriptors<'a> {
+    ptr: *mut *const AvbDescriptor,
+    all: &'a [*const AvbDescriptor],
+}
+
+// Wrapped pointer should be freed with avb_free.
+impl Drop for Descriptors<'_> {
+    fn drop(&mut self) {
+        // SAFETY: ptr is allocated by avb_descriptor_get_all
+        unsafe { avb_free(self.ptr as *mut c_void) }
+    }
+}
+
+impl<'a> Deref for Descriptors<'a> {
+    type Target = &'a [*const AvbDescriptor];
+    fn deref(&self) -> &Self::Target {
+        &self.all
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    fn to_hex_string(buf: &[u8]) -> String {
+        buf.iter().map(|b| format!("{:02x}", b)).collect()
+    }
+    #[test]
+    fn test_open_apex() {
+        let res = verify("tests/data/test.apex").unwrap();
+        assert_eq!(
+            to_hex_string(&res.root_digest),
+            "fe11ab17da0a3a738b54bdc3a13f6139cbdf91ec32f001f8d4bbbf8938e04e39"
+        );
+    }
+}
diff --git a/microdroid_manager/tests/data/README.md b/microdroid_manager/tests/data/README.md
new file mode 100644
index 0000000..82ebec6
--- /dev/null
+++ b/microdroid_manager/tests/data/README.md
@@ -0,0 +1,3 @@
+# Test data
+
+- test.apex: copied from system/apexshim/prebuilts/x86/com.android.apex.cts.shim.v1.apex
\ No newline at end of file
diff --git a/microdroid_manager/tests/data/test.apex b/microdroid_manager/tests/data/test.apex
new file mode 100644
index 0000000..fd79365
--- /dev/null
+++ b/microdroid_manager/tests/data/test.apex
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 97cd426..4c8f5eb 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -61,6 +61,9 @@
 
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
+
+        // remove any leftover files under test root
+        android.tryRun("rm", "-rf", TEST_ROOT + "*");
     }
 
     public static void cleanUpVirtualizationTestSetup(ITestDevice androidDevice)
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 8b9d0fa..54b32ec 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -34,6 +34,7 @@
         "libmicrodroid_metadata",
         "libmicrodroid_payload_config",
         "libonce_cell",
+        "librustutils",
         "libserde_json",
         "libserde_xml_rs",
         "libserde",
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index 571cc5d..30a4b03 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -39,7 +39,10 @@
     backend: {
         rust: {
             enabled: true,
-            apex_available: ["com.android.virt"],
+            apex_available: [
+                "com.android.virt",
+                "com.android.compos",
+            ],
         },
     },
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index b0ea0ba..6679da6 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -17,7 +17,7 @@
 use crate::composite::make_composite_image;
 use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmInstance, VmState};
 use crate::payload::add_microdroid_images;
-use crate::{Cid, FIRST_GUEST_CID};
+use crate::{Cid, FIRST_GUEST_CID, SYSPROP_LAST_CID};
 
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualMachine::{
@@ -35,17 +35,18 @@
     VirtualMachineState::VirtualMachineState,
 };
 use android_system_virtualizationservice::binder::{
-    self, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor, Status, Strong, ThreadState,
+    self, force_lazy_services_persist, 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 anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use ::binder::unstable_api::AsNative;
 use disk::QcowFile;
 use idsig::{V4Signature, HashAlgorithm};
 use log::{debug, error, warn, info};
 use microdroid_payload_config::VmPayloadConfig;
+use rustutils::system_properties;
 use std::convert::TryInto;
 use std::ffi::CString;
 use std::fs::{File, OpenOptions, create_dir};
@@ -100,7 +101,7 @@
         let requester_uid = ThreadState::get_calling_uid();
         let requester_sid = get_calling_sid()?;
         let requester_debug_pid = ThreadState::get_calling_pid();
-        let cid = state.allocate_cid()?;
+        let cid = next_cid().or(Err(ExceptionCode::ILLEGAL_STATE))?;
 
         // Counter to generate unique IDs for temporary image files.
         let mut next_temporary_image_id = 0;
@@ -678,9 +679,6 @@
 /// struct.
 #[derive(Debug)]
 struct State {
-    /// The next available unused CID.
-    next_cid: Cid,
-
     /// The VMs which have been started. When VMs are started a weak reference is added to this list
     /// while a strong reference is returned to the caller over Binder. Once all copies of the
     /// Binder client are dropped the weak reference here will become invalid, and will be removed
@@ -716,29 +714,50 @@
     /// 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))?;
-        Some(self.debug_held_vms.swap_remove(pos))
-    }
-
-    /// Get the next available CID, or an error if we have run out.
-    fn allocate_cid(&mut self) -> binder::Result<Cid> {
-        // TODO(qwandor): keep track of which CIDs are currently in use so that we can reuse them.
-        let cid = self.next_cid;
-        self.next_cid = self.next_cid.checked_add(1).ok_or(ExceptionCode::ILLEGAL_STATE)?;
-        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)
     }
 }
 
 impl Default for State {
     fn default() -> Self {
-        State { next_cid: FIRST_GUEST_CID, vms: vec![], debug_held_vms: vec![] }
+        State { vms: vec![], debug_held_vms: vec![] }
     }
 }
 
+/// Get the next available CID, or an error if we have run out. The last CID used is stored in
+/// a system property so that restart of virtualizationservice doesn't reuse CID while the host
+/// Android is up.
+fn next_cid() -> Result<Cid> {
+    let next = if let Ok(val) = system_properties::read(SYSPROP_LAST_CID) {
+        if let Ok(num) = val.parse::<u32>() {
+            num.checked_add(1).ok_or_else(|| anyhow!("run out of CID"))?
+        } else {
+            error!("Invalid last CID {}. Using {}", &val, FIRST_GUEST_CID);
+            FIRST_GUEST_CID
+        }
+    } else {
+        // First VM since the boot
+        FIRST_GUEST_CID
+    };
+    // Persist the last value for next use
+    let str_val = format!("{}", next);
+    system_properties::write(SYSPROP_LAST_CID, &str_val)?;
+    Ok(next)
+}
+
 /// Gets the `VirtualMachineState` of the given `VmInstance`.
 fn get_state(instance: &VmInstance) -> VirtualMachineState {
     match &*instance.vm_state.lock().unwrap() {
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index 8628c01..0e1e974 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -21,7 +21,7 @@
 
 use crate::aidl::{VirtualizationService, BINDER_SERVICE_IDENTIFIER, TEMPORARY_DIRECTORY};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::BnVirtualizationService;
-use android_system_virtualizationservice::binder::{add_service, BinderFeatures, ProcessState};
+use android_system_virtualizationservice::binder::{register_lazy_service, BinderFeatures, ProcessState};
 use anyhow::Error;
 use log::{info, Level};
 use std::fs::{remove_dir_all, remove_file, read_dir};
@@ -30,6 +30,8 @@
 /// are reserved for the host or other usage.
 const FIRST_GUEST_CID: Cid = 10;
 
+const SYSPROP_LAST_CID: &str = "virtualizationservice.state.last_cid";
+
 const LOG_TAG: &str = "VirtualizationService";
 
 /// The unique ID of a VM used (together with a port number) for vsock communication.
@@ -47,7 +49,7 @@
         service,
         BinderFeatures { set_requesting_sid: true, ..BinderFeatures::default() },
     );
-    add_service(BINDER_SERVICE_IDENTIFIER, service.as_binder()).unwrap();
+    register_lazy_service(BINDER_SERVICE_IDENTIFIER, service.as_binder()).unwrap();
     info!("Registered Binder service, joining threadpool.");
     ProcessState::join_thread_pool();
 }