diff --git a/TEST_MAPPING b/TEST_MAPPING
index 4f879b4..5a422df 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -13,7 +13,16 @@
       "name": "art_standalone_dexpreopt_tests"
     },
     {
+      "name": "composd_cmd.test"
+    },
+    {
       "name": "compos_key_tests"
+    },
+    {
+      "name": "composd_verify.test"
+    },
+    {
+      "name": "initrd_bootconfig.test"
     }
   ],
   "avf-postsubmit": [
diff --git a/authfs/TEST_MAPPING b/authfs/TEST_MAPPING
index 5e144c7..450f133 100644
--- a/authfs/TEST_MAPPING
+++ b/authfs/TEST_MAPPING
@@ -4,6 +4,9 @@
       "name": "authfs_device_test_src_lib"
     },
     {
+      "name": "fd_server.test"
+    },
+    {
       "name": "open_then_run.test"
     },
     {
diff --git a/authfs/fd_server/Android.bp b/authfs/fd_server/Android.bp
index 44407a2..f7cb5e3 100644
--- a/authfs/fd_server/Android.bp
+++ b/authfs/fd_server/Android.bp
@@ -23,3 +23,25 @@
     ],
     apex_available: ["com.android.virt"],
 }
+
+rust_test {
+    name: "fd_server.test",
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "authfs_aidl_interface-rust",
+        "libandroid_logger",
+        "libanyhow",
+        "libauthfs_fsverity_metadata",
+        "libbinder_rs",
+        "libclap",
+        "liblibc",
+        "liblog_rust",
+        "libnix",
+        "librpcbinder_rs",
+    ],
+    prefer_rlib: true,
+    shared_libs: [
+        "libbinder_rpc_unstable",
+    ],
+    test_suites: ["general-tests"],
+}
diff --git a/authfs/fd_server/src/main.rs b/authfs/fd_server/src/main.rs
index 9d97423..f91ebec 100644
--- a/authfs/fd_server/src/main.rs
+++ b/authfs/fd_server/src/main.rs
@@ -148,3 +148,15 @@
     server.join();
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn verify_args() {
+        // Check that the command parsing has been configured in a valid way.
+        Args::command().debug_assert();
+    }
+}
diff --git a/compos/composd_cmd/Android.bp b/compos/composd_cmd/Android.bp
index 61df328..54b0bad 100644
--- a/compos/composd_cmd/Android.bp
+++ b/compos/composd_cmd/Android.bp
@@ -18,3 +18,18 @@
         "com.android.compos",
     ],
 }
+
+rust_test {
+    name: "composd_cmd.test",
+    srcs: ["composd_cmd.rs"],
+    edition: "2021",
+    rustlibs: [
+        "android.system.composd-rust",
+        "libanyhow",
+        "libbinder_rs",
+        "libclap",
+        "libcompos_common",
+    ],
+    prefer_rlib: true,
+    test_suites: ["general-tests"],
+}
diff --git a/compos/composd_cmd/composd_cmd.rs b/compos/composd_cmd/composd_cmd.rs
index b6d82aa..19c3720 100644
--- a/compos/composd_cmd/composd_cmd.rs
+++ b/compos/composd_cmd/composd_cmd.rs
@@ -163,3 +163,15 @@
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn verify_actions() {
+        // Check that the command parsing has been configured in a valid way.
+        Actions::command().debug_assert();
+    }
+}
diff --git a/compos/verify/Android.bp b/compos/verify/Android.bp
index f68cc1b..9e30b0d 100644
--- a/compos/verify/Android.bp
+++ b/compos/verify/Android.bp
@@ -22,3 +22,22 @@
         "com.android.compos",
     ],
 }
+
+rust_test {
+    name: "compos_verify.test",
+    srcs: ["verify.rs"],
+    edition: "2021",
+    rustlibs: [
+        "compos_aidl_interface-rust",
+        "libandroid_logger",
+        "libanyhow",
+        "libbinder_rs",
+        "libclap",
+        "libcompos_common",
+        "libcompos_verify_native_rust",
+        "liblog_rust",
+        "libvmclient",
+    ],
+    prefer_rlib: true,
+    test_suites: ["general-tests"],
+}
diff --git a/compos/verify/verify.rs b/compos/verify/verify.rs
index 745d5e9..71d8bcc 100644
--- a/compos/verify/verify.rs
+++ b/compos/verify/verify.rs
@@ -138,3 +138,15 @@
     file.read_to_end(&mut data)?;
     Ok(data)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn verify_args() {
+        // Check that the command parsing has been configured in a valid way.
+        Args::command().debug_assert();
+    }
+}
diff --git a/libs/dice/src/bcc.rs b/libs/dice/src/bcc.rs
index 849dfa0..6dc0cc3 100644
--- a/libs/dice/src/bcc.rs
+++ b/libs/dice/src/bcc.rs
@@ -16,19 +16,29 @@
 
 //! Wrapper around dice/android/bcc.h.
 
+use core::ffi::CStr;
 use core::mem;
 use core::ptr;
 
+use open_dice_bcc_bindgen::BccConfigValues;
+use open_dice_bcc_bindgen::BccFormatConfigDescriptor;
+use open_dice_bcc_bindgen::BccHandoverMainFlow;
 use open_dice_bcc_bindgen::BccHandoverParse;
+use open_dice_bcc_bindgen::DiceInputValues;
+use open_dice_bcc_bindgen::BCC_INPUT_COMPONENT_NAME;
+use open_dice_bcc_bindgen::BCC_INPUT_COMPONENT_VERSION;
+use open_dice_bcc_bindgen::BCC_INPUT_RESETTABLE;
 
 use crate::check_call;
 use crate::Cdi;
 use crate::Error;
+use crate::InputValues;
 use crate::Result;
 
 /// Boot Chain Certificate handover format combining the BCC and CDIs in a single CBOR object.
 #[derive(Clone, Debug)]
 pub struct Handover<'a> {
+    buffer: &'a [u8],
     /// Attestation CDI.
     pub cdi_attest: &'a Cdi,
     /// Sealing CDI.
@@ -75,8 +85,80 @@
             Some(buffer.get(i..(i + bcc_size)).ok_or(Error::PlatformError)?)
         };
 
-        Ok(Self { cdi_attest, cdi_seal, bcc })
+        Ok(Self { buffer, cdi_attest, cdi_seal, bcc })
     }
+
+    /// Executes the main BCC handover flow.
+    pub fn main_flow(&self, input_values: &InputValues, buffer: &mut [u8]) -> Result<usize> {
+        let context = ptr::null_mut();
+        let mut size: usize = 0;
+        // SAFETY - The function only reads `self.buffer`, writes to `buffer` within its bounds,
+        // reads `input_values` as a constant input and doesn't store any pointer.
+        check_call(unsafe {
+            BccHandoverMainFlow(
+                context,
+                self.buffer.as_ptr(),
+                self.buffer.len(),
+                input_values as *const _ as *const DiceInputValues,
+                buffer.len(),
+                buffer.as_mut_ptr(),
+                &mut size as *mut usize,
+            )
+        })?;
+
+        Ok(size)
+    }
+}
+
+/// Formats a configuration descriptor following the BCC's specification.
+///
+/// ```
+/// BccConfigDescriptor = {
+///   ? -70002 : tstr,     ; Component name
+///   ? -70003 : int,      ; Component version
+///   ? -70004 : null,     ; Resettable
+/// }
+/// ```
+pub fn format_config_descriptor(
+    buffer: &mut [u8],
+    name: Option<&CStr>,
+    version: Option<u64>,
+    resettable: bool,
+) -> Result<usize> {
+    let mut inputs = 0;
+
+    if name.is_some() {
+        inputs |= BCC_INPUT_COMPONENT_NAME;
+    }
+
+    if version.is_some() {
+        inputs |= BCC_INPUT_COMPONENT_VERSION;
+    }
+
+    if resettable {
+        inputs |= BCC_INPUT_RESETTABLE;
+    }
+
+    let values = BccConfigValues {
+        inputs,
+        component_name: name.map_or(ptr::null(), |p| p.as_ptr()),
+        component_version: version.unwrap_or(0),
+    };
+
+    let mut buffer_size = 0;
+
+    // SAFETY - The function writes to the buffer, within the given bounds, and only reads the
+    // input values. It writes its result to buffer_size.
+    check_call(unsafe {
+        BccFormatConfigDescriptor(
+            &values as *const _,
+            buffer.len(),
+            buffer.as_mut_ptr(),
+            &mut buffer_size as *mut _,
+        )
+    })?;
+
+    Ok(buffer_size)
 }
 
 fn index_from_ptr(slice: &[u8], pointer: *const u8) -> Option<usize> {
diff --git a/libs/dice/src/lib.rs b/libs/dice/src/lib.rs
index 43d167f..9bbacc6 100644
--- a/libs/dice/src/lib.rs
+++ b/libs/dice/src/lib.rs
@@ -19,9 +19,19 @@
 #![no_std]
 
 use core::fmt;
+use core::mem;
+use core::ptr;
 use core::result;
 
+use open_dice_cbor_bindgen::DiceConfigType_kDiceConfigTypeDescriptor as DICE_CONFIG_TYPE_DESCRIPTOR;
+use open_dice_cbor_bindgen::DiceConfigType_kDiceConfigTypeInline as DICE_CONFIG_TYPE_INLINE;
 use open_dice_cbor_bindgen::DiceHash;
+use open_dice_cbor_bindgen::DiceInputValues;
+use open_dice_cbor_bindgen::DiceMode;
+use open_dice_cbor_bindgen::DiceMode_kDiceModeDebug as DICE_MODE_DEBUG;
+use open_dice_cbor_bindgen::DiceMode_kDiceModeMaintenance as DICE_MODE_MAINTENANCE;
+use open_dice_cbor_bindgen::DiceMode_kDiceModeNormal as DICE_MODE_NORMAL;
+use open_dice_cbor_bindgen::DiceMode_kDiceModeNotInitialized as DICE_MODE_NOT_INITIALIZED;
 use open_dice_cbor_bindgen::DiceResult;
 use open_dice_cbor_bindgen::DiceResult_kDiceResultBufferTooSmall as DICE_RESULT_BUFFER_TOO_SMALL;
 use open_dice_cbor_bindgen::DiceResult_kDiceResultInvalidInput as DICE_RESULT_INVALID_INPUT;
@@ -32,11 +42,17 @@
 
 const CDI_SIZE: usize = open_dice_cbor_bindgen::DICE_CDI_SIZE as usize;
 const HASH_SIZE: usize = open_dice_cbor_bindgen::DICE_HASH_SIZE as usize;
+const HIDDEN_SIZE: usize = open_dice_cbor_bindgen::DICE_HIDDEN_SIZE as usize;
+const INLINE_CONFIG_SIZE: usize = open_dice_cbor_bindgen::DICE_INLINE_CONFIG_SIZE as usize;
 
 /// Array type of CDIs.
 pub type Cdi = [u8; CDI_SIZE];
 /// Array type of hashes used by DICE.
 pub type Hash = [u8; HASH_SIZE];
+/// Array type of additional input.
+pub type Hidden = [u8; HIDDEN_SIZE];
+/// Array type of inline configuration values.
+pub type InlineConfig = [u8; INLINE_CONFIG_SIZE];
 
 /// Error type used by DICE.
 pub enum Error {
@@ -74,6 +90,79 @@
     }
 }
 
+/// DICE mode values.
+#[derive(Clone, Copy, Debug)]
+pub enum Mode {
+    /// At least one security mechanism has not been configured. Also acts as a catch-all.
+    /// Invalid mode values should be treated like this mode.
+    NotInitialized = DICE_MODE_NOT_INITIALIZED as _,
+    /// Indicates the device is operating normally under secure configuration.
+    Normal = DICE_MODE_NORMAL as _,
+    /// Indicates at least one criteria for Normal mode is not met.
+    Debug = DICE_MODE_DEBUG as _,
+    /// Indicates a recovery or maintenance mode of some kind.
+    Maintenance = DICE_MODE_MAINTENANCE as _,
+}
+
+impl From<Mode> for DiceMode {
+    fn from(mode: Mode) -> Self {
+        mode as Self
+    }
+}
+
+/// DICE configuration input type.
+#[derive(Debug)]
+pub enum ConfigType<'a> {
+    /// Uses the formatted 64-byte configuration input value (See the Open Profile for DICE).
+    Inline(InlineConfig),
+    /// Uses the 64-byte hash of more configuration data.
+    Descriptor(&'a [u8]),
+}
+
+/// Set of DICE inputs.
+#[repr(transparent)]
+#[derive(Clone, Debug)]
+pub struct InputValues(DiceInputValues);
+
+impl InputValues {
+    /// Wrap the DICE inputs in a InputValues, expected by bcc::main_flow().
+    pub fn new(
+        code_hash: &Hash,
+        code_descriptor: Option<&[u8]>,
+        config: &ConfigType,
+        auth_hash: Option<&Hash>,
+        auth_descriptor: Option<&[u8]>,
+        mode: Mode,
+        hidden: Option<&Hidden>,
+    ) -> Self {
+        const ZEROED_INLINE_CONFIG: InlineConfig = [0; INLINE_CONFIG_SIZE];
+        let (config_type, config_value, config_descriptor) = match config {
+            ConfigType::Inline(value) => (DICE_CONFIG_TYPE_INLINE, *value, None),
+            ConfigType::Descriptor(desc) => {
+                (DICE_CONFIG_TYPE_DESCRIPTOR, ZEROED_INLINE_CONFIG, Some(*desc))
+            }
+        };
+        let (code_descriptor, code_descriptor_size) = as_raw_parts(code_descriptor);
+        let (config_descriptor, config_descriptor_size) = as_raw_parts(config_descriptor);
+        let (authority_descriptor, authority_descriptor_size) = as_raw_parts(auth_descriptor);
+
+        Self(DiceInputValues {
+            code_hash: *code_hash,
+            code_descriptor,
+            code_descriptor_size,
+            config_type,
+            config_value,
+            config_descriptor,
+            config_descriptor_size,
+            authority_hash: auth_hash.map_or([0; mem::size_of::<Hash>()], |h| *h),
+            authority_descriptor,
+            authority_descriptor_size,
+            mode: mode.into(),
+            hidden: hidden.map_or([0; mem::size_of::<Hidden>()], |h| *h),
+        })
+    }
+}
+
 fn ctx() -> *mut core::ffi::c_void {
     core::ptr::null_mut()
 }
@@ -85,3 +174,10 @@
     check_call(unsafe { DiceHash(ctx(), bytes.as_ptr(), bytes.len(), output.as_mut_ptr()) })?;
     Ok(output)
 }
+
+fn as_raw_parts<T: Sized>(s: Option<&[T]>) -> (*const T, usize) {
+    match s {
+        Some(s) => (s.as_ptr(), s.len()),
+        None => (ptr::null(), 0),
+    }
+}
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 64e6746..7c72fab 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -557,6 +557,11 @@
         Ok(self.path_offset(path)?.map(|offset| FdtNodeMut { fdt: self, offset }))
     }
 
+    /// Return the device tree as a slice (may be smaller than the containing buffer).
+    pub fn as_slice(&self) -> &[u8] {
+        &self.buffer[..self.totalsize()]
+    }
+
     fn path_offset(&self, path: &CStr) -> Result<Option<c_int>> {
         let len = path.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?;
         // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
@@ -591,4 +596,13 @@
     fn capacity(&self) -> usize {
         self.buffer.len()
     }
+
+    fn header(&self) -> &libfdt_bindgen::fdt_header {
+        // SAFETY - A valid FDT (verified by constructor) must contain a valid fdt_header.
+        unsafe { &*(&self as *const _ as *const libfdt_bindgen::fdt_header) }
+    }
+
+    fn totalsize(&self) -> usize {
+        u32::from_be(self.header().totalsize) as usize
+    }
 }
diff --git a/microdroid/initrd/Android.bp b/microdroid/initrd/Android.bp
index 4531583..7a95ce6 100644
--- a/microdroid/initrd/Android.bp
+++ b/microdroid/initrd/Android.bp
@@ -12,6 +12,17 @@
     prefer_rlib: true,
 }
 
+rust_test_host {
+    name: "initrd_bootconfig.test",
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "libanyhow",
+        "libclap",
+    ],
+    prefer_rlib: true,
+    test_suites: ["general-tests"],
+}
+
 python_binary_host {
     name: "gen_vbmeta_bootconfig",
     srcs: ["gen_vbmeta_bootconfig.py"],
diff --git a/microdroid/initrd/src/main.rs b/microdroid/initrd/src/main.rs
index 3f1fab5..74e4ba6 100644
--- a/microdroid/initrd/src/main.rs
+++ b/microdroid/initrd/src/main.rs
@@ -73,3 +73,15 @@
 fn main() {
     try_main().unwrap()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn verify_args() {
+        // Check that the command parsing has been configured in a valid way.
+        Args::command().debug_assert();
+    }
+}
diff --git a/pvmfw/avb/src/error.rs b/pvmfw/avb/src/error.rs
new file mode 100644
index 0000000..8b06150
--- /dev/null
+++ b/pvmfw/avb/src/error.rs
@@ -0,0 +1,87 @@
+// Copyright 2022, 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.
+
+//! This module contains the error thrown by the payload verification API.
+
+use avb_bindgen::AvbSlotVerifyResult;
+
+use core::fmt;
+
+/// This error is the error part of `AvbSlotVerifyResult`.
+/// It is the error thrown by the payload verification API `verify_payload()`.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum AvbSlotVerifyError {
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_ARGUMENT
+    InvalidArgument,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA
+    InvalidMetadata,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_IO
+    Io,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_OOM
+    Oom,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED
+    PublicKeyRejected,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX
+    RollbackIndex,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION
+    UnsupportedVersion,
+    /// AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION
+    Verification,
+}
+
+impl fmt::Display for AvbSlotVerifyError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::InvalidArgument => write!(f, "Invalid parameters."),
+            Self::InvalidMetadata => write!(f, "Invalid metadata."),
+            Self::Io => write!(f, "I/O error while trying to load data or get a rollback index."),
+            Self::Oom => write!(f, "Unable to allocate memory."),
+            Self::PublicKeyRejected => write!(f, "Public key rejected or data not signed."),
+            Self::RollbackIndex => write!(f, "Rollback index is less than its stored value."),
+            Self::UnsupportedVersion => write!(
+                f,
+                "Some of the metadata requires a newer version of libavb than what is in use."
+            ),
+            Self::Verification => write!(f, "Data does not verify."),
+        }
+    }
+}
+
+pub(crate) fn slot_verify_result_to_verify_payload_result(
+    result: AvbSlotVerifyResult,
+) -> Result<(), AvbSlotVerifyError> {
+    match result {
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_OK => Ok(()),
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_ARGUMENT => {
+            Err(AvbSlotVerifyError::InvalidArgument)
+        }
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA => {
+            Err(AvbSlotVerifyError::InvalidMetadata)
+        }
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_IO => Err(AvbSlotVerifyError::Io),
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_OOM => Err(AvbSlotVerifyError::Oom),
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED => {
+            Err(AvbSlotVerifyError::PublicKeyRejected)
+        }
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX => {
+            Err(AvbSlotVerifyError::RollbackIndex)
+        }
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION => {
+            Err(AvbSlotVerifyError::UnsupportedVersion)
+        }
+        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION => {
+            Err(AvbSlotVerifyError::Verification)
+        }
+    }
+}
diff --git a/pvmfw/avb/src/lib.rs b/pvmfw/avb/src/lib.rs
index 1f39076..6a5b16d 100644
--- a/pvmfw/avb/src/lib.rs
+++ b/pvmfw/avb/src/lib.rs
@@ -18,6 +18,8 @@
 // For usize.checked_add_signed(isize), available in Rust 1.66.0
 #![feature(mixed_integer_ops)]
 
+mod error;
 mod verify;
 
-pub use verify::{verify_payload, AvbImageVerifyError};
+pub use error::AvbSlotVerifyError;
+pub use verify::verify_payload;
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index 1e3f8ed..fb18626 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -14,84 +14,16 @@
 
 //! This module handles the pvmfw payload verification.
 
-use avb_bindgen::{
-    avb_slot_verify, AvbHashtreeErrorMode, AvbIOResult, AvbOps, AvbSlotVerifyFlags,
-    AvbSlotVerifyResult,
-};
+use crate::error::{slot_verify_result_to_verify_payload_result, AvbSlotVerifyError};
+use avb_bindgen::{avb_slot_verify, AvbHashtreeErrorMode, AvbIOResult, AvbOps, AvbSlotVerifyFlags};
 use core::{
     ffi::{c_char, c_void, CStr},
-    fmt,
     ptr::{self, NonNull},
     slice,
 };
 
 static NULL_BYTE: &[u8] = b"\0";
 
-/// Error code from AVB image verification.
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum AvbImageVerifyError {
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_ARGUMENT
-    InvalidArgument,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA
-    InvalidMetadata,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_IO
-    Io,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_OOM
-    Oom,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED
-    PublicKeyRejected,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX
-    RollbackIndex,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION
-    UnsupportedVersion,
-    /// AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION
-    Verification,
-}
-
-impl fmt::Display for AvbImageVerifyError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            Self::InvalidArgument => write!(f, "Invalid parameters."),
-            Self::InvalidMetadata => write!(f, "Invalid metadata."),
-            Self::Io => write!(f, "I/O error while trying to load data or get a rollback index."),
-            Self::Oom => write!(f, "Unable to allocate memory."),
-            Self::PublicKeyRejected => write!(f, "Public key rejected or data not signed."),
-            Self::RollbackIndex => write!(f, "Rollback index is less than its stored value."),
-            Self::UnsupportedVersion => write!(
-                f,
-                "Some of the metadata requires a newer version of libavb than what is in use."
-            ),
-            Self::Verification => write!(f, "Data does not verify."),
-        }
-    }
-}
-
-fn to_avb_verify_result(result: AvbSlotVerifyResult) -> Result<(), AvbImageVerifyError> {
-    match result {
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_OK => Ok(()),
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_ARGUMENT => {
-            Err(AvbImageVerifyError::InvalidArgument)
-        }
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA => {
-            Err(AvbImageVerifyError::InvalidMetadata)
-        }
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_IO => Err(AvbImageVerifyError::Io),
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_OOM => Err(AvbImageVerifyError::Oom),
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED => {
-            Err(AvbImageVerifyError::PublicKeyRejected)
-        }
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX => {
-            Err(AvbImageVerifyError::RollbackIndex)
-        }
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION => {
-            Err(AvbImageVerifyError::UnsupportedVersion)
-        }
-        AvbSlotVerifyResult::AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION => {
-            Err(AvbImageVerifyError::Verification)
-        }
-    }
-}
-
 enum AvbIOError {
     /// AVB_IO_RESULT_ERROR_OOM,
     #[allow(dead_code)]
@@ -421,9 +353,9 @@
         }
     }
 
-    fn verify_partitions(&mut self, partition_names: &[&CStr]) -> Result<(), AvbImageVerifyError> {
+    fn verify_partitions(&mut self, partition_names: &[&CStr]) -> Result<(), AvbSlotVerifyError> {
         if partition_names.len() > Self::MAX_NUM_OF_HASH_DESCRIPTORS {
-            return Err(AvbImageVerifyError::InvalidArgument);
+            return Err(AvbSlotVerifyError::InvalidArgument);
         }
         let mut requested_partitions = [ptr::null(); Self::MAX_NUM_OF_HASH_DESCRIPTORS + 1];
         partition_names
@@ -464,7 +396,7 @@
                 out_data,
             )
         };
-        to_avb_verify_result(result)
+        slot_verify_result_to_verify_payload_result(result)
     }
 }
 
@@ -473,7 +405,7 @@
     kernel: &[u8],
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
-) -> Result<(), AvbImageVerifyError> {
+) -> Result<(), AvbSlotVerifyError> {
     let mut payload = Payload { kernel, initrd, trusted_public_key };
     let requested_partitions = [PartitionName::Kernel.as_cstr()];
     payload.verify_partitions(&requested_partitions)
@@ -525,7 +457,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             /*trusted_public_key=*/ &[0u8; 0],
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -535,7 +467,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             /*trusted_public_key=*/ &[0u8; 512],
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -545,7 +477,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA2048_PATH)?,
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -555,7 +487,7 @@
             &fs::read(UNSIGNED_TEST_IMG_PATH)?,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::Io,
+            AvbSlotVerifyError::Io,
         )
     }
 
@@ -568,7 +500,7 @@
             &kernel,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::Verification,
+            AvbSlotVerifyError::Verification,
         )
     }
 
@@ -582,7 +514,7 @@
             &kernel,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::InvalidMetadata,
+            AvbSlotVerifyError::InvalidMetadata,
         )
     }
 
@@ -590,7 +522,7 @@
         kernel: &[u8],
         initrd: &[u8],
         trusted_public_key: &[u8],
-        expected_error: AvbImageVerifyError,
+        expected_error: AvbSlotVerifyError,
     ) -> Result<()> {
         assert_eq!(Err(expected_error), verify_payload(kernel, Some(initrd), trusted_public_key));
         Ok(())
diff --git a/pvmfw/src/dice.rs b/pvmfw/src/dice.rs
new file mode 100644
index 0000000..b322850
--- /dev/null
+++ b/pvmfw/src/dice.rs
@@ -0,0 +1,58 @@
+// Copyright 2022, 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.
+
+//! Support for DICE derivation and BCC generation.
+
+use core::ffi::CStr;
+
+use dice::bcc::format_config_descriptor;
+use dice::bcc::Handover;
+use dice::hash;
+use dice::ConfigType;
+use dice::InputValues;
+
+/// Derive the VM-specific secrets and certificate through DICE.
+pub fn derive_next_bcc(
+    bcc: &Handover,
+    next_bcc: &mut [u8],
+    code: &[u8],
+    debug_mode: bool,
+    authority: &[u8],
+) -> dice::Result<usize> {
+    let code_hash = hash(code)?;
+    let auth_hash = hash(authority)?;
+    let mode = if debug_mode { dice::Mode::Debug } else { dice::Mode::Normal };
+    let component_name = CStr::from_bytes_with_nul(b"vm_entry\0").unwrap();
+    let mut config_descriptor_buffer = [0; 128];
+    let config_descriptor_size = format_config_descriptor(
+        &mut config_descriptor_buffer,
+        Some(component_name),
+        None,  // component_version
+        false, // resettable
+    )?;
+    let config = &config_descriptor_buffer[..config_descriptor_size];
+    let config = ConfigType::Descriptor(config);
+
+    let input_values = InputValues::new(
+        &code_hash,
+        None, // code_descriptor
+        &config,
+        Some(&auth_hash),
+        None, // auth_descriptor
+        mode,
+        None, // TODO(b/249723852): Get salt from instance.img (virtio-blk) and/or TRNG.
+    );
+
+    bcc.main_flow(&input_values, next_bcc)
+}
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 1b35c79..bfcb423 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -48,6 +48,8 @@
     InvalidRamdisk,
     /// Failed to verify the payload.
     PayloadVerificationError,
+    /// DICE layering process failed.
+    SecretDerivationError,
 }
 
 main!(start);
@@ -248,6 +250,7 @@
     crate::main(slices.fdt, slices.kernel, slices.ramdisk, &bcc, &mut memory)?;
 
     helpers::flushed_zeroize(bcc_slice);
+    helpers::flush(slices.fdt.as_slice());
 
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
     memory.mmio_unmap_all().map_err(|e| {
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index dcd17b7..b735b9c 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -48,3 +48,29 @@
 
     Ok(None)
 }
+
+/// Add a "google,open-dice"-compatible reserved-memory node to the tree.
+pub fn add_dice_node(fdt: &mut libfdt::Fdt, addr: usize, size: usize) -> libfdt::Result<()> {
+    fdt.unpack()?;
+
+    let reserved_memory = CStr::from_bytes_with_nul(b"/reserved-memory\0").unwrap();
+    // We reject DTs with missing reserved-memory node as validation should have checked that the
+    // "swiotlb" subnode (compatible = "restricted-dma-pool") was present.
+    let mut reserved_memory = fdt.node_mut(reserved_memory)?.ok_or(libfdt::FdtError::NotFound)?;
+
+    let dice = CStr::from_bytes_with_nul(b"dice\0").unwrap();
+    let mut dice = reserved_memory.add_subnode(dice)?;
+
+    let compatible = CStr::from_bytes_with_nul(b"compatible\0").unwrap();
+    dice.appendprop(compatible, b"google,open-dice\0")?;
+
+    let no_map = CStr::from_bytes_with_nul(b"no-map\0").unwrap();
+    dice.appendprop(no_map, &[])?;
+
+    let reg = CStr::from_bytes_with_nul(b"reg\0").unwrap();
+    dice.appendprop_addrrange(reg, addr as u64, size as u64)?;
+
+    fdt.pack()?;
+
+    Ok(())
+}
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index eab3bc4..e412f69 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -14,8 +14,11 @@
 
 //! Heap implementation.
 
+use alloc::alloc::alloc;
+use alloc::alloc::Layout;
+use alloc::boxed::Box;
+
 use core::alloc::GlobalAlloc as _;
-use core::alloc::Layout;
 use core::ffi::c_void;
 use core::mem;
 use core::num::NonZeroUsize;
@@ -33,6 +36,19 @@
     HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
 }
 
+/// Allocate an aligned but uninitialized slice of heap.
+pub fn aligned_boxed_slice(size: usize, align: usize) -> Option<Box<[u8]>> {
+    let size = NonZeroUsize::new(size)?.get();
+    let layout = Layout::from_size_align(size, align).ok()?;
+    // SAFETY - We verify that `size` and the returned `ptr` are non-null.
+    let ptr = unsafe { alloc(layout) };
+    let ptr = NonNull::new(ptr)?.as_ptr();
+    let slice_ptr = ptr::slice_from_raw_parts_mut(ptr, size);
+
+    // SAFETY - The memory was allocated using the proper layout by our global_allocator.
+    Some(unsafe { Box::from_raw(slice_ptr) })
+}
+
 #[no_mangle]
 unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
     malloc_(size).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index e8a20a8..40266f7 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -20,6 +20,8 @@
 pub const SIZE_4KB: usize = 4 << 10;
 pub const SIZE_2MB: usize = 2 << 20;
 
+pub const GUEST_PAGE_SIZE: usize = SIZE_4KB;
+
 /// Computes the largest multiple of the provided alignment smaller or equal to the address.
 ///
 /// Note: the result is undefined if alignment isn't a power of two.
@@ -89,8 +91,14 @@
 }
 
 #[inline]
+/// Flushes the slice to the point of unification.
+pub fn flush(reg: &[u8]) {
+    flush_region(reg.as_ptr() as usize, reg.len())
+}
+
+#[inline]
 /// Overwrites the slice with zeroes, to the point of unification.
 pub fn flushed_zeroize(reg: &mut [u8]) {
     reg.zeroize();
-    flush_region(reg.as_ptr() as usize, reg.len())
+    flush(reg)
 }
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 9b14644..d0fdd5a 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -19,8 +19,11 @@
 #![feature(default_alloc_error_handler)]
 #![feature(ptr_const_cast)] // Stabilized in 1.65.0
 
+extern crate alloc;
+
 mod avb;
 mod config;
+mod dice;
 mod entry;
 mod exceptions;
 mod fdt;
@@ -33,20 +36,28 @@
 mod pci;
 mod smccc;
 
+use alloc::boxed::Box;
+
 use crate::{
     avb::PUBLIC_KEY,
+    dice::derive_next_bcc,
     entry::RebootReason,
+    fdt::add_dice_node,
+    helpers::flush,
+    helpers::GUEST_PAGE_SIZE,
     memory::MemoryTracker,
     pci::{find_virtio_devices, map_mmio},
 };
-use dice::bcc;
+use ::dice::bcc;
 use fdtpci::{PciError, PciInfo};
 use libfdt::Fdt;
 use log::{debug, error, info, trace};
 use pvmfw_avb::verify_payload;
 
+const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
+
 fn main(
-    fdt: &Fdt,
+    fdt: &mut Fdt,
     signed_kernel: &[u8],
     ramdisk: Option<&[u8]>,
     bcc: &bcc::Handover,
@@ -77,6 +88,43 @@
         RebootReason::PayloadVerificationError
     })?;
 
+    let debug_mode = false; // TODO(b/256148034): Derive the DICE mode from the received initrd.
+    const HASH_SIZE: usize = 64;
+    let mut hashes = [0; HASH_SIZE * 2]; // TODO(b/256148034): Extract AvbHashDescriptor digests.
+    hashes[..HASH_SIZE].copy_from_slice(&::dice::hash(signed_kernel).map_err(|_| {
+        error!("Failed to hash the kernel");
+        RebootReason::InternalError
+    })?);
+    // Note: Using signed_kernel currently makes the DICE code input depend on its VBMeta fields.
+    let code_hash = if let Some(rd) = ramdisk {
+        hashes[HASH_SIZE..].copy_from_slice(&::dice::hash(rd).map_err(|_| {
+            error!("Failed to hash the ramdisk");
+            RebootReason::InternalError
+        })?);
+        &hashes[..]
+    } else {
+        &hashes[..HASH_SIZE]
+    };
+    let next_bcc = heap::aligned_boxed_slice(NEXT_BCC_SIZE, GUEST_PAGE_SIZE).ok_or_else(|| {
+        error!("Failed to allocate the next-stage BCC");
+        RebootReason::InternalError
+    })?;
+    // By leaking the slice, its content will be left behind for the next stage.
+    let next_bcc = Box::leak(next_bcc);
+    let next_bcc_size =
+        derive_next_bcc(bcc, next_bcc, code_hash, debug_mode, PUBLIC_KEY).map_err(|e| {
+            error!("Failed to derive next-stage DICE secrets: {e:?}");
+            RebootReason::SecretDerivationError
+        })?;
+    trace!("Next BCC: {:x?}", bcc::Handover::new(&next_bcc[..next_bcc_size]));
+
+    flush(next_bcc);
+
+    add_dice_node(fdt, next_bcc.as_ptr() as usize, NEXT_BCC_SIZE).map_err(|e| {
+        error!("Failed to add DICE node to device tree: {e}");
+        RebootReason::InternalError
+    })?;
+
     info!("Starting payload...");
     Ok(())
 }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 617c300..2ee33e6 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -652,6 +652,8 @@
     @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
     public void testMicrodroidBoots() throws Exception {
+        CommandRunner android = new CommandRunner(getDevice());
+
         final String configPath = "assets/vm_config.json"; // path inside the APK
         mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
@@ -662,6 +664,9 @@
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
 
+        String vmList = android.run("/apex/com.android.virt/bin/vm list");
+        assertThat(vmList).contains("requesterUid: " + android.run("id -u"));
+
         // Test writing to /data partition
         microdroid.run("echo MicrodroidTest > /data/local/tmp/test.txt");
         assertThat(microdroid.run("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
@@ -680,7 +685,6 @@
         assertThat(abis).hasLength(1);
 
         // Check that no denials have happened so far
-        CommandRunner android = new CommandRunner(getDevice());
         assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH)).isNull();
         assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", CONSOLE_PATH)).isNull();
 
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index f028c0f..91d91aa 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -36,7 +36,10 @@
 aidl_interface {
     name: "android.system.virtualizationservice_internal",
     srcs: ["android/system/virtualizationservice_internal/**/*.aidl"],
-    imports: ["android.system.virtualizationcommon"],
+    imports: [
+        "android.system.virtualizationcommon",
+        "android.system.virtualizationservice",
+    ],
     unstable: true,
     backend: {
         java: {
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index e417ec4..fc4c9e7 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -53,17 +53,4 @@
      * and as such is only permitted from the shell user.
      */
     VirtualMachineDebugInfo[] debugListVms();
-
-    /**
-     * Hold a strong reference to a VM in VirtualizationService. This method is only intended for
-     * debug purposes, and as such is only permitted from the shell user.
-     */
-    void debugHoldVmRef(IVirtualMachine vm);
-
-    /**
-     * Drop reference to a VM that is being held by VirtualizationService. Returns the reference if
-     * VM was found and null otherwise. This method is only intended for debug purposes, and as such
-     * is only permitted from the shell user.
-     */
-    @nullable IVirtualMachine debugDropVmRef(int cid);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
index 424eec1..870a342 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
@@ -33,7 +33,4 @@
      * the PID may have been reused for a different process, so this should not be trusted.
      */
     int requesterPid;
-
-    /** The current lifecycle state of the VM. */
-    VirtualMachineState state = VirtualMachineState.NOT_STARTED;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index d6b3536..5422a48 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -15,6 +15,7 @@
  */
 package android.system.virtualizationservice_internal;
 
+import android.system.virtualizationservice.VirtualMachineDebugInfo;
 import android.system.virtualizationservice_internal.AtomVmBooted;
 import android.system.virtualizationservice_internal.AtomVmCreationRequested;
 import android.system.virtualizationservice_internal.AtomVmExited;
@@ -35,7 +36,7 @@
      * The resources will not be recycled as long as there is a strong reference
      * to the returned object.
      */
-    IGlobalVmContext allocateGlobalVmContext();
+    IGlobalVmContext allocateGlobalVmContext(int requesterDebugPid);
 
     /** Forwards a VmBooted atom to statsd. */
     void atomVmBooted(in AtomVmBooted atom);
@@ -45,4 +46,7 @@
 
     /** Forwards a VmExited atom to statsd. */
     void atomVmExited(in AtomVmExited atom);
+
+    /** Get a list of all currently running VMs. */
+    VirtualMachineDebugInfo[] debugListVms();
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 0859a76..374b90f 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -74,6 +74,7 @@
 use std::num::NonZeroU32;
 use std::os::unix::fs::PermissionsExt;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
+use std::os::unix::raw::{pid_t, uid_t};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
 use tombstoned_client::{DebuggerdDumpType, TombstonedConnection};
@@ -123,15 +124,6 @@
     (GUEST_CID_MIN..=GUEST_CID_MAX).contains(&cid)
 }
 
-fn next_guest_cid(cid: Cid) -> Cid {
-    assert!(is_valid_guest_cid(cid));
-    if cid == GUEST_CID_MAX {
-        GUEST_CID_MIN
-    } else {
-        cid + 1
-    }
-}
-
 fn create_or_update_idsig_file(
     input_fd: &ParcelFileDescriptor,
     idsig_fd: &ParcelFileDescriptor,
@@ -197,10 +189,14 @@
         }
     }
 
-    fn allocateGlobalVmContext(&self) -> binder::Result<Strong<dyn IGlobalVmContext>> {
-        let client_uid = Uid::from_raw(get_calling_uid());
+    fn allocateGlobalVmContext(
+        &self,
+        requester_debug_pid: i32,
+    ) -> binder::Result<Strong<dyn IGlobalVmContext>> {
+        let requester_uid = get_calling_uid();
+        let requester_debug_pid = requester_debug_pid as pid_t;
         let state = &mut *self.state.lock().unwrap();
-        state.allocate_vm_context(client_uid).map_err(|e| {
+        state.allocate_vm_context(requester_uid, requester_debug_pid).map_err(|e| {
             Status::new_exception_str(ExceptionCode::ILLEGAL_STATE, Some(e.to_string()))
         })
     }
@@ -219,25 +215,55 @@
         forward_vm_exited_atom(atom);
         Ok(())
     }
+
+    fn debugListVms(&self) -> binder::Result<Vec<VirtualMachineDebugInfo>> {
+        let state = &mut *self.state.lock().unwrap();
+        let cids = state
+            .held_contexts
+            .iter()
+            .filter_map(|(_, inst)| Weak::upgrade(inst))
+            .map(|vm| VirtualMachineDebugInfo {
+                cid: vm.cid as i32,
+                temporaryDirectory: vm.get_temp_dir().to_string_lossy().to_string(),
+                requesterUid: vm.requester_uid as i32,
+                requesterPid: vm.requester_debug_pid as i32,
+            })
+            .collect();
+        Ok(cids)
+    }
+}
+
+#[derive(Debug, Default)]
+struct GlobalVmInstance {
+    /// The unique CID assigned to the VM for vsock communication.
+    cid: Cid,
+    /// UID of the client who requested this VM instance.
+    requester_uid: uid_t,
+    /// PID of the client who requested this VM instance.
+    requester_debug_pid: pid_t,
+}
+
+impl GlobalVmInstance {
+    fn get_temp_dir(&self) -> PathBuf {
+        let cid = self.cid;
+        format!("{TEMPORARY_DIRECTORY}/{cid}").into()
+    }
 }
 
 /// The mutable state of the VirtualizationServiceInternal. There should only be one instance
 /// of this struct.
 #[derive(Debug, Default)]
 struct GlobalState {
-    /// CIDs currently allocated to running VMs. A CID is never recycled as long
+    /// VM contexts currently allocated to running VMs. A CID is never recycled as long
     /// as there is a strong reference held by a GlobalVmContext.
-    held_cids: HashMap<Cid, Weak<Cid>>,
+    held_contexts: HashMap<Cid, Weak<GlobalVmInstance>>,
 }
 
 impl GlobalState {
     /// 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 allocate_cid(&mut self) -> Result<Arc<Cid>> {
-        // Garbage collect unused CIDs.
-        self.held_cids.retain(|_, cid| cid.strong_count() > 0);
-
+    fn get_next_available_cid(&mut self) -> Result<Cid> {
         // Start trying to find a CID from the last used CID + 1. This ensures
         // that we do not eagerly recycle CIDs. It makes debugging easier but
         // also means that retrying to allocate a CID, eg. because it is
@@ -259,55 +285,62 @@
             });
 
         let first_cid = if let Some(last_cid) = last_cid_prop {
-            next_guest_cid(last_cid)
+            if last_cid == GUEST_CID_MAX {
+                GUEST_CID_MIN
+            } else {
+                last_cid + 1
+            }
         } else {
             GUEST_CID_MIN
         };
 
         let cid = self
             .find_available_cid(first_cid..=GUEST_CID_MAX)
-            .or_else(|| self.find_available_cid(GUEST_CID_MIN..first_cid));
+            .or_else(|| self.find_available_cid(GUEST_CID_MIN..first_cid))
+            .ok_or_else(|| anyhow!("Could not find an available CID."))?;
 
-        if let Some(cid) = cid {
-            let cid_arc = Arc::new(cid);
-            self.held_cids.insert(cid, Arc::downgrade(&cid_arc));
-            system_properties::write(SYSPROP_LAST_CID, &format!("{}", cid))?;
-            Ok(cid_arc)
-        } else {
-            Err(anyhow!("Could not find an available CID."))
-        }
+        system_properties::write(SYSPROP_LAST_CID, &format!("{}", cid))?;
+        Ok(cid)
     }
 
     fn find_available_cid<I>(&self, mut range: I) -> Option<Cid>
     where
         I: Iterator<Item = Cid>,
     {
-        range.find(|cid| !self.held_cids.contains_key(cid))
+        range.find(|cid| !self.held_contexts.contains_key(cid))
     }
 
-    fn allocate_vm_context(&mut self, client_uid: Uid) -> Result<Strong<dyn IGlobalVmContext>> {
-        let cid = self.allocate_cid()?;
-        let temp_dir = create_vm_directory(client_uid, *cid)?;
-        let binder = GlobalVmContext { cid, temp_dir, ..Default::default() };
+    fn allocate_vm_context(
+        &mut self,
+        requester_uid: uid_t,
+        requester_debug_pid: pid_t,
+    ) -> Result<Strong<dyn IGlobalVmContext>> {
+        // Garbage collect unused VM contexts.
+        self.held_contexts.retain(|_, instance| instance.strong_count() > 0);
+
+        let cid = self.get_next_available_cid()?;
+        let instance = Arc::new(GlobalVmInstance { cid, requester_uid, requester_debug_pid });
+        create_temporary_directory(&instance.get_temp_dir(), requester_uid)?;
+
+        self.held_contexts.insert(cid, Arc::downgrade(&instance));
+        let binder = GlobalVmContext { instance, ..Default::default() };
         Ok(BnGlobalVmContext::new_binder(binder, BinderFeatures::default()))
     }
 }
 
-fn create_vm_directory(client_uid: Uid, cid: Cid) -> Result<PathBuf> {
-    let path: PathBuf = format!("{}/{}", TEMPORARY_DIRECTORY, cid).into();
+fn create_temporary_directory(path: &PathBuf, requester_uid: uid_t) -> Result<()> {
     if path.as_path().exists() {
-        remove_temporary_dir(&path).unwrap_or_else(|e| {
+        remove_temporary_dir(path).unwrap_or_else(|e| {
             warn!("Could not delete temporary directory {:?}: {}", path, e);
         });
     }
     // Create a directory that is owned by client's UID but system's GID, and permissions 0700.
     // If the chown() fails, this will leave behind an empty directory that will get removed
     // at the next attempt, or if virtualizationservice is restarted.
-    create_dir(&path)
-        .with_context(|| format!("Could not create temporary directory {:?}", path))?;
-    chown(&path, Some(client_uid), None)
+    create_dir(path).with_context(|| format!("Could not create temporary directory {:?}", path))?;
+    chown(path, Some(Uid::from_raw(requester_uid)), None)
         .with_context(|| format!("Could not set ownership of temporary directory {:?}", path))?;
-    Ok(path)
+    Ok(())
 }
 
 /// Removes a directory owned by a different user by first changing its owner back
@@ -333,10 +366,8 @@
 /// Implementation of the AIDL `IGlobalVmContext` interface.
 #[derive(Debug, Default)]
 struct GlobalVmContext {
-    /// The unique CID assigned to the VM for vsock communication.
-    cid: Arc<Cid>,
-    /// The temporary folder created for the VM and owned by the creator's UID.
-    temp_dir: PathBuf,
+    /// Strong reference to the context's instance data structure.
+    instance: Arc<GlobalVmInstance>,
     /// Keeps our service process running as long as this VM context exists.
     #[allow(dead_code)]
     lazy_service_guard: LazyServiceGuard,
@@ -346,11 +377,11 @@
 
 impl IGlobalVmContext for GlobalVmContext {
     fn getCid(&self) -> binder::Result<i32> {
-        Ok(*self.cid as i32)
+        Ok(self.instance.cid as i32)
     }
 
     fn getTemporaryDirectory(&self) -> binder::Result<String> {
-        Ok(self.temp_dir.to_string_lossy().to_string())
+        Ok(self.instance.get_temp_dir().to_string_lossy().to_string())
     }
 }
 
@@ -469,40 +500,7 @@
     /// and as such is only permitted from the shell user.
     fn debugListVms(&self) -> binder::Result<Vec<VirtualMachineDebugInfo>> {
         check_debug_access()?;
-
-        let state = &mut *self.state.lock().unwrap();
-        let vms = state.vms();
-        let cids = vms
-            .into_iter()
-            .map(|vm| VirtualMachineDebugInfo {
-                cid: vm.cid as i32,
-                temporaryDirectory: vm.temporary_directory.to_string_lossy().to_string(),
-                requesterUid: vm.requester_uid as i32,
-                requesterPid: vm.requester_debug_pid,
-                state: get_state(&vm),
-            })
-            .collect();
-        Ok(cids)
-    }
-
-    /// Hold a strong reference to a VM in VirtualizationService. This method is only intended for
-    /// debug purposes, and as such is only permitted from the shell user.
-    fn debugHoldVmRef(&self, vmref: &Strong<dyn IVirtualMachine>) -> binder::Result<()> {
-        check_debug_access()?;
-
-        let state = &mut *self.state.lock().unwrap();
-        state.debug_hold_vm(vmref.clone());
-        Ok(())
-    }
-
-    /// Drop reference to a VM that is being held by VirtualizationService. Returns the reference if
-    /// the VM was found and None otherwise. This method is only intended for debug purposes, and as
-    /// such is only permitted from the shell user.
-    fn debugDropVmRef(&self, cid: i32) -> binder::Result<Option<Strong<dyn IVirtualMachine>>> {
-        check_debug_access()?;
-
-        let state = &mut *self.state.lock().unwrap();
-        Ok(state.debug_drop_vm(cid))
+        GLOBAL_SERVICE.debugListVms()
     }
 }
 
@@ -564,13 +562,13 @@
         VirtualizationService::default()
     }
 
-    fn create_vm_context(&self) -> Result<(VmContext, Cid, PathBuf)> {
+    fn create_vm_context(&self, requester_debug_pid: pid_t) -> Result<(VmContext, Cid, PathBuf)> {
         const NUM_ATTEMPTS: usize = 5;
 
         for _ in 0..NUM_ATTEMPTS {
-            let global_context = GLOBAL_SERVICE.allocateGlobalVmContext()?;
-            let cid = global_context.getCid()? as Cid;
-            let temp_dir: PathBuf = global_context.getTemporaryDirectory()?.into();
+            let vm_context = GLOBAL_SERVICE.allocateGlobalVmContext(requester_debug_pid as i32)?;
+            let cid = vm_context.getCid()? as Cid;
+            let temp_dir: PathBuf = vm_context.getTemporaryDirectory()?.into();
             let service = VirtualMachineService::new_binder(self.state.clone(), cid).as_binder();
 
             // Start VM service listening for connections from the new CID on port=CID.
@@ -578,7 +576,7 @@
             match RpcServer::new_vsock(service, cid, port) {
                 Ok(vm_server) => {
                     vm_server.start();
-                    return Ok((VmContext::new(global_context, vm_server), cid, temp_dir));
+                    return Ok((VmContext::new(vm_context, vm_server), cid, temp_dir));
                 }
                 Err(err) => {
                     warn!("Could not start RpcServer on port {}: {}", port, err);
@@ -611,19 +609,21 @@
             check_use_custom_virtual_machine()?;
         }
 
-        let (vm_context, cid, temporary_directory) = self.create_vm_context().map_err(|e| {
-            error!("Failed to create VmContext: {:?}", e);
-            Status::new_service_specific_error_str(
-                -1,
-                Some(format!("Failed to create VmContext: {:?}", e)),
-            )
-        })?;
+        let requester_uid = get_calling_uid();
+        let requester_debug_pid = get_calling_pid();
+
+        let (vm_context, cid, temporary_directory) =
+            self.create_vm_context(requester_debug_pid).map_err(|e| {
+                error!("Failed to create VmContext: {:?}", e);
+                Status::new_service_specific_error_str(
+                    -1,
+                    Some(format!("Failed to create VmContext: {:?}", e)),
+                )
+            })?;
 
         let state = &mut *self.state.lock().unwrap();
         let console_fd = console_fd.map(clone_file).transpose()?;
         let log_fd = log_fd.map(clone_file).transpose()?;
-        let requester_uid = get_calling_uid();
-        let requester_debug_pid = get_calling_pid();
 
         // Counter to generate unique IDs for temporary image files.
         let mut next_temporary_image_id = 0;
@@ -664,6 +664,16 @@
             .try_for_each(check_label_for_partition)
             .map_err(|e| Status::new_service_specific_error_str(-1, Some(format!("{:?}", e))))?;
 
+        let kernel = maybe_clone_file(&config.kernel)?;
+        let initrd = maybe_clone_file(&config.initrd)?;
+
+        // In a protected VM, we require custom kernels to come from a trusted source (b/237054515).
+        if config.protectedVm {
+            check_label_for_kernel_files(&kernel, &initrd).map_err(|e| {
+                Status::new_service_specific_error_str(-1, Some(format!("{:?}", e)))
+            })?;
+        }
+
         let zero_filler_path = temporary_directory.join("zero.img");
         write_zero_filler(&zero_filler_path).map_err(|e| {
             error!("Failed to make composite image: {:?}", e);
@@ -706,8 +716,8 @@
             cid,
             name: config.name.clone(),
             bootloader: maybe_clone_file(&config.bootloader)?,
-            kernel: maybe_clone_file(&config.kernel)?,
-            initrd: maybe_clone_file(&config.initrd)?,
+            kernel,
+            initrd,
             disks,
             params: config.params.to_owned(),
             protected: *is_protected,
@@ -971,14 +981,8 @@
     check_permission("android.permission.USE_CUSTOM_VIRTUAL_MACHINE")
 }
 
-/// Check if a partition has selinux labels that are not allowed
-fn check_label_for_partition(partition: &Partition) -> Result<()> {
-    let ctx = getfilecon(partition.image.as_ref().unwrap().as_ref())?;
-    check_label_is_allowed(&ctx).with_context(|| format!("Partition {} invalid", &partition.label))
-}
-
-// Return whether a partition is exempt from selinux label checks, because we know that it does
-// not contain code and is likely to be generated in an app-writable directory.
+/// Return whether a partition is exempt from selinux label checks, because we know that it does
+/// not contain code and is likely to be generated in an app-writable directory.
 fn is_safe_app_partition(label: &str) -> bool {
     // See add_microdroid_system_images & add_microdroid_payload_images in payload.rs.
     label == "vm-instance"
@@ -988,23 +992,46 @@
         || label.starts_with("extra-idsig-")
 }
 
-fn check_label_is_allowed(ctx: &SeContext) -> Result<()> {
-    // We only want to allow code in a VM payload to be sourced from places that apps, and the
-    // system, do not have write access to.
-    // (Note that sepolicy must also grant read access for these types to both virtualization
-    // service and crosvm.)
-    // App private data files are deliberately excluded, to avoid arbitrary payloads being run on
-    // user devices (W^X).
-    match ctx.selinux_type()? {
+/// Check that a file SELinux label is acceptable.
+///
+/// We only want to allow code in a VM to be sourced from places that apps, and the
+/// system, do not have write access to.
+///
+/// Note that sepolicy must also grant read access for these types to both virtualization
+/// service and crosvm.
+///
+/// App private data files are deliberately excluded, to avoid arbitrary payloads being run on
+/// user devices (W^X).
+fn check_label_is_allowed(context: &SeContext) -> Result<()> {
+    match context.selinux_type()? {
         | "system_file" // immutable dm-verity protected partition
         | "apk_data_file" // APKs of an installed app
         | "staging_data_file" // updated/staged APEX images
         | "shell_data_file" // test files created via adb shell
          => Ok(()),
-        _ => bail!("Label {} is not allowed", ctx),
+        _ => bail!("Label {} is not allowed", context),
     }
 }
 
+fn check_label_for_partition(partition: &Partition) -> Result<()> {
+    let file = partition.image.as_ref().unwrap().as_ref();
+    check_label_is_allowed(&getfilecon(file)?)
+        .with_context(|| format!("Partition {} invalid", &partition.label))
+}
+
+fn check_label_for_kernel_files(kernel: &Option<File>, initrd: &Option<File>) -> Result<()> {
+    if let Some(f) = kernel {
+        check_label_for_file(f, "kernel")?;
+    }
+    if let Some(f) = initrd {
+        check_label_for_file(f, "initrd")?;
+    }
+    Ok(())
+}
+fn check_label_for_file(file: &File, name: &str) -> Result<()> {
+    check_label_is_allowed(&getfilecon(file)?).with_context(|| format!("{} file invalid", name))
+}
+
 /// Implementation of the AIDL `IVirtualMachine` interface. Used as a handle to a VM.
 #[derive(Debug)]
 struct VirtualMachine {
@@ -1163,10 +1190,6 @@
     /// Binder client are dropped the weak reference here will become invalid, and will be removed
     /// from the list opportunistically the next time `add_vm` is called.
     vms: Vec<Weak<VmInstance>>,
-
-    /// Vector of strong VM references held on behalf of users that cannot hold them themselves.
-    /// This is only used for debugging purposes.
-    debug_held_vms: Vec<Strong<dyn IVirtualMachine>>,
 }
 
 impl State {
@@ -1189,18 +1212,6 @@
     fn get_vm(&self, cid: Cid) -> Option<Arc<VmInstance>> {
         self.vms().into_iter().find(|vm| vm.cid == cid)
     }
-
-    /// Store a strong VM reference.
-    fn debug_hold_vm(&mut self, vm: Strong<dyn IVirtualMachine>) {
-        self.debug_held_vms.push(vm);
-    }
-
-    /// 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);
-        Some(vm)
-    }
 }
 
 /// Gets the `VirtualMachineState` of the given `VmInstance`.
diff --git a/vm/src/main.rs b/vm/src/main.rs
index bfc7920..ea744f7 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -60,10 +60,6 @@
         #[clap(long)]
         name: Option<String>,
 
-        /// Detach VM from the terminal and run in the background
-        #[clap(short, long)]
-        daemonize: bool,
-
         /// Path to the file backing the storage.
         /// Created if the option is used but the path does not exist in the device.
         #[clap(long)]
@@ -119,10 +115,6 @@
         #[clap(long)]
         name: Option<String>,
 
-        /// Detach VM from the terminal and run in the background
-        #[clap(short, long)]
-        daemonize: bool,
-
         /// Path to the file backing the storage.
         /// Created if the option is used but the path does not exist in the device.
         #[clap(long)]
@@ -171,10 +163,6 @@
         #[clap(long)]
         name: Option<String>,
 
-        /// Detach VM from the terminal and run in the background
-        #[clap(short, long)]
-        daemonize: bool,
-
         /// Number of vCPUs in the VM. If unspecified, defaults to 1.
         #[clap(long)]
         cpus: Option<u32>,
@@ -191,11 +179,6 @@
         #[clap(long)]
         log: Option<PathBuf>,
     },
-    /// Stop a virtual machine running in the background
-    Stop {
-        /// CID of the virtual machine
-        cid: u32,
-    },
     /// List running virtual machines
     List,
     /// Print information about virtual machine support
@@ -260,7 +243,6 @@
             storage_size,
             config_path,
             payload_binary_name,
-            daemonize,
             console,
             log,
             debug,
@@ -279,7 +261,6 @@
             storage_size,
             config_path,
             payload_binary_name,
-            daemonize,
             console.as_deref(),
             log.as_deref(),
             debug,
@@ -294,7 +275,6 @@
             work_dir,
             storage,
             storage_size,
-            daemonize,
             console,
             log,
             debug,
@@ -308,7 +288,6 @@
             work_dir,
             storage.as_deref(),
             storage_size,
-            daemonize,
             console.as_deref(),
             log.as_deref(),
             debug,
@@ -317,12 +296,11 @@
             cpus,
             task_profiles,
         ),
-        Opt::Run { name, config, daemonize, cpus, task_profiles, console, log } => {
+        Opt::Run { name, config, cpus, task_profiles, console, log } => {
             command_run(
                 name,
                 service.as_ref(),
                 &config,
-                daemonize,
                 console.as_deref(),
                 log.as_deref(),
                 /* mem */ None,
@@ -330,7 +308,6 @@
                 task_profiles,
             )
         }
-        Opt::Stop { cid } => command_stop(service.as_ref(), cid),
         Opt::List => command_list(service.as_ref()),
         Opt::Info => command_info(),
         Opt::CreatePartition { path, size, partition_type } => {
@@ -340,15 +317,6 @@
     }
 }
 
-/// Retrieve reference to a previously daemonized VM and stop it.
-fn command_stop(service: &dyn IVirtualizationService, cid: u32) -> Result<(), Error> {
-    service
-        .debugDropVmRef(cid as i32)
-        .context("Failed to get VM from VirtualizationService")?
-        .context("CID does not correspond to a running background VM")?;
-    Ok(())
-}
-
 /// List the VMs currently running.
 fn command_list(service: &dyn IVirtualizationService) -> Result<(), Error> {
     let vms = service.debugListVms().context("Failed to get list of VMs")?;
diff --git a/vm/src/run.rs b/vm/src/run.rs
index b99328a..6c21dbc 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -49,7 +49,6 @@
     storage_size: Option<u64>,
     config_path: Option<String>,
     payload_binary_name: Option<String>,
-    daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
     debug_level: DebugLevel,
@@ -145,7 +144,7 @@
         numCpus: cpus.unwrap_or(1) as i32,
         taskProfiles: task_profiles,
     });
-    run(service, &config, &payload_config_str, daemonize, console_path, log_path)
+    run(service, &config, &payload_config_str, console_path, log_path)
 }
 
 const EMPTY_PAYLOAD_APK: &str = "com.android.microdroid.empty_payload";
@@ -180,7 +179,6 @@
     work_dir: Option<PathBuf>,
     storage: Option<&Path>,
     storage_size: Option<u64>,
-    daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
     debug_level: DebugLevel,
@@ -211,7 +209,6 @@
         storage_size,
         /* config_path= */ None,
         Some(payload_binary_name.to_owned()),
-        daemonize,
         console_path,
         log_path,
         debug_level,
@@ -229,7 +226,6 @@
     name: Option<String>,
     service: &dyn IVirtualizationService,
     config_path: &Path,
-    daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
     mem: Option<u32>,
@@ -255,7 +251,6 @@
         service,
         &VirtualMachineConfig::RawConfig(config),
         &format!("{:?}", config_path),
-        daemonize,
         console_path,
         log_path,
     )
@@ -277,7 +272,6 @@
     service: &dyn IVirtualizationService,
     config: &VirtualMachineConfig,
     payload_config: &str,
-    daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
 ) -> Result<(), Error> {
@@ -286,8 +280,6 @@
             File::create(console_path)
                 .with_context(|| format!("Failed to open console file {:?}", console_path))?,
         )
-    } else if daemonize {
-        None
     } else {
         Some(duplicate_stdout()?)
     };
@@ -296,8 +288,6 @@
             File::create(log_path)
                 .with_context(|| format!("Failed to open log file {:?}", log_path))?,
         )
-    } else if daemonize {
-        None
     } else {
         Some(duplicate_stdout()?)
     };
@@ -314,17 +304,10 @@
         state_to_str(vm.state()?)
     );
 
-    if daemonize {
-        // Pass the VM reference back to VirtualizationService and have it hold it in the
-        // background.
-        service.debugHoldVmRef(&vm.vm).context("Failed to pass VM to VirtualizationService")?;
-    } else {
-        // Wait until the VM or VirtualizationService dies. If we just returned immediately then the
-        // IVirtualMachine Binder object would be dropped and the VM would be killed.
-        let death_reason = vm.wait_for_death();
-        println!("VM ended: {:?}", death_reason);
-    }
-
+    // Wait until the VM or VirtualizationService dies. If we just returned immediately then the
+    // IVirtualMachine Binder object would be dropped and the VM would be killed.
+    let death_reason = vm.wait_for_death();
+    println!("VM ended: {:?}", death_reason);
     Ok(())
 }
 
diff --git a/vm/vm_shell.sh b/vm/vm_shell.sh
index 29cc7da..3db7003 100755
--- a/vm/vm_shell.sh
+++ b/vm/vm_shell.sh
@@ -92,7 +92,8 @@
         shift
     done
     if [[ "${auto_connect}" == true ]]; then
-        adb shell /apex/com.android.virt/bin/vm run-microdroid -d "${passthrough_args}"
+        adb shell /apex/com.android.virt/bin/vm run-microdroid "${passthrough_args}" &
+        trap "kill $!" EXIT
         sleep 2
         handle_connect_cmd
     else