Support amending vendor hashtree digest when virtmgr constructs DTBO based on host DT

Bug: 325555638
Test: atest MicrodroidTests#configuringVendorDiskImageRequiresCustomPermission
Test: atest MicrodroidTests#bootsWithVendorPartition
Test: atest MicrodroidTests#bootsWithCustomVendorPartitionForNonPvm
Test: atest MicrodroidTests#bootFailsWithCustomVendorPartitionForPvm
Test: atest MicrodroidTests#creationFailsWithUnsignedVendorPartition
Test: atest virtualizationmanager_device_test

Change-Id: I3a008c4bbb34d7673d843c18c0b479212b925065
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 2c72561..2c92f04 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -21,6 +21,7 @@
 import static com.google.common.truth.TruthJUnit.assume;
 
 import static org.junit.Assume.assumeTrue;
+import static org.junit.Assume.assumeFalse;
 
 import android.app.Instrumentation;
 import android.app.UiAutomation;
@@ -611,4 +612,8 @@
     protected void assumeProtectedVM() {
         assumeTrue("Skip on non-protected VM", mProtectedVm);
     }
+
+    protected void assumeNonProtectedVM() {
+        assumeFalse("Skip on protected VM", mProtectedVm);
+    }
 }
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 0132b0d..3b8b4ac 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -2133,29 +2133,34 @@
         }
     }
 
-    @Test
-    public void configuringVendorDiskImageRequiresCustomPermission() throws Exception {
+    private VirtualMachineConfig buildVmConfigWithVendor(File vendorDiskImage) throws Exception {
         assumeSupportedDevice();
+        // TODO(b/325094712): Boot fails with vendor partition in Cuttlefish.
         assumeFalse(
                 "Cuttlefish doesn't support device tree under /proc/device-tree", isCuttlefish());
-        // TODO(b/317567210): Boots fails with vendor partition in HWASAN enabled microdroid
+        // TODO(b/317567210): Boot fails with vendor partition in HWASAN enabled microdroid
         // after introducing verification based on DT and fstab in microdroid vendor partition.
         assumeFalse(
                 "boot with vendor partition is failing in HWASAN enabled Microdroid.", isHwasan());
         assumeFeatureEnabled(VirtualMachineManager.FEATURE_VENDOR_MODULES);
-        revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-
-        File vendorDiskImage =
-                new File("/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img");
         VirtualMachineConfig config =
                 newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
                         .setVendorDiskImage(vendorDiskImage)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        return config;
+    }
+
+    @Test
+    public void configuringVendorDiskImageRequiresCustomPermission() throws Exception {
+        File vendorDiskImage =
+                new File("/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img");
+        VirtualMachineConfig config = buildVmConfigWithVendor(vendorDiskImage);
+        revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
 
         VirtualMachine vm =
                 forceCreateNewVirtualMachine("test_vendor_image_req_custom_permission", config);
-
         SecurityException e =
                 assertThrows(
                         SecurityException.class, () -> runVmTestService(TAG, vm, (ts, tr) -> {}));
@@ -2166,27 +2171,11 @@
 
     @Test
     public void bootsWithVendorPartition() throws Exception {
-        assumeSupportedDevice();
-        assumeFalse(
-                "Cuttlefish doesn't support device tree under /proc/device-tree", isCuttlefish());
-        // TODO(b/317567210): Boots fails with vendor partition in HWASAN enabled microdroid
-        // after introducing verification based on DT and fstab in microdroid vendor partition.
-        assumeFalse(
-                "Boot with vendor partition is failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeFeatureEnabled(VirtualMachineManager.FEATURE_VENDOR_MODULES);
-
-        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-
         File vendorDiskImage = new File("/vendor/etc/avf/microdroid/microdroid_vendor.img");
         assumeTrue("Microdroid vendor image doesn't exist, skip", vendorDiskImage.exists());
-        VirtualMachineConfig config =
-                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
-                        .setVendorDiskImage(vendorDiskImage)
-                        .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .build();
+        VirtualMachineConfig config = buildVmConfigWithVendor(vendorDiskImage);
 
         VirtualMachine vm = forceCreateNewVirtualMachine("test_boot_with_vendor", config);
-
         TestResults testResults =
                 runVmTestService(
                         TAG,
@@ -2194,40 +2183,57 @@
                         (ts, tr) -> {
                             tr.mMountFlags = ts.getMountFlags("/vendor");
                         });
-
         assertThat(testResults.mException).isNull();
         int expectedFlags = MS_NOATIME | MS_RDONLY;
         assertThat(testResults.mMountFlags & expectedFlags).isEqualTo(expectedFlags);
     }
 
     @Test
-    public void creationFailsWithUnsignedVendorPartition() throws Exception {
-        assumeSupportedDevice();
-        assumeFalse(
-                "Cuttlefish doesn't support device tree under /proc/device-tree", isCuttlefish());
-        // TODO(b/317567210): Boots fails with vendor partition in HWASAN enabled microdroid
-        // after introducing verification based on DT and fstab in microdroid vendor partition.
-        assumeFalse(
-                "boot with vendor partition is failing in HWASAN enabled Microdroid.", isHwasan());
-        assumeFeatureEnabled(VirtualMachineManager.FEATURE_VENDOR_MODULES);
+    public void bootsWithCustomVendorPartitionForNonPvm() throws Exception {
+        assumeNonProtectedVM();
+        File vendorDiskImage =
+                new File("/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img");
+        VirtualMachineConfig config = buildVmConfigWithVendor(vendorDiskImage);
 
-        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        VirtualMachine vm =
+                forceCreateNewVirtualMachine("test_boot_with_custom_vendor_non_pvm", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (ts, tr) -> {
+                            tr.mMountFlags = ts.getMountFlags("/vendor");
+                        });
+        assertThat(testResults.mException).isNull();
+        int expectedFlags = MS_NOATIME | MS_RDONLY;
+        assertThat(testResults.mMountFlags & expectedFlags).isEqualTo(expectedFlags);
+    }
 
-        File unsignedVendorDiskImage =
-                new File(
-                        "/data/local/tmp/cts/microdroid/test_microdroid_vendor_image_unsigned.img");
-        VirtualMachineConfig config =
-                newVmConfigBuilderWithPayloadBinary("MicrodroidTestNativeLib.so")
-                        .setVendorDiskImage(unsignedVendorDiskImage)
-                        .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .build();
+    @Test
+    public void bootFailsWithCustomVendorPartitionForPvm() throws Exception {
+        assumeProtectedVM();
+        File vendorDiskImage =
+                new File("/data/local/tmp/cts/microdroid/test_microdroid_vendor_image.img");
+        VirtualMachineConfig config = buildVmConfigWithVendor(vendorDiskImage);
 
-        BootResult bootResult = tryBootVmWithConfig(config, "test_boot_with_unsigned_vendor");
+        BootResult bootResult = tryBootVmWithConfig(config, "test_boot_with_custom_vendor_pvm");
         assertThat(bootResult.payloadStarted).isFalse();
         assertThat(bootResult.deathReason).isEqualTo(VirtualMachineCallback.STOP_REASON_REBOOT);
     }
 
     @Test
+    public void creationFailsWithUnsignedVendorPartition() throws Exception {
+        File vendorDiskImage =
+                new File(
+                        "/data/local/tmp/cts/microdroid/test_microdroid_vendor_image_unsigned.img");
+        VirtualMachineConfig config = buildVmConfigWithVendor(vendorDiskImage);
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_boot_with_unsigned_vendor", config);
+        assertThrowsVmExceptionContaining(
+                () -> vm.run(), "Failed to extract vendor hashtree digest");
+    }
+
+    @Test
     public void systemPartitionMountFlags() throws Exception {
         assumeSupportedDevice();
 
diff --git a/virtualizationmanager/Android.bp b/virtualizationmanager/Android.bp
index bb6ccb7..d8f8209 100644
--- a/virtualizationmanager/Android.bp
+++ b/virtualizationmanager/Android.bp
@@ -42,6 +42,7 @@
         "libcommand_fds",
         "libdisk",
         "libglob",
+        "libhex",
         "libhypervisor_props",
         "liblazy_static",
         "liblibc",
@@ -62,6 +63,7 @@
         "libshared_child",
         "libstatslog_virtualization_rust",
         "libtombstoned_client_rust",
+        "libvbmeta_rust",
         "libvm_control",
         "libvmconfig",
         "libzip",
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 88700ec..c32960b 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -87,6 +87,7 @@
 use std::os::unix::raw::pid_t;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
+use vbmeta::VbMetaImage;
 use vmconfig::VmConfig;
 use vsock::VsockStream;
 use zip::ZipArchive;
@@ -381,6 +382,23 @@
             None
         };
 
+        let vendor_hashtree_digest = extract_vendor_hashtree_digest(config)
+            .context("Failed to extract vendor hashtree digest")
+            .or_service_specific_exception(-1)?;
+
+        let trusted_props = if let Some(ref vendor_hashtree_digest) = vendor_hashtree_digest {
+            info!(
+                "Passing vendor hashtree digest to pvmfw. This will be rejected if it doesn't \
+                match the trusted digest in the pvmfw config, causing the VM to fail to start."
+            );
+            vec![(
+                cstr!("vendor_hashtree_descriptor_root_digest"),
+                vendor_hashtree_digest.as_slice(),
+            )]
+        } else {
+            vec![]
+        };
+
         let untrusted_props = if cfg!(llpvm_changes) {
             // TODO(b/291213394): Replace this with a per-VM instance Id.
             let instance_id = b"sixtyfourbyteslonghardcoded_indeed_sixtyfourbyteslonghardcoded_h";
@@ -389,17 +407,23 @@
             vec![]
         };
 
-        let device_tree_overlay = if host_ref_dt.is_some() || !untrusted_props.is_empty() {
-            let dt_output = temporary_directory.join(VM_DT_OVERLAY_PATH);
-            let mut data = [0_u8; VM_DT_OVERLAY_MAX_SIZE];
-            let fdt = create_device_tree_overlay(&mut data, host_ref_dt, &untrusted_props)
+        let device_tree_overlay =
+            if host_ref_dt.is_some() || !untrusted_props.is_empty() || !trusted_props.is_empty() {
+                let dt_output = temporary_directory.join(VM_DT_OVERLAY_PATH);
+                let mut data = [0_u8; VM_DT_OVERLAY_MAX_SIZE];
+                let fdt = create_device_tree_overlay(
+                    &mut data,
+                    host_ref_dt,
+                    &untrusted_props,
+                    &trusted_props,
+                )
                 .map_err(|e| anyhow!("Failed to create DT overlay, {e:?}"))
                 .or_service_specific_exception(-1)?;
-            fs::write(&dt_output, fdt.as_slice()).or_service_specific_exception(-1)?;
-            Some(File::open(dt_output).or_service_specific_exception(-1)?)
-        } else {
-            None
-        };
+                fs::write(&dt_output, fdt.as_slice()).or_service_specific_exception(-1)?;
+                Some(File::open(dt_output).or_service_specific_exception(-1)?)
+            } else {
+                None
+            };
 
         let debug_level = match config {
             VirtualMachineConfig::AppConfig(config) => config.debugLevel,
@@ -601,6 +625,32 @@
     }
 }
 
+fn extract_vendor_hashtree_digest(config: &VirtualMachineConfig) -> Result<Option<Vec<u8>>> {
+    let VirtualMachineConfig::AppConfig(config) = config else {
+        return Ok(None);
+    };
+    let Some(custom_config) = &config.customConfig else {
+        return Ok(None);
+    };
+    let Some(file) = custom_config.vendorImage.as_ref() else {
+        return Ok(None);
+    };
+
+    let file = clone_file(file)?;
+    let size =
+        file.metadata().context("Failed to get metadata from microdroid vendor image")?.len();
+    let vbmeta = VbMetaImage::verify_reader_region(&file, 0, size)
+        .context("Failed to get vbmeta from microdroid-vendor.img")?;
+
+    for descriptor in vbmeta.descriptors()?.iter() {
+        if let vbmeta::Descriptor::Hashtree(_) = descriptor {
+            let root_digest = hex::encode(descriptor.to_hashtree()?.root_digest());
+            return Ok(Some(root_digest.as_bytes().to_vec()));
+        }
+    }
+    Err(anyhow!("No hashtree digest is extracted from microdroid vendor image"))
+}
+
 fn write_zero_filler(zero_filler_path: &Path) -> Result<()> {
     let file = OpenOptions::new()
         .create_new(true)
diff --git a/virtualizationmanager/src/dt_overlay.rs b/virtualizationmanager/src/dt_overlay.rs
index 71d3a26..b39ba3a 100644
--- a/virtualizationmanager/src/dt_overlay.rs
+++ b/virtualizationmanager/src/dt_overlay.rs
@@ -31,15 +31,18 @@
 /// * `dt_path` - (Optional) Path to (proc style) device tree to be included in the overlay.
 /// * `untrusted_props` - Include a property in /avf/untrusted node. This node is used to specify
 ///   host provided properties such as `instance-id`.
+/// * `trusted_props` - Include a property in /avf node. This overwrites nodes included with
+///   `dt_path`. In pVM, pvmfw will reject if it doesn't match the value in pvmfw config.
 ///
-/// Example: with `create_device_tree_overlay(_, _, [("instance-id", _),])`
+/// Example: with `create_device_tree_overlay(_, _, [("instance-id", _),], [("digest", _),])`
 /// ```
 ///   {
 ///     fragment@0 {
 ///         target-path = "/";
 ///         __overlay__ {
 ///             avf {
-///                 untrusted { instance-id = [0x01 0x23 .. ] }
+///                 digest = [ 0xaa 0xbb .. ]
+///                 untrusted { instance-id = [ 0x01 0x23 .. ] }
 ///               }
 ///             };
 ///         };
@@ -50,36 +53,54 @@
     buffer: &'a mut [u8],
     dt_path: Option<&'a Path>,
     untrusted_props: &[(&'a CStr, &'a [u8])],
+    trusted_props: &[(&'a CStr, &'a [u8])],
 ) -> Result<&'a mut Fdt> {
-    if dt_path.is_none() && untrusted_props.is_empty() {
+    if dt_path.is_none() && untrusted_props.is_empty() && trusted_props.is_empty() {
         return Err(anyhow!("Expected at least one device tree addition"));
     }
 
     let fdt =
         Fdt::create_empty_tree(buffer).map_err(|e| anyhow!("Failed to create empty Fdt: {e:?}"))?;
-    let root = fdt.root_mut().map_err(|e| anyhow!("Failed to get root: {e:?}"))?;
-    let mut node =
-        root.add_subnode(cstr!("fragment@0")).map_err(|e| anyhow!("Failed to fragment: {e:?}"))?;
-    node.setprop(cstr!("target-path"), b"/\0")
-        .map_err(|e| anyhow!("Failed to set target-path: {e:?}"))?;
-    let node = node
+    let root = fdt.root_mut().map_err(|e| anyhow!("Failed to get root node: {e:?}"))?;
+    let mut fragment = root
+        .add_subnode(cstr!("fragment@0"))
+        .map_err(|e| anyhow!("Failed to add fragment node: {e:?}"))?;
+    fragment
+        .setprop(cstr!("target-path"), b"/\0")
+        .map_err(|e| anyhow!("Failed to set target-path property: {e:?}"))?;
+    let overlay = fragment
         .add_subnode(cstr!("__overlay__"))
-        .map_err(|e| anyhow!("Failed to __overlay__ node: {e:?}"))?;
+        .map_err(|e| anyhow!("Failed to add __overlay__ node: {e:?}"))?;
+    let avf =
+        overlay.add_subnode(AVF_NODE_NAME).map_err(|e| anyhow!("Failed to add avf node: {e:?}"))?;
 
     if !untrusted_props.is_empty() {
-        let mut node = node
-            .add_subnode(AVF_NODE_NAME)
-            .map_err(|e| anyhow!("Failed to add avf node: {e:?}"))?
+        let mut untrusted = avf
             .add_subnode(UNTRUSTED_NODE_NAME)
-            .map_err(|e| anyhow!("Failed to add /avf/untrusted node: {e:?}"))?;
+            .map_err(|e| anyhow!("Failed to add untrusted node: {e:?}"))?;
         for (name, value) in untrusted_props {
-            node.setprop(name, value).map_err(|e| anyhow!("Failed to set property: {e:?}"))?;
+            untrusted
+                .setprop(name, value)
+                .map_err(|e| anyhow!("Failed to set untrusted property: {e:?}"))?;
         }
     }
 
+    // Read dt_path from host DT and overlay onto fdt.
     if let Some(path) = dt_path {
         fdt.overlay_onto(cstr!("/fragment@0/__overlay__"), path)?;
     }
+
+    if !trusted_props.is_empty() {
+        let mut avf = fdt
+            .node_mut(cstr!("/fragment@0/__overlay__/avf"))
+            .map_err(|e| anyhow!("Failed to search avf node: {e:?}"))?
+            .ok_or(anyhow!("Failed to get avf node"))?;
+        for (name, value) in trusted_props {
+            avf.setprop(name, value)
+                .map_err(|e| anyhow!("Failed to set trusted property: {e:?}"))?;
+        }
+    }
+
     fdt.pack().map_err(|e| anyhow!("Failed to pack DT overlay, {e:?}"))?;
 
     Ok(fdt)
@@ -92,7 +113,7 @@
     #[test]
     fn empty_overlays_not_allowed() {
         let mut buffer = vec![0_u8; VM_DT_OVERLAY_MAX_SIZE];
-        let res = create_device_tree_overlay(&mut buffer, None, &[]);
+        let res = create_device_tree_overlay(&mut buffer, None, &[], &[]);
         assert!(res.is_err());
     }
 
@@ -102,7 +123,8 @@
         let prop_name = cstr!("XOXO");
         let prop_val_input = b"OXOX";
         let fdt =
-            create_device_tree_overlay(&mut buffer, None, &[(prop_name, prop_val_input)]).unwrap();
+            create_device_tree_overlay(&mut buffer, None, &[(prop_name, prop_val_input)], &[])
+                .unwrap();
 
         let prop_value_dt = fdt
             .node(cstr!("/fragment@0/__overlay__/avf/untrusted"))
@@ -113,4 +135,23 @@
             .expect("Prop not found!");
         assert_eq!(prop_value_dt, prop_val_input, "Unexpected property value");
     }
+
+    #[test]
+    fn trusted_prop_test() {
+        let mut buffer = vec![0_u8; VM_DT_OVERLAY_MAX_SIZE];
+        let prop_name = cstr!("XOXOXO");
+        let prop_val_input = b"OXOXOX";
+        let fdt =
+            create_device_tree_overlay(&mut buffer, None, &[], &[(prop_name, prop_val_input)])
+                .unwrap();
+
+        let prop_value_dt = fdt
+            .node(cstr!("/fragment@0/__overlay__/avf"))
+            .unwrap()
+            .expect("/avf node doesn't exist")
+            .getprop(prop_name)
+            .unwrap()
+            .expect("Prop not found!");
+        assert_eq!(prop_value_dt, prop_val_input, "Unexpected property value");
+    }
 }