Merge "[attestation_ext] Check attestation extension in e2e test" into main
diff --git a/Android.bp b/Android.bp
index f50007f..2091a90 100644
--- a/Android.bp
+++ b/Android.bp
@@ -68,6 +68,7 @@
     module_type: "cc_defaults",
     config_namespace: "ANDROID",
     bool_variables: [
+        "release_avf_enable_dice_changes",
         "release_avf_enable_virt_cpufreq",
     ],
     properties: [
@@ -78,6 +79,9 @@
 avf_flag_aware_cc_defaults {
     name: "avf_build_flags_cc",
     soong_config_variables: {
+        release_avf_enable_dice_changes: {
+            cflags: ["-DAVF_OPEN_DICE_CHANGES=1"],
+        },
         release_avf_enable_virt_cpufreq: {
             cflags: ["-DAVF_ENABLE_VIRT_CPUFREQ=1"],
         },
diff --git a/avf_flags.aconfig b/avf_flags.aconfig
index 8abb9ee..589d227 100644
--- a/avf_flags.aconfig
+++ b/avf_flags.aconfig
@@ -2,6 +2,7 @@
 
 flag {
   name: "avf_v_test_apis"
+  is_exported: true
   namespace: "virtualization"
   description: "Only purpose of this flag is to be used in @FlaggedApi in our V test apis"
   bug: "325441024"
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/flags/cpp/Android.bp b/flags/cpp/Android.bp
new file mode 100644
index 0000000..da4158a
--- /dev/null
+++ b/flags/cpp/Android.bp
@@ -0,0 +1,13 @@
+cc_library_static {
+    name: "libavf_cc_flags",
+    defaults: ["avf_build_flags_cc"],
+    apex_available: [
+        "//apex_available:platform",
+        "com.android.virt",
+    ],
+
+    export_include_dirs: ["include"],
+    local_include_dirs: ["include"],
+    ramdisk_available: true,
+    recovery_available: true,
+}
diff --git a/flags/cpp/include/android/avf_cc_flags.h b/flags/cpp/include/android/avf_cc_flags.h
new file mode 100644
index 0000000..536ea9f
--- /dev/null
+++ b/flags/cpp/include/android/avf_cc_flags.h
@@ -0,0 +1,31 @@
+// 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.
+
+#pragma once
+
+// TODO(b/309090563): remove this file once build flags are exposed to aconfig.
+
+namespace android {
+namespace virtualization {
+
+inline bool IsOpenDiceChangesFlagEnabled() {
+#ifdef AVF_OPEN_DICE_CHANGES
+    return AVF_OPEN_DICE_CHANGES;
+#else
+    return false;
+#endif
+}
+
+} // namespace virtualization
+} // namespace android
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/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
index 693a7d7..12aeac8 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -829,12 +829,12 @@
         @SystemApi
         @NonNull
         public Builder setPayloadBinaryName(@NonNull String payloadBinaryName) {
+            requireNonNull(payloadBinaryName, "payloadBinaryName must not be null");
             if (payloadBinaryName.contains(File.separator)) {
                 throw new IllegalArgumentException(
                         "Invalid binary file name: " + payloadBinaryName);
             }
-            mPayloadBinaryName =
-                    requireNonNull(payloadBinaryName, "payloadBinaryName must not be null");
+            mPayloadBinaryName = payloadBinaryName;
             return this;
         }
 
diff --git a/libs/dice/driver/Android.bp b/libs/dice/driver/Android.bp
new file mode 100644
index 0000000..c93bd7d
--- /dev/null
+++ b/libs/dice/driver/Android.bp
@@ -0,0 +1,50 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_defaults {
+    name: "libdice_driver_defaults",
+    crate_name: "dice_driver",
+    defaults: [
+        "avf_build_flags_rust",
+    ],
+    srcs: ["src/lib.rs"],
+    edition: "2021",
+    prefer_rlib: true,
+    rustlibs: [
+        "libanyhow",
+        "libbyteorder",
+        "libcoset",
+        "libdice_policy_builder",
+        "libdiced_open_dice",
+        "libdiced_sample_inputs",
+        "libkeystore2_crypto_rust",
+        "liblibc",
+        "liblog_rust",
+        "libnix",
+        "libonce_cell",
+        "libopenssl",
+        "libthiserror",
+        "libserde_cbor",
+    ],
+    multilib: {
+        lib32: {
+            enabled: false,
+        },
+    },
+}
+
+rust_library {
+    name: "libdice_driver",
+    defaults: ["libdice_driver_defaults"],
+}
+
+rust_test {
+    name: "libdice_driver_test",
+    defaults: ["libdice_driver_defaults"],
+    test_suites: ["general-tests"],
+    rustlibs: [
+        "libhex",
+        "libtempfile",
+    ],
+}
diff --git a/microdroid_manager/src/dice_driver.rs b/libs/dice/driver/src/lib.rs
similarity index 70%
rename from microdroid_manager/src/dice_driver.rs
rename to libs/dice/driver/src/lib.rs
index 229f3e0..79edb51 100644
--- a/microdroid_manager/src/dice_driver.rs
+++ b/libs/dice/driver/src/lib.rs
@@ -32,13 +32,26 @@
 
 /// Artifacts that are mapped into the process address space from the driver.
 pub enum DiceDriver<'a> {
+    /// Implementation that reads bcc handover from the dice driver.
     Real {
+        /// Path to the driver character device (e.g. /dev/open-dice0).
         driver_path: PathBuf,
+        /// Address of the memory to mmap driver to.
         mmap_addr: *mut c_void,
+        /// Size of the mmap.
         mmap_size: usize,
+        /// BCC handover.
         bcc_handover: BccHandover<'a>,
     },
+    /// Fake implementation used in tests and non-protected VMs.
     Fake(OwnedDiceArtifacts),
+    /// Implementation that reads bcc handover from the file.
+    FromFile {
+        /// Path to the file to read dice chain from,
+        file_path: PathBuf,
+        /// Dice artifacts read from file_path,
+        dice_artifacts: OwnedDiceArtifacts,
+    },
 }
 
 impl DiceDriver<'_> {
@@ -46,13 +59,15 @@
         match self {
             Self::Real { bcc_handover, .. } => bcc_handover,
             Self::Fake(owned_dice_artifacts) => owned_dice_artifacts,
+            Self::FromFile { dice_artifacts, .. } => dice_artifacts,
         }
     }
 
-    pub fn new(driver_path: &Path) -> Result<Self> {
+    /// Creates a new dice driver from the given driver_path.
+    pub fn new(driver_path: &Path, is_strict_boot: bool) -> Result<Self> {
         if driver_path.exists() {
             log::info!("Using DICE values from driver");
-        } else if super::is_strict_boot() {
+        } else if is_strict_boot {
             bail!("Strict boot requires DICE value from driver but none were found");
         } else {
             log::warn!("Using sample DICE values");
@@ -90,6 +105,15 @@
         })
     }
 
+    /// Create a new dice driver that reads dice_artifacts from the given file.
+    pub fn from_file(file_path: &Path) -> Result<Self> {
+        let file =
+            fs::File::open(file_path).map_err(|error| Error::new(error).context("open file"))?;
+        let dice_artifacts = serde_cbor::from_reader(file)
+            .map_err(|error| Error::new(error).context("read file"))?;
+        Ok(Self::FromFile { file_path: file_path.to_path_buf(), dice_artifacts })
+    }
+
     /// Derives a sealing key of `key_length` bytes from the DICE sealing CDI.
     pub fn get_sealing_key(&self, identifier: &[u8], key_length: usize) -> Result<ZVec> {
         // Deterministically derive a key to use for sealing data, rather than using the CDI
@@ -101,6 +125,7 @@
         Ok(key)
     }
 
+    /// Derives a new dice chain.
     pub fn derive(
         self,
         code_hash: Hash,
@@ -147,3 +172,36 @@
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn assert_eq_bytes(expected: &[u8], actual: &[u8]) {
+        assert_eq!(
+            expected,
+            actual,
+            "Expected {}, got {}",
+            hex::encode(expected),
+            hex::encode(actual)
+        )
+    }
+
+    #[test]
+    fn test_write_bcc_to_file_read_from_file() -> Result<()> {
+        let dice_artifacts = diced_sample_inputs::make_sample_bcc_and_cdis()?;
+
+        let test_file = tempfile::NamedTempFile::new()?;
+        serde_cbor::to_writer(test_file.as_file(), &dice_artifacts)?;
+        test_file.as_file().sync_all()?;
+
+        let dice = DiceDriver::from_file(test_file.as_ref())?;
+
+        let dice_artifacts2 = dice.dice_artifacts();
+        assert_eq_bytes(dice_artifacts.cdi_attest(), dice_artifacts2.cdi_attest());
+        assert_eq_bytes(dice_artifacts.cdi_seal(), dice_artifacts2.cdi_seal());
+        assert_eq_bytes(dice_artifacts.bcc().expect("bcc"), dice_artifacts2.bcc().expect("bcc"));
+
+        Ok(())
+    }
+}
diff --git a/libs/dice/open_dice/Android.bp b/libs/dice/open_dice/Android.bp
index 79d0b96..ab3220e 100644
--- a/libs/dice/open_dice/Android.bp
+++ b/libs/dice/open_dice/Android.bp
@@ -39,12 +39,15 @@
     rustlibs: [
         "libopen_dice_android_bindgen",
         "libopen_dice_cbor_bindgen",
+        "libserde",
         "libzeroize",
     ],
     features: [
         "alloc",
+        "serde_derive",
         "std",
     ],
+    proc_macros: ["libserde_derive"],
     shared_libs: [
         "libcrypto",
     ],
@@ -69,7 +72,7 @@
 rust_defaults {
     name: "libdiced_open_dice_test_defaults",
     crate_name: "diced_open_dice_test",
-    srcs: ["tests/*.rs"],
+    srcs: ["tests/api_test.rs"],
     test_suites: ["general-tests"],
 }
 
diff --git a/libs/dice/open_dice/src/dice.rs b/libs/dice/open_dice/src/dice.rs
index e42e373..e330e00 100644
--- a/libs/dice/open_dice/src/dice.rs
+++ b/libs/dice/open_dice/src/dice.rs
@@ -23,6 +23,8 @@
     DICE_INLINE_CONFIG_SIZE, DICE_PRIVATE_KEY_SEED_SIZE, DICE_PRIVATE_KEY_SIZE,
     DICE_PUBLIC_KEY_SIZE, DICE_SIGNATURE_SIZE,
 };
+#[cfg(feature = "serde_derive")]
+use serde_derive::{Deserialize, Serialize};
 use std::{marker::PhantomData, ptr};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
@@ -82,6 +84,7 @@
 /// for sensitive data like CDI values and private key.
 /// CDI Values.
 #[derive(Debug, Zeroize, ZeroizeOnDrop, Default)]
+#[cfg_attr(feature = "serde_derive", derive(Serialize, Deserialize))]
 pub struct CdiValues {
     /// Attestation CDI.
     pub cdi_attest: [u8; CDI_SIZE],
diff --git a/libs/dice/open_dice/src/retry.rs b/libs/dice/open_dice/src/retry.rs
index a6303bd..d9551f3 100644
--- a/libs/dice/open_dice/src/retry.rs
+++ b/libs/dice/open_dice/src/retry.rs
@@ -25,12 +25,15 @@
 use crate::ops::generate_certificate;
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
+#[cfg(feature = "serde_derive")]
+use serde_derive::{Deserialize, Serialize};
 
 /// Artifacts stores a set of dice artifacts comprising CDI_ATTEST, CDI_SEAL,
 /// and the BCC formatted attestation certificate chain.
 /// As we align with the DICE standards today, this is the certificate chain
 /// is also called DICE certificate chain.
 #[derive(Debug)]
+#[cfg_attr(feature = "serde_derive", derive(Serialize, Deserialize))]
 pub struct OwnedDiceArtifacts {
     /// CDI Values.
     cdi_values: CdiValues,
diff --git a/libs/dice/sample_inputs/Android.bp b/libs/dice/sample_inputs/Android.bp
index 013038c..36fe8c7 100644
--- a/libs/dice/sample_inputs/Android.bp
+++ b/libs/dice/sample_inputs/Android.bp
@@ -54,7 +54,7 @@
 rust_defaults {
     name: "libdiced_sample_inputs_test_defaults",
     crate_name: "diced_sample_inputs_test",
-    srcs: ["tests/*.rs"],
+    srcs: ["tests/api_test.rs"],
     test_suites: ["general-tests"],
     rustlibs: [
         "libanyhow",
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/libs/libfdt/Android.bp b/libs/libfdt/Android.bp
index 1bb5692..7a7d71f 100644
--- a/libs/libfdt/Android.bp
+++ b/libs/libfdt/Android.bp
@@ -53,7 +53,7 @@
     name: "liblibfdt.integration_test",
     crate_name: "libfdt_test",
     defaults: ["avf_build_flags_rust"],
-    srcs: ["tests/*.rs"],
+    srcs: ["tests/api_test.rs"],
     test_suites: ["general-tests"],
     data: [
         ":fdt_test_tree_one_memory_range_dtb",
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 999dc52..e19a343 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -49,10 +49,12 @@
     module_type: "android_system_image",
     config_namespace: "ANDROID",
     bool_variables: [
+        "release_avf_enable_dice_changes",
         "release_avf_enable_multi_tenant_microdroid_vm",
     ],
     properties: [
         "deps",
+        "dirs",
     ],
 }
 
@@ -154,6 +156,14 @@
 
     // Below are dependencies that are conditionally enabled depending on value of build flags.
     soong_config_variables: {
+        release_avf_enable_dice_changes: {
+            deps: [
+                "derive_microdroid_vendor_dice_node",
+            ],
+            dirs: [
+                "microdroid_resources",
+            ],
+        },
         release_avf_enable_multi_tenant_microdroid_vm: {
             deps: [
                 "microdroid_etc_passwd",
@@ -295,7 +305,19 @@
     },
 }
 
-android_filesystem {
+soong_config_module_type {
+    name: "flag_aware_microdroid_filesystem",
+    module_type: "android_filesystem",
+    config_namespace: "ANDROID",
+    bool_variables: [
+        "release_avf_enable_dice_changes",
+    ],
+    properties: [
+        "dirs",
+    ],
+}
+
+flag_aware_microdroid_filesystem {
     name: "microdroid_ramdisk",
     deps: [
         "init_first_stage.microdroid",
@@ -305,14 +327,23 @@
         "proc",
         "sys",
 
-        // TODO(jiyong): remove these
         "mnt",
         "debug_ramdisk",
         "second_stage_resources",
     ],
     type: "compressed_cpio",
+
+    // Below are dependencies that are conditionally enabled depending on value of build flags.
+    soong_config_variables: {
+        release_avf_enable_dice_changes: {
+            dirs: [
+                "microdroid_resources",
+            ],
+        },
+    },
 }
 
+// TODO(ioffe): rename to microdroid_first_stage_ramdisk
 android_filesystem {
     name: "microdroid_fstab_ramdisk",
     deps: [
diff --git a/microdroid/derive_microdroid_vendor_dice_node/Android.bp b/microdroid/derive_microdroid_vendor_dice_node/Android.bp
new file mode 100644
index 0000000..de1bef7
--- /dev/null
+++ b/microdroid/derive_microdroid_vendor_dice_node/Android.bp
@@ -0,0 +1,22 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_defaults {
+    name: "derive_microdroid_vendor_dice_node_defaults",
+    crate_name: "derive_microdroid_vendor_dice_node",
+    defaults: ["avf_build_flags_rust"],
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "libanyhow",
+        "libclap",
+    ],
+    bootstrap: true,
+    prefer_rlib: true,
+}
+
+rust_binary {
+    name: "derive_microdroid_vendor_dice_node",
+    defaults: ["derive_microdroid_vendor_dice_node_defaults"],
+    stem: "derive_microdroid_vendor_dice_node",
+}
diff --git a/microdroid/derive_microdroid_vendor_dice_node/src/main.rs b/microdroid/derive_microdroid_vendor_dice_node/src/main.rs
new file mode 100644
index 0000000..1d5db0d
--- /dev/null
+++ b/microdroid/derive_microdroid_vendor_dice_node/src/main.rs
@@ -0,0 +1,38 @@
+// 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.
+
+//! Derives microdroid vendor dice node.
+
+use anyhow::Error;
+use clap::Parser;
+use std::path::PathBuf;
+
+#[derive(Parser)]
+struct Args {
+    /// Path to the dice driver (e.g. /dev/open-dice0)
+    #[arg(long)]
+    dice_driver: PathBuf,
+    /// Path to the microdroid-vendor.img disk image.
+    #[arg(long)]
+    microdroid_vendor_disk_image: PathBuf,
+    /// File to save resulting dice chain to.
+    #[arg(long)]
+    output: PathBuf,
+}
+
+fn main() -> Result<(), Error> {
+    let args = Args::parse();
+    eprintln!("{:?} {:?} {:?}", args.dice_driver, args.microdroid_vendor_disk_image, args.output);
+    Ok(())
+}
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/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 81bb409..9c9a3d0 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -29,6 +29,7 @@
         "libclient_vm_csr",
         "libciborium",
         "libcoset",
+        "libdice_driver",
         "libdice_policy_builder",
         "libdiced_open_dice",
         "libdiced_sample_inputs",
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index 2469325..7f65159 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -12,13 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::dice_driver::DiceDriver;
 use crate::instance::{ApexData, ApkData};
-use crate::{is_debuggable, is_strict_boot, MicrodroidData};
+use crate::{is_debuggable, MicrodroidData};
 use anyhow::{bail, Context, Result};
 use ciborium::{cbor, Value};
 use coset::CborSerializable;
-use diced_open_dice::{Hidden, OwnedDiceArtifacts, HIDDEN_SIZE};
+use dice_driver::DiceDriver;
+use diced_open_dice::OwnedDiceArtifacts;
 use microdroid_metadata::PayloadMetadata;
 use openssl::sha::{sha512, Sha512};
 use std::iter::once;
@@ -53,37 +53,10 @@
     let debuggable = is_debuggable()?;
 
     // Send the details to diced
-    let hidden = if cfg!(llpvm_changes) {
-        hidden_input_from_instance_id()?
-    } else {
-        instance_data.salt.clone().try_into().unwrap()
-    };
+    let hidden = instance_data.salt.clone().try_into().unwrap();
     dice.derive(code_hash, &config_descriptor, authority_hash, debuggable, hidden)
 }
 
-// Get the "Hidden input" for DICE derivation.
-// This provides differentiation of secrets for different VM instances with same payload.
-fn hidden_input_from_instance_id() -> Result<Hidden> {
-    // For protected VM: this is all 0s, pvmfw ensures differentiation is added early in secrets.
-    // For non-protected VM: this is derived from instance_id of the VM instance.
-    let hidden_input = if !is_strict_boot() {
-        if let Some(id) = super::get_instance_id()? {
-            sha512(&id)
-        } else {
-            // TODO(b/325094712): Absence of instance_id occurs due to missing DT in some
-            // x86_64 test devices (such as Cuttlefish). From security perspective, this is
-            // acceptable for non-protected VM.
-            log::warn!(
-                "Instance Id missing, this may lead to 2 non protected VMs having same secrets"
-            );
-            [0u8; HIDDEN_SIZE]
-        }
-    } else {
-        [0u8; HIDDEN_SIZE]
-    };
-    Ok(hidden_input)
-}
-
 struct Subcomponent {
     name: String,
     version: u64,
diff --git a/microdroid_manager/src/instance.rs b/microdroid_manager/src/instance.rs
index f42b86d..888c451 100644
--- a/microdroid_manager/src/instance.rs
+++ b/microdroid_manager/src/instance.rs
@@ -33,11 +33,11 @@
 //! The payload of a partition is encrypted/signed by a key that is unique to the loader and to the
 //! VM as well. Failing to decrypt/authenticate a partition by a loader stops the boot process.
 
-use crate::dice_driver::DiceDriver;
 use crate::ioutil;
 
 use anyhow::{anyhow, bail, Context, Result};
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use dice_driver::DiceDriver;
 use openssl::symm::{decrypt_aead, encrypt_aead, Cipher};
 use serde::{Deserialize, Serialize};
 use std::fs::{File, OpenOptions};
@@ -273,8 +273,6 @@
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct MicrodroidData {
-    // `salt` is obsolete, it was used as a differentiator for non-protected VM instances running
-    // same payload. Instance-id (present in DT) is used for that now.
     pub salt: Vec<u8>, // Should be [u8; 64] but that isn't serializable.
     pub apk_data: ApkData,
     pub extra_apks_data: Vec<ApkData>,
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index e8017e8..8d2c629 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -15,7 +15,6 @@
 //! Microdroid Manager
 
 mod dice;
-mod dice_driver;
 mod instance;
 mod ioutil;
 mod payload;
@@ -33,16 +32,16 @@
 };
 
 use crate::dice::dice_derivation;
-use crate::dice_driver::DiceDriver;
 use crate::instance::{InstanceDisk, MicrodroidData};
 use crate::verify::verify_payload;
 use crate::vm_payload_service::register_vm_payload_service;
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use binder::Strong;
+use dice_driver::DiceDriver;
 use keystore2_crypto::ZVec;
 use libc::VMADDR_CID_HOST;
 use log::{error, info};
-use microdroid_metadata::{Metadata, PayloadMetadata};
+use microdroid_metadata::PayloadMetadata;
 use microdroid_payload_config::{ApkConfig, OsConfig, Task, TaskType, VmPayloadConfig};
 use nix::sys::signal::Signal;
 use payload::load_metadata;
@@ -236,12 +235,17 @@
     }
 }
 
-fn verify_payload_with_instance_img(
-    metadata: &Metadata,
-    dice: &DiceDriver,
-) -> Result<MicrodroidData> {
+fn try_run_payload(
+    service: &Strong<dyn IVirtualMachineService>,
+    vm_payload_service_fd: OwnedFd,
+) -> Result<i32> {
+    let metadata = load_metadata().context("Failed to load payload metadata")?;
+    let dice = DiceDriver::new(Path::new("/dev/open-dice0"), is_strict_boot())
+        .context("Failed to load DICE")?;
+
     let mut instance = InstanceDisk::new().context("Failed to load instance.img")?;
-    let saved_data = instance.read_microdroid_data(dice).context("Failed to read identity data")?;
+    let saved_data =
+        instance.read_microdroid_data(&dice).context("Failed to read identity data")?;
 
     if is_strict_boot() {
         // Provisioning must happen on the first boot and never again.
@@ -261,7 +265,7 @@
     }
 
     // Verify the payload before using it.
-    let extracted_data = verify_payload(metadata, saved_data.as_ref())
+    let extracted_data = verify_payload(&metadata, saved_data.as_ref())
         .context("Payload verification failed")
         .map_err(|e| MicrodroidError::PayloadVerificationFailed(e.to_string()))?;
 
@@ -285,28 +289,10 @@
     } else {
         info!("Saving verified data.");
         instance
-            .write_microdroid_data(&extracted_data, dice)
+            .write_microdroid_data(&extracted_data, &dice)
             .context("Failed to write identity data")?;
         extracted_data
     };
-    Ok(instance_data)
-}
-
-fn try_run_payload(
-    service: &Strong<dyn IVirtualMachineService>,
-    vm_payload_service_fd: OwnedFd,
-) -> Result<i32> {
-    let metadata = load_metadata().context("Failed to load payload metadata")?;
-    let dice = DiceDriver::new(Path::new("/dev/open-dice0")).context("Failed to load DICE")?;
-
-    // TODO(b/291306122): Checking with host about Secretkeeper support multiple times introduces
-    // a whole range of security vulnerability since host can give different answers. Guest should
-    // check only once and the same answer should be known to pVM Firmware and Microdroid.
-    let instance_data = if let Some(_sk) = vm_secret::is_sk_supported(service)? {
-        verify_payload(&metadata, None)?
-    } else {
-        verify_payload_with_instance_img(&metadata, &dice)?
-    };
 
     let payload_metadata = metadata.payload.ok_or_else(|| {
         MicrodroidError::PayloadInvalidConfig("No payload config in metadata".to_string())
diff --git a/microdroid_manager/src/verify.rs b/microdroid_manager/src/verify.rs
index 65c32b0..445c1ae 100644
--- a/microdroid_manager/src/verify.rs
+++ b/microdroid_manager/src/verify.rs
@@ -169,14 +169,13 @@
     // verified is consistent with the root hash) or because we have the saved APK data which will
     // be checked as identical to the data we have verified.
 
-    let salt = if cfg!(llpvm_changes) || is_strict_boot() {
-        // Salt is obsolete with llpvm_changes.
-        vec![0u8; 64]
-    } else if let Some(saved_data) = saved_data {
-        // Use the salt from a verified instance.
+    // Use the salt from a verified instance, or generate a salt for a new instance.
+    let salt = if let Some(saved_data) = saved_data {
         saved_data.salt.clone()
+    } else if is_strict_boot() {
+        // No need to add more entropy as a previous stage must have used a new, random salt.
+        vec![0u8; 64]
     } else {
-        // Generate a salt for a new instance.
         let mut salt = vec![0u8; 64];
         salt.as_mut_slice().try_fill(&mut rand::thread_rng())?;
         salt
diff --git a/microdroid_manager/src/vm_secret.rs b/microdroid_manager/src/vm_secret.rs
index 7b65491..5ceedea 100644
--- a/microdroid_manager/src/vm_secret.rs
+++ b/microdroid_manager/src/vm_secret.rs
@@ -279,9 +279,9 @@
     anyhow!("{:?}", err)
 }
 
-/// Get the secretkeeper connection if supported. Host can be consulted whether the device supports
-/// secretkeeper but that should be used with caution for protected VM.
-pub fn is_sk_supported(
+// Get the secretkeeper connection if supported. Host can be consulted whether the device supports
+// secretkeeper but that should be used with caution for protected VM.
+fn is_sk_supported(
     host: &Strong<dyn IVirtualMachineService>,
 ) -> Result<Option<Strong<dyn ISecretkeeper>>> {
     let sk = if cfg!(llpvm_changes) {
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 6a6d199..4ee02c1 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -74,15 +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",
@@ -91,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: {
@@ -135,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"],
@@ -142,6 +160,13 @@
 }
 
 genrule {
+    name: "test_pvmfw_devices_overlapping_pvmfw",
+    defaults: ["test_device_assignment_dts_to_dtb"],
+    srcs: ["testdata/test_pvmfw_devices_overlapping_pvmfw.dts"],
+    out: ["test_pvmfw_devices_overlapping_pvmfw.dtb"],
+}
+
+genrule {
     name: "test_pvmfw_devices_with_rng",
     defaults: ["test_device_assignment_dts_to_dtb"],
     srcs: ["testdata/test_pvmfw_devices_with_rng.dts"],
@@ -197,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/README.md b/pvmfw/README.md
index 2758a5d..053e4f7 100644
--- a/pvmfw/README.md
+++ b/pvmfw/README.md
@@ -30,11 +30,11 @@
 hypervisor, although trusted, is also validated.
 
 Once it has been determined that the platform can be trusted, pvmfw derives
-unique secrets for the guest through the [_Boot Certificate Chain_][BCC]
-("BCC", see [Open Profile for DICE][open-dice]) that can be used to prove the
-identity of the pVM to local and remote actors. If any operation or check fails,
-or in case of a missing prerequisite, pvmfw will abort the boot process of the
-pVM, effectively preventing non-compliant pVMs and/or guests from running.
+unique secrets for the guest through the [_DICE Chain_][android-dice] (see
+[Open Profile for DICE][open-dice]) that can be used to prove the identity of
+the pVM to local and remote actors. If any operation or check fails, or in case
+of a missing prerequisite, pvmfw will abort the boot process of the pVM,
+effectively preventing non-compliant pVMs and/or guests from running.
 Otherwise, it hands over the pVM to the guest kernel by jumping to its first
 instruction, similarly to a bootloader.
 
@@ -42,7 +42,7 @@
 
 [AVF]: https://source.android.com/docs/core/virtualization
 [why-avf]: https://source.android.com/docs/core/virtualization/whyavf
-[BCC]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/android.md
+[android-dice]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/android.md
 [pKVM]: https://source.android.com/docs/core/virtualization/architecture#hypervisor
 [open-dice]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md
 
@@ -153,7 +153,7 @@
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+
 | (Padding to 8-byte alignment) |
 +===============================+ <-- FIRST
-|       {First blob: BCC}       |
+|   {First blob: DICE chain}    |
 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ <-- FIRST_END
 | (Padding to 8-byte alignment) |
 +===============================+ <-- SECOND
@@ -189,7 +189,7 @@
 The header format itself is agnostic of the internal format of the individual
 blos it refers to. In version 1.0, it describes two blobs:
 
-- entry 0 must point to a valid BCC Handover (see below)
+- entry 0 must point to a valid DICE chain handover (see below)
 - entry 1 may point to a [DTBO] to be applied to the pVM device tree. See
   [debug policy][debug_policy] for an example.
 
@@ -230,39 +230,39 @@
 [secretkeeper_key]: https://android.googlesource.com/platform/system/secretkeeper/+/refs/heads/main/README.md#secretkeeper-public-key
 [vendor_hashtree_digest]: ../microdroid/README.md#verification-of-vendor-image
 
-#### Virtual Platform Boot Certificate Chain Handover
+#### Virtual Platform DICE Chain Handover
 
-The format of the BCC entry mentioned above, compatible with the
-[`BccHandover`][BccHandover] defined by the Open Profile for DICE reference
-implementation, is described by the following [CDDL][CDDL]:
+The format of the DICE chain entry mentioned above, compatible with the
+[`AndroidDiceHandover`][AndroidDiceHandover] defined by the Open Profile for
+DICE reference implementation, is described by the following [CDDL][CDDL]:
 ```
-PvmfwBccHandover = {
+PvmfwDiceHandover = {
   1 : bstr .size 32,     ; CDI_Attest
   2 : bstr .size 32,     ; CDI_Seal
-  3 : Bcc,               ; Certificate chain
+  3 : DiceCertChain,     ; Android DICE chain
 }
 ```
 
 and contains the _Compound Device Identifiers_ ("CDIs"), used to derive the
 next-stage secret, and a certificate chain, intended for pVM attestation. Note
-that it differs from the `BccHandover` defined by the specification in that its
-`Bcc` field is mandatory (while optional in the original).
+that it differs from the `AndroidDiceHandover` defined by the specification in
+that its `DiceCertChain` field is mandatory (while optional in the original).
 
 Devices that fully implement DICE should provide a certificate rooted at the
 Unique Device Secret (UDS) in a boot stage preceding the pvmfw loader (typically
-ABL), in such a way that it would receive a valid `BccHandover`, that can be
-passed to [`BccHandoverMainFlow`][BccHandoverMainFlow] along with the inputs
-described below.
+ABL), in such a way that it would receive a valid `AndroidDiceHandover`, that
+can be passed to [`DiceAndroidHandoverMainFlow`][DiceAndroidHandoverMainFlow] along with
+the inputs described below.
 
 Otherwise, as an intermediate step towards supporting DICE throughout the
-software stack of the device, incomplete implementations may root the BCC at the
-pvmfw loader, using an arbitrary constant as initial CDI. The pvmfw loader can
-easily do so by:
+software stack of the device, incomplete implementations may root the DICE chain
+at the pvmfw loader, using an arbitrary constant as initial CDI. The pvmfw
+loader can easily do so by:
 
-1. Building a BCC-less `BccHandover` using CBOR operations
-   ([example][Trusty-BCC]) and containing the constant CDIs
-1. Passing the resulting `BccHandover` to `BccHandoverMainFlow` as described
-   above
+1. Building an "empty" `AndroidDiceHandover` using CBOR operations only
+   containing constant CDIs ([example][Trusty-BCC])
+1. Passing the resulting `AndroidDiceHandover` to `DiceAndroidHandoverMainFlow`
+   as described above
 
 The recommended DICE inputs at this stage are:
 
@@ -278,13 +278,14 @@
   storage and changes during every factory reset) or similar that changes as
   part of the device lifecycle (_e.g._ reset)
 
-The resulting `BccHandover` is then used by pvmfw in a similar way to derive
-another [DICE layer][Layering], passed to the guest through a `/reserved-memory`
-device tree node marked as [`compatible=”google,open-dice”`][dice-dt].
+The resulting `AndroidDiceHandover` is then used by pvmfw in a similar way to
+derive another [DICE layer][Layering], passed to the guest through a
+`/reserved-memory` device tree node marked as
+[`compatible=”google,open-dice”`][dice-dt].
 
 [AVB]: https://source.android.com/docs/security/features/verifiedboot/boot-flow
-[BccHandover]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c/src/android/bcc.c#260
-[BccHandoverMainFlow]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c/src/android/bcc.c#199
+[AndroidDiceHandover]: https://pigweed.googlesource.com/open-dice/+/42ae7760023/src/android.c#212
+[DiceAndroidHandoverMainFlow]: https://pigweed.googlesource.com/open-dice/+/42ae7760023/src/android.c#221
 [CDDL]: https://datatracker.ietf.org/doc/rfc8610
 [dice-mode]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#Mode-Value-Details
 [dice-dt]: https://www.kernel.org/doc/Documentation/devicetree/bindings/reserved-memory/google%2Copen-dice.yaml
@@ -383,7 +384,7 @@
 After verifying the guest kernel, pvmfw boots it using the Linux ABI described
 above. It uses the device tree to pass the following:
 
-- a reserved memory node containing the produced BCC:
+- a reserved memory node containing the produced DICE chain:
 
     ```
     / {
@@ -462,15 +463,15 @@
 above must be replicated to produce a single file containing the pvmfw binary
 and its configuration data.
 
-As a quick prototyping solution, a valid BCC (such as the [bcc.dat] test file)
-can be appended to the `pvmfw.bin` image with `pvmfw-tool`.
+As a quick prototyping solution, a valid DICE chain (such as this [test
+file][bcc.dat]) can be appended to the `pvmfw.bin` image with `pvmfw-tool`.
 
 ```shell
 m pvmfw-tool pvmfw_bin
 PVMFW_BIN=${ANDROID_PRODUCT_OUT}/system/etc/pvmfw.bin
-BCC_DAT=${ANDROID_BUILD_TOP}/packages/modules/Virtualization/tests/pvmfw/assets/bcc.dat
+DICE=${ANDROID_BUILD_TOP}/packages/modules/Virtualization/tests/pvmfw/assets/bcc.dat
 
-pvmfw-tool custom_pvmfw ${PVMFW_BIN} ${BCC_DAT}
+pvmfw-tool custom_pvmfw ${PVMFW_BIN} ${DICE}
 ```
 
 The result can then be pushed to the device. Pointing the system property
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 6101a0c..558152d 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -25,7 +25,7 @@
     name: "libpvmfw_avb.integration_test",
     crate_name: "pvmfw_avb_test",
     defaults: ["avf_build_flags_rust"],
-    srcs: ["tests/*.rs"],
+    srcs: ["tests/api_test.rs"],
     test_suites: ["general-tests"],
     data: [
         ":avb_testkey_rsa2048_pub_bin",
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 9c3e566..885cd22 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -27,9 +27,12 @@
 use core::ffi::CStr;
 use core::iter::Iterator;
 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 {
@@ -55,8 +58,10 @@
     InvalidSymbols,
     /// Malformed <reg>. Can't parse.
     MalformedReg,
-    /// Invalid <reg>. Failed to validate with HVC.
-    InvalidReg,
+    /// Invalid physical <reg> of assigned device.
+    InvalidPhysReg(u64, u64),
+    /// Invalid virtual <reg> of assigned device.
+    InvalidReg(u64, u64),
     /// Invalid <interrupts>
     InvalidInterrupts,
     /// Malformed <iommus>
@@ -73,6 +78,8 @@
     DuplicatedIommuIds,
     /// Duplicated pvIOMMU IDs exist
     DuplicatedPvIommuIds,
+    /// Unsupported path format. Only supports full path.
+    UnsupportedPathFormat,
     /// Unsupported overlay target syntax. Only supports <target-path> with full path.
     UnsupportedOverlayTarget,
     /// Unsupported PhysIommu,
@@ -102,7 +109,12 @@
                 "Invalid property in /__symbols__. Must point to valid assignable device node."
             ),
             Self::MalformedReg => write!(f, "Malformed <reg>. Can't parse"),
-            Self::InvalidReg => write!(f, "Invalid <reg>. Failed to validate with hypervisor"),
+            Self::InvalidReg(addr, size) => {
+                write!(f, "Invalid guest MMIO region (addr: {addr:#x}, size: {size:#x})")
+            }
+            Self::InvalidPhysReg(addr, size) => {
+                write!(f, "Invalid physical MMIO region (addr: {addr:#x}, size: {size:#x})")
+            }
             Self::InvalidInterrupts => write!(f, "Invalid <interrupts>"),
             Self::MalformedIommus => write!(f, "Malformed <iommus>. Can't parse."),
             Self::InvalidIommus => {
@@ -120,6 +132,9 @@
             Self::DuplicatedPvIommuIds => {
                 write!(f, "Duplicated pvIOMMU IDs exist. IDs must unique among iommu node")
             }
+            Self::UnsupportedPathFormat => {
+                write!(f, "Unsupported UnsupportedPathFormat. Only supports full path")
+            }
             Self::UnsupportedOverlayTarget => {
                 write!(f, "Unsupported overlay target. Only supports 'target-path = \"/\"'")
             }
@@ -140,6 +155,133 @@
 
 pub type Result<T> = core::result::Result<T, DeviceAssignmentError>;
 
+#[derive(Clone, Default, Ord, PartialOrd, Eq, PartialEq)]
+pub struct DtPathTokens<'a> {
+    tokens: Vec<&'a [u8]>,
+}
+
+impl<'a> fmt::Debug for DtPathTokens<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut list = f.debug_list();
+        for token in &self.tokens {
+            let mut bytes = token.to_vec();
+            bytes.push(b'\0');
+            match CString::from_vec_with_nul(bytes) {
+                Ok(string) => list.entry(&string),
+                Err(_) => list.entry(token),
+            };
+        }
+        list.finish()
+    }
+}
+
+impl<'a> DtPathTokens<'a> {
+    fn new(path: &'a CStr) -> Result<Self> {
+        if path.to_bytes().first() != Some(&b'/') {
+            return Err(DeviceAssignmentError::UnsupportedPathFormat);
+        }
+        let tokens: Vec<_> = path
+            .to_bytes()
+            .split(|char| *char == b'/')
+            .filter(|&component| !component.is_empty())
+            .collect();
+        Ok(Self { tokens })
+    }
+
+    fn to_overlay_target_path(&self) -> Result<Self> {
+        if !self.is_overlayable_node() {
+            return Err(DeviceAssignmentError::InvalidDtbo);
+        }
+        Ok(Self { tokens: self.tokens.as_slice()[2..].to_vec() })
+    }
+
+    fn to_cstring(&self) -> CString {
+        if self.tokens.is_empty() {
+            return CString::new(*b"/\0").unwrap();
+        }
+
+        let size = self.tokens.iter().fold(0, |sum, token| sum + token.len() + 1);
+        let mut path = Vec::with_capacity(size + 1);
+        for token in &self.tokens {
+            path.push(b'/');
+            path.extend_from_slice(token);
+        }
+        path.push(b'\0');
+
+        CString::from_vec_with_nul(path).unwrap()
+    }
+
+    fn is_overlayable_node(&self) -> bool {
+        self.tokens.get(1) == Some(&&b"__overlay__"[..])
+    }
+}
+
+#[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);
@@ -179,14 +321,9 @@
     // node and/or properties. The enforcement is for compatibility reason.
     fn locate_overlay_target_path(
         &self,
-        dtbo_node_path: &CStr,
+        dtbo_node_path: &DtPathTokens,
         dtbo_node: &FdtNode,
     ) -> Result<CString> {
-        let dtbo_node_path_bytes = dtbo_node_path.to_bytes();
-        if dtbo_node_path_bytes.first() != Some(&b'/') {
-            return Err(DeviceAssignmentError::UnsupportedOverlayTarget);
-        }
-
         let fragment_node = dtbo_node.supernode_at_depth(1)?;
         let target_path = fragment_node
             .getprop_str(cstr!("target-path"))?
@@ -195,22 +332,8 @@
             return Err(DeviceAssignmentError::UnsupportedOverlayTarget);
         }
 
-        let mut components = dtbo_node_path_bytes
-            .split(|char| *char == b'/')
-            .filter(|&component| !component.is_empty())
-            .skip(1);
-        let overlay_node_name = components.next();
-        if overlay_node_name != Some(b"__overlay__") {
-            return Err(DeviceAssignmentError::InvalidDtbo);
-        }
-        let mut overlaid_path = Vec::with_capacity(dtbo_node_path_bytes.len());
-        for component in components {
-            overlaid_path.push(b'/');
-            overlaid_path.extend_from_slice(component);
-        }
-        overlaid_path.push(b'\0');
-
-        Ok(CString::from_vec_with_nul(overlaid_path).unwrap())
+        let overlaid_path = dtbo_node_path.to_overlay_target_path()?;
+        Ok(overlaid_path.to_cstring())
     }
 
     fn parse_physical_iommus(physical_node: &FdtNode) -> Result<BTreeMap<Phandle, PhysIommu>> {
@@ -281,15 +404,125 @@
         let phys_iommus = Self::parse_physical_iommus(&physical_node)?;
         Self::parse_physical_devices_with_iommus(&physical_node, &phys_iommus)
     }
-}
 
-fn is_overlayable_node(dtbo_path: &CStr) -> bool {
-    dtbo_path
-        .to_bytes()
-        .split(|char| *char == b'/')
-        .filter(|&component| !component.is_empty())
-        .nth(1)
-        .map_or(false, |name| name == b"__overlay__")
+    fn node(&self, path: &DtPathTokens) -> Result<Option<FdtNode>> {
+        let mut node = self.as_ref().root();
+        for token in &path.tokens {
+            let Some(subnode) = node.subnode_with_name_bytes(token)? else {
+                return Ok(None);
+            };
+            node = subnode;
+        }
+        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<()> {
@@ -324,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
@@ -363,6 +628,12 @@
     size: u64,
 }
 
+impl DeviceReg {
+    pub fn overlaps(&self, range: &Range<u64>) -> bool {
+        self.addr < range.end && range.start < self.addr.checked_add(self.size).unwrap()
+    }
+}
+
 impl TryFrom<Reg<u64>> for DeviceReg {
     type Error = DeviceAssignmentError;
 
@@ -458,8 +729,6 @@
 struct AssignedDeviceInfo {
     // Node path of assigned device (e.g. "/rng")
     node_path: CString,
-    // DTBO node path of the assigned device (e.g. "/fragment@rng/__overlay__/rng")
-    dtbo_node_path: CString,
     // <reg> property from the crosvm DT
     reg: Vec<DeviceReg>,
     // <interrupts> property from the crosvm DT
@@ -474,21 +743,35 @@
         physical_device_reg: &[DeviceReg],
         hypervisor: &dyn DeviceAssigningHypervisor,
     ) -> Result<()> {
-        if device_reg.len() != physical_device_reg.len() {
-            return Err(DeviceAssignmentError::InvalidReg);
-        }
+        let mut virt_regs = device_reg.iter();
+        let mut phys_regs = physical_device_reg.iter();
+        // TODO(b/308694211): Move this constant to vmbase::layout once vmbase is std-compatible.
+        const PVMFW_RANGE: Range<u64> = 0x7fc0_0000..0x8000_0000;
         // PV reg and physical reg should have 1:1 match in order.
-        for (reg, phys_reg) in device_reg.iter().zip(physical_device_reg.iter()) {
+        for (reg, phys_reg) in virt_regs.by_ref().zip(phys_regs.by_ref()) {
+            if reg.overlaps(&PVMFW_RANGE) {
+                return Err(DeviceAssignmentError::InvalidReg(reg.addr, reg.size));
+            }
+            // If this call returns successfully, hyp has mapped the MMIO region at `reg`.
             let addr = hypervisor.get_phys_mmio_token(reg.addr, reg.size).map_err(|e| {
-                error!("Failed to validate device <reg>, error={e:?}, reg={reg:x?}");
-                DeviceAssignmentError::InvalidReg
+                error!("Hypervisor error while requesting MMIO token: {e}");
+                DeviceAssignmentError::InvalidReg(reg.addr, reg.size)
             })?;
-            // Only check address because hypervisor guaranatees size match when success.
+            // Only check address because hypervisor guarantees size match when success.
             if phys_reg.addr != addr {
-                error!("Failed to validate device <reg>. No matching phys reg for reg={reg:x?}");
-                return Err(DeviceAssignmentError::InvalidReg);
+                error!("Assigned device {reg:x?} has unexpected physical address");
+                return Err(DeviceAssignmentError::InvalidPhysReg(addr, reg.size));
             }
         }
+
+        if let Some(DeviceReg { addr, size }) = virt_regs.next() {
+            return Err(DeviceAssignmentError::InvalidReg(*addr, *size));
+        }
+
+        if let Some(DeviceReg { addr, size }) = phys_regs.next() {
+            return Err(DeviceAssignmentError::InvalidPhysReg(*addr, *size));
+        }
+
         Ok(())
     }
 
@@ -545,11 +828,11 @@
         // So we need to mark what's matched or not.
         let mut physical_device_iommu = physical_device_iommu.to_vec();
         for (pviommu, vsid) in iommus {
-            let (id, sid) = hypervisor.get_phys_iommu_token(pviommu.id.into(), vsid.0.into())
-            .map_err(|e| {
-                error!("Failed to validate device <iommus>, error={e:?}, pviommu={pviommu:?}, vsid={vsid:?}");
-                DeviceAssignmentError::InvalidIommus
-            })?;
+            let (id, sid) =
+                hypervisor.get_phys_iommu_token(pviommu.id.into(), vsid.0.into()).map_err(|e| {
+                    error!("Hypervisor error while requesting IOMMU token ({pviommu:?}, {vsid:?}): {e}");
+                    DeviceAssignmentError::InvalidIommus
+                })?;
 
             let pos = physical_device_iommu
                 .iter()
@@ -568,13 +851,13 @@
     fn parse(
         fdt: &Fdt,
         vm_dtbo: &VmDtbo,
-        dtbo_node_path: &CStr,
+        dtbo_node_path: &DtPathTokens,
         physical_devices: &BTreeMap<Phandle, PhysicalDeviceInfo>,
         pviommus: &BTreeMap<Phandle, PvIommu>,
         hypervisor: &dyn DeviceAssigningHypervisor,
     ) -> Result<Option<Self>> {
         let dtbo_node =
-            vm_dtbo.as_ref().node(dtbo_node_path)?.ok_or(DeviceAssignmentError::InvalidSymbols)?;
+            vm_dtbo.node(dtbo_node_path)?.ok_or(DeviceAssignmentError::InvalidSymbols)?;
         let node_path = vm_dtbo.locate_overlay_target_path(dtbo_node_path, &dtbo_node)?;
 
         let Some(node) = fdt.node(&node_path)? else { return Ok(None) };
@@ -595,7 +878,7 @@
         let iommus = Self::parse_iommus(&node, pviommus)?;
         Self::validate_iommus(&iommus, &physical_device.iommus, hypervisor)?;
 
-        Ok(Some(Self { node_path, dtbo_node_path: dtbo_node_path.into(), reg, interrupts, iommus }))
+        Ok(Some(Self { node_path, reg, interrupts, iommus }))
     }
 
     fn patch(&self, fdt: &mut Fdt, pviommu_phandles: &BTreeMap<PvIommu, Phandle>) -> Result<()> {
@@ -614,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 {
@@ -676,26 +959,26 @@
         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)
                 .or(Err(DeviceAssignmentError::InvalidSymbols))?;
-            if !is_overlayable_node(dtbo_node_path) {
+            let dtbo_node_path = DtPathTokens::new(dtbo_node_path)?;
+            if !dtbo_node_path.is_overlayable_node() {
                 continue;
             }
             let assigned_device = AssignedDeviceInfo::parse(
                 fdt,
                 vm_dtbo,
-                dtbo_node_path,
+                &dtbo_node_path,
                 &physical_devices,
                 &pviommus,
                 hypervisor,
             )?;
             if let Some(assigned_device) = assigned_device {
                 assigned_devices.push(assigned_device);
-            } else {
-                filtered_dtbo_paths.push(dtbo_node_path.into());
+                assigned_device_paths.push(dtbo_node_path);
             }
         }
         if assigned_devices.is_empty() {
@@ -704,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)
     }
 
@@ -768,20 +1048,37 @@
     }
 }
 
+/// Cleans device trees not to contain any pre-populated nodes/props for device assignment.
+pub fn clean(fdt: &mut Fdt) -> Result<()> {
+    let mut compatible = fdt.root_mut().next_compatible(cstr!("pkvm,pviommu"))?;
+    // Filters pre-populated
+    while let Some(filtered_pviommu) = compatible {
+        compatible = filtered_pviommu.delete_and_next_compatible(cstr!("pkvm,pviommu"))?;
+    }
+
+    // Removes any dangling references in __symbols__ (e.g. removed pvIOMMUs)
+    filter_dangling_symbols(fdt)
+}
+
 #[cfg(test)]
 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";
+    const FDT_WITH_DEVICE_OVERLAPPING_PVMFW: &str = "test_pvmfw_devices_overlapping_pvmfw.dtb";
     const FDT_WITH_MULTIPLE_DEVICES_IOMMUS_FILE_PATH: &str =
         "test_pvmfw_devices_with_multiple_devices_iommus.dtb";
     const FDT_WITH_IOMMU_SHARING: &str = "test_pvmfw_devices_with_iommu_sharing.dtb";
@@ -790,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 {
@@ -922,7 +1229,6 @@
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/bus0/backlight").unwrap(),
-            dtbo_node_path: cstr!("/fragment@0/__overlay__/bus0/backlight").into(),
             reg: vec![[0x9, 0xFF].into()],
             interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
             iommus: vec![],
@@ -946,7 +1252,6 @@
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/rng").unwrap(),
-            dtbo_node_path: cstr!("/fragment@0/__overlay__/rng").into(),
             reg: vec![[0x9, 0xFF].into()],
             interrupts: into_fdt_prop(vec![0x0, 0xF, 0x4]),
             iommus: vec![(PvIommu { id: 0x4 }, Vsid(0xFF0))],
@@ -1250,7 +1555,7 @@
         };
         let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
 
-        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg));
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x9, 0xFF)));
     }
 
     #[test]
@@ -1266,7 +1571,7 @@
         };
         let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
 
-        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg));
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidPhysReg(0xF10000, 0x1000)));
     }
 
     #[test]
@@ -1332,4 +1637,127 @@
 
         assert_eq!(device_info, Err(DeviceAssignmentError::InvalidIommus));
     }
+
+    #[test]
+    fn device_info_overlaps_pvmfw() {
+        let mut fdt_data = fs::read(FDT_WITH_DEVICE_OVERLAPPING_PVMFW).unwrap();
+        let mut vm_dtbo_data = fs::read(VM_DTBO_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 hypervisor = MockHypervisor {
+            mmio_tokens: [((0x7fee0000, 0x1000), 0xF00000)].into(),
+            iommu_tokens: [((0xFF, 0xF), (0x40000, 0x4))].into(),
+        };
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x7fee0000, 0x1000)));
+    }
+
+    #[test]
+    fn device_assignment_clean() {
+        let mut platform_dt_data = pvmfw_fdt_template::RAW.to_vec();
+        let platform_dt = Fdt::from_mut_slice(&mut platform_dt_data).unwrap();
+
+        let compatible = platform_dt.root().next_compatible(cstr!("pkvm,pviommu"));
+        assert_ne!(None, compatible.unwrap());
+
+        clean(platform_dt).unwrap();
+
+        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/src/fdt.rs b/pvmfw/src/fdt.rs
index 51ba112..d847ca2 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -15,7 +15,7 @@
 //! High-level FDT functions.
 
 use crate::bootargs::BootArgsIterator;
-use crate::device_assignment::{DeviceAssignmentInfo, VmDtbo};
+use crate::device_assignment::{self, DeviceAssignmentInfo, VmDtbo};
 use crate::helpers::GUEST_PAGE_SIZE;
 use crate::Box;
 use crate::RebootReason;
@@ -1158,6 +1158,11 @@
             error!("Failed to patch device assignment info to DT: {e}");
             RebootReason::InvalidFdt
         })?;
+    } else {
+        device_assignment::clean(fdt).map_err(|e| {
+            error!("Failed to clean pre-polulated DT nodes for device assignment: {e}");
+            RebootReason::InvalidFdt
+        })?;
     }
     patch_untrusted_props(fdt, &info.untrusted_props).map_err(|e| {
         error!("Failed to patch untrusted properties: {e}");
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_crosvm_dt_base.dtsi b/pvmfw/testdata/test_crosvm_dt_base.dtsi
index 0c1a311..10d7e6d 100644
--- a/pvmfw/testdata/test_crosvm_dt_base.dtsi
+++ b/pvmfw/testdata/test_crosvm_dt_base.dtsi
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 // This is generated manually by removing unassigned pvIOMMU nodes
 // from patched platform.dts.
diff --git a/pvmfw/testdata/test_pvmfw_devices_overlapping_pvmfw.dts b/pvmfw/testdata/test_pvmfw_devices_overlapping_pvmfw.dts
new file mode 100644
index 0000000..2743dd8
--- /dev/null
+++ b/pvmfw/testdata/test_pvmfw_devices_overlapping_pvmfw.dts
@@ -0,0 +1,16 @@
+/dts-v1/;
+
+/include/ "test_crosvm_dt_base.dtsi"
+
+/ {
+    light {
+        reg = <0x0 0x7fee0000 0x0 0x1000>;
+        iommus = <&pviommu_0 0xF>;
+    };
+
+    pviommu_0: pviommu0 {
+        compatible = "pkvm,pviommu";
+        id = <0xFF>;
+        #iommu-cells = <1>;
+    };
+};
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_vm_dtbo_without_symbols.dts b/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_without_symbols.dts
index 20a8c1b..495a0eb 100644
--- a/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_without_symbols.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_vm_dtbo_without_symbols.dts
@@ -1,4 +1,5 @@
 /dts-v1/;
+// /plugin/ omitted as this DTBO has been written by hand as a DTB in this DTS.
 
 / {
     host {
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_duplicated_pviommus.dts b/pvmfw/testdata/test_pvmfw_devices_with_duplicated_pviommus.dts
index 4ebf034..5646c7f 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_duplicated_pviommus.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_duplicated_pviommus.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts b/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
index a9e30be..04052fa 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_iommu_id_conflict.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts b/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
index 2470725..32e5610 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_iommu_sharing.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
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/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts b/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
index 3aaafdd..3698c1d 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_multiple_devices_iommus.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 / {
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_multiple_reg_iommus.dts b/pvmfw/testdata/test_pvmfw_devices_with_multiple_reg_iommus.dts
index 0676aa3..94fe18e 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_multiple_reg_iommus.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_multiple_reg_iommus.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_with_rng.dts b/pvmfw/testdata/test_pvmfw_devices_with_rng.dts
index a987098..429771c 100644
--- a/pvmfw/testdata/test_pvmfw_devices_with_rng.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_with_rng.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_without_device.dts b/pvmfw/testdata/test_pvmfw_devices_without_device.dts
index ee0be3a..05daaa7 100644
--- a/pvmfw/testdata/test_pvmfw_devices_without_device.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_without_device.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/pvmfw/testdata/test_pvmfw_devices_without_iommus.dts b/pvmfw/testdata/test_pvmfw_devices_without_iommus.dts
index 1a12c87..96fb073 100644
--- a/pvmfw/testdata/test_pvmfw_devices_without_iommus.dts
+++ b/pvmfw/testdata/test_pvmfw_devices_without_iommus.dts
@@ -1,5 +1,4 @@
 /dts-v1/;
-/plugin/;
 
 /include/ "test_crosvm_dt_base.dtsi"
 
diff --git a/service_vm/comm/Android.bp b/service_vm/comm/Android.bp
index bf923a4..15b4ef8 100644
--- a/service_vm/comm/Android.bp
+++ b/service_vm/comm/Android.bp
@@ -50,7 +50,7 @@
 rust_defaults {
     name: "libservice_vm_comm_test_defaults",
     crate_name: "diced_open_dice_test",
-    srcs: ["tests/*.rs"],
+    srcs: ["tests/api_test.rs"],
     test_suites: ["general-tests"],
     prefer_rlib: true,
     rustlibs: [
diff --git a/service_vm/requests/src/rkp.rs b/service_vm/requests/src/rkp.rs
index 569ab01..08ee08e 100644
--- a/service_vm/requests/src/rkp.rs
+++ b/service_vm/requests/src/rkp.rs
@@ -21,7 +21,10 @@
 use alloc::vec;
 use alloc::vec::Vec;
 use bssl_avf::EcKey;
-use ciborium::{cbor, value::Value};
+use ciborium::{
+    cbor,
+    value::{CanonicalValue, Value},
+};
 use core::result;
 use coset::{iana, AsCborValue, CoseSign1, CoseSign1Builder, HeaderBuilder};
 use diced_open_dice::{derive_cdi_leaf_priv, kdf, sign, DiceArtifacts, PrivateKey};
@@ -106,18 +109,24 @@
 
 /// Generates the device info required by the RKP server as a temporary placeholder.
 /// More details in b/301592917.
-fn device_info() -> Value {
-    cbor!({"brand" => "aosp-avf",
-    "manufacturer" => "aosp-avf",
-    "product" => "avf",
-    "model" => "avf",
-    "device" => "avf",
-    "vbmeta_digest" => Value::Bytes(vec![0u8; 0]),
-    "system_patch_level" => 202402,
-    "boot_patch_level" => 20240202,
-    "vendor_patch_level" => 20240202,
-    "fused" => 1})
+///
+/// The keys of the map should be in the length-first core deterministic encoding order
+/// as per RFC8949.
+fn device_info() -> CanonicalValue {
+    cbor!({
+        "brand" => "aosp-avf",
+        "fused" => 1,
+        "model" => "avf",
+        "device" => "avf",
+        "product" => "avf",
+        "manufacturer" => "aosp-avf",
+        "vbmeta_digest" => Value::Bytes(vec![0u8; 0]),
+        "boot_patch_level" => 20240202,
+        "system_patch_level" => 202402,
+        "vendor_patch_level" => 20240202,
+    })
     .unwrap()
+    .into()
 }
 
 fn derive_hmac_key(dice_artifacts: &dyn DiceArtifacts) -> Result<Zeroizing<[u8; HMAC_KEY_LENGTH]>> {
@@ -153,3 +162,25 @@
         })?
         .to_vec())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The keys of device info map should be in the length-first core deterministic encoding
+    /// order as per RFC8949.
+    /// The CBOR ordering rules are:
+    /// 1. If two keys have different lengths, the shorter one sorts earlier;
+    /// 2. If two keys have the same length, the one with the lower value in
+    ///  (bytewise) lexical order sorts earlier.
+    #[test]
+    fn device_info_is_in_length_first_deterministic_order() {
+        let device_info = cbor!(device_info()).unwrap();
+        let device_info_map = device_info.as_map().unwrap();
+        let device_info_keys: Vec<&str> =
+            device_info_map.iter().map(|k| k.0.as_text().unwrap()).collect();
+        let mut sorted_keys = device_info_keys.clone();
+        sorted_keys.sort_by(|a, b| a.len().cmp(&b.len()).then(a.cmp(b)));
+        assert_eq!(device_info_keys, sorted_keys);
+    }
+}
diff --git a/service_vm/test_apk/AndroidTest.rkpd.xml b/service_vm/test_apk/AndroidTest.rkpd.xml
index 2c0380c..39eca32 100644
--- a/service_vm/test_apk/AndroidTest.rkpd.xml
+++ b/service_vm/test_apk/AndroidTest.rkpd.xml
@@ -28,4 +28,11 @@
     <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/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index 2abf110..14cc0ae 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -31,12 +31,13 @@
 import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
 
 import org.json.JSONArray;
 
 import java.io.File;
-import java.io.FileNotFoundException;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -143,15 +144,29 @@
     }
 
     public File findTestFile(String name) {
-        return findTestFile(getBuild(), name);
-    }
+        String moduleName = getInvocationContext().getConfigurationDescriptor().getModuleName();
+        IBuildInfo buildInfo = getBuild();
+        CompatibilityBuildHelper helper = new CompatibilityBuildHelper(buildInfo);
 
-    private static File findTestFile(IBuildInfo buildInfo, String name) {
+        // We're not using helper.getTestFile here because it sometimes picks a file
+        // from a different module, which may be old and/or wrong. See b/328779049.
         try {
-            return (new CompatibilityBuildHelper(buildInfo)).getTestFile(name);
-        } catch (FileNotFoundException e) {
-            throw new AssertionError("Missing test file: " + name, e);
+            File testsDir = helper.getTestsDir().getAbsoluteFile();
+
+            for (File subDir : FileUtil.findDirsUnder(testsDir, testsDir.getParentFile())) {
+                if (!subDir.getName().equals(moduleName)) {
+                    continue;
+                }
+                File testFile = FileUtil.findFile(subDir, name);
+                if (testFile != null) {
+                    return testFile;
+                }
+            }
+        } catch (IOException e) {
+            throw new AssertionError(
+                    "Failed to find test file " + name + " for module " + moduleName, e);
         }
+        throw new AssertionError("Failed to find test file " + name + " for module " + moduleName);
     }
 
     public String getPathForPackage(String packageName)
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 4f502ab..6dd3afe 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -793,12 +793,11 @@
         assertWithMessage("Incorrect ABI list").that(abis).hasLength(1);
 
         // Check that no denials have happened so far
+        String logText =
+                getDevice().pullFileContents(CONSOLE_PATH) + getDevice().pullFileContents(LOG_PATH);
         assertWithMessage("Unexpected denials during VM boot")
-                .that(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH))
-                .isNull();
-        assertWithMessage("Unexpected denials during VM boot")
-                .that(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", CONSOLE_PATH))
-                .isNull();
+                .that(logText)
+                .doesNotContainMatch("avc:\s+denied");
 
         assertThat(getDeviceNumCpus(microdroid)).isEqualTo(getDeviceNumCpus(android));
 
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..d9d425a
--- /dev/null
+++ b/tests/pvmfw/java/com/android/pvmfw/test/CustomPvmfwHostTestCaseBase.java
@@ -0,0 +1,165 @@
+/*
+ * 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 org.junit.Assume.assumeTrue;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.util.FileUtil;
+
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.File;
+import java.util.Objects;
+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 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";
+
+    @Nullable private static File mPvmfwBinFileOnHost;
+    @Nullable private static File mBccFileOnHost;
+
+    @Nullable private TestDevice mAndroidDevice;
+    @Nullable private ITestDevice mMicrodroidDevice;
+
+    @Nullable private File mCustomPvmfwFileOnHost;
+
+    @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 artifacts 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().
+        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 {
+        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(mCustomPvmfwFileOnHost);
+
+        cleanUpVirtualizationTestSetup(mAndroidDevice);
+    }
+
+    /** Returns pvmfw.bin file on host for building custom pvmfw with */
+    public File getPvmfwBinFile() {
+        return mPvmfwBinFileOnHost;
+    }
+
+    /** Returns BCC file on host for building custom pvmfw with */
+    public File getBccFile() {
+        return mBccFileOnHost;
+    }
+
+    /**
+     * 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.
+     */
+    public File getCustomPvmfwFile() {
+        return mCustomPvmfwFileOnHost;
+    }
+
+    /**
+     * 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;
+    }
+}
diff --git a/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java b/tests/pvmfw/java/com/android/pvmfw/test/DebugPolicyHostTests.java
index 26f5993..803405d 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 {
@@ -261,10 +193,10 @@
                         .getDependencyFile(debugPolicyFileName, /* targetFirst= */ false);
 
         Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost)
+                new Pvmfw.Builder(getPvmfwBinFile(), getBccFile())
                         .setDebugPolicyOverlay(mCustomDebugPolicyFileOnHost)
                         .build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        pvmfw.serialize(getCustomPvmfwFile());
     }
 
     private boolean hasConsoleOutput(@NonNull CommandResult result)
@@ -274,29 +206,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 +229,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 +252,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..b68316d 100644
--- a/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
+++ b/tests/pvmfw/java/com/android/pvmfw/test/PvmfwImgTest.java
@@ -16,124 +16,36 @@
 
 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);
+        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 0).build();
+        pvmfw.serialize(getCustomPvmfwFile());
 
         launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
     }
 
     @Test
     public void testConfigVersion1_1_boots() throws Exception {
-        Pvmfw pvmfw =
-                new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost).setVersion(1, 1).build();
-        pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+        Pvmfw pvmfw = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile()).setVersion(1, 1).build();
+        pvmfw.serialize(getCustomPvmfwFile());
 
         launchProtectedVmAndWaitForBootCompleted(BOOT_COMPLETE_TIMEOUT_MS);
     }
@@ -153,7 +65,7 @@
                         new int[] {0xFFFF, 1},
                         new int[] {0xFFFF, 0xFFFF});
 
-        Pvmfw.Builder builder = new Pvmfw.Builder(mPvmfwBinFileOnHost, mBccFileOnHost);
+        Pvmfw.Builder builder = new Pvmfw.Builder(getPvmfwBinFile(), getBccFile());
 
         for (int[] pair : invalid_versions) {
             int major = pair[0];
@@ -161,7 +73,7 @@
             String version = "v" + major + "." + minor;
 
             Pvmfw pvmfw = builder.setVersion(major, minor).build();
-            pvmfw.serialize(mCustomPvmfwBinFileOnHost);
+            pvmfw.serialize(getCustomPvmfwFile());
 
             assertThrows(
                     "pvmfw shouldn't boot with invalid version " + version,
@@ -170,17 +82,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/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/tests/testapk/src/java/com/android/microdroid/test/MicrodroidCapabilitiesTest.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidCapabilitiesTest.java
index eb23e21..61f4cba 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidCapabilitiesTest.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidCapabilitiesTest.java
@@ -15,11 +15,16 @@
  */
 package com.android.microdroid.test;
 
-import static com.google.common.truth.Truth.assertWithMessage;
+import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
 
+import static com.google.common.truth.Truth.assertWithMessage;
+import static com.google.common.truth.TruthJUnit.assume;
+
+import android.os.SystemProperties;
 import android.system.virtualmachine.VirtualMachineManager;
 
 import com.android.compatibility.common.util.CddTest;
+import com.android.compatibility.common.util.VsrTest;
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
 
 import org.junit.Ignore;
@@ -28,16 +33,16 @@
 import org.junit.runners.JUnit4;
 
 /**
- * Test the advertised AVF capabilities include the ability to start some type of VM.
+ * Test the device's AVF capabilities.
  *
  * <p>Tests in MicrodroidTests run on either protected or non-protected VMs, provided they are
  * supported. If neither is they are all skipped. So we need a separate test (that doesn't call
- * {@link #prepareTestSetup}) to make sure that at least one of these is available.
+ * {@link #prepareTestSetup}) when we need to run on such devices.
  */
 @RunWith(JUnit4.class)
 public class MicrodroidCapabilitiesTest extends MicrodroidDeviceTestBase {
     @Test
-    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    @CddTest(requirements = "9.17/C-1-6")
     @Ignore("b/326092480")
     public void supportForProtectedOrNonProtectedVms() {
         assumeSupportedDevice();
@@ -57,4 +62,16 @@
                 .that(vmCapabilities)
                 .isNotEqualTo(0);
     }
+
+    @Test
+    @VsrTest(requirements = "VSR-7.1-001.005")
+    public void avfIsRequired() {
+        int vendorApiLevel = SystemProperties.getInt("ro.vendor.api_level", 0);
+        assume().withMessage("Requirement doesn't apply due to vendor API level")
+                .that(vendorApiLevel)
+                .isAtLeast(202404);
+        boolean avfSupported =
+                getContext().getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK);
+        assertWithMessage("Device doesn't support AVF").that(avfSupported).isTrue();
+    }
 }
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index ea3a481..22bea58 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -49,7 +49,7 @@
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
         BnVirtualMachineService, IVirtualMachineService,
 };
-use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::ISecretkeeper::{BnSecretkeeper, ISecretkeeper};
+use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::ISecretkeeper::ISecretkeeper;
 use android_hardware_security_secretkeeper::aidl::android::hardware::security::secretkeeper::SecretId::SecretId;
 use android_hardware_security_authgraph::aidl::android::hardware::security::authgraph::{
     Arc::Arc as AuthgraphArc, IAuthGraphKeyExchange::IAuthGraphKeyExchange,
@@ -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)?)
@@ -1506,12 +1502,10 @@
     }
 
     fn getSecretkeeper(&self) -> binder::Result<Option<Strong<dyn ISecretkeeper>>> {
-        let sk = if is_secretkeeper_supported() {
-            Some(binder::wait_for_interface(SECRETKEEPER_IDENTIFIER)?)
-        } else {
-            None
-        };
-        Ok(sk.map(|s| BnSecretkeeper::new_binder(SecretkeeperProxy(s), BinderFeatures::default())))
+        // TODO(b/327526008): Session establishment wth secretkeeper is failing.
+        // Re-enable this when fixed.
+        let _sk_supported = is_secretkeeper_supported();
+        Ok(None)
     }
 
     fn requestAttestation(&self, csr: &[u8], test_mode: bool) -> binder::Result<Vec<Certificate>> {
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]]
         );