Merge "Revert "Mount /data with MS_NOEXEC""
diff --git a/encryptedstore/src/main.rs b/encryptedstore/src/main.rs
index 888485b..2f54534 100644
--- a/encryptedstore/src/main.rs
+++ b/encryptedstore/src/main.rs
@@ -63,7 +63,7 @@
 
 fn encryptedstore_init(blkdevice: &Path, key: &str, mountpoint: &Path) -> Result<()> {
     ensure!(
-        std::fs::metadata(&blkdevice)
+        std::fs::metadata(blkdevice)
             .context(format!("Failed to get metadata of {:?}", blkdevice))?
             .file_type()
             .is_block_device(),
diff --git a/javalib/32/public/api/android.system.virtualmachine-removed.txt b/javalib/32/public/api/android.system.virtualmachine-removed.txt
deleted file mode 100644
index d802177..0000000
--- a/javalib/32/public/api/android.system.virtualmachine-removed.txt
+++ /dev/null
@@ -1 +0,0 @@
-// Signature format: 2.0
diff --git a/javalib/32/public/api/android.system.virtualmachine.txt b/javalib/32/public/api/android.system.virtualmachine.txt
deleted file mode 100644
index d802177..0000000
--- a/javalib/32/public/api/android.system.virtualmachine.txt
+++ /dev/null
@@ -1 +0,0 @@
-// Signature format: 2.0
diff --git a/javalib/32/system/api/android.system.virtualmachine-removed.txt b/javalib/32/system/api/android.system.virtualmachine-removed.txt
deleted file mode 100644
index d802177..0000000
--- a/javalib/32/system/api/android.system.virtualmachine-removed.txt
+++ /dev/null
@@ -1 +0,0 @@
-// Signature format: 2.0
diff --git a/javalib/32/system/api/android.system.virtualmachine.txt b/javalib/32/system/api/android.system.virtualmachine.txt
deleted file mode 100644
index d802177..0000000
--- a/javalib/32/system/api/android.system.virtualmachine.txt
+++ /dev/null
@@ -1 +0,0 @@
-// Signature format: 2.0
diff --git a/javalib/Android.bp b/javalib/Android.bp
index 118b648..f5f91f2 100644
--- a/javalib/Android.bp
+++ b/javalib/Android.bp
@@ -52,11 +52,6 @@
     ],
 }
 
-prebuilt_apis {
-    name: "android-virtualization-framework-sdk",
-    api_dirs: ["32"],
-}
-
 java_api_contribution {
     name: "framework-virtualization-public-stubs",
     api_surface: "public",
diff --git a/libs/apkverify/src/v3.rs b/libs/apkverify/src/v3.rs
index fcd966b..e1b728d 100644
--- a/libs/apkverify/src/v3.rs
+++ b/libs/apkverify/src/v3.rs
@@ -110,7 +110,7 @@
 ) -> Result<(Signer, ApkSections<R>)> {
     let mut sections = ApkSections::new(apk)?;
     let mut block = sections.find_signature(APK_SIGNATURE_SCHEME_V3_BLOCK_ID).context(
-        "Fallback to v2 when v3 block not found is not yet implemented. See b/197052981.",
+        "Fallback to v2 when v3 block not found is not yet implemented.", // b/197052981
     )?;
     let mut supported = block
         .read::<Signers>()?
@@ -135,7 +135,7 @@
             .iter()
             .filter(|sig| sig.signature_algorithm_id.map_or(false, |algo| algo.is_supported()))
             .max_by_key(|sig| sig.signature_algorithm_id.unwrap().content_digest_algorithm())
-            .context("No supported signatures found")?)
+            .context("No supported APK signatures found; DSA is not supported")?)
     }
 
     pub(crate) fn find_digest_by_algorithm(
diff --git a/libs/apkverify/tests/apkverify_test.rs b/libs/apkverify/tests/apkverify_test.rs
index 047538c..baf7c42 100644
--- a/libs/apkverify/tests/apkverify_test.rs
+++ b/libs/apkverify/tests/apkverify_test.rs
@@ -44,7 +44,7 @@
     for key_name in KEY_NAMES_DSA.iter() {
         let res = verify(format!("tests/data/v3-only-with-dsa-sha256-{}.apk", key_name));
         assert!(res.is_err(), "DSA algorithm is not supported for verification. See b/197052981.");
-        assert_contains(&res.unwrap_err().to_string(), "No supported signatures found");
+        assert_contains(&res.unwrap_err().to_string(), "No supported APK signatures found");
     }
 }
 
@@ -151,7 +151,7 @@
 fn test_verify_v3_no_supported_sig_algs() {
     let res = verify("tests/data/v3-only-no-supported-sig-algs.apk");
     assert!(res.is_err());
-    assert_contains(&res.unwrap_err().to_string(), "No supported signatures found");
+    assert_contains(&res.unwrap_err().to_string(), "No supported APK signatures found");
 }
 
 #[test]
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 7c72fab..c6d6c2c 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -508,6 +508,19 @@
         fdt_err_expect_zero(ret)
     }
 
+    /// Applies a DT overlay on the base DT.
+    ///
+    /// # Safety
+    ///
+    /// On failure, the library corrupts the DT and overlay so both must be discarded.
+    pub unsafe fn apply_overlay<'a>(&'a mut self, overlay: &'a mut Fdt) -> Result<&'a mut Self> {
+        fdt_err_expect_zero(libfdt_bindgen::fdt_overlay_apply(
+            self.as_mut_ptr(),
+            overlay.as_mut_ptr(),
+        ))?;
+        Ok(self)
+    }
+
     /// Return an iterator of memory banks specified the "/memory" node.
     ///
     /// NOTE: This does not support individual "/memory@XXXX" banks.
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index f48611b..24a12f7 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -77,6 +77,7 @@
     "/sys/firmware/devicetree/base/virtualization/guest/debug-microdroid,no-verified-boot";
 
 const APEX_CONFIG_DONE_PROP: &str = "apex_config.done";
+const TOMBSTONE_TRANSMIT_DONE_PROP: &str = "tombstone_transmit.init_done";
 const DEBUGGABLE_PROP: &str = "ro.boot.microdroid.debuggable";
 
 // SYNC WITH virtualizationservice/src/crosvm.rs
@@ -435,6 +436,10 @@
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice_context)?;
 
+    if config.export_tombstones {
+        wait_for_tombstone_transmit_done()?;
+    }
+
     // Wait for encryptedstore to finish mounting the storage (if enabled) before setting
     // microdroid_manager.init_done. Reason is init stops uneventd after that.
     // Encryptedstore, however requires ueventd
@@ -723,6 +728,11 @@
     wait_for_property_true(APEX_CONFIG_DONE_PROP).context("Failed waiting for apex config done")
 }
 
+fn wait_for_tombstone_transmit_done() -> Result<()> {
+    wait_for_property_true(TOMBSTONE_TRANSMIT_DONE_PROP)
+        .context("Failed waiting for tombstone transmit done")
+}
+
 fn wait_for_property_true(property_name: &str) -> Result<()> {
     let mut prop = PropertyWatcher::new(property_name)?;
     loop {
diff --git a/pvmfw/README.md b/pvmfw/README.md
index e5ba88b..f46c718 100644
--- a/pvmfw/README.md
+++ b/pvmfw/README.md
@@ -1,12 +1,107 @@
 # Protected Virtual Machine Firmware
 
-## Configuration Data Format
+In the context of the [Android Virtualization Framework][AVF], a hypervisor
+(_e.g._ [pKVM]) enforces full memory isolation between its virtual machines
+(VMs) and the host.  As a result, the host is only allowed to access memory that
+has been explicitly shared back by a VM. Such _protected VMs_ (“pVMs”) are
+therefore able to manipulate secrets without being at risk of an attacker
+stealing them by compromising the Android host.
 
-pvmfw will expect a [header] to have been appended to its loaded binary image
-at the next 4KiB boundary. It describes the configuration data entries that
-pvmfw will use and, being loaded by the pvmfw loader, is necessarily trusted.
+As pVMs are started dynamically by a _virtual machine manager_ (“VMM”) running
+as a host process and as pVMs must not trust the host (see [_Why
+AVF?_][why-avf]), the virtual machine it configures can't be trusted either.
+Furthermore, even though the isolation mentioned above allows pVMs to protect
+their secrets from the host, it does not help with provisioning them during
+boot. In particular, the threat model would prohibit the host from ever having
+access to those secrets, preventing the VMM from passing them to the pVM.
 
-The layout of the configuration data is as follows:
+To address these concerns the hypervisor securely loads the pVM firmware
+(“pvmfw”) in the pVM from a protected memory region (this prevents the host or
+any pVM from tampering with it), setting it as the entry point of the virtual
+machine. As a result, pvmfw becomes the very first code that gets executed in
+the pVM, allowing it to validate the environment and abort the boot sequence if
+necessary. This process takes place whenever the VMM places a VM in protected
+mode and can’t be prevented by the host.
+
+Given the threat model, pvmfw is not allowed to trust the devices or device
+layout provided by the virtual platform it is running on as those are configured
+by the VMM. Instead, it performs all the necessary checks to ensure that the pVM
+was set up as expected. For functional purposes, the interface with the
+hypervisor, although trusted, is also validated.
+
+Once it has been determined that the platform can be trusted, pvmfw derives
+unique secrets for the guest through the [_Boot Certificate Chain_][BCC]
+("BCC", see [Open Profile for DICE][open-dice]) that can be used to prove the
+identity of the pVM to local and remote actors. If any operation or check fails,
+or in case of a missing prerequisite, pvmfw will abort the boot process of the
+pVM, effectively preventing non-compliant pVMs and/or guests from running.
+Otherwise, it hands over the pVM to the guest kernel by jumping to its first
+instruction, similarly to a bootloader.
+
+pvmfw currently only supports AArch64.
+
+[AVF]: https://source.android.com/docs/core/virtualization
+[why-avf]: https://source.android.com/docs/core/virtualization/whyavf
+[BCC]: https://pigweed.googlesource.com/open-dice/+/master/src/android/README.md
+[pKVM]: https://source.android.com/docs/core/virtualization/architecture#hypervisor
+[open-dice]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md
+
+## Integration
+
+### pvmfw Loading
+
+When running pKVM, the physical memory from which the hypervisor loads pvmfw
+into guest address space is not initially populated by the hypervisor itself.
+Instead, it receives a pre-loaded memory region from a trusted pvmfw loader and
+only then becomes responsible for protecting it. As a result, the hypervisor is
+kept generic (beyond AVF) and small as it is not expected (nor necessary) for it
+to know how to interpret or obtain the content of that region.
+
+#### Android Bootloader (ABL) Support
+
+Starting in Android T, the `PRODUCT_BUILD_PVMFW_IMAGE` build variable controls
+the generation of `pvmfw.img`, a new [ABL partition][ABL-part] containing the
+pvmfw binary and following the internal format of the [`boot`][boot-img]
+partition, intended to be verified and loaded by ABL on AVF-compatible devices.
+
+To support pKVM, ABL is expected to describe the region using a reserved memory
+device tree node where both address and size have been properly aligned to the
+page size used by the hypervisor. For example, the following node describes a
+region of size `0x40000` at address `0x80000000`:
+```
+reserved-memory {
+    ...
+    pkvm_guest_firmware {
+        compatible = "linux,pkvm-guest-firmware-memory";
+        reg = <0x0 0x80000000 0x40000>;
+        no-map;
+    }
+}
+```
+
+[ABL-part]: https://source.android.com/docs/core/architecture/bootloader/partitions
+[boot-img]: https://source.android.com/docs/core/architecture/bootloader/boot-image-header
+
+### Configuration Data
+
+As part of the process of loading pvmfw, the loader (typically the Android
+Bootloader, "ABL") is expected to pass device-specific pvmfw configuration data
+by appending it to the pvmfw binary and including it in the region passed to the
+hypervisor. As a result, the hypervisor will give the same protection to this
+data as it does to pvmfw and will transparently load it in guest memory, making
+it available to pvmfw at runtime. This enables pvmfw to be kept device-agnostic,
+simplifying its adoption and distribution as a centralized signed binary, while
+also being able to support device-specific details.
+
+The configuration data will be read by pvmfw at the next 4KiB boundary from the
+end of its loaded binary. Even if the pvmfw is position-independent, it will be
+expected for it to also have been loaded at a 4-KiB boundary. As a result, the
+location of the configuration data is implicitly passed to pvmfw and known to it
+at build time.
+
+#### Configuration Data Format
+
+The configuration data is described using the following [header]:
 
 ```
 +===============================+
@@ -64,9 +159,62 @@
 The header format itself is agnostic of the internal format of the individual
 blos it refers to. In version 1.0, it describes two blobs:
 
-- entry 0 must point to a valid [BCC Handover]
+- entry 0 must point to a valid BCC Handover (see below)
 - entry 1 may point to a [DTBO] to be applied to the pVM device tree
 
 [header]: src/config.rs
-[BCC Handover]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c6efcd8c35506d818c18d1e73b9834a/src/android/bcc.c#260
 [DTBO]: https://android.googlesource.com/platform/external/dtc/+/refs/heads/master/Documentation/dt-object-internal.txt
+
+#### Virtual Platform Boot Certificate Chain Handover
+
+The format of the BCC entry mentioned above, compatible with the
+[`BccHandover`][BccHandover] defined by the Open Profile for DICE reference
+implementation, is described by the following [CDDL][CDDL]:
+```
+PvmfwBccHandover = {
+  1 : bstr .size 32,     ; CDI_Attest
+  2 : bstr .size 32,     ; CDI_Seal
+  3 : Bcc,               ; Certificate chain
+}
+```
+
+and contains the _Compound Device Identifiers_ ("CDIs"), used to derive the
+next-stage secret, and a certificate chain, intended for pVM attestation. Note
+that it differs from the `BccHandover` defined by the specification in that its
+`Bcc` field is mandatory (while optional in the original).
+
+The handover expected by pvmfw can be generated as follows:
+
+- by passing a `BccHandover` received from a previous boot stage (_e.g._ Trusted
+  Firmware, ROM bootloader, ...) to
+  [`BccHandoverMainFlow`][BccHandoverMainFlow];
+
+- by generating a `BccHandover` (as an example, see [Trusty][Trusty-BCC]) with
+  both CDIs set to an arbitrary constant value and no `Bcc`, and pass it to
+  `BccHandoverMainFlow`, which will both derive the pvmfw CDIs and start a
+  valid certificate chain, making the pvmfw loader the root of the BCC.
+
+The recommended DICE inputs at this stage are:
+
+- **Code**: hash of the pvmfw image, hypervisor (`boot.img`), and other target
+  code relevant to the secure execution of pvmfw (_e.g._ `vendor_boot.img`)
+- **Configuration Data**: any extra input relevant to pvmfw security
+- **Authority Data**: must cover all the public keys used to sign and verify the
+  code contributing to the **Code** input
+- **Mode Decision**: Set according to the [specification][dice-mode]. In
+  particular, should only be `Normal` if secure boot is being properly enforced
+  (_e.g._ locked device in [Android Verified Boot][AVB])
+- **Hidden Inputs**: Factory Reset Secret (FRS, stored in a tamper evident
+  storage and changes during every factory reset) or similar that changes as
+  part of the device lifecycle (_e.g._ reset)
+
+The resulting `BccHandover` is then used by pvmfw in a similar way to derive
+another [DICE layer][Layering], passed to the guest through a `/reserved-memory`
+device tree node marked as [`compatible=”google,open-dice”`][dice-dt].
+
+[AVB]: https://source.android.com/docs/security/features/verifiedboot/boot-flow
+[BccHandover]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c/src/android/bcc.c#260
+[CDDL]: https://datatracker.ietf.org/doc/rfc8610
+[dice-dt]: https://www.kernel.org/doc/Documentation/devicetree/bindings/reserved-memory/google%2Copen-dice.yaml
+[Layering]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#layering-details
+[Trusty-BCC]: https://android.googlesource.com/trusty/lib/+/1696be0a8f3a7103/lib/hwbcc/common/swbcc.c#554
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 0527dfb..837f747 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -27,7 +27,7 @@
 rust_test {
     name: "libpvmfw_avb.integration_test",
     crate_name: "pvmfw_avb_test",
-    srcs: ["tests/*_test.rs"],
+    srcs: ["tests/*.rs"],
     test_suites: ["general-tests"],
     data: [
         ":avb_testkey_rsa2048_pub_bin",
diff --git a/pvmfw/avb/src/lib.rs b/pvmfw/avb/src/lib.rs
index 6a5b16d..f0ee4ed 100644
--- a/pvmfw/avb/src/lib.rs
+++ b/pvmfw/avb/src/lib.rs
@@ -19,7 +19,9 @@
 #![feature(mixed_integer_ops)]
 
 mod error;
+mod partition;
+mod utils;
 mod verify;
 
 pub use error::AvbSlotVerifyError;
-pub use verify::verify_payload;
+pub use verify::{verify_payload, DebugLevel};
diff --git a/pvmfw/avb/src/partition.rs b/pvmfw/avb/src/partition.rs
new file mode 100644
index 0000000..82d89f8
--- /dev/null
+++ b/pvmfw/avb/src/partition.rs
@@ -0,0 +1,86 @@
+// Copyright 2023, 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.
+
+//! Struct and functions relating to well-known partition names.
+
+use crate::error::AvbIOError;
+use crate::utils::is_not_null;
+use core::ffi::{c_char, CStr};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub(crate) 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";
+
+    pub(crate) fn as_cstr(&self) -> &CStr {
+        CStr::from_bytes_with_nul(self.as_bytes()).unwrap()
+    }
+
+    fn as_non_null_terminated_bytes(&self) -> &[u8] {
+        let partition_name = self.as_bytes();
+        &partition_name[..partition_name.len() - 1]
+    }
+
+    fn as_bytes(&self) -> &[u8] {
+        match self {
+            Self::Kernel => Self::KERNEL_PARTITION_NAME,
+            Self::InitrdNormal => Self::INITRD_NORMAL_PARTITION_NAME,
+            Self::InitrdDebug => Self::INITRD_DEBUG_PARTITION_NAME,
+        }
+    }
+}
+
+impl TryFrom<*const c_char> for PartitionName {
+    type Error = AvbIOError;
+
+    fn try_from(partition_name: *const c_char) -> Result<Self, Self::Error> {
+        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) };
+        partition_name.try_into()
+    }
+}
+
+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),
+        }
+    }
+}
+
+impl TryFrom<&[u8]> for PartitionName {
+    type Error = AvbIOError;
+
+    fn try_from(non_null_terminated_name: &[u8]) -> Result<Self, Self::Error> {
+        match non_null_terminated_name {
+            x if x == Self::Kernel.as_non_null_terminated_bytes() => Ok(Self::Kernel),
+            x if x == Self::InitrdNormal.as_non_null_terminated_bytes() => Ok(Self::InitrdNormal),
+            x if x == Self::InitrdDebug.as_non_null_terminated_bytes() => Ok(Self::InitrdDebug),
+            _ => Err(AvbIOError::NoSuchPartition),
+        }
+    }
+}
diff --git a/pvmfw/avb/src/utils.rs b/pvmfw/avb/src/utils.rs
new file mode 100644
index 0000000..1b35c22
--- /dev/null
+++ b/pvmfw/avb/src/utils.rs
@@ -0,0 +1,56 @@
+// Copyright 2023, 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.
+
+//! Common utility functions.
+
+use crate::error::AvbIOError;
+use core::ptr::NonNull;
+use core::result;
+
+pub(crate) type Result<T> = result::Result<T, AvbIOError>;
+
+pub(crate) fn write<T>(ptr: *mut T, value: T) -> Result<()> {
+    let ptr = to_nonnull(ptr)?;
+    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
+    unsafe {
+        *ptr.as_ptr() = value;
+    }
+    Ok(())
+}
+
+pub(crate) fn as_ref<'a, T>(ptr: *mut T) -> Result<&'a T> {
+    let ptr = to_nonnull(ptr)?;
+    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
+    unsafe { Ok(ptr.as_ref()) }
+}
+
+pub(crate) fn to_nonnull<T>(ptr: *mut T) -> Result<NonNull<T>> {
+    NonNull::new(ptr).ok_or(AvbIOError::NoSuchValue)
+}
+
+pub(crate) fn is_not_null<T>(ptr: *const T) -> Result<()> {
+    if ptr.is_null() {
+        Err(AvbIOError::NoSuchValue)
+    } else {
+        Ok(())
+    }
+}
+
+pub(crate) fn to_usize<T: TryInto<usize>>(num: T) -> Result<usize> {
+    num.try_into().map_err(|_| AvbIOError::InvalidValueSize)
+}
+
+pub(crate) fn usize_checked_add(x: usize, y: usize) -> Result<usize> {
+    x.checked_add(y).ok_or(AvbIOError::InvalidValueSize)
+}
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index 9d40075..b39bb5a 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -17,6 +17,8 @@
 use crate::error::{
     slot_verify_result_to_verify_payload_result, to_avb_io_result, AvbIOError, AvbSlotVerifyError,
 };
+use crate::partition::PartitionName;
+use crate::utils::{as_ref, is_not_null, to_nonnull, to_usize, usize_checked_add, write};
 use avb_bindgen::{
     avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, avb_slot_verify,
     avb_slot_verify_data_free, AvbDescriptor, AvbHashDescriptor, AvbHashtreeErrorMode, AvbIOResult,
@@ -25,12 +27,20 @@
 use core::{
     ffi::{c_char, c_void, CStr},
     mem::{size_of, MaybeUninit},
-    ptr::{self, NonNull},
-    slice,
+    ptr, slice,
 };
 
 const NULL_BYTE: &[u8] = b"\0";
 
+/// This enum corresponds to the `DebugLevel` in `VirtualMachineConfig`.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum DebugLevel {
+    /// Not debuggable at all.
+    None,
+    /// Fully debuggable.
+    Full,
+}
+
 extern "C" fn read_is_device_unlocked(
     _ops: *mut AvbOps,
     out_is_unlocked: *mut bool,
@@ -156,38 +166,31 @@
     AvbIOResult::AVB_IO_RESULT_OK
 }
 
-extern "C" fn validate_public_key_for_partition(
+extern "C" fn validate_vbmeta_public_key(
     ops: *mut AvbOps,
-    partition: *const c_char,
     public_key_data: *const u8,
     public_key_length: usize,
     public_key_metadata: *const u8,
     public_key_metadata_length: usize,
     out_is_trusted: *mut bool,
-    out_rollback_index_location: *mut u32,
 ) -> AvbIOResult {
-    to_avb_io_result(try_validate_public_key_for_partition(
+    to_avb_io_result(try_validate_vbmeta_public_key(
         ops,
-        partition,
         public_key_data,
         public_key_length,
         public_key_metadata,
         public_key_metadata_length,
         out_is_trusted,
-        out_rollback_index_location,
     ))
 }
 
-#[allow(clippy::too_many_arguments)]
-fn try_validate_public_key_for_partition(
+fn try_validate_vbmeta_public_key(
     ops: *mut AvbOps,
-    partition: *const c_char,
     public_key_data: *const u8,
     public_key_length: usize,
     _public_key_metadata: *const u8,
     _public_key_metadata_length: usize,
     out_is_trusted: *mut bool,
-    _out_rollback_index_location: *mut u32,
 ) -> Result<(), AvbIOError> {
     is_not_null(public_key_data)?;
     // SAFETY: It is safe to create a slice with the given pointer and length as
@@ -195,8 +198,6 @@
     // `public_key_length`.
     let public_key = unsafe { slice::from_raw_parts(public_key_data, public_key_length) };
     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;
     write(out_is_trusted, public_key == trusted_public_key)
 }
@@ -276,108 +277,6 @@
     }
 }
 
-fn to_usize<T: TryInto<usize>>(num: T) -> Result<usize, AvbIOError> {
-    num.try_into().map_err(|_| AvbIOError::InvalidValueSize)
-}
-
-fn usize_checked_add(x: usize, y: usize) -> Result<usize, AvbIOError> {
-    x.checked_add(y).ok_or(AvbIOError::InvalidValueSize)
-}
-
-fn write<T>(ptr: *mut T, value: T) -> Result<(), AvbIOError> {
-    let ptr = to_nonnull(ptr)?;
-    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
-    unsafe {
-        *ptr.as_ptr() = value;
-    }
-    Ok(())
-}
-
-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>(ptr: *mut T) -> Result<NonNull<T>, AvbIOError> {
-    NonNull::new(ptr).ok_or(AvbIOError::NoSuchValue)
-}
-
-fn is_not_null<T>(ptr: *const T) -> Result<(), AvbIOError> {
-    if ptr.is_null() {
-        Err(AvbIOError::NoSuchValue)
-    } else {
-        Ok(())
-    }
-}
-
-#[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 {
-        CStr::from_bytes_with_nul(self.as_bytes()).unwrap()
-    }
-
-    fn as_non_null_terminated_bytes(&self) -> &[u8] {
-        let partition_name = self.as_bytes();
-        &partition_name[..partition_name.len() - 1]
-    }
-
-    fn as_bytes(&self) -> &[u8] {
-        match self {
-            Self::Kernel => Self::KERNEL_PARTITION_NAME,
-            Self::InitrdNormal => Self::INITRD_NORMAL_PARTITION_NAME,
-            Self::InitrdDebug => Self::INITRD_DEBUG_PARTITION_NAME,
-        }
-    }
-}
-
-impl TryFrom<*const c_char> for PartitionName {
-    type Error = AvbIOError;
-
-    fn try_from(partition_name: *const c_char) -> Result<Self, Self::Error> {
-        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) };
-        partition_name.try_into()
-    }
-}
-
-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),
-        }
-    }
-}
-
-impl TryFrom<&[u8]> for PartitionName {
-    type Error = AvbIOError;
-
-    fn try_from(non_null_terminated_name: &[u8]) -> Result<Self, Self::Error> {
-        match non_null_terminated_name {
-            x if x == Self::Kernel.as_non_null_terminated_bytes() => Ok(Self::Kernel),
-            x if x == Self::InitrdNormal.as_non_null_terminated_bytes() => Ok(Self::InitrdNormal),
-            x if x == Self::InitrdDebug.as_non_null_terminated_bytes() => Ok(Self::InitrdDebug),
-            _ => Err(AvbIOError::NoSuchPartition),
-        }
-    }
-}
-
 struct AvbSlotVerifyDataWrap(*mut AvbSlotVerifyData);
 
 impl TryFrom<*mut AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
@@ -457,7 +356,7 @@
             read_from_partition: Some(read_from_partition),
             get_preloaded_partition: Some(get_preloaded_partition),
             write_to_partition: None,
-            validate_vbmeta_public_key: None,
+            validate_vbmeta_public_key: Some(validate_vbmeta_public_key),
             read_rollback_index: Some(read_rollback_index),
             write_rollback_index: None,
             read_is_device_unlocked: Some(read_is_device_unlocked),
@@ -465,7 +364,7 @@
             get_size_of_partition: Some(get_size_of_partition),
             read_persistent_value: None,
             write_persistent_value: None,
-            validate_public_key_for_partition: Some(validate_public_key_for_partition),
+            validate_public_key_for_partition: None,
         };
         let ab_suffix = CStr::from_bytes_with_nul(NULL_BYTE).unwrap();
         let mut out_data = MaybeUninit::uninit();
@@ -478,7 +377,7 @@
                 &mut avb_ops,
                 requested_partitions.as_ptr(),
                 ab_suffix.as_ptr(),
-                AvbSlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION,
+                AvbSlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
                 AvbHashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
                 out_data.as_mut_ptr(),
             )
@@ -528,7 +427,7 @@
     kernel: &[u8],
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
-) -> Result<(), AvbSlotVerifyError> {
+) -> Result<DebugLevel, AvbSlotVerifyError> {
     let mut payload = Payload { kernel, initrd, trusted_public_key };
     let kernel_verify_result = payload.verify_partition(PartitionName::Kernel.as_cstr())?;
     let vbmeta_images = kernel_verify_result.vbmeta_images()?;
@@ -540,8 +439,17 @@
     verify_vbmeta_is_from_kernel_partition(&vbmeta_image)?;
     if payload.initrd.is_none() {
         verify_vbmeta_has_no_initrd_descriptor(&vbmeta_image)?;
+        return Ok(DebugLevel::None);
     }
     // TODO(b/256148034): Check the vbmeta doesn't have hash descriptors other than
     // boot, initrd_normal, initrd_debug.
-    Ok(())
+
+    let debug_level = if payload.verify_partition(PartitionName::InitrdNormal.as_cstr()).is_ok() {
+        DebugLevel::None
+    } else if payload.verify_partition(PartitionName::InitrdDebug.as_cstr()).is_ok() {
+        DebugLevel::Full
+    } else {
+        return Err(AvbSlotVerifyError::Verification);
+    };
+    Ok(debug_level)
 }
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index 872ad63..41ead59 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -14,72 +14,68 @@
  * limitations under the License.
  */
 
-use anyhow::Result;
-use avb_bindgen::{
-    avb_footer_validate_and_byteswap, avb_vbmeta_image_header_to_host_byte_order, AvbFooter,
-    AvbVBMetaImageHeader,
-};
-use pvmfw_avb::{verify_payload, AvbSlotVerifyError};
-use std::{
-    fs,
-    mem::{size_of, transmute, MaybeUninit},
-    ptr,
-};
+mod utils;
 
-const MICRODROID_KERNEL_IMG_PATH: &str = "microdroid_kernel";
-const INITRD_NORMAL_IMG_PATH: &str = "microdroid_initrd_normal.img";
-const INITRD_DEBUG_IMG_PATH: &str = "microdroid_initrd_debuggable.img";
+use anyhow::Result;
+use avb_bindgen::{AvbFooter, AvbVBMetaImageHeader};
+use pvmfw_avb::{AvbSlotVerifyError, DebugLevel};
+use std::{fs, mem::size_of, ptr};
+use utils::*;
+
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
 const TEST_IMG_WITH_PROP_DESC_PATH: &str = "test_image_with_prop_desc.img";
 const TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH: &str = "test_image_with_non_initrd_hashdesc.img";
 const UNSIGNED_TEST_IMG_PATH: &str = "unsigned_test.img";
 
-const PUBLIC_KEY_RSA2048_PATH: &str = "data/testkey_rsa2048_pub.bin";
-const PUBLIC_KEY_RSA4096_PATH: &str = "data/testkey_rsa4096_pub.bin";
 const RANDOM_FOOTER_POS: usize = 30;
 
 /// This test uses the Microdroid payload compiled on the fly to check that
 /// the latest payload can be verified successfully.
 #[test]
 fn latest_normal_payload_passes_verification() -> Result<()> {
-    assert_payload_verification_succeeds(
+    assert_payload_verification_with_initrd_eq(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
+        Ok(DebugLevel::None),
     )
 }
 
 #[test]
 fn latest_debug_payload_passes_verification() -> Result<()> {
-    assert_payload_verification_succeeds(
+    assert_payload_verification_with_initrd_eq(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_debug()?,
         &load_trusted_public_key()?,
+        Ok(DebugLevel::Full),
     )
 }
 
 #[test]
 fn payload_expecting_no_initrd_passes_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_with_no_initrd_eq(
+    assert_payload_verification_eq(
         &fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH)?,
+        /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Ok(()),
+        Ok(DebugLevel::None),
     )
 }
 
 #[test]
 fn payload_with_non_initrd_descriptor_passes_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_with_no_initrd_eq(
+    assert_payload_verification_eq(
         &fs::read(TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH)?,
+        /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Ok(()),
+        Ok(DebugLevel::None),
     )
 }
 
 #[test]
 fn payload_with_prop_descriptor_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_with_no_initrd_eq(
+    assert_payload_verification_eq(
         &fs::read(TEST_IMG_WITH_PROP_DESC_PATH)?,
+        /*initrd=*/ None,
         &load_trusted_public_key()?,
         Err(AvbSlotVerifyError::InvalidMetadata),
     )
@@ -87,8 +83,9 @@
 
 #[test]
 fn payload_expecting_initrd_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_with_no_initrd_eq(
+    assert_payload_verification_eq(
         &load_latest_signed_kernel()?,
+        /*initrd=*/ None,
         &load_trusted_public_key()?,
         Err(AvbSlotVerifyError::InvalidMetadata),
     )
@@ -96,41 +93,51 @@
 
 #[test]
 fn payload_with_empty_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         /*trusted_public_key=*/ &[0u8; 0],
-        AvbSlotVerifyError::PublicKeyRejected,
+        Err(AvbSlotVerifyError::PublicKeyRejected),
     )
 }
 
 #[test]
 fn payload_with_an_invalid_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         /*trusted_public_key=*/ &[0u8; 512],
-        AvbSlotVerifyError::PublicKeyRejected,
+        Err(AvbSlotVerifyError::PublicKeyRejected),
     )
 }
 
 #[test]
 fn payload_with_a_different_valid_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         &fs::read(PUBLIC_KEY_RSA2048_PATH)?,
-        AvbSlotVerifyError::PublicKeyRejected,
+        Err(AvbSlotVerifyError::PublicKeyRejected),
+    )
+}
+
+#[test]
+fn payload_with_an_invalid_initrd_fails_verification() -> Result<()> {
+    assert_payload_verification_with_initrd_eq(
+        &load_latest_signed_kernel()?,
+        /*initrd=*/ &fs::read(UNSIGNED_TEST_IMG_PATH)?,
+        &load_trusted_public_key()?,
+        Err(AvbSlotVerifyError::Verification),
     )
 }
 
 #[test]
 fn unsigned_kernel_fails_verification() -> Result<()> {
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &fs::read(UNSIGNED_TEST_IMG_PATH)?,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::Io,
+        Err(AvbSlotVerifyError::Io),
     )
 }
 
@@ -139,25 +146,58 @@
     let mut kernel = load_latest_signed_kernel()?;
     kernel[1] = !kernel[1]; // Flip the bits
 
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::Verification,
+        Err(AvbSlotVerifyError::Verification),
     )
 }
 
 #[test]
+fn kernel_footer_with_vbmeta_offset_overwritten_fails_verification() -> Result<()> {
+    // Arrange.
+    let mut kernel = load_latest_signed_kernel()?;
+    let total_len = kernel.len() as u64;
+    let footer = extract_avb_footer(&kernel)?;
+    assert!(footer.vbmeta_offset < total_len);
+    let vbmeta_offset_addr = ptr::addr_of!(footer.vbmeta_offset) as *const u8;
+    // SAFETY: It is safe as both raw pointers `vbmeta_offset_addr` and `footer` are not null.
+    let vbmeta_offset_start =
+        unsafe { vbmeta_offset_addr.offset_from(ptr::addr_of!(footer) as *const u8) };
+    let footer_start = kernel.len() - size_of::<AvbFooter>();
+    let vbmeta_offset_start = footer_start + usize::try_from(vbmeta_offset_start)?;
+
+    let wrong_offsets = [total_len, u64::MAX];
+    for &wrong_offset in wrong_offsets.iter() {
+        // Act.
+        kernel[vbmeta_offset_start..(vbmeta_offset_start + size_of::<u64>())]
+            .copy_from_slice(&wrong_offset.to_be_bytes());
+
+        // Assert.
+        let footer = extract_avb_footer(&kernel)?;
+        assert_eq!(wrong_offset, footer.vbmeta_offset as u64);
+        assert_payload_verification_with_initrd_eq(
+            &kernel,
+            &load_latest_initrd_normal()?,
+            &load_trusted_public_key()?,
+            Err(AvbSlotVerifyError::Io),
+        )?;
+    }
+    Ok(())
+}
+
+#[test]
 fn tampered_kernel_footer_fails_verification() -> Result<()> {
     let mut kernel = load_latest_signed_kernel()?;
     let avb_footer_index = kernel.len() - size_of::<AvbFooter>() + RANDOM_FOOTER_POS;
     kernel[avb_footer_index] = !kernel[avb_footer_index];
 
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::InvalidMetadata,
+        Err(AvbSlotVerifyError::InvalidMetadata),
     )
 }
 
@@ -169,11 +209,11 @@
 
     kernel[vbmeta_index] = !kernel[vbmeta_index]; // Flip the bits
 
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::InvalidMetadata,
+        Err(AvbSlotVerifyError::InvalidMetadata),
     )
 }
 
@@ -192,17 +232,17 @@
     kernel[public_key_offset..(public_key_offset + public_key_size)]
         .copy_from_slice(&empty_public_key);
 
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &empty_public_key,
-        AvbSlotVerifyError::Verification,
+        Err(AvbSlotVerifyError::Verification),
     )?;
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::Verification,
+        Err(AvbSlotVerifyError::Verification),
     )
 }
 
@@ -234,81 +274,10 @@
         AVB_VBMETA_IMAGE_FLAGS_VERIFICATION_DISABLED, vbmeta_header.flags as u32,
         "VBMeta verification flag should be disabled now."
     );
-    assert_payload_verification_fails(
+    assert_payload_verification_with_initrd_eq(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        AvbSlotVerifyError::Verification,
+        Err(AvbSlotVerifyError::Verification),
     )
 }
-
-fn extract_avb_footer(kernel: &[u8]) -> Result<AvbFooter> {
-    let footer_start = kernel.len() - size_of::<AvbFooter>();
-    // SAFETY: The slice is the same size as the struct which only contains simple data types.
-    let mut footer = unsafe {
-        transmute::<[u8; size_of::<AvbFooter>()], AvbFooter>(kernel[footer_start..].try_into()?)
-    };
-    // SAFETY: The function updates the struct in-place.
-    unsafe {
-        avb_footer_validate_and_byteswap(&footer, &mut footer);
-    }
-    Ok(footer)
-}
-
-fn extract_vbmeta_header(kernel: &[u8], footer: &AvbFooter) -> Result<AvbVBMetaImageHeader> {
-    let vbmeta_offset: usize = footer.vbmeta_offset.try_into()?;
-    let vbmeta_size: usize = footer.vbmeta_size.try_into()?;
-    let vbmeta_src = &kernel[vbmeta_offset..(vbmeta_offset + vbmeta_size)];
-    // SAFETY: The latest kernel has a valid VBMeta header at the position specified in footer.
-    let vbmeta_header = unsafe {
-        let mut header = MaybeUninit::uninit();
-        let src = vbmeta_src.as_ptr() as *const _ as *const AvbVBMetaImageHeader;
-        avb_vbmeta_image_header_to_host_byte_order(src, header.as_mut_ptr());
-        header.assume_init()
-    };
-    Ok(vbmeta_header)
-}
-
-fn assert_payload_verification_with_no_initrd_eq(
-    kernel: &[u8],
-    trusted_public_key: &[u8],
-    expected_result: Result<(), AvbSlotVerifyError>,
-) -> Result<()> {
-    assert_eq!(expected_result, verify_payload(kernel, /*initrd=*/ None, trusted_public_key));
-    Ok(())
-}
-
-fn assert_payload_verification_fails(
-    kernel: &[u8],
-    initrd: &[u8],
-    trusted_public_key: &[u8],
-    expected_error: AvbSlotVerifyError,
-) -> Result<()> {
-    assert_eq!(Err(expected_error), verify_payload(kernel, Some(initrd), trusted_public_key));
-    Ok(())
-}
-
-fn assert_payload_verification_succeeds(
-    kernel: &[u8],
-    initrd: &[u8],
-    trusted_public_key: &[u8],
-) -> Result<()> {
-    assert_eq!(Ok(()), verify_payload(kernel, Some(initrd), trusted_public_key));
-    Ok(())
-}
-
-fn load_latest_signed_kernel() -> Result<Vec<u8>> {
-    Ok(fs::read(MICRODROID_KERNEL_IMG_PATH)?)
-}
-
-fn load_latest_initrd_normal() -> Result<Vec<u8>> {
-    Ok(fs::read(INITRD_NORMAL_IMG_PATH)?)
-}
-
-fn load_latest_initrd_debug() -> Result<Vec<u8>> {
-    Ok(fs::read(INITRD_DEBUG_IMG_PATH)?)
-}
-
-fn load_trusted_public_key() -> Result<Vec<u8>> {
-    Ok(fs::read(PUBLIC_KEY_RSA4096_PATH)?)
-}
diff --git a/pvmfw/avb/tests/utils.rs b/pvmfw/avb/tests/utils.rs
new file mode 100644
index 0000000..0a2eac6
--- /dev/null
+++ b/pvmfw/avb/tests/utils.rs
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2023 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.
+ */
+
+//! Utility functions used by API tests.
+
+use anyhow::Result;
+use avb_bindgen::{
+    avb_footer_validate_and_byteswap, avb_vbmeta_image_header_to_host_byte_order, AvbFooter,
+    AvbVBMetaImageHeader,
+};
+use pvmfw_avb::{verify_payload, AvbSlotVerifyError, DebugLevel};
+use std::{
+    fs,
+    mem::{size_of, transmute, MaybeUninit},
+};
+
+const MICRODROID_KERNEL_IMG_PATH: &str = "microdroid_kernel";
+const INITRD_NORMAL_IMG_PATH: &str = "microdroid_initrd_normal.img";
+const INITRD_DEBUG_IMG_PATH: &str = "microdroid_initrd_debuggable.img";
+const PUBLIC_KEY_RSA4096_PATH: &str = "data/testkey_rsa4096_pub.bin";
+
+pub const PUBLIC_KEY_RSA2048_PATH: &str = "data/testkey_rsa2048_pub.bin";
+
+pub fn assert_payload_verification_with_initrd_eq(
+    kernel: &[u8],
+    initrd: &[u8],
+    trusted_public_key: &[u8],
+    expected_result: Result<DebugLevel, AvbSlotVerifyError>,
+) -> Result<()> {
+    assert_payload_verification_eq(kernel, Some(initrd), trusted_public_key, expected_result)
+}
+
+pub fn assert_payload_verification_eq(
+    kernel: &[u8],
+    initrd: Option<&[u8]>,
+    trusted_public_key: &[u8],
+    expected_result: Result<DebugLevel, AvbSlotVerifyError>,
+) -> Result<()> {
+    assert_eq!(expected_result, verify_payload(kernel, initrd, trusted_public_key));
+    Ok(())
+}
+
+pub fn load_latest_signed_kernel() -> Result<Vec<u8>> {
+    Ok(fs::read(MICRODROID_KERNEL_IMG_PATH)?)
+}
+
+pub fn load_latest_initrd_normal() -> Result<Vec<u8>> {
+    Ok(fs::read(INITRD_NORMAL_IMG_PATH)?)
+}
+
+pub fn load_latest_initrd_debug() -> Result<Vec<u8>> {
+    Ok(fs::read(INITRD_DEBUG_IMG_PATH)?)
+}
+
+pub fn load_trusted_public_key() -> Result<Vec<u8>> {
+    Ok(fs::read(PUBLIC_KEY_RSA4096_PATH)?)
+}
+
+pub fn extract_avb_footer(kernel: &[u8]) -> Result<AvbFooter> {
+    let footer_start = kernel.len() - size_of::<AvbFooter>();
+    // SAFETY: The slice is the same size as the struct which only contains simple data types.
+    let mut footer = unsafe {
+        transmute::<[u8; size_of::<AvbFooter>()], AvbFooter>(kernel[footer_start..].try_into()?)
+    };
+    // SAFETY: The function updates the struct in-place.
+    unsafe {
+        avb_footer_validate_and_byteswap(&footer, &mut footer);
+    }
+    Ok(footer)
+}
+
+pub fn extract_vbmeta_header(kernel: &[u8], footer: &AvbFooter) -> Result<AvbVBMetaImageHeader> {
+    let vbmeta_offset: usize = footer.vbmeta_offset.try_into()?;
+    let vbmeta_size: usize = footer.vbmeta_size.try_into()?;
+    let vbmeta_src = &kernel[vbmeta_offset..(vbmeta_offset + vbmeta_size)];
+    // SAFETY: The latest kernel has a valid VBMeta header at the position specified in footer.
+    let vbmeta_header = unsafe {
+        let mut header = MaybeUninit::uninit();
+        let src = vbmeta_src.as_ptr() as *const _ as *const AvbVBMetaImageHeader;
+        avb_vbmeta_image_header_to_host_byte_order(src, header.as_mut_ptr());
+        header.assume_init()
+    };
+    Ok(vbmeta_header)
+}
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index bfcb423..4f30902 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -178,6 +178,37 @@
     }
 }
 
+/// Applies the debug policy device tree overlay to the pVM DT.
+///
+/// # Safety
+///
+/// When an error is returned by this function, the input `Fdt` should be discarded as it may have
+/// have been partially corrupted during the overlay application process.
+unsafe fn apply_debug_policy(
+    fdt: &mut libfdt::Fdt,
+    debug_policy: &mut [u8],
+) -> Result<(), RebootReason> {
+    let overlay = libfdt::Fdt::from_mut_slice(debug_policy).map_err(|e| {
+        error!("Failed to load the debug policy overlay: {e}");
+        RebootReason::InvalidConfig
+    })?;
+
+    fdt.unpack().map_err(|e| {
+        error!("Failed to unpack DT for debug policy: {e}");
+        RebootReason::InternalError
+    })?;
+
+    let fdt = fdt.apply_overlay(overlay).map_err(|e| {
+        error!("Failed to apply the debug policy overlay: {e}");
+        RebootReason::InvalidConfig
+    })?;
+
+    fdt.pack().map_err(|e| {
+        error!("Failed to re-pack DT after debug policy: {e}");
+        RebootReason::InternalError
+    })
+}
+
 /// Sets up the environment for main() and wraps its result for start().
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
@@ -252,6 +283,11 @@
     helpers::flushed_zeroize(bcc_slice);
     helpers::flush(slices.fdt.as_slice());
 
+    if let Some(debug_policy) = appended.get_debug_policy() {
+        // SAFETY - As we `?` the result, there is no risk of re-using a bad `slices.fdt`.
+        unsafe { apply_debug_policy(slices.fdt, debug_policy) }?;
+    }
+
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
     memory.mmio_unmap_all().map_err(|e| {
         error!("Failed to unshare MMIO ranges: {e}");
diff --git a/vm/vm_shell.sh b/vm/vm_shell.sh
index 3db7003..b73a9dc 100755
--- a/vm/vm_shell.sh
+++ b/vm/vm_shell.sh
@@ -66,12 +66,16 @@
     fi
 
     if [ ! -n "${selected_cid}" ]; then
-        PS3="Select CID of VM to adb-shell into: "
-        select cid in ${available_cids}
-        do
-            selected_cid=${cid}
-            break
-        done
+        if [ ${#selected_cid[@]} -eq 1 ]; then
+            selected_cid=${available_cids[0]}
+        else
+            PS3="Select CID of VM to adb-shell into: "
+            select cid in ${available_cids}
+            do
+                selected_cid=${cid}
+                break
+            done
+        fi
     fi
 
     if [[ ! " ${available_cids[*]} " =~ " ${selected_cid} " ]]; then