Merge "Add more clap tests."
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index cb4e98c..d9fc70c 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -352,6 +352,7 @@
      * that would alter the identity of the VM (e.g. using a different payload or changing the debug
      * mode) are considered incompatible.
      *
+     * @see VirtualMachine#setConfig
      * @hide
      */
     @SystemApi
@@ -536,6 +537,14 @@
         /**
          * Sets the debug level. Defaults to {@link #DEBUG_LEVEL_NONE}.
          *
+         * <p>If {@link #DEBUG_LEVEL_FULL} is set then logs from inside the VM are exported to the
+         * host and adb connections from the host are possible. This is convenient for debugging but
+         * may compromise the integrity of the VM - including bypassing the protections offered by a
+         * {@linkplain #setProtectedVm protected VM}.
+         *
+         * <p>Note that it isn't possible to {@linkplain #isCompatibleWith change} the debug level
+         * of a VM instance; debug and non-debug VMs always have different secrets.
+         *
          * @hide
          */
         @SystemApi
@@ -552,6 +561,13 @@
          * Sets whether to protect the VM memory from the host. No default is provided, this must be
          * set explicitly.
          *
+         * <p>Note that if debugging is {@linkplain #setDebugLevel enabled} for a protected VM, the
+         * VM is not truly protected - direct memory access by the host is prevented, but e.g. the
+         * debugger can be used to access the VM's internals.
+         *
+         * <p>It isn't possible to {@linkplain #isCompatibleWith change} the protected status of a
+         * VM instance; protected and non-protected VMs always have different secrets.
+         *
          * @see VirtualMachineManager#getCapabilities
          * @hide
          */
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index b46f112..1fc44a2 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -544,7 +544,7 @@
 avb_gen_vbmeta_image {
     name: "microdroid_initrd_normal_hashdesc",
     src: ":microdroid_initrd_normal",
-    partition_name: "vbmeta_initrd_normal",
+    partition_name: "initrd_normal",
     salt: initrd_normal_salt,
     enabled: false,
     arch: {
@@ -564,7 +564,7 @@
 avb_gen_vbmeta_image {
     name: "microdroid_initrd_debug_hashdesc",
     src: ":microdroid_initrd_debuggable",
-    partition_name: "vbmeta_initrd_debug",
+    partition_name: "initrd_debug",
     salt: initrd_debug_salt,
     enabled: false,
     arch: {
@@ -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/src/main.rs b/microdroid/initrd/src/main.rs
index 0edb650..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())?;
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 a465b12..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)]
@@ -169,7 +101,7 @@
     out_pointer: *mut *mut u8,
     out_num_bytes_preloaded: *mut usize,
 ) -> Result<(), AvbIOError> {
-    let ops = as_avbops_ref(ops)?;
+    let ops = as_ref(ops)?;
     let partition = ops.as_ref().get_partition(partition)?;
     let out_pointer = to_nonnull(out_pointer)?;
     // SAFETY: It is safe as the raw pointer `out_pointer` is a nonnull pointer.
@@ -211,7 +143,7 @@
     buffer: *mut c_void,
     out_num_read: *mut usize,
 ) -> Result<(), AvbIOError> {
-    let ops = as_avbops_ref(ops)?;
+    let ops = as_ref(ops)?;
     let partition = ops.as_ref().get_partition(partition)?;
     let buffer = to_nonnull(buffer)?;
     // SAFETY: It is safe to copy the requested number of bytes to `buffer` as `buffer`
@@ -253,7 +185,7 @@
     partition: *const c_char,
     out_size_num_bytes: *mut u64,
 ) -> Result<(), AvbIOError> {
-    let ops = as_avbops_ref(ops)?;
+    let ops = as_ref(ops)?;
     let partition = ops.as_ref().get_partition(partition)?;
     let partition_size =
         u64::try_from(partition.len()).map_err(|_| AvbIOError::InvalidValueSize)?;
@@ -324,7 +256,7 @@
     // `public_key_data` is a valid pointer and it points to an array of length
     // `public_key_length`.
     let public_key = unsafe { slice::from_raw_parts(public_key_data, public_key_length) };
-    let ops = as_avbops_ref(ops)?;
+    let ops = as_ref(ops)?;
     // Verifies the public key for the known partitions only.
     ops.as_ref().get_partition(partition)?;
     let trusted_public_key = ops.as_ref().trusted_public_key;
@@ -336,14 +268,14 @@
     Ok(())
 }
 
-fn as_avbops_ref<'a>(ops: *mut AvbOps) -> Result<&'a AvbOps, AvbIOError> {
-    let ops = to_nonnull(ops)?;
-    // SAFETY: It is safe as the raw pointer `ops` is a nonnull pointer.
-    unsafe { Ok(ops.as_ref()) }
+fn as_ref<'a, T>(ptr: *mut T) -> Result<&'a T, AvbIOError> {
+    let ptr = to_nonnull(ptr)?;
+    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
+    unsafe { Ok(ptr.as_ref()) }
 }
 
-fn to_nonnull<T>(p: *mut T) -> Result<NonNull<T>, AvbIOError> {
-    NonNull::new(p).ok_or(AvbIOError::NoSuchValue)
+fn to_nonnull<T>(ptr: *mut T) -> Result<NonNull<T>, AvbIOError> {
+    NonNull::new(ptr).ok_or(AvbIOError::NoSuchValue)
 }
 
 fn is_not_null<T>(ptr: *const T) -> Result<(), AvbIOError> {
@@ -354,6 +286,41 @@
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+enum PartitionName {
+    Kernel,
+    InitrdNormal,
+    InitrdDebug,
+}
+
+impl PartitionName {
+    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";
+
+    fn as_cstr(&self) -> &CStr {
+        let partition_name = match self {
+            Self::Kernel => Self::KERNEL_PARTITION_NAME,
+            Self::InitrdNormal => Self::INITRD_NORMAL_PARTITION_NAME,
+            Self::InitrdDebug => Self::INITRD_DEBUG_PARTITION_NAME,
+        };
+        CStr::from_bytes_with_nul(partition_name).unwrap()
+    }
+}
+
+impl TryFrom<&CStr> for PartitionName {
+    type Error = AvbIOError;
+
+    fn try_from(partition_name: &CStr) -> Result<Self, Self::Error> {
+        match partition_name.to_bytes_with_nul() {
+            Self::KERNEL_PARTITION_NAME => Ok(Self::Kernel),
+            Self::INITRD_NORMAL_PARTITION_NAME => Ok(Self::InitrdNormal),
+            Self::INITRD_DEBUG_PARTITION_NAME => Ok(Self::InitrdDebug),
+            _ => Err(AvbIOError::NoSuchPartition),
+        }
+    }
+}
+
 struct Payload<'a> {
     kernel: &'a [u8],
     initrd: Option<&'a [u8]>,
@@ -372,28 +339,23 @@
 }
 
 impl<'a> Payload<'a> {
-    const KERNEL_PARTITION_NAME: &[u8] = b"bootloader\0";
-    const INITRD_NORMAL_PARTITION_NAME: &[u8] = b"initrd_normal\0";
-    const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
-
     const MAX_NUM_OF_HASH_DESCRIPTORS: usize = 3;
 
     fn get_partition(&self, partition_name: *const c_char) -> Result<&[u8], AvbIOError> {
         is_not_null(partition_name)?;
         // SAFETY: It is safe as the raw pointer `partition_name` is a nonnull pointer.
         let partition_name = unsafe { CStr::from_ptr(partition_name) };
-        match partition_name.to_bytes_with_nul() {
-            Self::KERNEL_PARTITION_NAME => Ok(self.kernel),
-            Self::INITRD_NORMAL_PARTITION_NAME | Self::INITRD_DEBUG_PARTITION_NAME => {
+        match partition_name.try_into()? {
+            PartitionName::Kernel => Ok(self.kernel),
+            PartitionName::InitrdNormal | PartitionName::InitrdDebug => {
                 self.initrd.ok_or(AvbIOError::NoSuchPartition)
             }
-            _ => Err(AvbIOError::NoSuchPartition),
         }
     }
 
-    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
@@ -434,7 +396,7 @@
                 out_data,
             )
         };
-        to_avb_verify_result(result)
+        slot_verify_result_to_verify_payload_result(result)
     }
 }
 
@@ -443,10 +405,9 @@
     kernel: &[u8],
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
-) -> Result<(), AvbImageVerifyError> {
+) -> Result<(), AvbSlotVerifyError> {
     let mut payload = Payload { kernel, initrd, trusted_public_key };
-    let kernel = CStr::from_bytes_with_nul(Payload::KERNEL_PARTITION_NAME).unwrap();
-    let requested_partitions = [kernel];
+    let requested_partitions = [PartitionName::Kernel.as_cstr()];
     payload.verify_partitions(&requested_partitions)
 }
 
@@ -496,7 +457,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             /*trusted_public_key=*/ &[0u8; 0],
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -506,7 +467,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             /*trusted_public_key=*/ &[0u8; 512],
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -516,7 +477,7 @@
             &load_latest_signed_kernel()?,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA2048_PATH)?,
-            AvbImageVerifyError::PublicKeyRejected,
+            AvbSlotVerifyError::PublicKeyRejected,
         )
     }
 
@@ -526,7 +487,7 @@
             &fs::read(UNSIGNED_TEST_IMG_PATH)?,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::Io,
+            AvbSlotVerifyError::Io,
         )
     }
 
@@ -539,7 +500,7 @@
             &kernel,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::Verification,
+            AvbSlotVerifyError::Verification,
         )
     }
 
@@ -553,7 +514,7 @@
             &kernel,
             &load_latest_initrd_normal()?,
             &fs::read(PUBLIC_KEY_RSA4096_PATH)?,
-            AvbImageVerifyError::InvalidMetadata,
+            AvbSlotVerifyError::InvalidMetadata,
         )
     }
 
@@ -561,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/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/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..1c82155 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,20 +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)
+        GLOBAL_SERVICE.debugListVms()
     }
 
     /// Hold a strong reference to a VM in VirtualizationService. This method is only intended for
@@ -564,13 +582,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 +596,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 +629,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 +684,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 +736,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 +1001,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 +1012,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 {
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