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/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index c3fa55f..aab4148 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -246,6 +246,7 @@
                 builder.setProtectedVm(true);
                 if (debug) {
                     builder.setDebugLevel(VirtualMachineConfig.DEBUG_LEVEL_FULL);
+                    builder.setVmOutputCaptured(true);
                 }
                 VirtualMachineConfig config = builder.build();
                 VirtualMachineManager vmm =
@@ -261,10 +262,12 @@
                 mVirtualMachine.setCallback(Executors.newSingleThreadExecutor(), callback);
                 mStatus.postValue(mVirtualMachine.getStatus());
 
-                InputStream console = mVirtualMachine.getConsoleOutput();
-                InputStream log = mVirtualMachine.getLogOutput();
-                mExecutorService.execute(new Reader("console", mConsoleOutput, console));
-                mExecutorService.execute(new Reader("log", mLogOutput, log));
+                if (debug) {
+                    InputStream console = mVirtualMachine.getConsoleOutput();
+                    InputStream log = mVirtualMachine.getLogOutput();
+                    mExecutorService.execute(new Reader("console", mConsoleOutput, console));
+                    mExecutorService.execute(new Reader("log", mLogOutput, log));
+                }
             } catch (VirtualMachineException e) {
                 throw new RuntimeException(e);
             }
diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index 71bdf13..1977321 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -65,6 +65,7 @@
     method public boolean isCompatibleWith(@NonNull android.system.virtualmachine.VirtualMachineConfig);
     method public boolean isEncryptedStorageEnabled();
     method public boolean isProtectedVm();
+    method public boolean isVmOutputCaptured();
     field public static final int DEBUG_LEVEL_FULL = 1; // 0x1
     field public static final int DEBUG_LEVEL_NONE = 0; // 0x0
   }
@@ -79,6 +80,7 @@
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setNumCpus(@IntRange(from=1) int);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadBinaryName(@NonNull String);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setProtectedVm(boolean);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setVmOutputCaptured(boolean);
   }
 
   public final class VirtualMachineDescriptor implements android.os.Parcelable {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index b57cb5e..1f0c8ea 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -293,6 +293,7 @@
     /** Lock protecting callbacks. */
     private final Object mCallbackLock = new Object();
 
+    private final boolean mVmOutputCaptured;
 
     /** The configuration that is currently associated with this VM. */
     @GuardedBy("mLock")
@@ -371,6 +372,7 @@
                         ? new File(thisVmDir, ENCRYPTED_STORE_FILE)
                         : null;
 
+        mVmOutputCaptured = config.isVmOutputCaptured();
     }
 
     /**
@@ -772,7 +774,9 @@
             IVirtualizationService service = mVirtualizationService.connect();
 
             try {
-                createVmPipes();
+                if (mVmOutputCaptured) {
+                    createVmPipes();
+                }
 
                 VirtualMachineAppConfig appConfig = getConfig().toVsConfig();
                 appConfig.name = mName;
@@ -890,17 +894,26 @@
     }
 
     /**
-     * Returns the stream object representing the console output from the virtual machine.
+     * Returns the stream object representing the console output from the virtual machine. The
+     * console output is only available if the {@link VirtualMachineConfig} specifies that it should
+     * be {@linkplain VirtualMachineConfig#isVmOutputCaptured captured}.
+     *
+     * <p>If you turn on output capture, you must consume data from {@code getConsoleOutput} -
+     * because otherwise the code in the VM may get blocked when the pipe buffer fills up.
      *
      * <p>NOTE: This method may block and should not be called on the main thread.
      *
-     * @throws VirtualMachineException if the stream could not be created.
+     * @throws VirtualMachineException if the stream could not be created, or capturing is turned
+     *     off.
      * @hide
      */
     @SystemApi
     @WorkerThread
     @NonNull
     public InputStream getConsoleOutput() throws VirtualMachineException {
+        if (!mVmOutputCaptured) {
+            throw new VirtualMachineException("Capturing vm outputs is turned off");
+        }
         synchronized (mLock) {
             createVmPipes();
             return new FileInputStream(mConsoleReader.getFileDescriptor());
@@ -908,17 +921,26 @@
     }
 
     /**
-     * Returns the stream object representing the log output from the virtual machine.
+     * Returns the stream object representing the log output from the virtual machine. The log
+     * output is only available if the VirtualMachineConfig specifies that it should be {@linkplain
+     * VirtualMachineConfig#isVmOutputCaptured captured}.
+     *
+     * <p>If you turn on output capture, you must consume data from {@code getLogOutput} - because
+     * otherwise the code in the VM may get blocked when the pipe buffer fills up.
      *
      * <p>NOTE: This method may block and should not be called on the main thread.
      *
-     * @throws VirtualMachineException if the stream could not be created.
+     * @throws VirtualMachineException if the stream could not be created, or capturing is turned
+     *     off.
      * @hide
      */
     @SystemApi
     @WorkerThread
     @NonNull
     public InputStream getLogOutput() throws VirtualMachineException {
+        if (!mVmOutputCaptured) {
+            throw new VirtualMachineException("Capturing vm outputs is turned off");
+        }
         synchronized (mLock) {
             createVmPipes();
             return new FileInputStream(mLogReader.getFileDescriptor());
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index d9fc70c..f5c3cd2 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -58,7 +58,7 @@
     private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
-    private static final int VERSION = 2;
+    private static final int VERSION = 3;
     private static final String KEY_VERSION = "version";
     private static final String KEY_APKPATH = "apkPath";
     private static final String KEY_PAYLOADCONFIGPATH = "payloadConfigPath";
@@ -68,6 +68,7 @@
     private static final String KEY_MEMORY_MIB = "memoryMib";
     private static final String KEY_NUM_CPUS = "numCpus";
     private static final String KEY_ENCRYPTED_STORAGE_KIB = "encryptedStorageKib";
+    private static final String KEY_VM_OUTPUT_CAPTURED = "vmOutputCaptured";
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -124,6 +125,9 @@
     /** The size of storage in KiB. 0 indicates that encryptedStorage is not required */
     private final long mEncryptedStorageKib;
 
+    /** Whether the app can read console and log output. */
+    private final boolean mVmOutputCaptured;
+
     private VirtualMachineConfig(
             @NonNull String apkPath,
             @Nullable String payloadConfigPath,
@@ -132,7 +136,8 @@
             boolean protectedVm,
             int memoryMib,
             int numCpus,
-            long encryptedStorageKib) {
+            long encryptedStorageKib,
+            boolean vmOutputCaptured) {
         // This is only called from Builder.build(); the builder handles parameter validation.
         mApkPath = apkPath;
         mPayloadConfigPath = payloadConfigPath;
@@ -142,6 +147,7 @@
         mMemoryMib = memoryMib;
         mNumCpus = numCpus;
         mEncryptedStorageKib = encryptedStorageKib;
+        mVmOutputCaptured = vmOutputCaptured;
     }
 
     /** Loads a config from a file. */
@@ -210,6 +216,7 @@
         if (encryptedStorageKib != 0) {
             builder.setEncryptedStorageKib(encryptedStorageKib);
         }
+        builder.setVmOutputCaptured(b.getBoolean(KEY_VM_OUTPUT_CAPTURED));
 
         return builder.build();
     }
@@ -239,6 +246,7 @@
         if (mEncryptedStorageKib > 0) {
             b.putLong(KEY_ENCRYPTED_STORAGE_KIB, mEncryptedStorageKib);
         }
+        b.putBoolean(KEY_VM_OUTPUT_CAPTURED, mVmOutputCaptured);
         b.writeToStream(output);
     }
 
@@ -346,6 +354,18 @@
     }
 
     /**
+     * Returns whether the app can read the VM console or log output. If not, the VM output is
+     * automatically forwarded to the host logcat.
+     *
+     * @see #setVmOutputCaptured
+     * @hide
+     */
+    @SystemApi
+    public boolean isVmOutputCaptured() {
+        return mVmOutputCaptured;
+    }
+
+    /**
      * Tests if this config is compatible with other config. Being compatible means that the configs
      * can be interchangeably used for the same virtual machine; they do not change the VM identity
      * or secrets. Such changes include varying the number of CPUs or the size of the RAM. Changes
@@ -360,6 +380,7 @@
         return this.mDebugLevel == other.mDebugLevel
                 && this.mProtectedVm == other.mProtectedVm
                 && this.mEncryptedStorageKib == other.mEncryptedStorageKib
+                && this.mVmOutputCaptured == other.mVmOutputCaptured
                 && Objects.equals(this.mPayloadConfigPath, other.mPayloadConfigPath)
                 && Objects.equals(this.mPayloadBinaryName, other.mPayloadBinaryName)
                 && this.mApkPath.equals(other.mApkPath);
@@ -417,6 +438,7 @@
         private int mMemoryMib;
         private int mNumCpus = 1;
         private long mEncryptedStorageKib;
+        private boolean mVmOutputCaptured = false;
 
         /**
          * Creates a builder for the given context.
@@ -469,6 +491,10 @@
                 throw new IllegalStateException("setProtectedVm must be called explicitly");
             }
 
+            if (mVmOutputCaptured && mDebugLevel != DEBUG_LEVEL_FULL) {
+                throw new IllegalStateException("debug level must be FULL to capture output");
+            }
+
             return new VirtualMachineConfig(
                     apkPath,
                     mPayloadConfigPath,
@@ -477,7 +503,8 @@
                     mProtectedVm,
                     mMemoryMib,
                     mNumCpus,
-                    mEncryptedStorageKib);
+                    mEncryptedStorageKib,
+                    mVmOutputCaptured);
         }
 
         /**
@@ -654,5 +681,31 @@
             mEncryptedStorageKib = encryptedStorageKib;
             return this;
         }
+
+        /**
+         * Sets whether to allow the app to read the VM outputs (console / log). Default is {@code
+         * false}.
+         *
+         * <p>By default, console and log outputs of a {@linkplain #setDebugLevel debuggable} VM are
+         * automatically forwarded to the host logcat. Setting this as {@code true} will allow the
+         * app to directly read {@linkplain VirtualMachine#getConsoleOutput console output} and
+         * {@linkplain VirtualMachine#getLogOutput log output}, instead of forwarding them to the
+         * host logcat.
+         *
+         * <p>If you turn on output capture, you must consume data from {@link
+         * VirtualMachine#getConsoleOutput} and {@link VirtualMachine#getLogOutput} - because
+         * otherwise the code in the VM may get blocked when the pipe buffer fills up.
+         *
+         * <p>The {@linkplain #setDebugLevel debug level} must be {@link #DEBUG_LEVEL_FULL} to be
+         * set as true.
+         *
+         * @hide
+         */
+        @SystemApi
+        @NonNull
+        public Builder setVmOutputCaptured(boolean captured) {
+            mVmOutputCaptured = captured;
+            return this;
+        }
     }
 }
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/Android.bp b/microdroid/Android.bp
index b898097..529686d 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -582,7 +582,7 @@
     name: "microdroid_kernel_signed",
     src: "empty_kernel",
     filename: "microdroid_kernel",
-    partition_name: "bootloader",
+    partition_name: "boot",
     private_key: ":microdroid_sign_key",
     salt: bootloader_salt,
     enabled: false,
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 7402481..bc42791 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -165,7 +165,6 @@
     class core
     critical
     seclabel u:r:ueventd:s0
-    shutdown critical
     capabilities CHOWN DAC_OVERRIDE DAC_READ_SEARCH FOWNER FSETID MKNOD NET_ADMIN SETGID SETUID SYS_MODULE SYS_RAWIO
 
 service console /system/bin/sh
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 69c6ae4..74e4ba6 100644
--- a/microdroid/initrd/src/main.rs
+++ b/microdroid/initrd/src/main.rs
@@ -54,7 +54,8 @@
         checksum += get_checksum(&bootconfig)?;
     }
 
-    let padding_size: usize = FOOTER_ALIGNMENT - (initrd_size + bootconfig_size) % FOOTER_ALIGNMENT;
+    let padding_size: usize =
+        (FOOTER_ALIGNMENT - (initrd_size + bootconfig_size) % FOOTER_ALIGNMENT) % FOOTER_ALIGNMENT;
     output_file.write_all(&ZEROS[..padding_size])?;
     output_file.write_all(&((padding_size + bootconfig_size) as u32).to_le_bytes())?;
     output_file.write_all(&checksum.to_le_bytes())?;
@@ -72,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/Android.bp b/pvmfw/avb/Android.bp
index cbec235..d3a5e4e 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -62,7 +62,7 @@
 avb_add_hash_footer {
     name: "test_image_with_one_hashdesc",
     src: ":unsigned_test_image",
-    partition_name: "bootloader",
+    partition_name: "boot",
     private_key: ":pvmfw_sign_key",
     salt: "1111",
 }
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 0bab960..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)]
@@ -362,7 +294,7 @@
 }
 
 impl PartitionName {
-    const KERNEL_PARTITION_NAME: &[u8] = b"bootloader\0";
+    const KERNEL_PARTITION_NAME: &[u8] = b"boot\0";
     const INITRD_NORMAL_PARTITION_NAME: &[u8] = b"initrd_normal\0";
     const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
 
@@ -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/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 9e4b228..40114fd 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -211,6 +211,7 @@
                     newVmConfigBuilder()
                             .setPayloadBinaryName("MicrodroidIdleNativeLib.so")
                             .setDebugLevel(DEBUG_LEVEL_FULL)
+                            .setVmOutputCaptured(true)
                             .setMemoryMib(256)
                             .build();
             forceCreateNewVirtualMachine("test_vm_boot_time_debug", normalConfig);
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index d762310..e8a36ce 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -212,9 +212,11 @@
                 throws VirtualMachineException, InterruptedException {
             vm.setCallback(mExecutorService, this);
             vm.run();
-            logVmOutputAndMonitorBootEvents(
-                    logTag, vm.getConsoleOutput(), "Console", mConsoleOutput);
-            logVmOutput(logTag, vm.getLogOutput(), "Log", mLogOutput);
+            if (vm.getConfig().isVmOutputCaptured()) {
+                logVmOutputAndMonitorBootEvents(
+                        logTag, vm.getConsoleOutput(), "Console", mConsoleOutput);
+                logVmOutput(logTag, vm.getLogOutput(), "Log", mLogOutput);
+            }
             mExecutorService.awaitTermination(300, TimeUnit.SECONDS);
         }
 
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/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index d8e74f7..5c3be5f 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -153,6 +153,7 @@
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setMemoryMib(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setVmOutputCaptured(false)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
 
@@ -332,18 +333,21 @@
         assertThat(minimal.isProtectedVm()).isEqualTo(isProtectedVm());
         assertThat(minimal.isEncryptedStorageEnabled()).isFalse();
         assertThat(minimal.getEncryptedStorageKib()).isEqualTo(0);
+        assertThat(minimal.isVmOutputCaptured()).isEqualTo(false);
 
         // Maximal has everything that can be set to some non-default value. (And has different
         // values than minimal for the required fields.)
         int maxCpus = Runtime.getRuntime().availableProcessors();
         VirtualMachineConfig.Builder maximalBuilder =
-                newVmConfigBuilder()
+                new VirtualMachineConfig.Builder(getContext())
+                        .setProtectedVm(mProtectedVm)
                         .setPayloadConfigPath("config/path")
                         .setApkPath("/apk/path")
                         .setNumCpus(maxCpus)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .setMemoryMib(42)
-                        .setEncryptedStorageKib(1024);
+                        .setEncryptedStorageKib(1024)
+                        .setVmOutputCaptured(true);
         VirtualMachineConfig maximal = maximalBuilder.build();
 
         assertThat(maximal.getApkPath()).isEqualTo("/apk/path");
@@ -355,6 +359,7 @@
         assertThat(maximal.isProtectedVm()).isEqualTo(isProtectedVm());
         assertThat(maximal.isEncryptedStorageEnabled()).isTrue();
         assertThat(maximal.getEncryptedStorageKib()).isEqualTo(1024);
+        assertThat(maximal.isVmOutputCaptured()).isEqualTo(true);
 
         assertThat(minimal.isCompatibleWith(maximal)).isFalse();
         assertThat(minimal.isCompatibleWith(minimal)).isTrue();
@@ -392,6 +397,14 @@
                 new VirtualMachineConfig.Builder(getContext()).setPayloadBinaryName("binary.so");
         e = assertThrows(IllegalStateException.class, () -> protectedNotSet.build());
         assertThat(e).hasMessageThat().contains("setProtectedVm must be called");
+
+        VirtualMachineConfig.Builder captureOutputOnNonDebuggable =
+                newVmConfigBuilder()
+                        .setPayloadBinaryName("binary.so")
+                        .setDebugLevel(VirtualMachineConfig.DEBUG_LEVEL_NONE)
+                        .setVmOutputCaptured(true);
+        e = assertThrows(IllegalStateException.class, () -> captureOutputOnNonDebuggable.build());
+        assertThat(e).hasMessageThat().contains("debug level must be FULL to capture output");
     }
 
     private VirtualMachineConfig.Builder newBaselineBuilder() {
@@ -431,6 +444,11 @@
         assertConfigCompatible(baseline, newBaselineBuilder().setApkPath("/different")).isFalse();
         assertConfigCompatible(baseline, newBaselineBuilder().setEncryptedStorageKib(100))
                 .isFalse();
+
+        VirtualMachineConfig.Builder debuggableBuilder =
+                newBaselineBuilder().setDebugLevel(DEBUG_LEVEL_FULL);
+        VirtualMachineConfig debuggable = debuggableBuilder.build();
+        assertConfigCompatible(debuggable, debuggableBuilder.setVmOutputCaptured(true)).isFalse();
     }
 
     private BooleanSubject assertConfigCompatible(
@@ -662,6 +680,7 @@
                             .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                             .setMemoryMib(memMib)
                             .setDebugLevel(DEBUG_LEVEL_NONE)
+                            .setVmOutputCaptured(false)
                             .build();
             VirtualMachine vm = forceCreateNewVirtualMachine("low_mem", lowMemConfig);
             final CompletableFuture<Boolean> onPayloadReadyExecuted = new CompletableFuture<>();
@@ -704,7 +723,8 @@
         VirtualMachineConfig.Builder builder =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setDebugLevel(fromLevel);
+                        .setDebugLevel(fromLevel)
+                        .setVmOutputCaptured(false);
         VirtualMachineConfig normalConfig = builder.build();
         forceCreateNewVirtualMachine("test_vm", normalConfig);
         assertThat(tryBootVm(TAG, "test_vm").payloadStarted).isTrue();
@@ -1031,9 +1051,12 @@
 
     @Test
     public void bootFailsWhenBinaryIsMissingEntryFunction() throws Exception {
-        VirtualMachineConfig.Builder builder =
-                newVmConfigBuilder().setPayloadBinaryName("MicrodroidEmptyNativeLib.so");
-        VirtualMachineConfig normalConfig = builder.setDebugLevel(DEBUG_LEVEL_FULL).build();
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadBinaryName("MicrodroidEmptyNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setVmOutputCaptured(true)
+                        .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_missing_entry", normalConfig);
 
         assertThatPayloadFailsDueTo(vm, "Failed to find entrypoint");
@@ -1041,9 +1064,12 @@
 
     @Test
     public void bootFailsWhenBinaryTriesToLinkAgainstPrivateLibs() throws Exception {
-        VirtualMachineConfig.Builder builder =
-                newVmConfigBuilder().setPayloadBinaryName("MicrodroidPrivateLinkingNativeLib.so");
-        VirtualMachineConfig normalConfig = builder.setDebugLevel(DEBUG_LEVEL_FULL).build();
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadBinaryName("MicrodroidPrivateLinkingNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setVmOutputCaptured(true)
+                        .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_private_linking", normalConfig);
 
         assertThatPayloadFailsDueTo(vm, "Failed to dlopen");
@@ -1233,6 +1259,29 @@
         assertThat(testResults.mFileContent).isEqualTo("Hello, I am a file!");
     }
 
+    @Test
+    public void outputsShouldBeExplicitlyForwarded() throws Exception {
+        assumeSupportedKernel();
+
+        final VirtualMachineConfig vmConfig =
+                new VirtualMachineConfig.Builder(ApplicationProvider.getApplicationContext())
+                        .setProtectedVm(mProtectedVm)
+                        .setPayloadBinaryName("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        final VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_forward_log", vmConfig);
+        vm.run();
+
+        try {
+            assertThrowsVmExceptionContaining(
+                    () -> vm.getConsoleOutput(), "Capturing vm outputs is turned off");
+            assertThrowsVmExceptionContaining(
+                    () -> vm.getLogOutput(), "Capturing vm outputs is turned off");
+        } finally {
+            vm.stop();
+        }
+    }
+
     private void assertFileContentsAreEqualInTwoVms(String fileName, String vmName1, String vmName2)
             throws IOException {
         File file1 = getVmFile(vmName1, fileName);
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