Merge "pvmfw: Refactor get_or_generate_instance_salt" into main
diff --git a/TEST_MAPPING b/TEST_MAPPING
index ec9042c..5b0c000 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -62,6 +62,9 @@
       // TODO(b/325610326): Add this target to presubmit once there is enough
       // SLO data for it.
       "name": "AvfRkpdAppIntegrationTests"
+    },
+    {
+      "name": "AvfRkpdVmAttestationTestApp"
     }
   ],
   "postsubmit": [
diff --git a/compos/common/Android.bp b/compos/common/Android.bp
index 01ab7c9..72cb5e1 100644
--- a/compos/common/Android.bp
+++ b/compos/common/Android.bp
@@ -20,6 +20,7 @@
         "libnum_traits",
         "librustutils",
         "libvmclient",
+        "libplatformproperties_rust",
     ],
     proc_macros: ["libnum_derive"],
     apex_available: [
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 077a0ef..ffdd0ea 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -35,6 +35,7 @@
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
 use glob::glob;
 use log::{info, warn};
+use platformproperties::hypervisorproperties;
 use rustutils::system_properties;
 use std::fs::File;
 use std::path::{Path, PathBuf};
@@ -232,7 +233,7 @@
 
 fn want_protected_vm() -> Result<bool> {
     let have_protected_vm =
-        system_properties::read_bool("ro.boot.hypervisor.protected_vm.supported", false)?;
+        hypervisorproperties::hypervisor_protected_vm_supported()?.unwrap_or(false);
     if have_protected_vm {
         info!("Starting protected VM");
         return Ok(true);
@@ -243,8 +244,7 @@
         bail!("Protected VM not supported, unable to start VM");
     }
 
-    let have_non_protected_vm =
-        system_properties::read_bool("ro.boot.hypervisor.vm.supported", false)?;
+    let have_non_protected_vm = hypervisorproperties::hypervisor_vm_supported()?.unwrap_or(false);
     if have_non_protected_vm {
         warn!("Protected VM not supported, falling back to non-protected on debuggable build");
         return Ok(false);
diff --git a/java/framework/Android.bp b/java/framework/Android.bp
index 32b2aee..26ea214 100644
--- a/java/framework/Android.bp
+++ b/java/framework/Android.bp
@@ -43,4 +43,7 @@
     impl_library_visibility: [
         "//packages/modules/Virtualization:__subpackages__",
     ],
+    aconfig_declarations: [
+        "avf_aconfig_flags",
+    ],
 }
diff --git a/libs/hypervisor_props/Android.bp b/libs/hypervisor_props/Android.bp
index af08b01..af6d417 100644
--- a/libs/hypervisor_props/Android.bp
+++ b/libs/hypervisor_props/Android.bp
@@ -9,7 +9,7 @@
     edition: "2021",
     rustlibs: [
         "libanyhow",
-        "librustutils",
+        "libplatformproperties_rust",
     ],
     apex_available: [
         "com.android.compos",
diff --git a/libs/hypervisor_props/src/lib.rs b/libs/hypervisor_props/src/lib.rs
index 120a48c..14614fd 100644
--- a/libs/hypervisor_props/src/lib.rs
+++ b/libs/hypervisor_props/src/lib.rs
@@ -14,18 +14,17 @@
 
 //! Access to hypervisor capabilities via system properties set by the bootloader.
 
-use anyhow::{Error, Result};
-use rustutils::system_properties;
+use anyhow::Result;
+use platformproperties::hypervisorproperties;
 
 /// Returns whether there is a hypervisor present that supports non-protected VMs.
 pub fn is_vm_supported() -> Result<bool> {
-    system_properties::read_bool("ro.boot.hypervisor.vm.supported", false).map_err(Error::new)
+    Ok(hypervisorproperties::hypervisor_vm_supported()?.unwrap_or(false))
 }
 
 /// Returns whether there is a hypervisor present that supports protected VMs.
 pub fn is_protected_vm_supported() -> Result<bool> {
-    system_properties::read_bool("ro.boot.hypervisor.protected_vm.supported", false)
-        .map_err(Error::new)
+    Ok(hypervisorproperties::hypervisor_protected_vm_supported()?.unwrap_or(false))
 }
 
 /// Returns whether there is a hypervisor present that supports any sort of VM, either protected
@@ -36,5 +35,5 @@
 
 /// Returns the version of the hypervisor, if there is one.
 pub fn version() -> Result<Option<String>> {
-    system_properties::read("ro.boot.hypervisor.version").map_err(Error::new)
+    Ok(hypervisorproperties::hypervisor_version()?)
 }
diff --git a/microdroid/init_debug_policy/src/init_debug_policy.rs b/microdroid/init_debug_policy/src/init_debug_policy.rs
index 90d04ac..c443088 100644
--- a/microdroid/init_debug_policy/src/init_debug_policy.rs
+++ b/microdroid/init_debug_policy/src/init_debug_policy.rs
@@ -15,7 +15,7 @@
 //! Applies debug policies when booting microdroid
 
 use rustutils::system_properties;
-use rustutils::system_properties::PropertyWatcherError;
+use rustutils::system_properties::error::PropertyWatcherError;
 use std::fs::File;
 use std::io::Read;
 
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index cce0e73..4ee02c1 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -74,16 +74,19 @@
     srcs: ["src/device_assignment.rs"],
     defaults: ["libpvmfw.test.defaults"],
     rustlibs: [
+        "libdts",
         "libhyp",
         "liblibfdt",
         "liblog_rust",
         "libpvmfw_fdt_template",
+        "libzerocopy",
     ],
     data: [
         ":test_pvmfw_devices_vm_dtbo",
         ":test_pvmfw_devices_vm_dtbo_without_symbols",
         ":test_pvmfw_devices_vm_dtbo_with_duplicated_iommus",
         ":test_pvmfw_devices_overlapping_pvmfw",
+        ":test_pvmfw_devices_vm_dtbo_with_dependencies",
         ":test_pvmfw_devices_with_rng",
         ":test_pvmfw_devices_with_multiple_devices_iommus",
         ":test_pvmfw_devices_with_iommu_sharing",
@@ -92,7 +95,13 @@
         ":test_pvmfw_devices_without_iommus",
         ":test_pvmfw_devices_with_duplicated_pviommus",
         ":test_pvmfw_devices_with_multiple_reg_iommus",
+        ":test_pvmfw_devices_with_dependency",
+        ":test_pvmfw_devices_with_dependency_loop",
+        ":test_pvmfw_devices_with_multiple_dependencies",
+        ":test_pvmfw_devices_expected_dt",
     ],
+    data_bins: ["dtc_static"],
+    compile_multilib: "first",
     // To use libpvmfw_fdt_template for testing
     enabled: false,
     target: {
@@ -136,6 +145,14 @@
     out: ["test_pvmfw_devices_vm_dtbo_with_duplicated_iommus.dtbo"],
 }
 
+genrule {
+    name: "test_pvmfw_devices_vm_dtbo_with_dependencies",
+    tools: ["dtc"],
+    cmd: "$(location dtc) -@ -I dts -O dtb $(in) -o $(out)",
+    srcs: ["testdata/test_pvmfw_devices_vm_dtbo_with_dependencies.dts"],
+    out: ["test_pvmfw_devices_vm_dtbo_with_dependencies.dtbo"],
+}
+
 genrule_defaults {
     name: "test_device_assignment_dts_to_dtb",
     defaults: ["dts_to_dtb"],
@@ -205,6 +222,53 @@
     out: ["test_pvmfw_devices_with_multiple_reg_iommus.dtb"],
 }
 
+genrule {
+    name: "test_pvmfw_devices_with_dependency",
+    defaults: ["test_device_assignment_dts_to_dtb"],
+    srcs: ["testdata/test_pvmfw_devices_with_dependency.dts"],
+    out: ["test_pvmfw_devices_with_dependency.dtb"],
+}
+
+genrule {
+    name: "test_pvmfw_devices_with_multiple_dependencies",
+    defaults: ["test_device_assignment_dts_to_dtb"],
+    srcs: ["testdata/test_pvmfw_devices_with_multiple_dependencies.dts"],
+    out: ["test_pvmfw_devices_with_multiple_dependencies.dtb"],
+}
+
+genrule {
+    name: "test_pvmfw_devices_with_dependency_loop",
+    defaults: ["test_device_assignment_dts_to_dtb"],
+    srcs: ["testdata/test_pvmfw_devices_with_dependency_loop.dts"],
+    out: ["test_pvmfw_devices_with_dependency_loop.dtb"],
+}
+
+// We can't use genrule because preprocessed platform DT is built with cc_object.
+// cc_genrule doesn't support default, so we'll build all expected DTs in
+// a single build rule.
+cc_genrule {
+    name: "test_pvmfw_devices_expected_dt",
+    srcs: [
+        ":pvmfw_platform.dts.preprocessed",
+        "testdata/expected_dt_with_dependency.dts",
+        "testdata/expected_dt_with_multiple_dependencies.dts",
+        "testdata/expected_dt_with_dependency_loop.dts",
+    ],
+    out: [
+        "expected_dt_with_dependency.dtb",
+        "expected_dt_with_multiple_dependencies.dtb",
+        "expected_dt_with_dependency_loop.dtb",
+    ],
+    tools: ["dtc"],
+    cmd: "FILES=($(in));" +
+        "cp $${FILES[0]} $(genDir)/platform_preprocessed.dts;" +
+        "for DTS in $${FILES[@]:1}; do" +
+        "  DTB=$$(basename -s .dts $${DTS}).dtb;" +
+        "  $(location dtc) -@ -i $(genDir) -I dts -O dtb $${DTS} -o $(genDir)/$${DTB};" +
+        "done",
+    visibility: ["//visibility:private"],
+}
+
 cc_binary {
     name: "pvmfw",
     defaults: ["vmbase_elf_defaults"],
diff --git a/pvmfw/platform.dts b/pvmfw/platform.dts
index 275a1c9..8074188 100644
--- a/pvmfw/platform.dts
+++ b/pvmfw/platform.dts
@@ -706,6 +706,14 @@
 		timeout-sec = <8>;
 	};
 
+	cpufreq {
+		compatible = "virtual,android-v-only-cpufreq";
+		reg = <0x0 0x1040000 PLACEHOLDER2>;
+	};
+
+	// Keep pvIOMMUs at the last for making test happy.
+	// Otherwise, phandle of other nodes are changed when unused pvIOMMU nodes
+	// are removed, so hardcoded phandles in test data would mismatch.
 	pviommu_0: pviommu0 {
 		compatible = "pkvm,pviommu";
 		id = <PLACEHOLDER>;
@@ -766,8 +774,5 @@
 		#iommu-cells = <1>;
 	};
 
-	cpufreq {
-		compatible = "virtual,android-v-only-cpufreq";
-		reg = <0x0 0x1040000 PLACEHOLDER2>;
-	};
+	// Do not add new node below
 };
diff --git a/pvmfw/src/device_assignment.rs b/pvmfw/src/device_assignment.rs
index c3ccf96..885cd22 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -29,8 +29,10 @@
 use core::mem;
 use core::ops::Range;
 use hyp::DeviceAssigningHypervisor;
-use libfdt::{Fdt, FdtError, FdtNode, Phandle, Reg};
+use libfdt::{Fdt, FdtError, FdtNode, FdtNodeMut, Phandle, Reg};
 use log::error;
+use zerocopy::byteorder::big_endian::U32;
+use zerocopy::FromBytes as _;
 
 // TODO(b/308694211): Use cstr! from vmbase instead.
 macro_rules! cstr {
@@ -214,6 +216,72 @@
     }
 }
 
+#[derive(Debug, Eq, PartialEq)]
+enum DeviceTreeChildrenMask {
+    Partial(Vec<DeviceTreeMask>),
+    All,
+}
+
+#[derive(Eq, PartialEq)]
+struct DeviceTreeMask {
+    name_bytes: Vec<u8>,
+    children: DeviceTreeChildrenMask,
+}
+
+impl fmt::Debug for DeviceTreeMask {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let name_bytes = [self.name_bytes.as_slice(), b"\0"].concat();
+
+        f.debug_struct("DeviceTreeMask")
+            .field("name", &CStr::from_bytes_with_nul(&name_bytes).unwrap())
+            .field("children", &self.children)
+            .finish()
+    }
+}
+
+impl DeviceTreeMask {
+    fn new() -> Self {
+        Self { name_bytes: b"/".to_vec(), children: DeviceTreeChildrenMask::Partial(Vec::new()) }
+    }
+
+    fn mask_internal(&mut self, path: &DtPathTokens, leaf_mask: DeviceTreeChildrenMask) -> bool {
+        let mut iter = self;
+        let mut newly_masked = false;
+        'next_token: for path_token in &path.tokens {
+            let DeviceTreeChildrenMask::Partial(ref mut children) = &mut iter.children else {
+                return false;
+            };
+
+            // Note: Can't use iterator for 'get or insert'. (a.k.a. polonius Rust)
+            #[allow(clippy::needless_range_loop)]
+            for i in 0..children.len() {
+                if children[i].name_bytes.as_slice() == *path_token {
+                    iter = &mut children[i];
+                    newly_masked = false;
+                    continue 'next_token;
+                }
+            }
+            let child = Self {
+                name_bytes: path_token.to_vec(),
+                children: DeviceTreeChildrenMask::Partial(Vec::new()),
+            };
+            children.push(child);
+            newly_masked = true;
+            iter = children.last_mut().unwrap()
+        }
+        iter.children = leaf_mask;
+        newly_masked
+    }
+
+    fn mask(&mut self, path: &DtPathTokens) -> bool {
+        self.mask_internal(path, DeviceTreeChildrenMask::Partial(Vec::new()))
+    }
+
+    fn mask_all(&mut self, path: &DtPathTokens) {
+        self.mask_internal(path, DeviceTreeChildrenMask::All);
+    }
+}
+
 /// Represents VM DTBO
 #[repr(transparent)]
 pub struct VmDtbo(Fdt);
@@ -347,6 +415,114 @@
         }
         Ok(Some(node))
     }
+
+    fn collect_overlayable_nodes_with_phandle(&self) -> Result<BTreeMap<Phandle, DtPathTokens>> {
+        let mut paths = BTreeMap::new();
+        let mut path: DtPathTokens = Default::default();
+        let root = self.as_ref().root();
+        for (node, depth) in root.descendants() {
+            path.tokens.truncate(depth - 1);
+            path.tokens.push(node.name()?.to_bytes());
+            if !path.is_overlayable_node() {
+                continue;
+            }
+            if let Some(phandle) = node.get_phandle()? {
+                paths.insert(phandle, path.clone());
+            }
+        }
+        Ok(paths)
+    }
+
+    fn collect_phandle_references_from_overlayable_nodes(
+        &self,
+    ) -> Result<BTreeMap<DtPathTokens, Vec<Phandle>>> {
+        const CELL_SIZE: usize = core::mem::size_of::<u32>();
+
+        let vm_dtbo = self.as_ref();
+
+        let mut phandle_map = BTreeMap::new();
+        let Some(local_fixups) = vm_dtbo.node(cstr!("/__local_fixups__"))? else {
+            return Ok(phandle_map);
+        };
+
+        let mut path: DtPathTokens = Default::default();
+        for (fixup_node, depth) in local_fixups.descendants() {
+            let node_name = fixup_node.name()?;
+            path.tokens.truncate(depth - 1);
+            path.tokens.push(node_name.to_bytes());
+            if path.tokens.len() != depth {
+                return Err(DeviceAssignmentError::Internal);
+            }
+            if !path.is_overlayable_node() {
+                continue;
+            }
+            let target_node = self.node(&path)?.ok_or(DeviceAssignmentError::InvalidDtbo)?;
+
+            let mut phandles = vec![];
+            for fixup_prop in fixup_node.properties()? {
+                let target_prop = target_node
+                    .getprop(fixup_prop.name()?)
+                    .or(Err(DeviceAssignmentError::InvalidDtbo))?
+                    .ok_or(DeviceAssignmentError::InvalidDtbo)?;
+                let fixup_prop_values = fixup_prop.value()?;
+                if fixup_prop_values.is_empty() || fixup_prop_values.len() % CELL_SIZE != 0 {
+                    return Err(DeviceAssignmentError::InvalidDtbo);
+                }
+
+                for fixup_prop_cell in fixup_prop_values.chunks(CELL_SIZE) {
+                    let phandle_offset: usize = u32::from_be_bytes(
+                        fixup_prop_cell.try_into().or(Err(DeviceAssignmentError::InvalidDtbo))?,
+                    )
+                    .try_into()
+                    .or(Err(DeviceAssignmentError::InvalidDtbo))?;
+                    if phandle_offset % CELL_SIZE != 0 {
+                        return Err(DeviceAssignmentError::InvalidDtbo);
+                    }
+                    let phandle_value = target_prop
+                        .get(phandle_offset..phandle_offset + CELL_SIZE)
+                        .ok_or(DeviceAssignmentError::InvalidDtbo)?;
+                    let phandle: Phandle = U32::ref_from(phandle_value)
+                        .unwrap()
+                        .get()
+                        .try_into()
+                        .or(Err(DeviceAssignmentError::InvalidDtbo))?;
+
+                    phandles.push(phandle);
+                }
+            }
+            if !phandles.is_empty() {
+                phandle_map.insert(path.clone(), phandles);
+            }
+        }
+
+        Ok(phandle_map)
+    }
+
+    fn build_mask(&self, assigned_devices: Vec<DtPathTokens>) -> Result<DeviceTreeMask> {
+        if assigned_devices.is_empty() {
+            return Err(DeviceAssignmentError::Internal);
+        }
+
+        let dependencies = self.collect_phandle_references_from_overlayable_nodes()?;
+        let paths = self.collect_overlayable_nodes_with_phandle()?;
+
+        let mut mask = DeviceTreeMask::new();
+        let mut stack = assigned_devices;
+        while let Some(path) = stack.pop() {
+            if !mask.mask(&path) {
+                continue;
+            }
+            let Some(dst_phandles) = dependencies.get(&path) else {
+                continue;
+            };
+            for dst_phandle in dst_phandles {
+                let dst_path = paths.get(dst_phandle).ok_or(DeviceAssignmentError::Internal)?;
+                stack.push(dst_path.clone());
+            }
+        }
+
+        Ok(mask)
+    }
 }
 
 fn filter_dangling_symbols(fdt: &mut Fdt) -> Result<()> {
@@ -381,6 +557,38 @@
     }
 }
 
+// Filter any node that isn't masked by DeviceTreeMask.
+fn filter_with_mask(anchor: FdtNodeMut, mask: &DeviceTreeMask) -> Result<()> {
+    let mut stack = vec![mask];
+    let mut iter = anchor.next_node(0)?;
+    while let Some((node, depth)) = iter {
+        stack.truncate(depth);
+        let parent_mask = stack.last().unwrap();
+        let DeviceTreeChildrenMask::Partial(parent_mask_children) = &parent_mask.children else {
+            // Shouldn't happen. We only step-in if parent has DeviceTreeChildrenMask::Partial.
+            return Err(DeviceAssignmentError::Internal);
+        };
+
+        let name = node.as_node().name()?.to_bytes();
+        let mask = parent_mask_children.iter().find(|child_mask| child_mask.name_bytes == name);
+        if let Some(masked) = mask {
+            if let DeviceTreeChildrenMask::Partial(_) = &masked.children {
+                // This node is partially masked. Stepping-in.
+                stack.push(masked);
+                iter = node.next_node(depth)?;
+            } else {
+                // This node is fully masked. Stepping-out.
+                iter = node.next_node_skip_subnodes(depth)?;
+            }
+        } else {
+            // This node isn't masked.
+            iter = node.delete_and_next_node(depth)?;
+        }
+    }
+
+    Ok(())
+}
+
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
 struct PvIommu {
     // ID from pvIOMMU node
@@ -689,11 +897,11 @@
     }
 }
 
-#[derive(Debug, Default, Eq, PartialEq)]
+#[derive(Debug, Eq, PartialEq)]
 pub struct DeviceAssignmentInfo {
     pviommus: BTreeSet<PvIommu>,
     assigned_devices: Vec<AssignedDeviceInfo>,
-    filtered_dtbo_paths: Vec<CString>,
+    vm_dtbo_mask: DeviceTreeMask,
 }
 
 impl DeviceAssignmentInfo {
@@ -751,7 +959,7 @@
         let physical_devices = vm_dtbo.parse_physical_devices()?;
 
         let mut assigned_devices = vec![];
-        let mut filtered_dtbo_paths = vec![];
+        let mut assigned_device_paths = vec![];
         for symbol_prop in symbols_node.properties()? {
             let symbol_prop_value = symbol_prop.value()?;
             let dtbo_node_path = CStr::from_bytes_with_nul(symbol_prop_value)
@@ -770,8 +978,7 @@
             )?;
             if let Some(assigned_device) = assigned_device {
                 assigned_devices.push(assigned_device);
-            } else {
-                filtered_dtbo_paths.push(dtbo_node_path.to_cstring());
+                assigned_device_paths.push(dtbo_node_path);
             }
         }
         if assigned_devices.is_empty() {
@@ -780,32 +987,29 @@
 
         Self::validate_pviommu_topology(&assigned_devices)?;
 
-        // Clean up any nodes that wouldn't be overlaid but may contain reference to filtered nodes.
-        // Otherwise, `fdt_apply_overlay()` would fail because of missing phandle reference.
-        // TODO(b/277993056): Also filter other unused nodes/props in __local_fixups__
-        filtered_dtbo_paths.push(CString::new("/__local_fixups__/host").unwrap());
+        let mut vm_dtbo_mask = vm_dtbo.build_mask(assigned_device_paths)?;
+        vm_dtbo_mask.mask_all(&DtPathTokens::new(cstr!("/__local_fixups__"))?);
+        vm_dtbo_mask.mask_all(&DtPathTokens::new(cstr!("/__symbols__"))?);
 
         // Note: Any node without __overlay__ will be ignored by fdt_apply_overlay,
         // so doesn't need to be filtered.
 
-        Ok(Some(Self { pviommus: unique_pviommus, assigned_devices, filtered_dtbo_paths }))
+        Ok(Some(Self { pviommus: unique_pviommus, assigned_devices, vm_dtbo_mask }))
     }
 
     /// Filters VM DTBO to only contain necessary information for booting pVM
-    /// In detail, this will remove followings by setting nop node / nop property.
-    ///   - Removes unassigned devices
-    // TODO(b/277993056): remove unused dependencies in VM DTBO.
-    // TODO(b/277993056): remove supernodes' properties.
-    // TODO(b/277993056): remove unused alises.
     pub fn filter(&self, vm_dtbo: &mut VmDtbo) -> Result<()> {
         let vm_dtbo = vm_dtbo.as_mut();
 
-        // Filters unused node in assigned devices
-        for filtered_dtbo_path in &self.filtered_dtbo_paths {
-            let node = vm_dtbo.node_mut(filtered_dtbo_path).unwrap().unwrap();
-            node.nop()?;
+        // Filter unused references in /__local_fixups__
+        if let Some(local_fixups) = vm_dtbo.node_mut(cstr!("/__local_fixups__"))? {
+            filter_with_mask(local_fixups, &self.vm_dtbo_mask)?;
         }
 
+        // Filter unused nodes in rest of tree
+        let root = vm_dtbo.root_mut();
+        filter_with_mask(root, &self.vm_dtbo_mask)?;
+
         filter_dangling_symbols(vm_dtbo)
     }
 
@@ -860,13 +1064,17 @@
 mod tests {
     use super::*;
     use alloc::collections::{BTreeMap, BTreeSet};
+    use dts::Dts;
     use std::fs;
+    use std::path::Path;
 
     const VM_DTBO_FILE_PATH: &str = "test_pvmfw_devices_vm_dtbo.dtbo";
     const VM_DTBO_WITHOUT_SYMBOLS_FILE_PATH: &str =
         "test_pvmfw_devices_vm_dtbo_without_symbols.dtbo";
     const VM_DTBO_WITH_DUPLICATED_IOMMUS_FILE_PATH: &str =
         "test_pvmfw_devices_vm_dtbo_with_duplicated_iommus.dtbo";
+    const VM_DTBO_WITH_DEPENDENCIES_FILE_PATH: &str =
+        "test_pvmfw_devices_vm_dtbo_with_dependencies.dtbo";
     const FDT_WITHOUT_IOMMUS_FILE_PATH: &str = "test_pvmfw_devices_without_iommus.dtb";
     const FDT_WITHOUT_DEVICE_FILE_PATH: &str = "test_pvmfw_devices_without_device.dtb";
     const FDT_FILE_PATH: &str = "test_pvmfw_devices_with_rng.dtb";
@@ -879,6 +1087,16 @@
         "test_pvmfw_devices_with_duplicated_pviommus.dtb";
     const FDT_WITH_MULTIPLE_REG_IOMMU_FILE_PATH: &str =
         "test_pvmfw_devices_with_multiple_reg_iommus.dtb";
+    const FDT_WITH_DEPENDENCY_FILE_PATH: &str = "test_pvmfw_devices_with_dependency.dtb";
+    const FDT_WITH_MULTIPLE_DEPENDENCIES_FILE_PATH: &str =
+        "test_pvmfw_devices_with_multiple_dependencies.dtb";
+    const FDT_WITH_DEPENDENCY_LOOP_FILE_PATH: &str = "test_pvmfw_devices_with_dependency_loop.dtb";
+
+    const EXPECTED_FDT_WITH_DEPENDENCY_FILE_PATH: &str = "expected_dt_with_dependency.dtb";
+    const EXPECTED_FDT_WITH_MULTIPLE_DEPENDENCIES_FILE_PATH: &str =
+        "expected_dt_with_multiple_dependencies.dtb";
+    const EXPECTED_FDT_WITH_DEPENDENCY_LOOP_FILE_PATH: &str =
+        "expected_dt_with_dependency_loop.dtb";
 
     #[derive(Debug, Default)]
     struct MockHypervisor {
@@ -1449,4 +1667,97 @@
         let compatible = platform_dt.root().next_compatible(cstr!("pkvm,pviommu"));
         assert_eq!(Ok(None), compatible);
     }
+
+    #[test]
+    fn device_info_dependency() {
+        let mut fdt_data = fs::read(FDT_WITH_DEPENDENCY_FILE_PATH).unwrap();
+        let mut vm_dtbo_data = fs::read(VM_DTBO_WITH_DEPENDENCIES_FILE_PATH).unwrap();
+        let fdt = Fdt::from_mut_slice(&mut fdt_data).unwrap();
+        let vm_dtbo = VmDtbo::from_mut_slice(&mut vm_dtbo_data).unwrap();
+        let mut platform_dt_data = pvmfw_fdt_template::RAW.to_vec();
+        platform_dt_data.resize(pvmfw_fdt_template::RAW.len() * 2, 0);
+        let platform_dt = Fdt::from_mut_slice(&mut platform_dt_data).unwrap();
+        platform_dt.unpack().unwrap();
+
+        let hypervisor = MockHypervisor {
+            mmio_tokens: [((0xFF000, 0x1), 0xF000)].into(),
+            iommu_tokens: Default::default(),
+        };
+
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        device_info.filter(vm_dtbo).unwrap();
+
+        // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
+        unsafe {
+            platform_dt.apply_overlay(vm_dtbo.as_mut()).unwrap();
+        }
+        device_info.patch(platform_dt).unwrap();
+
+        let expected = Dts::from_dtb(Path::new(EXPECTED_FDT_WITH_DEPENDENCY_FILE_PATH)).unwrap();
+        let platform_dt = Dts::from_fdt(platform_dt).unwrap();
+
+        assert_eq!(expected, platform_dt);
+    }
+
+    #[test]
+    fn device_info_multiple_dependencies() {
+        let mut fdt_data = fs::read(FDT_WITH_MULTIPLE_DEPENDENCIES_FILE_PATH).unwrap();
+        let mut vm_dtbo_data = fs::read(VM_DTBO_WITH_DEPENDENCIES_FILE_PATH).unwrap();
+        let fdt = Fdt::from_mut_slice(&mut fdt_data).unwrap();
+        let vm_dtbo = VmDtbo::from_mut_slice(&mut vm_dtbo_data).unwrap();
+        let mut platform_dt_data = pvmfw_fdt_template::RAW.to_vec();
+        platform_dt_data.resize(pvmfw_fdt_template::RAW.len() * 2, 0);
+        let platform_dt = Fdt::from_mut_slice(&mut platform_dt_data).unwrap();
+        platform_dt.unpack().unwrap();
+
+        let hypervisor = MockHypervisor {
+            mmio_tokens: [((0xFF000, 0x1), 0xF000), ((0xFF100, 0x1), 0xF100)].into(),
+            iommu_tokens: Default::default(),
+        };
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        device_info.filter(vm_dtbo).unwrap();
+
+        // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
+        unsafe {
+            platform_dt.apply_overlay(vm_dtbo.as_mut()).unwrap();
+        }
+        device_info.patch(platform_dt).unwrap();
+
+        let expected =
+            Dts::from_dtb(Path::new(EXPECTED_FDT_WITH_MULTIPLE_DEPENDENCIES_FILE_PATH)).unwrap();
+        let platform_dt = Dts::from_fdt(platform_dt).unwrap();
+
+        assert_eq!(expected, platform_dt);
+    }
+
+    #[test]
+    fn device_info_dependency_loop() {
+        let mut fdt_data = fs::read(FDT_WITH_DEPENDENCY_LOOP_FILE_PATH).unwrap();
+        let mut vm_dtbo_data = fs::read(VM_DTBO_WITH_DEPENDENCIES_FILE_PATH).unwrap();
+        let fdt = Fdt::from_mut_slice(&mut fdt_data).unwrap();
+        let vm_dtbo = VmDtbo::from_mut_slice(&mut vm_dtbo_data).unwrap();
+        let mut platform_dt_data = pvmfw_fdt_template::RAW.to_vec();
+        platform_dt_data.resize(pvmfw_fdt_template::RAW.len() * 2, 0);
+        let platform_dt = Fdt::from_mut_slice(&mut platform_dt_data).unwrap();
+        platform_dt.unpack().unwrap();
+
+        let hypervisor = MockHypervisor {
+            mmio_tokens: [((0xFF200, 0x1), 0xF200)].into(),
+            iommu_tokens: Default::default(),
+        };
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        device_info.filter(vm_dtbo).unwrap();
+
+        // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
+        unsafe {
+            platform_dt.apply_overlay(vm_dtbo.as_mut()).unwrap();
+        }
+        device_info.patch(platform_dt).unwrap();
+
+        let expected =
+            Dts::from_dtb(Path::new(EXPECTED_FDT_WITH_DEPENDENCY_LOOP_FILE_PATH)).unwrap();
+        let platform_dt = Dts::from_fdt(platform_dt).unwrap();
+
+        assert_eq!(expected, platform_dt);
+    }
 }
diff --git a/pvmfw/testdata/expected_dt_with_dependency.dts b/pvmfw/testdata/expected_dt_with_dependency.dts
new file mode 100644
index 0000000..7e0ad20
--- /dev/null
+++ b/pvmfw/testdata/expected_dt_with_dependency.dts
@@ -0,0 +1,47 @@
+/dts-v1/;
+
+/include/ "platform_preprocessed.dts"
+
+// Note: This uses manually written __symbols__ so we don't
+
+/ {
+    node_a: node_a {
+        phandle = <0x2E>;
+        val = <0x6>;
+        dep = <&node_a_dep &common>;
+        reg = <0x0 0xFF000 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        iommus;
+    };
+
+    node_a_dep: node_a_dep {
+        phandle = <0x31>;
+        val = <0xFF>;
+        dep = <&node_aa_nested_dep>;
+    };
+
+    node_aa {
+        should_be_preserved = <0xFF>;
+
+        node_aa_nested_dep: node_aa_nested_dep {
+            phandle = <0x33>;
+            tag = <0x9>;
+        };
+    };
+
+    common: common {
+        phandle = <0x32>;
+        id = <0x9>;
+    };
+
+    /delete-node/ pviommu0;
+    /delete-node/ pviommu1;
+    /delete-node/ pviommu2;
+    /delete-node/ pviommu3;
+    /delete-node/ pviommu4;
+    /delete-node/ pviommu5;
+    /delete-node/ pviommu6;
+    /delete-node/ pviommu7;
+    /delete-node/ pviommu8;
+    /delete-node/ pviommu9;
+};
diff --git a/pvmfw/testdata/expected_dt_with_dependency_loop.dts b/pvmfw/testdata/expected_dt_with_dependency_loop.dts
new file mode 100644
index 0000000..61031ab
--- /dev/null
+++ b/pvmfw/testdata/expected_dt_with_dependency_loop.dts
@@ -0,0 +1,29 @@
+/dts-v1/;
+
+/include/ "platform_preprocessed.dts"
+
+/ {
+    node_c: node_c {
+        phandle = <0x30>;
+        loop_dep = <&node_c_loop>;
+        reg = <0x0 0xFF200 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        iommus;
+    };
+
+    node_c_loop: node_c_loop {
+        phandle = <0x36>;
+        loop_dep = <&node_c>;
+    };
+
+    /delete-node/ pviommu0;
+    /delete-node/ pviommu1;
+    /delete-node/ pviommu2;
+    /delete-node/ pviommu3;
+    /delete-node/ pviommu4;
+    /delete-node/ pviommu5;
+    /delete-node/ pviommu6;
+    /delete-node/ pviommu7;
+    /delete-node/ pviommu8;
+    /delete-node/ pviommu9;
+};
diff --git a/pvmfw/testdata/expected_dt_with_multiple_dependencies.dts b/pvmfw/testdata/expected_dt_with_multiple_dependencies.dts
new file mode 100644
index 0000000..dc8c357
--- /dev/null
+++ b/pvmfw/testdata/expected_dt_with_multiple_dependencies.dts
@@ -0,0 +1,70 @@
+/dts-v1/;
+
+// Note: We can't use label syntax here.
+// Implementation applies overlay after removing /__symbols__,
+// so using label syntax here wouldn't match with the actual reasult.
+
+/include/ "platform_preprocessed.dts"
+
+/ {
+    node_a: node_a {
+        phandle = <0x2E>;
+        val = <0x6>;
+        dep = <&node_a_dep &common>;
+        reg = <0x0 0xFF000 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        iommus;
+    };
+
+    node_a_dep: node_a_dep {
+        phandle = <0x31>;
+        val = <0xFF>;
+        dep = <&node_aa_nested_dep>;
+    };
+
+    node_aa {
+        should_be_preserved = <0xFF>;
+
+        node_aa_nested_dep: node_aa_nested_dep {
+            phandle = <0x33>;
+            tag = <0x9>;
+        };
+    };
+
+    node_b: node_b {
+        phandle = <0x2F>;
+        tag = <0x33>;
+        version = <0x1 0x2>;
+        dep = <&node_b_dep1 &node_b_dep2>;
+        reg = <0x00 0xFF100 0x00 0x01>;
+        interrupts = <0x00 0x0F 0x04>;
+        iommus;
+    };
+
+    node_b_dep1: node_b_dep1 {
+        phandle = <0x34>;
+        placeholder;
+    };
+
+    node_b_dep2: node_b_dep2 {
+        phandle = <0x35>;
+        placeholder;
+        dep = <&common>;
+    };
+
+    common: common {
+        phandle = <0x32>;
+        id = <0x9>;
+    };
+
+    /delete-node/ pviommu0;
+    /delete-node/ pviommu1;
+    /delete-node/ pviommu2;
+    /delete-node/ pviommu3;
+    /delete-node/ pviommu4;
+    /delete-node/ pviommu5;
+    /delete-node/ pviommu6;
+    /delete-node/ pviommu7;
+    /delete-node/ pviommu8;
+    /delete-node/ pviommu9;
+};
diff --git a/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_with_dependencies.dts b/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_with_dependencies.dts
new file mode 100644
index 0000000..21075e7
--- /dev/null
+++ b/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_with_dependencies.dts
@@ -0,0 +1,77 @@
+/dts-v1/;
+/plugin/;
+
+/ {
+    host {
+        #address-cells = <0x2>;
+        #size-cells = <0x1>;
+        node_a {
+            reg = <0x0 0xF000 0x1>;
+            android,pvmfw,target = <&node_a>;
+        };
+        node_b {
+            reg = <0x0 0xF100 0x1>;
+            android,pvmfw,target = <&node_b>;
+        };
+        node_c {
+            reg = <0x0 0xF200 0x1>;
+            android,pvmfw,target = <&node_c>;
+        };
+    };
+};
+
+&{/} {
+    node_a: node_a {
+        val = <0x6>;
+        dep = <&node_a_dep &common>;
+    };
+
+    node_a_dep: node_a_dep {
+        val = <0xFF>;
+        dep = <&node_aa_nested_dep>;
+
+        node_a_internal {
+            val;
+        };
+    };
+
+    node_aa {
+        should_be_preserved = <0xFF>;
+        node_aa_nested_dep: node_aa_nested_dep {
+            tag = <0x9>;
+        };
+    };
+};
+
+&{/} {
+    node_b: node_b {
+        tag = <0x33>;
+        version = <0x1 0x2>;
+        dep = <&node_b_dep1 &node_b_dep2>;
+    };
+
+    node_b_dep1: node_b_dep1 {
+        placeholder;
+    };
+
+    node_b_dep2: node_b_dep2 {
+        placeholder;
+        dep = <&common>;
+    };
+};
+
+&{/} {
+    node_c: node_c {
+        loop_dep = <&node_c_loop>;
+    };
+
+    node_c_loop: node_c_loop {
+        loop_dep = <&node_c>;
+    };
+};
+
+&{/} {
+    common: common {
+        id = <0x9>;
+    };
+};
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_dependency.dts b/pvmfw/testdata/test_pvmfw_devices_with_dependency.dts
new file mode 100644
index 0000000..b1cf6c7
--- /dev/null
+++ b/pvmfw/testdata/test_pvmfw_devices_with_dependency.dts
@@ -0,0 +1,36 @@
+/dts-v1/;
+
+/include/ "test_crosvm_dt_base.dtsi"
+
+/ {
+    node_a: node_a {
+        reg = <0x0 0xFF000 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        val = <0x6>;
+        dep = <&node_a_dep &common>;
+
+        node_a_internal {
+            parent = <&node_a>;
+        };
+    };
+
+    node_a_dep: node_a_dep {
+        val = <0xFF>;
+        dep = <&node_aa_nested_dep>;
+
+        node_a_dep_internal {
+            val;
+        };
+    };
+
+    node_aa {
+        should_be_preserved = <0xFF>;
+        node_aa_nested_dep: node_aa_nested_dep {
+            tag = <0x9>;
+        };
+    };
+
+    common: common {
+        id = <0x9>;
+    };
+};
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_dependency_loop.dts b/pvmfw/testdata/test_pvmfw_devices_with_dependency_loop.dts
new file mode 100644
index 0000000..9a62cb5
--- /dev/null
+++ b/pvmfw/testdata/test_pvmfw_devices_with_dependency_loop.dts
@@ -0,0 +1,15 @@
+/dts-v1/;
+
+/include/ "test_crosvm_dt_base.dtsi"
+
+/ {
+    node_c: node_c {
+        reg = <0x0 0xFF200 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        loop_dep = <&node_c_loop>;
+    };
+
+    node_c_loop: node_c_loop {
+        loop_dep = <&node_c>;
+    };
+};
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_multiple_dependencies.dts b/pvmfw/testdata/test_pvmfw_devices_with_multiple_dependencies.dts
new file mode 100644
index 0000000..573bdcf
--- /dev/null
+++ b/pvmfw/testdata/test_pvmfw_devices_with_multiple_dependencies.dts
@@ -0,0 +1,50 @@
+/dts-v1/;
+
+/include/ "test_crosvm_dt_base.dtsi"
+
+/ {
+    node_a: node_a {
+        reg = <0x0 0xFF000 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        val = <0x6>;
+        dep = <&node_a_dep &common>;
+    };
+
+    node_a_dep: node_a_dep {
+        val = <0xFF>;
+        dep = <&node_nested_dep>;
+
+        node_a_internal {
+            val;
+        };
+    };
+
+    node_aa {
+        should_be_preserved = <0xFF>;
+        node_nested_dep: node_aa_nested_dep {
+            tag = <0x9>;
+        };
+    };
+
+    node_b: node_b {
+        reg = <0x0 0xFF100 0x0 0x1>;
+        interrupts = <0x0 0xF 0x4>;
+        tag = <0x33>;
+        version = <0x1 0x2>;
+        phandle = <0x5>;
+        dep = <&node_b_dep1 &node_b_dep2>;
+    };
+
+    node_b_dep1: node_b_dep1 {
+        placeholder;
+    };
+
+    node_b_dep2: node_b_dep2 {
+        placeholder;
+        dep = <&common>;
+    };
+
+    common: common {
+        id = <0x9>;
+    };
+};
diff --git a/service_vm/test_apk/Android.bp b/service_vm/test_apk/Android.bp
index 8f5fb41..cd992db 100644
--- a/service_vm/test_apk/Android.bp
+++ b/service_vm/test_apk/Android.bp
@@ -2,12 +2,11 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
-android_test {
-    name: "VmAttestationTestApp",
+java_defaults {
+    name: "vm_attestation_testapp_defaults",
     test_suites: [
         "general-tests",
     ],
-    srcs: ["src/java/**/*.java"],
     static_libs: [
         "MicrodroidDeviceTestHelper",
         "androidx.test.runner",
@@ -19,7 +18,12 @@
     jni_uses_platform_apis: true,
     use_embedded_native_libs: true,
     sdk_version: "test_current",
-    compile_multilib: "first",
+}
+
+android_test {
+    name: "VmAttestationTestApp",
+    srcs: ["src/java/com/android/virt/vm_attestation/testapp/*.java"],
+    defaults: ["vm_attestation_testapp_defaults"],
 }
 
 rust_defaults {
@@ -41,4 +45,22 @@
 rust_ffi {
     name: "libvm_attestation_test_payload",
     defaults: ["vm_attestation_test_payload_defaults"],
+    visibility: [":__subpackages__"],
+}
+
+android_test {
+    name: "AvfRkpdVmAttestationTestApp",
+    srcs: ["src/java/com/android/virt/rkpd/vm_attestation/testapp/*.java"],
+    defaults: ["vm_attestation_testapp_defaults"],
+    manifest: "AndroidManifest.rkpd.xml",
+    test_config: "AndroidTest.rkpd.xml",
+    static_libs: [
+        "RkpdAppTestUtil",
+        "androidx.work_work-testing",
+        "bouncycastle-unbundled",
+    ],
+    instrumentation_for: "rkpdapp",
+    // This app is a variation of rkpdapp, with additional permissions to run
+    // a VM. It is defined in packages/modules/RemoteKeyProvisioning.
+    data: [":avf-rkpdapp"],
 }
diff --git a/service_vm/test_apk/AndroidManifest.rkpd.xml b/service_vm/test_apk/AndroidManifest.rkpd.xml
new file mode 100644
index 0000000..6ecc5a9
--- /dev/null
+++ b/service_vm/test_apk/AndroidManifest.rkpd.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.virt.rkpd.vm_attestation.testapp">
+
+    <instrumentation
+            android:name="androidx.test.runner.AndroidJUnitRunner"
+            android:targetPackage="com.android.rkpdapp"
+            android:label="AVF rkpd app integration tests" />
+</manifest>
diff --git a/service_vm/test_apk/AndroidTest.rkpd.xml b/service_vm/test_apk/AndroidTest.rkpd.xml
new file mode 100644
index 0000000..39eca32
--- /dev/null
+++ b/service_vm/test_apk/AndroidTest.rkpd.xml
@@ -0,0 +1,38 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+<configuration description="VM attestation integration tests with the rkpd app.">
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <!-- Need to disable SELinux policy to allow com.android.rkpdapp to run a VM. -->
+    <target_preparer class="com.android.tradefed.targetprep.DisableSELinuxTargetPreparer"/>
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="test-file-name" value="AvfRkpdVmAttestationTestApp.apk" />
+        <option name="test-file-name" value="avf-rkpdapp.apk" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.virt.rkpd.vm_attestation.testapp" />
+    </test>
+
+    <!-- Only run if RKPD mainline module is installed -->
+    <object type="module_controller"
+            class="com.android.tradefed.testtype.suite.module.MainlineTestModuleController">
+        <option name="enable" value="true" />
+        <option name="mainline-module-package-name" value="com.android.rkpd" />
+    </object>
+</configuration>
diff --git a/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl b/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl
index 94a7b8d..34c8549 100644
--- a/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl
+++ b/service_vm/test_apk/aidl/com/android/virt/vm_attestation/testservice/IAttestationService.aidl
@@ -21,7 +21,31 @@
     const int PORT = 5679;
 
     /**
-     * Requests attestation for testing.
+     * The result of signing a message with the attested key.
+     */
+    parcelable SigningResult {
+        /** The DER-encoded ECDSA signature of the message. */
+        byte[] signature;
+
+        /** The DER-encoded attestation X509 certificate chain. */
+        byte[] certificateChain;
+    }
+
+    /**
+     * Requests attestation with {@link AVmPayload_requestAttestation} API and signs the
+     * given message with the attested key.
+     *
+     * The remotely provisioned keys are retrieved from RKPD and are provisioned from the
+     * real RKP server.
+     *
+     * @param challenge the challenge to include in the attestation output.
+     * @param message the message to sign.
+     * @return the result of signing the message with the attested key.
+     */
+    SigningResult signWithAttestationKey(in byte[] challenge, in byte[] message);
+
+    /**
+     * Requests attestation for testing with {@link AVmPayload_requestAttestationForTesting} API.
      *
      * A fake key pair should be provisioned with the call to
      * {@link VirtualMachine#enableTestAttestation()} before calling this method.
diff --git a/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
new file mode 100644
index 0000000..2a771f3
--- /dev/null
+++ b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
@@ -0,0 +1,266 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.virt.rkpd.vm_attestation.testapp;
+
+import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_FULL;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+import static com.google.common.truth.TruthJUnit.assume;
+
+import android.content.Context;
+import android.hardware.security.keymint.IRemotelyProvisionedComponent;
+import android.os.SystemProperties;
+import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineConfig;
+
+import androidx.work.ListenableWorker;
+import androidx.work.testing.TestWorkerBuilder;
+
+import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
+import com.android.rkpdapp.database.ProvisionedKeyDao;
+import com.android.rkpdapp.database.RkpdDatabase;
+import com.android.rkpdapp.interfaces.ServerInterface;
+import com.android.rkpdapp.interfaces.ServiceManagerInterface;
+import com.android.rkpdapp.interfaces.SystemInterface;
+import com.android.rkpdapp.provisioner.PeriodicProvisioner;
+import com.android.rkpdapp.testutil.SystemInterfaceSelector;
+import com.android.rkpdapp.utils.Settings;
+import com.android.rkpdapp.utils.X509Utils;
+import com.android.virt.vm_attestation.testservice.IAttestationService;
+import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
+
+import org.bouncycastle.asn1.ASN1Boolean;
+import org.bouncycastle.asn1.ASN1Encodable;
+import org.bouncycastle.asn1.ASN1OctetString;
+import org.bouncycastle.asn1.ASN1Sequence;
+import org.bouncycastle.asn1.DEROctetString;
+import org.bouncycastle.asn1.DERUTF8String;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.security.Signature;
+import java.security.cert.X509Certificate;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
+
+/**
+ * End-to-end test for the pVM remote attestation.
+ *
+ * <p>The test checks the two major steps of the pVM remote attestation:
+ *
+ * <p>1. Key provisioning: The test provisions AVF keys from the RKP server and verifies that the
+ * keys are for AVF.
+ *
+ * <p>2. VM attestation: The test creates a VM with a payload binary that requests to attest the VM,
+ * and then signs a message with the attestation key.
+ *
+ * <p>To run this test, you need to:
+ *
+ * <p>- Have an arm64 device supporting protected VMs.
+ *
+ * <p>- Have a stable network connection on the device.
+ *
+ * <p>- Have the RKP server hostname configured in the device. If not, you can set it using: $ adb
+ * shell setprop remote_provisioning.hostname remoteprovisioning.googleapis.com
+ */
+@RunWith(Parameterized.class)
+public class RkpdVmAttestationTest extends MicrodroidDeviceTestBase {
+    private static final String TAG = "RkpdVmAttestationTest";
+    private static final String AVF_ATTESTATION_EXTENSION_OID = "1.3.6.1.4.1.11129.2.1.29.1";
+    private static final String SERVICE_NAME = IRemotelyProvisionedComponent.DESCRIPTOR + "/avf";
+    private static final String VM_PAYLOAD_PATH = "libvm_attestation_test_payload.so";
+    private static final String MESSAGE = "Hello RKP from AVF!";
+    private static final String TEST_APP_PACKAGE_NAME =
+            "com.android.virt.rkpd.vm_attestation.testapp";
+
+    private ProvisionedKeyDao mKeyDao;
+    private PeriodicProvisioner mProvisioner;
+
+    @Parameterized.Parameter(0)
+    public String mGki;
+
+    @Parameterized.Parameters(name = "gki={0}")
+    public static Collection<Object[]> params() {
+        List<Object[]> ret = new ArrayList<>();
+        ret.add(new Object[] {null /* use microdroid kernel */});
+        for (String gki : SUPPORTED_GKI_VERSIONS) {
+            ret.add(new Object[] {gki});
+        }
+        return ret;
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        assume().withMessage("The RKP server hostname is not configured -- assume RKP disabled.")
+                .that(SystemProperties.get("remote_provisioning.hostname"))
+                .isNotEmpty();
+        assume().withMessage("RKP Integration tests rely on network availability.")
+                .that(ServerInterface.isNetworkConnected(getContext()))
+                .isTrue();
+        // TODO(b/329652894): Assume that pVM remote attestation feature is supported.
+
+        prepareTestSetup(true /* protectedVm */, mGki);
+
+        Settings.clearPreferences(getContext());
+        mKeyDao = RkpdDatabase.getDatabase(getContext()).provisionedKeyDao();
+        mKeyDao.deleteAllKeys();
+
+        mProvisioner =
+                TestWorkerBuilder.from(
+                                getContext(),
+                                PeriodicProvisioner.class,
+                                Executors.newSingleThreadExecutor())
+                        .build();
+
+        SystemInterface systemInterface =
+                SystemInterfaceSelector.getSystemInterfaceForServiceName(SERVICE_NAME);
+        ServiceManagerInterface.setInstances(new SystemInterface[] {systemInterface});
+
+        setMaxPerformanceTaskProfile();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        ServiceManagerInterface.setInstances(null);
+        if (mKeyDao != null) {
+            mKeyDao.deleteAllKeys();
+        }
+        Settings.clearPreferences(getContext());
+    }
+
+    @Test
+    public void usingProvisionedKeyForVmAttestationSucceeds() throws Exception {
+        // Provision keys.
+        assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
+        assertThat(mKeyDao.getTotalUnassignedKeysForIrpc(SERVICE_NAME)).isGreaterThan(0);
+
+        // Arrange.
+        Context ctx = getContext();
+        Context otherAppCtx = ctx.createPackageContext(TEST_APP_PACKAGE_NAME, 0);
+        VirtualMachineConfig config =
+                new VirtualMachineConfig.Builder(otherAppCtx)
+                        .setProtectedVm(true)
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setPayloadBinaryName(VM_PAYLOAD_PATH)
+                        .setVmOutputCaptured(true)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("attestation_with_rkpd_client", config);
+        byte[] challenge = new byte[32];
+        Arrays.fill(challenge, (byte) 0xab);
+
+        // Act.
+        CompletableFuture<Exception> exception = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
+        CompletableFuture<SigningResult> signingResultFuture = new CompletableFuture<>();
+        VmEventListener listener =
+                new VmEventListener() {
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        payloadReady.complete(true);
+                        try {
+                            IAttestationService service =
+                                    IAttestationService.Stub.asInterface(
+                                            vm.connectToVsockServer(IAttestationService.PORT));
+                            signingResultFuture.complete(
+                                    service.signWithAttestationKey(challenge, MESSAGE.getBytes()));
+                        } catch (Exception e) {
+                            exception.complete(e);
+                        } finally {
+                            forceStop(vm);
+                        }
+                    }
+                };
+        listener.runToFinish(TAG, vm);
+
+        // Assert.
+        assertThat(payloadReady.getNow(false)).isTrue();
+        assertThat(exception.getNow(null)).isNull();
+        SigningResult signingResult = signingResultFuture.getNow(null);
+        assertThat(signingResult).isNotNull();
+
+        // Parsing the certificate chain successfully indicates that the certificate
+        // chain is valid, that each certificate is signed by the next one and the last
+        // one is self-signed.
+        X509Certificate[] certs = X509Utils.formatX509Certs(signingResult.certificateChain);
+        assertThat(certs.length).isGreaterThan(2);
+        assertWithMessage("The first certificate should be generated in the RKP VM")
+                .that(certs[0].getSubjectX500Principal().getName())
+                .isEqualTo("CN=Android Protected Virtual Machine Key");
+        checkAvfAttestationExtension(certs[0], challenge);
+        assertWithMessage("The second certificate should contain AVF in the subject")
+                .that(certs[1].getSubjectX500Principal().getName())
+                .contains("O=AVF");
+
+        // Verify the signature using the public key from the leaf certificate generated
+        // in the RKP VM.
+        Signature sig = Signature.getInstance("SHA256withECDSA");
+        sig.initVerify(certs[0].getPublicKey());
+        sig.update(MESSAGE.getBytes());
+        assertThat(sig.verify(signingResult.signature)).isTrue();
+    }
+
+    private void checkAvfAttestationExtension(X509Certificate cert, byte[] challenge)
+            throws Exception {
+        byte[] extensionValue = cert.getExtensionValue(AVF_ATTESTATION_EXTENSION_OID);
+        ASN1OctetString extString = ASN1OctetString.getInstance(extensionValue);
+        ASN1Sequence seq = ASN1Sequence.getInstance(extString.getOctets());
+        // AVF attestation extension should contain 3 elements in the following format:
+        //
+        //  AttestationExtension ::= SEQUENCE {
+        //     attestationChallenge       OCTET_STRING,
+        //     isVmSecure                 BOOLEAN,
+        //     vmComponents               SEQUENCE OF VmComponent,
+        //  }
+        //   VmComponent ::= SEQUENCE {
+        //     name               UTF8String,
+        //     securityVersion    INTEGER,
+        //     codeHash           OCTET STRING,
+        //     authorityHash      OCTET STRING,
+        //  }
+        assertThat(seq).hasSize(3);
+
+        ASN1OctetString expectedChallenge = new DEROctetString(challenge);
+        assertThat(seq.getObjectAt(0)).isEqualTo(expectedChallenge);
+        assertWithMessage("The VM should be unsecure as it is debuggable.")
+                .that(seq.getObjectAt(1))
+                .isEqualTo(ASN1Boolean.FALSE);
+        ASN1Sequence vmComponents = ASN1Sequence.getInstance(seq.getObjectAt(2));
+        assertExtensionContainsPayloadApk(vmComponents);
+    }
+
+    private void assertExtensionContainsPayloadApk(ASN1Sequence vmComponents) throws Exception {
+        DERUTF8String payloadApkName = new DERUTF8String("apk:" + TEST_APP_PACKAGE_NAME);
+        boolean found = false;
+        for (ASN1Encodable encodable : vmComponents) {
+            ASN1Sequence vmComponent = ASN1Sequence.getInstance(encodable);
+            assertThat(vmComponent).hasSize(4);
+            if (payloadApkName.equals(vmComponent.getObjectAt(0))) {
+                assertWithMessage("Payload APK should not be found twice.").that(found).isFalse();
+                found = true;
+            }
+        }
+        assertWithMessage("vmComponents should contain the payload APK.").that(found).isTrue();
+    }
+}
diff --git a/service_vm/test_apk/src/native/main.rs b/service_vm/test_apk/src/native/main.rs
index 199b45c..a04fb1f 100644
--- a/service_vm/test_apk/src/native/main.rs
+++ b/service_vm/test_apk/src/native/main.rs
@@ -18,7 +18,7 @@
 use avflog::LogResult;
 use com_android_virt_vm_attestation_testservice::{
     aidl::com::android::virt::vm_attestation::testservice::IAttestationService::{
-        BnAttestationService, IAttestationService, PORT,
+        BnAttestationService, IAttestationService, SigningResult::SigningResult, PORT,
     },
     binder::{self, unstable_api::AsNative, BinderFeatures, Interface, IntoBinderResult, Strong},
 };
@@ -34,7 +34,7 @@
     AIBinder, AVmAttestationResult, AVmAttestationResult_free,
     AVmAttestationResult_getCertificateAt, AVmAttestationResult_getCertificateCount,
     AVmAttestationResult_getPrivateKey, AVmAttestationResult_sign, AVmAttestationStatus,
-    AVmAttestationStatus_toString, AVmPayload_notifyPayloadReady,
+    AVmAttestationStatus_toString, AVmPayload_notifyPayloadReady, AVmPayload_requestAttestation,
     AVmPayload_requestAttestationForTesting, AVmPayload_runVsockRpcServer,
 };
 
@@ -89,13 +89,8 @@
 
 impl IAttestationService for AttestationService {
     fn requestAttestationForTesting(&self) -> binder::Result<()> {
-        // The data below is only a placeholder generated randomly with urandom
-        let challenge = &[
-            0x6c, 0xad, 0x52, 0x50, 0x15, 0xe7, 0xf4, 0x1d, 0xa5, 0x60, 0x7e, 0xd2, 0x7d, 0xf1,
-            0x51, 0x67, 0xc3, 0x3e, 0x73, 0x9b, 0x30, 0xbd, 0x04, 0x20, 0x2e, 0xde, 0x3b, 0x1d,
-            0xc8, 0x07, 0x11, 0x7b,
-        ];
-        let res = AttestationResult::request_attestation(challenge)
+        const CHALLENGE: &[u8] = &[0xaa; 32];
+        let res = AttestationResult::request_attestation_for_testing(CHALLENGE)
             .map_err(|e| anyhow!("Unexpected status: {:?}", status_to_cstr(e)))
             .with_log()
             .or_service_specific_exception(-1)?;
@@ -103,6 +98,21 @@
         Ok(())
     }
 
+    fn signWithAttestationKey(
+        &self,
+        challenge: &[u8],
+        message: &[u8],
+    ) -> binder::Result<SigningResult> {
+        let res = AttestationResult::request_attestation(challenge)
+            .map_err(|e| anyhow!("Unexpected status: {:?}", status_to_cstr(e)))
+            .with_log()
+            .or_service_specific_exception(-1)?;
+        let certificate_chain =
+            res.certificate_chain().with_log().or_service_specific_exception(-1)?;
+        let signature = res.sign(message).with_log().or_service_specific_exception(-1)?;
+        Ok(SigningResult { certificateChain: certificate_chain, signature })
+    }
+
     fn validateAttestationResult(&self) -> binder::Result<()> {
         // TODO(b/191073073): Returns the attestation result to the host for validation.
         self.res.lock().unwrap().as_ref().unwrap().log().or_service_specific_exception(-1)
@@ -116,7 +126,9 @@
 unsafe impl Send for AttestationResult {}
 
 impl AttestationResult {
-    fn request_attestation(challenge: &[u8]) -> result::Result<Self, AVmAttestationStatus> {
+    fn request_attestation_for_testing(
+        challenge: &[u8],
+    ) -> result::Result<Self, AVmAttestationStatus> {
         let mut res: *mut AVmAttestationResult = ptr::null_mut();
         // SAFETY: It is safe as we only read the challenge within its bounds and the
         // function does not retain any reference to it.
@@ -136,11 +148,31 @@
         }
     }
 
-    fn certificate_chain(&self) -> Result<Vec<Box<[u8]>>> {
+    fn request_attestation(challenge: &[u8]) -> result::Result<Self, AVmAttestationStatus> {
+        let mut res: *mut AVmAttestationResult = ptr::null_mut();
+        // SAFETY: It is safe as we only read the challenge within its bounds and the
+        // function does not retain any reference to it.
+        let status = unsafe {
+            AVmPayload_requestAttestation(
+                challenge.as_ptr() as *const c_void,
+                challenge.len(),
+                &mut res,
+            )
+        };
+        if status == AVmAttestationStatus::ATTESTATION_OK {
+            info!("Attestation succeeds. Status: {:?}", status_to_cstr(status));
+            let res = NonNull::new(res).expect("The attestation result is null");
+            Ok(Self(res))
+        } else {
+            Err(status)
+        }
+    }
+
+    fn certificate_chain(&self) -> Result<Vec<u8>> {
         let num_certs = get_certificate_count(self.as_ref());
-        let mut certs = Vec::with_capacity(num_certs);
+        let mut certs = Vec::new();
         for i in 0..num_certs {
-            certs.push(get_certificate_at(self.as_ref(), i)?);
+            certs.extend(get_certificate_at(self.as_ref(), i)?.iter());
         }
         Ok(certs)
     }
@@ -149,7 +181,7 @@
         get_private_key(self.as_ref())
     }
 
-    fn sign(&self, message: &[u8]) -> Result<Box<[u8]>> {
+    fn sign(&self, message: &[u8]) -> Result<Vec<u8>> {
         sign_with_attested_key(self.as_ref(), message)
     }
 
@@ -231,7 +263,7 @@
     Ok(private_key.into_boxed_slice())
 }
 
-fn sign_with_attested_key(res: &AVmAttestationResult, message: &[u8]) -> Result<Box<[u8]>> {
+fn sign_with_attested_key(res: &AVmAttestationResult, message: &[u8]) -> Result<Vec<u8>> {
     // SAFETY: The result is returned by `AVmPayload_requestAttestation` and should be valid
     // before getting freed.
     let size = unsafe {
@@ -258,7 +290,7 @@
     };
     ensure!(size <= signature.len());
     signature.truncate(size);
-    Ok(signature.into_boxed_slice())
+    Ok(signature)
 }
 
 fn status_to_cstr(status: AVmAttestationStatus) -> &'static CStr {
diff --git a/tests/pvmfw/Android.bp b/tests/pvmfw/Android.bp
index c12f67a..0483066 100644
--- a/tests/pvmfw/Android.bp
+++ b/tests/pvmfw/Android.bp
@@ -53,4 +53,5 @@
         ":test_avf_debug_policy_without_adb",
         "assets/bcc.dat",
     ],
+    data_device_bins_first: ["dtc_static"],
 }
diff --git a/tests/pvmfw/AndroidTest.xml b/tests/pvmfw/AndroidTest.xml
index 6ff7b6f..5784f26 100644
--- a/tests/pvmfw/AndroidTest.xml
+++ b/tests/pvmfw/AndroidTest.xml
@@ -22,6 +22,24 @@
         <option name="force-root" value="true"/>
     </target_preparer>
 
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <!-- Prepare test directories. -->
+        <option name="run-command" value="mkdir -p /data/local/tmp/pvmfw" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/pvmfw" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+        <option name="cleanup" value="true" />
+        <option name="abort-on-push-failure" value="true" />
+        <option name="push-file" key="dtc_static" value="/data/local/tmp/pvmfw/dtc_static" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <option name="run-command" value="[ ! -d /proc/device-tree/avf/reference ] || /data/local/tmp/pvmfw/dtc_static -f -qqq /proc/device-tree/avf/reference -o /data/local/tmp/pvmfw/reference_dt.dtb" />
+    </target_preparer>
+
     <test class="com.android.compatibility.common.tradefed.testtype.JarHostTest" >
         <option name="jar" value="CustomPvmfwHostTestCases.jar" />
     </test>
diff --git a/tests/pvmfw/helper/Android.bp b/tests/pvmfw/helper/Android.bp
index 1b96842..90ca03e 100644
--- a/tests/pvmfw/helper/Android.bp
+++ b/tests/pvmfw/helper/Android.bp
@@ -5,5 +5,8 @@
 java_library_host {
     name: "PvmfwHostTestHelper",
     srcs: ["java/**/*.java"],
-    libs: ["androidx.annotation_annotation"],
+    libs: [
+        "androidx.annotation_annotation",
+        "truth",
+    ],
 }
diff --git a/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java b/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
index b0c1207..a77ba40 100644
--- a/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
+++ b/tests/pvmfw/helper/java/com/android/pvmfw/test/host/Pvmfw.java
@@ -16,6 +16,8 @@
 
 package com.android.pvmfw.test.host;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static java.nio.ByteOrder.LITTLE_ENDIAN;
 
 import androidx.annotation.NonNull;
@@ -34,22 +36,52 @@
     private static final int SIZE_4K = 4 << 10; // 4 KiB, PAGE_SIZE
     private static final int BUFFER_SIZE = 1024;
     private static final int HEADER_MAGIC = 0x666d7670;
-    private static final int HEADER_DEFAULT_VERSION = getVersion(1, 0);
+    private static final int HEADER_DEFAULT_VERSION = makeVersion(1, 2);
     private static final int HEADER_FLAGS = 0;
 
+    private static final int PVMFW_ENTRY_BCC = 0;
+    private static final int PVMFW_ENTRY_DP = 1;
+    private static final int PVMFW_ENTRY_VM_DTBO = 2;
+    private static final int PVMFW_ENTRY_VM_REFERENCE_DT = 3;
+    private static final int PVMFW_ENTRY_MAX = 4;
+
     @NonNull private final File mPvmfwBinFile;
-    @NonNull private final File mBccFile;
-    @Nullable private final File mDebugPolicyFile;
+    private final File[] mEntries;
+    private final int mEntryCnt;
     private final int mVersion;
 
+    public static int makeVersion(int major, int minor) {
+        return ((major & 0xFFFF) << 16) | (minor & 0xFFFF);
+    }
+
     private Pvmfw(
             @NonNull File pvmfwBinFile,
             @NonNull File bccFile,
             @Nullable File debugPolicyFile,
+            @Nullable File vmDtboFile,
+            @Nullable File vmReferenceDtFile,
             int version) {
         mPvmfwBinFile = Objects.requireNonNull(pvmfwBinFile);
-        mBccFile = Objects.requireNonNull(bccFile);
-        mDebugPolicyFile = debugPolicyFile;
+
+        if (version >= makeVersion(1, 2)) {
+            mEntryCnt = PVMFW_ENTRY_VM_REFERENCE_DT + 1;
+        } else if (version >= makeVersion(1, 1)) {
+            mEntryCnt = PVMFW_ENTRY_VM_DTBO + 1;
+        } else {
+            mEntryCnt = PVMFW_ENTRY_DP + 1;
+        }
+
+        mEntries = new File[PVMFW_ENTRY_MAX];
+        mEntries[PVMFW_ENTRY_BCC] = Objects.requireNonNull(bccFile);
+        mEntries[PVMFW_ENTRY_DP] = debugPolicyFile;
+
+        if (PVMFW_ENTRY_VM_DTBO < mEntryCnt) {
+            mEntries[PVMFW_ENTRY_VM_DTBO] = vmDtboFile;
+        }
+        if (PVMFW_ENTRY_VM_REFERENCE_DT < mEntryCnt) {
+            mEntries[PVMFW_ENTRY_VM_REFERENCE_DT] = Objects.requireNonNull(vmReferenceDtFile);
+        }
+
         mVersion = version;
     }
 
@@ -60,62 +92,54 @@
     public void serialize(@NonNull File outFile) throws IOException {
         Objects.requireNonNull(outFile);
 
-        int headerSize = alignTo(getHeaderSize(mVersion), SIZE_8B);
-        int bccOffset = headerSize;
-        int bccSize = (int) mBccFile.length();
+        int headerSize = alignTo(getHeaderSize(), SIZE_8B);
+        int[] entryOffsets = new int[mEntryCnt];
+        int[] entrySizes = new int[mEntryCnt];
 
-        int debugPolicyOffset = alignTo(bccOffset + bccSize, SIZE_8B);
-        int debugPolicySize = mDebugPolicyFile == null ? 0 : (int) mDebugPolicyFile.length();
+        entryOffsets[PVMFW_ENTRY_BCC] = headerSize;
+        entrySizes[PVMFW_ENTRY_BCC] = (int) mEntries[PVMFW_ENTRY_BCC].length();
 
-        int totalSize = debugPolicyOffset + debugPolicySize;
-        if (hasVmDtbo(mVersion)) {
-            // Add VM DTBO size as well.
-            totalSize += Integer.BYTES * 2;
+        for (int i = 1; i < mEntryCnt; i++) {
+            entryOffsets[i] = alignTo(entryOffsets[i - 1] + entrySizes[i - 1], SIZE_8B);
+            entrySizes[i] = mEntries[i] == null ? 0 : (int) mEntries[i].length();
         }
 
+        int totalSize = alignTo(entryOffsets[mEntryCnt - 1] + entrySizes[mEntryCnt - 1], SIZE_8B);
+
         ByteBuffer header = ByteBuffer.allocate(headerSize).order(LITTLE_ENDIAN);
         header.putInt(HEADER_MAGIC);
         header.putInt(mVersion);
         header.putInt(totalSize);
         header.putInt(HEADER_FLAGS);
-        header.putInt(bccOffset);
-        header.putInt(bccSize);
-        header.putInt(debugPolicyOffset);
-        header.putInt(debugPolicySize);
-
-        if (hasVmDtbo(mVersion)) {
-            // Add placeholder entry for VM DTBO.
-            // TODO(b/291191157): Add a real DTBO and test.
-            header.putInt(0);
-            header.putInt(0);
+        for (int i = 0; i < mEntryCnt; i++) {
+            header.putInt(entryOffsets[i]);
+            header.putInt(entrySizes[i]);
         }
 
         try (FileOutputStream pvmfw = new FileOutputStream(outFile)) {
             appendFile(pvmfw, mPvmfwBinFile);
             padTo(pvmfw, SIZE_4K);
+
+            int baseOffset = (int) pvmfw.getChannel().size();
             pvmfw.write(header.array());
-            padTo(pvmfw, SIZE_8B);
-            appendFile(pvmfw, mBccFile);
-            if (mDebugPolicyFile != null) {
+
+            for (int i = 0; i < mEntryCnt; i++) {
                 padTo(pvmfw, SIZE_8B);
-                appendFile(pvmfw, mDebugPolicyFile);
+                if (mEntries[i] != null) {
+                    assertThat((int) pvmfw.getChannel().size() - baseOffset)
+                            .isEqualTo(entryOffsets[i]);
+                    appendFile(pvmfw, mEntries[i]);
+                }
             }
+
             padTo(pvmfw, SIZE_4K);
         }
     }
 
     private void appendFile(@NonNull FileOutputStream out, @NonNull File inFile)
             throws IOException {
-        byte buffer[] = new byte[BUFFER_SIZE];
         try (FileInputStream in = new FileInputStream(inFile)) {
-            int size;
-            while (true) {
-                size = in.read(buffer);
-                if (size < 0) {
-                    return;
-                }
-                out.write(buffer, /* offset= */ 0, size);
-            }
+            in.transferTo(out);
         }
     }
 
@@ -126,27 +150,15 @@
         }
     }
 
-    private static int getHeaderSize(int version) {
-        if (version == getVersion(1, 0)) {
-            return Integer.BYTES * 8; // Header has 8 integers.
-        }
-        return Integer.BYTES * 10; // Default + VM DTBO (offset, size)
-    }
-
-    private static boolean hasVmDtbo(int version) {
-        int major = getMajorVersion(version);
-        int minor = getMinorVersion(version);
-        return major > 1 || (major == 1 && minor >= 1);
+    private int getHeaderSize() {
+        // Header + (entry offset, entry, size) * mEntryCnt
+        return Integer.BYTES * (4 + mEntryCnt * 2);
     }
 
     private static int alignTo(int x, int size) {
         return (x + size - 1) & ~(size - 1);
     }
 
-    private static int getVersion(int major, int minor) {
-        return ((major & 0xFFFF) << 16) | (minor & 0xFFFF);
-    }
-
     private static int getMajorVersion(int version) {
         return (version >> 16) & 0xFFFF;
     }
@@ -160,6 +172,8 @@
         @NonNull private final File mPvmfwBinFile;
         @NonNull private final File mBccFile;
         @Nullable private File mDebugPolicyFile;
+        @Nullable private File mVmDtboFile;
+        @Nullable private File mVmReferenceDtFile;
         private int mVersion;
 
         public Builder(@NonNull File pvmfwBinFile, @NonNull File bccFile) {
@@ -175,14 +189,32 @@
         }
 
         @NonNull
+        public Builder setVmDtbo(@Nullable File vmDtboFile) {
+            mVmDtboFile = vmDtboFile;
+            return this;
+        }
+
+        @NonNull
+        public Builder setVmReferenceDt(@Nullable File vmReferenceDtFile) {
+            mVmReferenceDtFile = vmReferenceDtFile;
+            return this;
+        }
+
+        @NonNull
         public Builder setVersion(int major, int minor) {
-            mVersion = getVersion(major, minor);
+            mVersion = makeVersion(major, minor);
             return this;
         }
 
         @NonNull
         public Pvmfw build() {
-            return new Pvmfw(mPvmfwBinFile, mBccFile, mDebugPolicyFile, mVersion);
+            return new Pvmfw(
+                    mPvmfwBinFile,
+                    mBccFile,
+                    mDebugPolicyFile,
+                    mVmDtboFile,
+                    mVmReferenceDtFile,
+                    mVersion);
         }
     }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
new file mode 100644
index 0000000..a3216c2
--- /dev/null
+++ b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
@@ -0,0 +1,229 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.pvmfw.test;
+
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.junit.Assume.assumeTrue;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.microdroid.test.host.CommandRunner;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.CommandStatus;
+import com.android.tradefed.util.FileUtil;
+
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.File;
+import java.util.Map;
+
+/** Base class for testing custom pvmfw */
+public class CustomPvmfwHostTestCaseBase extends MicrodroidHostTestCaseBase {
+    @NonNull public static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
+    @NonNull public static final String BCC_FILE_NAME = "bcc.dat";
+    @NonNull public static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
+    @NonNull public static final String PACKAGE_NAME = "com.android.microdroid.test";
+    @NonNull public static final String MICRODROID_DEBUG_FULL = "full";
+    @NonNull public static final String MICRODROID_DEBUG_NONE = "none";
+
+    @NonNull
+    public static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
+
+    @NonNull
+    public static final String VM_REFERENCE_DT_PATH = "/data/local/tmp/pvmfw/reference_dt.dtb";
+
+    @NonNull public static final String MICRODROID_LOG_PATH = TEST_ROOT + "log.txt";
+    public static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
+    public static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
+    public static final int CONSOLE_OUTPUT_WAIT_MS = 5000; // 5 seconds
+
+    @NonNull public static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
+    @NonNull public static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
+    @NonNull public static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
+    @NonNull public static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
+
+    @NonNull private static final String DUMPSYS = "dumpsys";
+
+    @NonNull
+    private static final String DUMPSYS_MISSING_SERVICE_MSG_PREFIX = "Can't find service: ";
+
+    @NonNull
+    private static final String SECRET_KEEPER_AIDL =
+            "android.hardware.security.secretkeeper.ISecretkeeper/default";
+
+    @Nullable private File mPvmfwBinFileOnHost;
+    @Nullable private File mBccFileOnHost;
+    @Nullable private File mVmReferenceDtFile;
+    private boolean mSecretKeeperSupported;
+
+    @NonNull private TestDevice mAndroidDevice;
+    @Nullable private ITestDevice mMicrodroidDevice;
+
+    @Nullable private File mCustomPvmfwFileOnHost;
+
+    @Before
+    public void setUp() throws Exception {
+        mAndroidDevice = (TestDevice) getDevice();
+
+        // Check device capabilities
+        assumeDeviceIsCapable(mAndroidDevice);
+        assumeTrue(
+                "Skip if protected VMs are not supported",
+                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
+
+        mPvmfwBinFileOnHost = findTestFile(PVMFW_FILE_NAME);
+        mBccFileOnHost = findTestFile(BCC_FILE_NAME);
+
+        // This is prepared by AndroidTest.xml
+        mVmReferenceDtFile = mAndroidDevice.pullFile(VM_REFERENCE_DT_PATH);
+
+        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandResult result = runner.runForResult(DUMPSYS, SECRET_KEEPER_AIDL);
+
+        // dumpsys prints 'Can't find service: ~' to stderr if secret keeper HAL is missing,
+        // but it doesn't return any error code for it.
+        // Read stderr to know whether secret keeper is supported, and stop test for any other case.
+        assertWithMessage("Failed to run " + DUMPSYS + ", result=" + result)
+                .that(result.getStatus() == CommandStatus.SUCCESS && result.getExitCode() == 0)
+                .isTrue();
+        if (result.getStderr() != null && !result.getStderr().trim().isEmpty()) {
+            assertWithMessage(
+                            "Unexpected stderr from " + DUMPSYS + ", stderr=" + result.getStderr())
+                    .that(result.getStderr().trim().startsWith(DUMPSYS_MISSING_SERVICE_MSG_PREFIX))
+                    .isTrue();
+        } else {
+            mSecretKeeperSupported = true;
+        }
+
+        // Prepare for system properties for custom pvmfw.img.
+        // File will be prepared later in individual test and then pushed to device
+        // when launching with launchProtectedVmAndWaitForBootCompleted().
+        mCustomPvmfwFileOnHost =
+                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
+        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
+
+        // Prepare for launching microdroid
+        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
+        prepareVirtualizationTestSetup(mAndroidDevice);
+        mMicrodroidDevice = null;
+    }
+
+    @After
+    public void shutdown() throws Exception {
+        shutdownMicrodroid();
+
+        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
+
+        FileUtil.deleteFile(mVmReferenceDtFile);
+
+        // Cleanup for custom pvmfw.img
+        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
+        FileUtil.deleteFile(mCustomPvmfwFileOnHost);
+
+        cleanUpVirtualizationTestSetup(mAndroidDevice);
+    }
+
+    /** Returns android device */
+    @NonNull
+    public TestDevice getAndroidDevice() {
+        return mAndroidDevice;
+    }
+
+    /** Returns pvmfw.bin file on host for building custom pvmfw with */
+    @NonNull
+    public File getPvmfwBinFile() {
+        return mPvmfwBinFileOnHost;
+    }
+
+    /** Returns BCC file on host for building custom pvmfw with */
+    @NonNull
+    public File getBccFile() {
+        return mBccFileOnHost;
+    }
+
+    /** Returns VM reference DT, generated from DUT, on host for building custom pvmfw with. */
+    @Nullable
+    public File getVmReferenceDtFile() {
+        return mVmReferenceDtFile;
+    }
+
+    /**
+     * Returns a custom pvmfw file.
+     *
+     * <p>This is a temporary file on host. The file should been prepared as a custom pvmfw because
+     * calling {@link #launchProtectedVmAndWaitForBootCompleted}, so virtualization manager can read
+     * the file path from sysprop and boot pVM with it.
+     */
+    @NonNull
+    public File getCustomPvmfwFile() {
+        return mCustomPvmfwFileOnHost;
+    }
+
+    /**
+     * Returns whether a secretkeeper is supported.
+     *
+     * <p>If {@code true}, then VM reference DT must exist. (i.e. {@link #getVmReferenceDtFile} must
+     * exist {@code null}).
+     */
+    public boolean isSecretKeeperSupported() {
+        return mSecretKeeperSupported;
+    }
+
+    /**
+     * Launches protected VM with custom pvmfw ({@link #getCustomPvmfwFile}) and wait for boot
+     * completed. Throws exception when boot failed.
+     */
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(
+            String debugLevel, long adbTimeoutMs, @NonNull Map<String, File> bootFiles)
+            throws DeviceNotAvailableException {
+        MicrodroidBuilder builder =
+                MicrodroidBuilder.fromDevicePath(
+                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
+                        .debugLevel(debugLevel)
+                        .protectedVm(/* protectedVm= */ true)
+                        .addBootFile(mCustomPvmfwFileOnHost, PVMFW_FILE_NAME)
+                        .setAdbConnectTimeoutMs(adbTimeoutMs);
+        for (String name : bootFiles.keySet()) {
+            File file = bootFiles.get(name);
+            builder.addBootFile(file, name);
+        }
+
+        mMicrodroidDevice = builder.build(mAndroidDevice);
+
+        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
+        assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
+        return mMicrodroidDevice;
+    }
+
+    /** Shuts down microdroid if it's running */
+    public void shutdownMicrodroid() throws Exception {
+        if (mMicrodroidDevice != null) {
+            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
+            mMicrodroidDevice = null;
+        }
+    }
+}
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
index 26f5993..223f93f 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
@@ -16,28 +16,22 @@
 
 package com.android.pvmfw.test;
 
-import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
-
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import static org.junit.Assume.assumeTrue;
 import static org.junit.Assert.assertThrows;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.microdroid.test.host.CommandRunner;
-import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.pvmfw.test.host.Pvmfw;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceRuntimeException;
 import com.android.tradefed.device.ITestDevice;
-import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.CommandResult;
-import com.android.tradefed.util.FileUtil;
 
 import org.junit.After;
 import org.junit.Before;
@@ -45,32 +39,13 @@
 import org.junit.runner.RunWith;
 
 import java.io.File;
-import java.util.Objects;
+import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /** Tests debug policy */
 @RunWith(DeviceJUnit4ClassRunner.class)
-public class DebugPolicyHostTests extends MicrodroidHostTestCaseBase {
-    @NonNull private static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
-    @NonNull private static final String BCC_FILE_NAME = "bcc.dat";
-    @NonNull private static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
-    @NonNull private static final String PACKAGE_NAME = "com.android.microdroid.test";
-    @NonNull private static final String MICRODROID_DEBUG_FULL = "full";
-    @NonNull private static final String MICRODROID_DEBUG_NONE = "none";
-
-    @NonNull
-    private static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
-
-    @NonNull private static final String MICRODROID_LOG_PATH = TEST_ROOT + "log.txt";
-    private static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
-    private static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
-    private static final int CONSOLE_OUTPUT_WAIT_MS = 5000; // 5 seconds
-
-    @NonNull private static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
-    @NonNull private static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
-
+public class DebugPolicyHostTests extends CustomPvmfwHostTestCaseBase {
     @NonNull private static final String CUSTOM_DEBUG_POLICY_FILE_NAME = "debug_policy.dtb";
 
     @NonNull
@@ -98,63 +73,22 @@
     @NonNull private static final String HEX_STRING_ZERO = "00000000";
     @NonNull private static final String HEX_STRING_ONE = "00000001";
 
-    @Nullable private static File mPvmfwBinFileOnHost;
-    @Nullable private static File mBccFileOnHost;
-
-    @Nullable private TestDevice mAndroidDevice;
-    @Nullable private ITestDevice mMicrodroidDevice;
-    @Nullable private File mCustomPvmfwBinFileOnHost;
     @Nullable private File mCustomDebugPolicyFileOnHost;
 
     @Before
     public void setUp() throws Exception {
-        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
+        super.setUp();
 
-        // Check device capabilities
-        assumeDeviceIsCapable(mAndroidDevice);
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
-
-        // tradefed copies the test artfacts under /tmp when running tests,
-        // so we should *find* the artifacts with the file name.
-        mPvmfwBinFileOnHost =
-                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
-        mBccFileOnHost =
-                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
-
-        // Prepare for system properties for custom debug policy.
-        // File will be prepared later in individual test by setupCustomDebugPolicy()
-        // and then pushed to device when launching with launchProtectedVmAndWaitForBootCompleted()
-        // or tryLaunchProtectedNonDebuggableVm().
-        mCustomPvmfwBinFileOnHost =
-                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_DEBUG_POLICY_PATH_PROP, CUSTOM_DEBUG_POLICY_PATH);
-
-        // Prepare for launching microdroid
-        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
-        prepareVirtualizationTestSetup(mAndroidDevice);
-        mMicrodroidDevice = null;
+        // Prepare system properties for custom debug policy.
+        setPropertyOrThrow(getDevice(), CUSTOM_DEBUG_POLICY_PATH_PROP, CUSTOM_DEBUG_POLICY_PATH);
     }
 
     @After
     public void shutdown() throws Exception {
-        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
-            return;
-        }
-        if (mMicrodroidDevice != null) {
-            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
-            mMicrodroidDevice = null;
-        }
-        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
+        super.shutdown();
 
         // Cleanup for custom debug policies
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_DEBUG_POLICY_PATH_PROP, "");
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
-        FileUtil.deleteFile(mCustomPvmfwBinFileOnHost);
-
-        cleanUpVirtualizationTestSetup(mAndroidDevice);
+        setPropertyOrThrow(getDevice(), CUSTOM_DEBUG_POLICY_PATH_PROP, "");
     }
 
     @Test
@@ -198,43 +132,41 @@
     @Test
     public void testRamdumpInDebugPolicy_withDebugLevelNone_hasRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_with_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
-                .contains("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH)).contains("crashkernel=");
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ONE);
     }
 
     @Test
     public void testNoRamdumpInDebugPolicy_withDebugLevelNone_noRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_without_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_NONE);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH))
                 .doesNotContain("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH))
                 .doesNotContain("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ZERO);
     }
 
     @Test
     public void testNoRamdumpInDebugPolicy_withDebugLevelFull_hasRamdumpArgs() throws Exception {
         prepareCustomDebugPolicy("avf_debug_policy_without_ramdump.dtbo");
-        mMicrodroidDevice = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_FULL);
+        ITestDevice device = launchProtectedVmAndWaitForBootCompleted(MICRODROID_DEBUG_FULL);
 
-        assertThat(readMicrodroidFileAsString(MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
-        assertThat(readMicrodroidFileAsString(MICRODROID_DT_BOOTARGS_PATH))
-                .contains("crashkernel=");
-        assertThat(readMicrodroidFileAsHexString(MICRODROID_DT_RAMDUMP_PATH))
+        assertThat(readFileAsString(device, MICRODROID_CMDLINE_PATH)).contains("crashkernel=");
+        assertThat(readFileAsString(device, MICRODROID_DT_BOOTARGS_PATH)).contains("crashkernel=");
+        assertThat(readFileAsHexString(device, MICRODROID_DT_RAMDUMP_PATH))
                 .isEqualTo(HEX_STRING_ZERO);
     }
 
     private boolean isDebugPolicyEnabled(@NonNull String dtPropertyPath)
             throws DeviceNotAvailableException {
-        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandRunner runner = new CommandRunner(getDevice());
         CommandResult result =
                 runner.runForResult("xxd", "-p", "/proc/device-tree" + dtPropertyPath);
         if (result.getStatus() == CommandStatus.SUCCESS) {
@@ -244,15 +176,15 @@
     }
 
     @NonNull
-    private String readMicrodroidFileAsString(@NonNull String path)
+    private String readFileAsString(@NonNull ITestDevice device, @NonNull String path)
             throws DeviceNotAvailableException {
-        return new CommandRunner(mMicrodroidDevice).run("cat", path);
+        return new CommandRunner(device).run("cat", path);
     }
 
     @NonNull
-    private String readMicrodroidFileAsHexString(@NonNull String path)
+    private String readFileAsHexString(@NonNull ITestDevice device, @NonNull String path)
             throws DeviceNotAvailableException {
-        return new CommandRunner(mMicrodroidDevice).run("xxd", "-p", path);
+        return new CommandRunner(device).run("xxd", "-p", path);
     }
 
     private void prepareCustomDebugPolicy(@NonNull String debugPolicyFileName) throws Exception {
@@ -260,11 +192,16 @@
                 getTestInformation()
                         .getDependencyFile(debugPolicyFileName, /* targetFirst= */ false);
 
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost)
-                        .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost)
-                        .build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        Pvmfw.Builder builder =
+                new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
+                        .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost);
+        if (isSecretKeeperSupported()) {
+            builder.setVmReferenceDt(getVmReferenceDtFile());
+        } else {
+            builder.setVersion(1, 1);
+        }
+        Pvmfw pvmfw = builder.build();
+        pvmfw.serialize(getCustomPvmfwFile());
     }
 
     private boolean hasConsoleOutput(@NonNull CommandResult result)
@@ -274,29 +211,22 @@
 
     private boolean hasMicrodroidLogcatOutput() throws DeviceNotAvailableException {
         CommandResult result =
-                new CommandRunner(mAndroidDevice).runForResult("test", "-s", MICRODROID_LOG_PATH);
+                new CommandRunner(getDevice()).runForResult("test", "-s", MICRODROID_LOG_PATH);
         return result.getExitCode() == 0;
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(String debugLevel)
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(String debugLevel)
             throws DeviceNotAvailableException {
         return launchProtectedVmAndWaitForBootCompleted(debugLevel, BOOT_COMPLETE_TIMEOUT_MS);
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(
             String debugLevel, long adbTimeoutMs) throws DeviceNotAvailableException {
-        mMicrodroidDevice =
-                MicrodroidBuilder.fromDevicePath(
-                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
-                        .debugLevel(debugLevel)
-                        .protectedVm(/* protectedVm= */ true)
-                        .addBootFile(mCustomPvmfwBinFileOnHost, PVMFW_FILE_NAME)
-                        .addBootFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_FILE_NAME)
-                        .setAdbConnectTimeoutMs(adbTimeoutMs)
-                        .build(mAndroidDevice);
-        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
-        assertThat(mMicrodroidDevice.enableAdbRoot()).isTrue();
-        return mMicrodroidDevice;
+        Map<String, File> bootFiles =
+                Collections.singletonMap(
+                        CUSTOM_DEBUG_POLICY_FILE_NAME, mCustomDebugPolicyFileOnHost);
+
+        return launchProtectedVmAndWaitForBootCompleted(debugLevel, adbTimeoutMs, bootFiles);
     }
 
     // Try to launch protected non-debuggable VM for a while and quit.
@@ -304,10 +234,10 @@
     private CommandResult tryLaunchProtectedNonDebuggableVm() throws Exception {
         // Can't use MicrodroidBuilder because it expects adb connection
         // but non-debuggable VM may not enable adb.
-        CommandRunner runner = new CommandRunner(mAndroidDevice);
+        CommandRunner runner = new CommandRunner(getDevice());
         runner.run("mkdir", "-p", TEST_ROOT);
-        mAndroidDevice.pushFile(mCustomPvmfwBinFileOnHost, CUSTOM_PVMFW_IMG_PATH);
-        mAndroidDevice.pushFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_PATH);
+        getDevice().pushFile(getCustomPvmfwFile(), CUSTOM_PVMFW_IMG_PATH);
+        getDevice().pushFile(mCustomDebugPolicyFileOnHost, CUSTOM_DEBUG_POLICY_PATH);
 
         // This will fail because app wouldn't finish itself.
         // But let's run the app once and get logs.
@@ -327,7 +257,11 @@
         if (isFeatureEnabled("com.android.kvm.LLPVM_CHANGES")) {
             command = String.join(" ", command, "--instance-id-file", TEST_ROOT + "instance_id");
         }
-        return mAndroidDevice.executeShellV2Command(
-                command, CONSOLE_OUTPUT_WAIT_MS, TimeUnit.MILLISECONDS, /* retryAttempts= */ 0);
+        return getDevice()
+                .executeShellV2Command(
+                        command,
+                        CONSOLE_OUTPUT_WAIT_MS,
+                        TimeUnit.MILLISECONDS,
+                        /* retryAttempts= */ 0);
     }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
index 9fbbd87..19334d6 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
@@ -16,126 +16,50 @@
 
 package com.android.pvmfw.test;
 
-import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.junit.Assume.assumeTrue;
 import static org.junit.Assert.assertThrows;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.pvmfw.test.host.Pvmfw;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceRuntimeException;
 import com.android.tradefed.device.ITestDevice;
-import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
-import com.android.tradefed.util.FileUtil;
 
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.File;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
-import java.util.Objects;
 
 /** Tests pvmfw.img and pvmfw */
 @RunWith(DeviceJUnit4ClassRunner.class)
-public class PvmfwImgTest extends MicrodroidHostTestCaseBase {
-    @NonNull private static final String PVMFW_FILE_NAME = "pvmfw_test.bin";
-    @NonNull private static final String BCC_FILE_NAME = "bcc.dat";
-    @NonNull private static final String PACKAGE_FILE_NAME = "MicrodroidTestApp.apk";
-    @NonNull private static final String PACKAGE_NAME = "com.android.microdroid.test";
-    @NonNull private static final String MICRODROID_DEBUG_FULL = "full";
-
-    @NonNull
-    private static final String MICRODROID_CONFIG_PATH = "assets/microdroid/vm_config_apex.json";
-
-    private static final int BOOT_COMPLETE_TIMEOUT_MS = 30000; // 30 seconds
-    private static final int BOOT_FAILURE_WAIT_TIME_MS = 10000; // 10 seconds
-
-    @NonNull private static final String CUSTOM_PVMFW_FILE_PREFIX = "pvmfw";
-    @NonNull private static final String CUSTOM_PVMFW_FILE_SUFFIX = ".bin";
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH = TEST_ROOT + PVMFW_FILE_NAME;
-    @NonNull private static final String CUSTOM_PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
-
-    @Nullable private static File mPvmfwBinFileOnHost;
-    @Nullable private static File mBccFileOnHost;
-
-    @Nullable private TestDevice mAndroidDevice;
-    @Nullable private ITestDevice mMicrodroidDevice;
-    @Nullable private File mCustomPvmfwBinFileOnHost;
-
-    @Before
-    public void setUp() throws Exception {
-        mAndroidDevice = (TestDevice) Objects.requireNonNull(getDevice());
-
-        // Check device capabilities
-        assumeDeviceIsCapable(mAndroidDevice);
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true));
-
-        // tradefed copies the test artfacts under /tmp when running tests,
-        // so we should *find* the artifacts with the file name.
-        mPvmfwBinFileOnHost =
-                getTestInformation().getDependencyFile(PVMFW_FILE_NAME, /* targetFirst= */ false);
-        mBccFileOnHost =
-                getTestInformation().getDependencyFile(BCC_FILE_NAME, /* targetFirst= */ false);
-
-        // Prepare for system properties for custom pvmfw.img.
-        // File will be prepared later in individual test and then pushed to device
-        // when launching with launchProtectedVmAndWaitForBootCompleted().
-        mCustomPvmfwBinFileOnHost =
-                FileUtil.createTempFile(CUSTOM_PVMFW_FILE_PREFIX, CUSTOM_PVMFW_FILE_SUFFIX);
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, CUSTOM_PVMFW_IMG_PATH);
-
-        // Prepare for launching microdroid
-        mAndroidDevice.installPackage(findTestFile(PACKAGE_FILE_NAME), /* reinstall */ false);
-        prepareVirtualizationTestSetup(mAndroidDevice);
-        mMicrodroidDevice = null;
-    }
-
-    @After
-    public void shutdown() throws Exception {
-        if (!mAndroidDevice.supportsMicrodroid(/* protectedVm= */ true)) {
-            return;
-        }
-        if (mMicrodroidDevice != null) {
-            mAndroidDevice.shutdownMicrodroid(mMicrodroidDevice);
-            mMicrodroidDevice = null;
-        }
-        mAndroidDevice.uninstallPackage(PACKAGE_NAME);
-
-        // Cleanup for custom pvmfw.img
-        setPropertyOrThrow(mAndroidDevice, CUSTOM_PVMFW_IMG_PATH_PROP, "");
-        FileUtil.deleteFile(mCustomPvmfwBinFileOnHost);
-
-        cleanUpVirtualizationTestSetup(mAndroidDevice);
-    }
-
+public class PvmfwImgTest extends CustomPvmfwHostTestCaseBase {
     @Test
-    public void testConfigVersion1_0_boots() throws Exception {
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 0).build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+    public void testPvmfw_beforeVmReferenceDt_whenSecretKeeperExists() throws Exception {
+        // VM reference DT is added since version 1.2
+        List<int[]> earlyVersions = Arrays.asList(new int[] {1, 0}, new int[] {1, 1});
+        Pvmfw.Builder builder = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile());
 
-        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
-    }
+        for (int[] pair : earlyVersions) {
+            int major = pair[0];
+            int minor = pair[1];
+            String version = "v" + major + "." + minor;
 
-    @Test
-    public void testConfigVersion1_1_boots() throws Exception {
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 1).build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+            // Pvmfw config before v1.2 can't have secret keeper key in VM reference DT.
+            Pvmfw pvmfw = builder.setVersion(major, minor).build();
+            pvmfw.serialize(getCustomPvmfwFile());
 
-        launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+            if (isSecretKeeperSupported()) {
+                // If secret keeper is supported, we can't boot with early version
+                assertThrows(
+                        "pvmfw shouldn't boot without VM reference DT, version=" + version,
+                        DeviceRuntimeException.class,
+                        () -> launchProtectedVmAndWaitForBootCompleted(BOOT_FAILURE_WAIT_TIME_MS));
+            } else {
+                launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
+                shutdownMicrodroid();
+            }
+        }
     }
 
     @Test
@@ -153,15 +77,23 @@
                         new int[] {0xFFFF, 1},
                         new int[] {0xFFFF, 0xFFFF});
 
-        Pvmfw.Builder builder = new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost);
+        Pvmfw.Builder builder =
+                new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
+                        .setVmReferenceDt(getVmReferenceDtFile());
 
         for (int[] pair : invalid_versions) {
             int major = pair[0];
             int minor = pair[1];
             String version = "v" + major + "." + minor;
 
+            if (Pvmfw.makeVersion(major, minor) >= Pvmfw.makeVersion(1, 2)
+                    && getVmReferenceDtFile() == null) {
+                // VM reference DT is unavailable, so we can't even build Pvmfw.
+                continue;
+            }
+
             Pvmfw pvmfw = builder.setVersion(major, minor).build();
-            pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+            pvmfw.serialize(getCustomPvmfwFile());
 
             assertThrows(
                     "pvmfw shouldn't boot with invalid version " + version,
@@ -170,17 +102,9 @@
         }
     }
 
-    private ITestDevice launchProtectedVmAndWaitForBootCompleted(long adbTimeoutMs)
+    public ITestDevice launchProtectedVmAndWaitForBootCompleted(long adbTimeoutMs)
             throws DeviceNotAvailableException {
-        mMicrodroidDevice =
-                MicrodroidBuilder.fromDevicePath(
-                                getPathForPackage(PACKAGE_NAME), MICRODROID_CONFIG_PATH)
-                        .debugLevel(MICRODROID_DEBUG_FULL)
-                        .protectedVm(true)
-                        .addBootFile(mCustomPvmfwBinFileOnHost, PVMFW_FILE_NAME)
-                        .setAdbConnectTimeoutMs(adbTimeoutMs)
-                        .build(mAndroidDevice);
-        assertThat(mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS)).isTrue();
-        return mMicrodroidDevice;
+        return launchProtectedVmAndWaitForBootCompleted(
+                MICRODROID_DEBUG_FULL, adbTimeoutMs, Collections.emptyMap());
     }
 }
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java
new file mode 100644
index 0000000..bea72eb
--- /dev/null
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwTest.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.pvmfw.test;
+
+import static org.junit.Assert.assertThrows;
+
+import com.android.pvmfw.test.host.Pvmfw;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Test test helper */
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class PvmfwTest extends CustomPvmfwHostTestCaseBase {
+    @Test
+    public void testPvmfw_withConfig1_2_requiresReferenceDt() {
+        assertThrows(
+                "pvmfw config 1.2 must require VM reference DT",
+                NullPointerException.class,
+                () -> {
+                    new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 2).build();
+                });
+    }
+
+    @Test
+    public void testPvmfw_before1_2_doesNotRequiresReferenceDt() {
+        new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 1).build();
+    }
+}
diff --git a/tests/pvmfw/tools/PvmfwTool.java b/tests/pvmfw/tools/PvmfwTool.java
index 62c641b..e150ec4 100644
--- a/tests/pvmfw/tools/PvmfwTool.java
+++ b/tests/pvmfw/tools/PvmfwTool.java
@@ -25,28 +25,41 @@
 public class PvmfwTool {
     public static void printUsage() {
         System.out.println("pvmfw-tool: Appends pvmfw.bin and config payloads.");
-        System.out.println("Requires BCC and optional debug policy dtbo files");
-        System.out.println("");
-        System.out.println("Usage: pvmfw-tool <out> <pvmfw.bin> <bcc.dat> [<dp.dtbo>]");
+        System.out.println("            Requires BCC and VM reference DT.");
+        System.out.println("            VM DTBO and Debug policy can optionally be specified");
+        System.out.println(
+                "Usage: pvmfw-tool <out> <pvmfw.bin> <bcc.dat> <VM reference DT> [VM DTBO] [debug"
+                        + " policy]");
     }
 
     public static void main(String[] args) {
-        if (args.length != 4 && args.length != 3) {
+        if (args.length < 3 || args.length > 6) {
             printUsage();
             System.exit(1);
         }
 
         File out = new File(args[0]);
-        File pvmfw_bin = new File(args[1]);
-        File bcc_dat = new File(args[2]);
+        File pvmfwBin = new File(args[1]);
+        File bccData = new File(args[2]);
+        File vmReferenceDt = new File(args[3]);
+
+        File vmDtbo = null;
+        File dp = null;
+        if (args.length > 4) {
+            vmDtbo = new File(args[4]);
+        }
+        if (args.length > 5) {
+            dp = new File(args[5]);
+        }
 
         try {
-            Pvmfw.Builder builder = new Pvmfw.Builder(pvmfw_bin, bcc_dat);
-            if (args.length == 4) {
-                File dtbo = new File(args[3]);
-                builder.setDebugPolicyOverlay(dtbo);
-            }
-            builder.build().serialize(out);
+            Pvmfw pvmfw =
+                    new Pvmfw.Builder(pvmfwBin, bccData)
+                            .setVmReferenceDt(vmReferenceDt)
+                            .setDebugPolicyOverlay(dp)
+                            .setVmDtbo(vmDtbo)
+                            .build();
+            pvmfw.serialize(out);
         } catch (IOException e) {
             e.printStackTrace();
             printUsage();
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 2a04103..732be94 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -17,6 +17,7 @@
     name: "MicrodroidTestAppsDefaults",
     test_suites: [
         "cts",
+        "vts",
         "general-tests",
     ],
     static_libs: [
diff --git a/tests/testapk/AndroidTest.xml b/tests/testapk/AndroidTest.xml
index 8a4c367..22cd0dc 100644
--- a/tests/testapk/AndroidTest.xml
+++ b/tests/testapk/AndroidTest.xml
@@ -15,6 +15,7 @@
 -->
 <configuration description="Runs Microdroid device-side tests.">
     <option name="test-suite-tag" value="cts" />
+    <option name="test-suite-tag" value="vts" />
     <option name="config-descriptor:metadata" key="component" value="security" />
     <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
     <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 278365c..22bea58 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -434,11 +434,7 @@
                 None
             };
 
-        let debug_level = match config {
-            VirtualMachineConfig::AppConfig(config) => config.debugLevel,
-            _ => DebugLevel::NONE,
-        };
-        let debug_config = DebugConfig::new(debug_level);
+        let debug_config = DebugConfig::new(config);
 
         let ramdump = if debug_config.is_ramdump_needed() {
             Some(prepare_ramdump_file(&temporary_directory)?)
diff --git a/virtualizationmanager/src/debug_config.rs b/virtualizationmanager/src/debug_config.rs
index e2b657a..451d1c6 100644
--- a/virtualizationmanager/src/debug_config.rs
+++ b/virtualizationmanager/src/debug_config.rs
@@ -15,17 +15,17 @@
 //! Functions for AVF debug policy and debug level
 
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    VirtualMachineAppConfig::DebugLevel::DebugLevel,
+    VirtualMachineAppConfig::DebugLevel::DebugLevel, VirtualMachineConfig::VirtualMachineConfig,
 };
 use anyhow::{anyhow, Context, Error, Result};
+use lazy_static::lazy_static;
+use libfdt::{Fdt, FdtError};
+use log::{info, warn};
+use rustutils::system_properties;
+use std::ffi::{CString, NulError};
 use std::fs;
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
-use std::ffi::{CString, NulError};
-use log::{warn, info};
-use rustutils::system_properties;
-use libfdt::{Fdt, FdtError};
-use lazy_static::lazy_static;
 
 const CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP: &str =
     "hypervisor.virtualizationmanager.debug_policy.path";
@@ -156,7 +156,12 @@
 }
 
 impl DebugConfig {
-    pub fn new(debug_level: DebugLevel) -> Self {
+    pub fn new(config: &VirtualMachineConfig) -> Self {
+        let debug_level = match config {
+            VirtualMachineConfig::AppConfig(config) => config.debugLevel,
+            _ => DebugLevel::NONE,
+        };
+
         match system_properties::read(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP).unwrap_or_default() {
             Some(path) if !path.is_empty() => {
                 match Self::from_custom_debug_overlay_policy(debug_level, Path::new(&path)) {
@@ -179,6 +184,11 @@
         }
 
         info!("Debug policy is disabled");
+        Self::new_with_debug_level(debug_level)
+    }
+
+    /// Creates a new DebugConfig with debug level. Only use this for test purpose.
+    pub fn new_with_debug_level(debug_level: DebugLevel) -> Self {
         Self {
             debug_level,
             debug_policy_log: false,
@@ -203,7 +213,6 @@
         self.debug_level != DebugLevel::NONE || self.debug_policy_ramdump
     }
 
-    // TODO: Remove this code path in user build for removing libfdt depenency.
     fn from_custom_debug_overlay_policy(debug_level: DebugLevel, path: &Path) -> Result<Self> {
         match OwnedFdt::from_overlay_onto_new_fdt(path) {
             Ok(fdt) => Ok(Self {
@@ -229,14 +238,6 @@
 #[cfg(test)]
 mod tests {
     use super::*;
-    use anyhow::ensure;
-
-    fn can_set_sysprop() -> bool {
-        if let Ok(Some(value)) = system_properties::read("ro.build.type") {
-            return "user".eq(&value);
-        }
-        false // if we're in doubt, skip test.
-    }
 
     #[test]
     fn test_read_avf_debug_policy_with_ramdump() -> Result<()> {
@@ -317,40 +318,4 @@
 
         Ok(())
     }
-
-    fn test_new_with_custom_policy_internal() -> Result<()> {
-        let debug_config = DebugConfig::new(DebugLevel::NONE);
-
-        ensure!(debug_config.debug_level == DebugLevel::NONE);
-        ensure!(!debug_config.debug_policy_log);
-        ensure!(!debug_config.debug_policy_ramdump);
-        ensure!(debug_config.debug_policy_adb);
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_new_with_custom_policy() -> Result<()> {
-        if !can_set_sysprop() {
-            // Skip test if we can't override sysprop.
-            return Ok(());
-        }
-
-        // Setup
-        let old_sysprop = system_properties::read(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP)
-            .context("Failed to read existing sysprop")?
-            .unwrap_or_default();
-        let file_name = "avf_debug_policy_with_adb.dtbo";
-        system_properties::write(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP, file_name)
-            .context("Failed to set sysprop")?;
-
-        // Run test
-        let test_result = test_new_with_custom_policy_internal();
-
-        // Clean up.
-        system_properties::write(CUSTOM_DEBUG_POLICY_OVERLAY_SYSPROP, &old_sysprop)
-            .context("Failed to restore sysprop")?;
-
-        test_result
-    }
 }
diff --git a/virtualizationmanager/src/payload.rs b/virtualizationmanager/src/payload.rs
index 05626d3..9d0c7d6 100644
--- a/virtualizationmanager/src/payload.rs
+++ b/virtualizationmanager/src/payload.rs
@@ -631,7 +631,7 @@
             collect_apex_infos(
                 &apex_info_list,
                 &apex_configs,
-                &DebugConfig::new(DebugLevel::FULL)
+                &DebugConfig::new_with_debug_level(DebugLevel::FULL)
             )?,
             vec![
                 // Pass active/required APEXes
@@ -660,8 +660,11 @@
         };
         let apex_configs = vec![ApexConfig { name: "apex-vendor".to_string() }];
 
-        let ret =
-            collect_apex_infos(&apex_info_list, &apex_configs, &DebugConfig::new(DebugLevel::NONE));
+        let ret = collect_apex_infos(
+            &apex_info_list,
+            &apex_configs,
+            &DebugConfig::new_with_debug_level(DebugLevel::NONE),
+        );
         assert!(ret
             .is_err_and(|ret| ret.to_string()
                 == "Non-system APEX apex-vendor is not supported in Microdroid"));
@@ -687,7 +690,7 @@
             collect_apex_infos(
                 &apex_info_list,
                 &apex_configs,
-                &DebugConfig::new(DebugLevel::NONE)
+                &DebugConfig::new_with_debug_level(DebugLevel::NONE)
             )?,
             vec![&apex_info_list.list[0]]
         );
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index fc7fcd2..0c39501 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -77,6 +77,9 @@
         "virtualizationservice_defaults",
     ],
     test_suites: ["general-tests"],
+    rustlibs: [
+        "libtempfile",
+    ],
     data: [
         ":test_rkp_cert_chain",
     ],
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index bbfb220..2fe14c0 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -39,7 +39,10 @@
 use openssl::x509::X509;
 use rand::Fill;
 use rkpd_client::get_rkpd_attestation_key;
-use rustutils::system_properties;
+use rustutils::{
+    system_properties,
+    users::{multiuser_get_app_id, multiuser_get_user_id},
+};
 use serde::Deserialize;
 use service_vm_comm::Response;
 use std::collections::{HashMap, HashSet};
@@ -385,7 +388,6 @@
         Ok(ParcelFileDescriptor::new(file))
     }
 
-    // TODO(b/294177871) Persist this Id, along with client uuid.
     fn allocateInstanceId(&self) -> binder::Result<[u8; 64]> {
         let mut id = [0u8; 64];
         id.try_fill(&mut rand::thread_rng())
@@ -393,6 +395,16 @@
             .or_service_specific_exception(-1)?;
         let uid = get_calling_uid();
         info!("Allocated a VM's instance_id: {:?}, for uid: {:?}", hex::encode(id), uid);
+        let state = &mut *self.state.lock().unwrap();
+        if let Some(sk_state) = &mut state.sk_state {
+            let user_id = multiuser_get_user_id(uid);
+            let app_id = multiuser_get_app_id(uid);
+            info!("Recording potential existence of state for (user_id={user_id}, app_id={app_id}");
+            if let Err(e) = sk_state.add_id(&id, user_id, app_id) {
+                error!("Failed to record the instance_id: {e:?}");
+            }
+        }
+
         Ok(id)
     }
 
diff --git a/virtualizationservice/src/maintenance.rs b/virtualizationservice/src/maintenance.rs
index 7fc2f37..0a367c5 100644
--- a/virtualizationservice/src/maintenance.rs
+++ b/virtualizationservice/src/maintenance.rs
@@ -15,7 +15,7 @@
 use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::{
     ISecretkeeper::ISecretkeeper, SecretId::SecretId,
 };
-use anyhow::Result;
+use anyhow::{Context, Result};
 use log::{error, info, warn};
 
 mod vmdb;
@@ -88,6 +88,13 @@
         }
     }
 
+    /// Record a new VM ID.
+    pub fn add_id(&mut self, vm_id: &VmId, user_id: u32, app_id: u32) -> Result<()> {
+        let user_id: i32 = user_id.try_into().context(format!("user_id {user_id} out of range"))?;
+        let app_id: i32 = app_id.try_into().context(format!("app_id {app_id} out of range"))?;
+        self.vm_id_db.add_vm_id(vm_id, user_id, app_id)
+    }
+
     /// Delete the VM IDs associated with Android user ID `user_id`.
     pub fn delete_ids_for_user(&mut self, user_id: i32) -> Result<()> {
         let vm_ids = self.vm_id_db.vm_ids_for_user(user_id)?;
diff --git a/virtualizationservice/src/maintenance/vmdb.rs b/virtualizationservice/src/maintenance/vmdb.rs
index bdff034..ce1e1e7 100644
--- a/virtualizationservice/src/maintenance/vmdb.rs
+++ b/virtualizationservice/src/maintenance/vmdb.rs
@@ -14,7 +14,7 @@
 
 //! Database of VM IDs.
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use log::{debug, error, info, warn};
 use rusqlite::{params, params_from_iter, Connection, OpenFlags, Rows};
 use std::path::PathBuf;
@@ -29,6 +29,15 @@
 /// (Default value of `SQLITE_LIMIT_VARIABLE_NUMBER` for <= 3.32.0)
 const MAX_VARIABLES: usize = 999;
 
+/// Return the current time as milliseconds since epoch.
+fn db_now() -> u64 {
+    let now = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap_or(std::time::Duration::ZERO)
+        .as_millis();
+    now.try_into().unwrap_or(u64::MAX)
+}
+
 /// Identifier for a VM and its corresponding secret.
 pub type VmId = [u8; 64];
 
@@ -37,6 +46,8 @@
     conn: Connection,
 }
 
+struct RetryOnFailure(bool);
+
 impl VmIdDb {
     /// Connect to the VM ID database file held in the given directory, creating it if necessary.
     /// The second return value indicates whether a new database file was created.
@@ -49,8 +60,11 @@
             std::fs::create_dir(&db_path).context("failed to create {db_path:?}")?;
             info!("created persistent db dir {db_path:?}");
         }
-
         db_path.push(DB_FILENAME);
+        Self::new_at_path(db_path, RetryOnFailure(true))
+    }
+
+    fn new_at_path(db_path: PathBuf, retry: RetryOnFailure) -> Result<(Self, bool)> {
         let (flags, created) = if db_path.exists() {
             debug!("connecting to existing database {db_path:?}");
             (
@@ -69,15 +83,42 @@
                 true,
             )
         };
-        let mut result = Self {
-            conn: Connection::open_with_flags(db_path, flags)
+        let mut db = Self {
+            conn: Connection::open_with_flags(&db_path, flags)
                 .context(format!("failed to open/create DB with {flags:?}"))?,
         };
 
         if created {
-            result.init_tables().context("failed to create tables")?;
+            db.init_tables().context("failed to create tables")?;
+        } else {
+            // An existing .sqlite file may have an earlier schema.
+            match db.schema_version() {
+                Err(e) => {
+                    // Couldn't determine a schema version, so wipe and try again.
+                    error!("failed to determine VM DB schema: {e:?}");
+                    if retry.0 {
+                        // This is the first attempt, so wipe and retry.
+                        error!("resetting database file {db_path:?}");
+                        let _ = std::fs::remove_file(&db_path);
+                        return Self::new_at_path(db_path, RetryOnFailure(false));
+                    } else {
+                        // An earlier attempt at wiping/retrying has failed, so give up.
+                        return Err(anyhow!("failed to reset database file {db_path:?}"));
+                    }
+                }
+                Ok(0) => db.upgrade_tables_v0_v1().context("failed to upgrade schema v0 -> v1")?,
+                Ok(1) => {
+                    // Current version, no action needed.
+                }
+                Ok(version) => {
+                    // If the database looks like it's from a future version, leave it alone and
+                    // fail to connect to it.
+                    error!("database from the future (v{version})");
+                    return Err(anyhow!("database from the future (v{version})"));
+                }
+            }
         }
-        Ok((result, created))
+        Ok((db, created))
     }
 
     /// Delete the associated database file.
@@ -94,8 +135,63 @@
         }
     }
 
-    /// Create the database table and indices.
+    fn schema_version(&mut self) -> Result<i32> {
+        let version: i32 = self
+            .conn
+            .query_row("PRAGMA main.user_version", (), |row| row.get(0))
+            .context("failed to read pragma")?;
+        Ok(version)
+    }
+
+    /// Create the database table and indices using the current schema.
     fn init_tables(&mut self) -> Result<()> {
+        self.init_tables_v1()
+    }
+
+    /// Create the database table and indices using the v1 schema.
+    fn init_tables_v1(&mut self) -> Result<()> {
+        info!("creating v1 database schema");
+        self.conn
+            .execute(
+                "CREATE TABLE IF NOT EXISTS main.vmids (
+                     vm_id BLOB PRIMARY KEY,
+                     user_id INTEGER,
+                     app_id INTEGER,
+                     created INTEGER
+                 ) WITHOUT ROWID;",
+                (),
+            )
+            .context("failed to create table")?;
+        self.conn
+            .execute("CREATE INDEX IF NOT EXISTS main.vmids_user_index ON vmids(user_id);", [])
+            .context("Failed to create user index")?;
+        self.conn
+            .execute(
+                "CREATE INDEX IF NOT EXISTS main.vmids_app_index ON vmids(user_id, app_id);",
+                [],
+            )
+            .context("Failed to create app index")?;
+        self.conn
+            .execute("PRAGMA main.user_version = 1;", ())
+            .context("failed to declare version")?;
+        Ok(())
+    }
+
+    fn upgrade_tables_v0_v1(&mut self) -> Result<()> {
+        let _rows = self
+            .conn
+            .execute("ALTER TABLE main.vmids ADD COLUMN created INTEGER;", ())
+            .context("failed to alter table v0->v1")?;
+        self.conn
+            .execute("PRAGMA main.user_version = 1;", ())
+            .context("failed to set schema version")?;
+        Ok(())
+    }
+
+    /// Create the database table and indices using the v0 schema.
+    #[cfg(test)]
+    fn init_tables_v0(&mut self) -> Result<()> {
+        info!("creating v0 database schema");
         self.conn
             .execute(
                 "CREATE TABLE IF NOT EXISTS main.vmids (
@@ -119,13 +215,13 @@
     }
 
     /// Add the given VM ID into the database.
-    #[allow(dead_code)] // TODO(b/294177871): connect this up
     pub fn add_vm_id(&mut self, vm_id: &VmId, user_id: i32, app_id: i32) -> Result<()> {
+        let now = db_now();
         let _rows = self
             .conn
             .execute(
-                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
-                params![vm_id, &user_id, &app_id],
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
+                params![vm_id, &user_id, &app_id, &now],
             )
             .context("failed to add VM ID")?;
         Ok(())
@@ -177,16 +273,21 @@
     }
 }
 
+/// Current schema version.
+#[cfg(test)]
+const SCHEMA_VERSION: usize = 1;
+
+/// Create a new in-memory database for testing.
 #[cfg(test)]
 pub fn new_test_db() -> VmIdDb {
-    let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
-    db.init_tables().unwrap();
-    db
+    tests::new_test_db_version(SCHEMA_VERSION)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::io::Write;
+
     const VM_ID1: VmId = [1u8; 64];
     const VM_ID2: VmId = [2u8; 64];
     const VM_ID3: VmId = [3u8; 64];
@@ -201,6 +302,113 @@
     const APP_C: i32 = 70;
     const APP_UNKNOWN: i32 = 99;
 
+    pub fn new_test_db_version(version: usize) -> VmIdDb {
+        let mut db = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        match version {
+            0 => db.init_tables_v0().unwrap(),
+            1 => db.init_tables_v1().unwrap(),
+            _ => panic!("unexpected version {version}"),
+        }
+        db
+    }
+
+    fn show_contents(db: &VmIdDb) {
+        let mut stmt = db.conn.prepare("SELECT * FROM main.vmids;").unwrap();
+        let mut rows = stmt.query(()).unwrap();
+        while let Some(row) = rows.next().unwrap() {
+            println!("  {row:?}");
+        }
+    }
+
+    #[test]
+    fn test_schema_version0() {
+        let mut db0 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        db0.init_tables_v0().unwrap();
+        let version = db0.schema_version().unwrap();
+        assert_eq!(0, version);
+    }
+
+    #[test]
+    fn test_schema_version1() {
+        let mut db1 = VmIdDb { conn: Connection::open_in_memory().unwrap() };
+        db1.init_tables_v1().unwrap();
+        let version = db1.schema_version().unwrap();
+        assert_eq!(1, version);
+    }
+
+    #[test]
+    fn test_schema_upgrade_v0_v1() {
+        let mut db = new_test_db_version(0);
+        let version = db.schema_version().unwrap();
+        assert_eq!(0, version);
+
+        // Manually insert a row before upgrade.
+        db.conn
+            .execute(
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
+                params![&VM_ID1, &USER1, APP_A],
+            )
+            .unwrap();
+
+        db.upgrade_tables_v0_v1().unwrap();
+        let version = db.schema_version().unwrap();
+        assert_eq!(1, version);
+
+        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
+        show_contents(&db);
+    }
+
+    #[test]
+    fn test_corrupt_database_file() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        let mut db_path = db_dir.path().to_owned();
+        db_path.push(DB_FILENAME);
+        {
+            let mut file = std::fs::File::create(db_path).unwrap();
+            let _ = file.write_all(b"This is not an SQLite file!");
+        }
+
+        // Non-DB file should be wiped and start over.
+        let (mut db, created) =
+            VmIdDb::new(&db_dir.path().to_string_lossy()).expect("failed to replace bogus DB");
+        assert!(created);
+        db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
+        assert_eq!(vec![VM_ID1], db.vm_ids_for_user(USER1).unwrap());
+    }
+
+    #[test]
+    fn test_non_upgradable_database_file() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        let mut db_path = db_dir.path().to_owned();
+        db_path.push(DB_FILENAME);
+        {
+            // Create an unrelated database that happens to apparently have a schema version of 0.
+            let (db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
+            assert!(created);
+            db.conn.execute("DROP TABLE main.vmids", ()).unwrap();
+            db.conn.execute("PRAGMA main.user_version = 0;", ()).unwrap();
+        }
+
+        // Should fail to open a database because the upgrade fails.
+        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_database_from_the_future() {
+        let db_dir = tempfile::Builder::new().prefix("vmdb-test-").tempdir().unwrap();
+        {
+            let (mut db, created) = VmIdDb::new(&db_dir.path().to_string_lossy()).unwrap();
+            assert!(created);
+            db.add_vm_id(&VM_ID1, USER1, APP_A).unwrap();
+            // Make the database look like it's from a future version.
+            db.conn.execute("PRAGMA main.user_version = 99;", ()).unwrap();
+        }
+        // Should fail to open a database from the future.
+        let result = VmIdDb::new(&db_dir.path().to_string_lossy());
+        assert!(result.is_err());
+    }
+
     #[test]
     fn test_add_remove() {
         let mut db = new_test_db();
@@ -239,6 +447,7 @@
         assert_eq!(vec![VM_ID5], db.vm_ids_for_user(USER3).unwrap());
         assert_eq!(empty, db.vm_ids_for_user(USER_UNKNOWN).unwrap());
         assert_eq!(empty, db.vm_ids_for_app(USER1, APP_UNKNOWN).unwrap());
+        show_contents(&db);
     }
 
     #[test]
@@ -254,12 +463,13 @@
         // Manually insert a row with a VM ID that's the wrong size.
         db.conn
             .execute(
-                "REPLACE INTO main.vmids (vm_id, user_id, app_id) VALUES (?1, ?2, ?3);",
-                params![&[99u8; 60], &USER1, APP_A],
+                "REPLACE INTO main.vmids (vm_id, user_id, app_id, created) VALUES (?1, ?2, ?3, ?4);",
+                params![&[99u8; 60], &USER1, APP_A, &db_now()],
             )
             .unwrap();
 
         // Invalid row is skipped and remainder returned.
         assert_eq!(vec![VM_ID1, VM_ID2, VM_ID3], db.vm_ids_for_user(USER1).unwrap());
+        show_contents(&db);
     }
 }