Merge "Set ttyd web page's title as an app name" into main
diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
index 246c28e..af41c85 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/MainActivity.java
@@ -45,6 +45,7 @@
 import android.widget.Toast;
 
 import com.android.virtualization.vmlauncher.InstallUtils;
+import com.android.virtualization.vmlauncher.VmLauncherService;
 import com.android.virtualization.vmlauncher.VmLauncherServices;
 
 import com.google.android.material.appbar.MaterialToolbar;
@@ -209,6 +210,8 @@
                                     public void onComplete(long requestId) {
                                         if (requestId == mRequestId) {
                                             android.os.Trace.endAsyncSection("executeTerminal", 0);
+                                            findViewById(R.id.boot_progress)
+                                                    .setVisibility(View.GONE);
                                             view.setVisibility(View.VISIBLE);
                                         }
                                     }
@@ -430,22 +433,46 @@
         Intent intent = new Intent();
         PendingIntent pendingIntent = PendingIntent.getActivity(this, 0, intent,
                 PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
-
+        Intent stopIntent = new Intent();
+        stopIntent.setClass(this, VmLauncherService.class);
+        stopIntent.setAction(VmLauncherServices.ACTION_STOP_VM_LAUNCHER_SERVICE);
+        PendingIntent stopPendingIntent =
+                PendingIntent.getService(
+                        this,
+                        0,
+                        stopIntent,
+                        PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
         Icon icon = Icon.createWithResource(getResources(), R.drawable.ic_launcher_foreground);
-        Notification notification = new Notification.Builder(this, TAG)
-                .setChannelId(TAG)
-                .setSmallIcon(R.drawable.ic_launcher_foreground)
-                .setContentTitle(getResources().getString(R.string.service_notification_title))
-                .setContentText(getResources().getString(R.string.service_notification_content))
-                .setContentIntent(pendingIntent)
-                .setOngoing(true)
-                .addAction(new Notification.Action.Builder(icon,
-                        getResources().getString(R.string.service_notification_settings),
-                        pendingIntent).build())
-                .addAction(new Notification.Action.Builder(icon,
-                        getResources().getString(R.string.service_notification_quit_action),
-                        pendingIntent).build())
-                .build();
+        Notification notification =
+                new Notification.Builder(this, TAG)
+                        .setChannelId(TAG)
+                        .setSmallIcon(R.drawable.ic_launcher_foreground)
+                        .setContentTitle(
+                                getResources().getString(R.string.service_notification_title))
+                        .setContentText(
+                                getResources().getString(R.string.service_notification_content))
+                        .setContentIntent(pendingIntent)
+                        .setOngoing(true)
+                        .addAction(
+                                new Notification.Action.Builder(
+                                                icon,
+                                                getResources()
+                                                        .getString(
+                                                                R.string
+                                                                        .service_notification_settings),
+                                                pendingIntent)
+                                        .build())
+                        .addAction(
+                                new Notification.Action.Builder(
+                                                icon,
+                                                getResources()
+                                                        .getString(
+                                                                R.string
+                                                                        .service_notification_quit_action),
+                                                stopPendingIntent)
+                                        .build())
+                        .setDeleteIntent(stopPendingIntent)
+                        .build();
 
         android.os.Trace.beginAsyncSection("executeTerminal", 0);
         VmLauncherServices.startVmLauncherService(this, this, notification);
diff --git a/android/TerminalApp/res/layout/activity_headless.xml b/android/TerminalApp/res/layout/activity_headless.xml
index 7baaf5c..7da87af 100644
--- a/android/TerminalApp/res/layout/activity_headless.xml
+++ b/android/TerminalApp/res/layout/activity_headless.xml
@@ -16,6 +16,7 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent">
         <LinearLayout
+            android:id="@+id/boot_progress"
             android:orientation="vertical"
             android:gravity="center"
             android:layout_gravity="center"
diff --git a/android/TerminalApp/res/values/config.xml b/android/TerminalApp/res/values/config.xml
new file mode 100644
index 0000000..055abb7
--- /dev/null
+++ b/android/TerminalApp/res/values/config.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--  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.
+ -->
+
+<resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="preference_file_key" translatable="false">com.android.virtualization.terminal.PREFERENCE_FILE_KEY</string>
+    <string name="preference_disk_size_key" translatable="false">PREFERENCE_DISK_SIZE_KEY</string>
+    <string name="preference_min_disk_size_key" translatable="false">PREFERENCE_MIN_DISK_SIZE_KEY</string>
+</resources>
\ No newline at end of file
diff --git a/android/TerminalApp/res/values/strings.xml b/android/TerminalApp/res/values/strings.xml
index dfe7b95..f8350a0 100644
--- a/android/TerminalApp/res/values/strings.xml
+++ b/android/TerminalApp/res/values/strings.xml
@@ -98,9 +98,4 @@
     <string name="service_notification_content">Click to open the terminal.</string>
     <!-- Notification action button for closing the virtual machine [CHAR LIMIT=none] -->
     <string name="service_notification_quit_action">Close</string>
-
-    <!-- Preference Keys -->
-    <string name="preference_file_key">com.android.virtualization.terminal.PREFERENCE_FILE_KEY</string>
-    <string name="preference_disk_size_key">PREFERENCE_DISK_SIZE_KEY</string>
-    <string name="preference_min_disk_size_key">PREFERENCE_MIN_DISK_SIZE_KEY</string>
 </resources>
diff --git a/build/apex/Android.bp b/build/apex/Android.bp
index e485aa7..b0ecdde 100644
--- a/build/apex/Android.bp
+++ b/build/apex/Android.bp
@@ -47,7 +47,6 @@
         "android.system.virtualmachine.res",
     ] + select(release_flag("RELEASE_AVF_SUPPORT_CUSTOM_VM_WITH_PARAVIRTUALIZED_DEVICES"), {
         true: [
-            "VmLauncherApp",
             "VmTerminalApp",
         ],
         default: [],
diff --git a/guest/pvmfw/src/device_assignment.rs b/guest/pvmfw/src/device_assignment.rs
index 5edfe97..da9462b 100644
--- a/guest/pvmfw/src/device_assignment.rs
+++ b/guest/pvmfw/src/device_assignment.rs
@@ -30,6 +30,7 @@
 use core::ops::Range;
 use libfdt::{Fdt, FdtError, FdtNode, FdtNodeMut, Phandle, Reg};
 use log::error;
+use log::warn;
 // TODO(b/308694211): Use vmbase::hyp::{DeviceAssigningHypervisor, Error} proper for tests.
 #[cfg(not(test))]
 use vmbase::hyp::DeviceAssigningHypervisor;
@@ -60,10 +61,16 @@
     InvalidSymbols,
     /// Malformed <reg>. Can't parse.
     MalformedReg,
-    /// Invalid physical <reg> of assigned device.
-    InvalidPhysReg(u64, u64),
+    /// Missing physical <reg> of assigned device.
+    MissingReg(u64, u64),
+    /// Extra <reg> of assigned device.
+    ExtraReg(u64, u64),
     /// Invalid virtual <reg> of assigned device.
-    InvalidReg(u64, u64),
+    InvalidReg(u64),
+    /// Token for <reg> of assigned device does not match expected value.
+    InvalidRegToken(u64, u64),
+    /// Invalid virtual <reg> size of assigned device.
+    InvalidRegSize(u64, u64),
     /// Invalid <interrupts>
     InvalidInterrupts,
     /// Malformed <iommus>
@@ -111,11 +118,20 @@
                 "Invalid property in /__symbols__. Must point to valid assignable device node."
             ),
             Self::MalformedReg => write!(f, "Malformed <reg>. Can't parse"),
-            Self::InvalidReg(addr, size) => {
-                write!(f, "Invalid guest MMIO region (addr: {addr:#x}, size: {size:#x})")
+            Self::MissingReg(addr, size) => {
+                write!(f, "Missing physical 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::ExtraReg(addr, size) => {
+                write!(f, "Unexpected extra MMIO region: addr:{addr:#x}), size:{size:#x}")
+            }
+            Self::InvalidReg(addr) => {
+                write!(f, "Invalid guest MMIO granule (addr: {addr:#x})")
+            }
+            Self::InvalidRegSize(size, expected) => {
+                write!(f, "Unexpected MMIO size ({size:#x}), should be {expected:#x}")
+            }
+            Self::InvalidRegToken(token, expected) => {
+                write!(f, "Unexpected MMIO token ({token:#x}), should be {expected:#x}")
             }
             Self::InvalidInterrupts => write!(f, "Invalid <interrupts>"),
             Self::MalformedIommus => write!(f, "Malformed <iommus>. Can't parse."),
@@ -634,6 +650,10 @@
     pub fn overlaps(&self, range: &Range<u64>) -> bool {
         self.addr < range.end && range.start < self.addr.checked_add(self.size).unwrap()
     }
+
+    pub fn is_aligned(&self, granule: u64) -> bool {
+        self.addr % granule == 0 && self.size % granule == 0
+    }
 }
 
 impl TryFrom<Reg<u64>> for DeviceReg {
@@ -744,34 +764,43 @@
         device_reg: &[DeviceReg],
         physical_device_reg: &[DeviceReg],
         hypervisor: &dyn DeviceAssigningHypervisor,
+        granule: usize,
     ) -> Result<()> {
         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 virt_regs.by_ref().zip(phys_regs.by_ref()) {
-            if reg.overlaps(&PVMFW_RANGE) {
-                return Err(DeviceAssignmentError::InvalidReg(reg.addr, reg.size));
+            if !reg.is_aligned(granule.try_into().unwrap()) {
+                let DeviceReg { addr, size } = reg;
+                warn!("Assigned region ({addr:#x}, {size:#x}) not aligned to {granule:#x}");
+                // TODO(ptosi): Fix our test data so that we can return Err(...);
             }
+            if reg.overlaps(&PVMFW_RANGE) {
+                return Err(DeviceAssignmentError::InvalidReg(reg.addr));
+            }
+            if reg.size != phys_reg.size {
+                return Err(DeviceAssignmentError::InvalidRegSize(reg.size, phys_reg.size));
+            }
+            let expected_token = phys_reg.addr;
             // 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| {
+            let token = hypervisor.get_phys_mmio_token(reg.addr, reg.size).map_err(|e| {
                 error!("Hypervisor error while requesting MMIO token: {e}");
-                DeviceAssignmentError::InvalidReg(reg.addr, reg.size)
+                DeviceAssignmentError::InvalidReg(reg.addr)
             })?;
-            // Only check address because hypervisor guarantees size match when success.
-            if phys_reg.addr != addr {
-                error!("Assigned device {reg:x?} has unexpected physical address");
-                return Err(DeviceAssignmentError::InvalidPhysReg(addr, reg.size));
+            if token != expected_token {
+                return Err(DeviceAssignmentError::InvalidRegToken(token, expected_token));
             }
         }
 
         if let Some(DeviceReg { addr, size }) = virt_regs.next() {
-            return Err(DeviceAssignmentError::InvalidReg(*addr, *size));
+            return Err(DeviceAssignmentError::ExtraReg(*addr, *size));
         }
 
         if let Some(DeviceReg { addr, size }) = phys_regs.next() {
-            return Err(DeviceAssignmentError::InvalidPhysReg(*addr, *size));
+            return Err(DeviceAssignmentError::MissingReg(*addr, *size));
         }
 
         Ok(())
@@ -857,6 +886,7 @@
         physical_devices: &BTreeMap<Phandle, PhysicalDeviceInfo>,
         pviommus: &BTreeMap<Phandle, PvIommu>,
         hypervisor: &dyn DeviceAssigningHypervisor,
+        granule: usize,
     ) -> Result<Option<Self>> {
         let dtbo_node =
             vm_dtbo.node(dtbo_node_path)?.ok_or(DeviceAssignmentError::InvalidSymbols)?;
@@ -873,7 +903,7 @@
         };
 
         let reg = parse_node_reg(&node)?;
-        Self::validate_reg(&reg, &physical_device.reg, hypervisor)?;
+        Self::validate_reg(&reg, &physical_device.reg, hypervisor, granule)?;
 
         let interrupts = Self::parse_interrupts(&node)?;
 
@@ -945,8 +975,9 @@
         fdt: &Fdt,
         vm_dtbo: &VmDtbo,
         hypervisor: &dyn DeviceAssigningHypervisor,
+        granule: usize,
     ) -> Result<Option<Self>> {
-        Self::internal_parse(fdt, vm_dtbo, hypervisor)
+        Self::internal_parse(fdt, vm_dtbo, hypervisor, granule)
     }
 
     #[cfg(not(test))]
@@ -957,14 +988,16 @@
         fdt: &Fdt,
         vm_dtbo: &VmDtbo,
         hypervisor: &dyn DeviceAssigningHypervisor,
+        granule: usize,
     ) -> Result<Option<Self>> {
-        Self::internal_parse(fdt, vm_dtbo, hypervisor)
+        Self::internal_parse(fdt, vm_dtbo, hypervisor, granule)
     }
 
     fn internal_parse(
         fdt: &Fdt,
         vm_dtbo: &VmDtbo,
         hypervisor: &dyn DeviceAssigningHypervisor,
+        granule: usize,
     ) -> Result<Option<Self>> {
         let Some(symbols_node) = vm_dtbo.as_ref().symbols()? else {
             // /__symbols__ should contain all assignable devices.
@@ -997,6 +1030,7 @@
                 &physical_devices,
                 &pviommus,
                 hypervisor,
+                granule,
             )?;
             if let Some(assigned_device) = assigned_device {
                 assigned_devices.push(assigned_device);
@@ -1153,15 +1187,27 @@
     const EXPECTED_FDT_WITH_DEPENDENCY_LOOP_FILE_PATH: &str =
         "expected_dt_with_dependency_loop.dtb";
 
+    // TODO(b/308694211): Use vmbase::SIZE_4KB.
+    const SIZE_4KB: usize = 4 << 10;
+
     #[derive(Debug, Default)]
     struct MockHypervisor {
         mmio_tokens: BTreeMap<(u64, u64), u64>,
         iommu_tokens: BTreeMap<(u64, u64), (u64, u64)>,
     }
 
+    impl MockHypervisor {
+        // TODO(ptosi): Improve these tests to cover multi-page devices.
+        fn get_mmio_token(&self, addr: u64) -> Option<&u64> {
+            // We currently only have single (or sub-) page MMIO test data so can ignore sizes.
+            let key = self.mmio_tokens.keys().find(|(virt, _)| *virt == addr)?;
+            self.mmio_tokens.get(key)
+        }
+    }
+
     impl DeviceAssigningHypervisor for MockHypervisor {
-        fn get_phys_mmio_token(&self, base_ipa: u64, size: u64) -> MockHypervisorResult<u64> {
-            let token = self.mmio_tokens.get(&(base_ipa, size));
+        fn get_phys_mmio_token(&self, base_ipa: u64, _size: u64) -> MockHypervisorResult<u64> {
+            let token = self.get_mmio_token(base_ipa);
 
             Ok(*token.ok_or(MockHypervisorError::FailedGetPhysMmioToken)?)
         }
@@ -1247,6 +1293,8 @@
         }
     }
 
+    // TODO(ptosi): Add tests with varying HYP_GRANULE values.
+
     #[test]
     fn device_info_new_without_symbols() {
         let mut fdt_data = fs::read(FDT_FILE_PATH).unwrap();
@@ -1255,7 +1303,9 @@
         let vm_dtbo = VmDtbo::from_mut_slice(&mut vm_dtbo_data).unwrap();
 
         let hypervisor: MockHypervisor = Default::default();
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap();
         assert_eq!(device_info, None);
     }
 
@@ -1267,7 +1317,9 @@
         let vm_dtbo = VmDtbo::from_mut_slice(&mut vm_dtbo_data).unwrap();
 
         let hypervisor: MockHypervisor = Default::default();
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap();
         assert_eq!(device_info, None);
     }
 
@@ -1282,7 +1334,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x300)].into(),
             iommu_tokens: BTreeMap::new(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/bus0/backlight").unwrap(),
@@ -1305,7 +1359,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x12F00000)].into(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 0x3))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
 
         let expected = [AssignedDeviceInfo {
             node_path: CString::new("/rng").unwrap(),
@@ -1328,7 +1384,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x12F00000)].into(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 0x3))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         let vm_dtbo = vm_dtbo.as_mut();
@@ -1369,7 +1427,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x300)].into(),
             iommu_tokens: BTreeMap::new(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1420,7 +1480,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x300)].into(),
             iommu_tokens: BTreeMap::new(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1455,7 +1517,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x12F00000)].into(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 0x3))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1503,7 +1567,9 @@
             ]
             .into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1550,7 +1616,9 @@
             mmio_tokens: [((0x9, 0xFF), 0x12F00000), ((0x1000, 0x9), 0x12000000)].into(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 3)), ((0x4, 0xFF1), (0x12E40000, 9))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1594,7 +1662,8 @@
             mmio_tokens: [((0x9, 0xFF), 0x300)].into(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 0x3))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
         assert_eq!(device_info, Err(DeviceAssignmentError::DuplicatedPvIommuIds));
     }
@@ -1610,9 +1679,10 @@
             mmio_tokens: BTreeMap::new(),
             iommu_tokens: [((0x4, 0xFF0), (0x12E40000, 0x3))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
-        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x9, 0xFF)));
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x9)));
     }
 
     #[test]
@@ -1626,9 +1696,10 @@
             mmio_tokens: [((0xF000, 0x1000), 0xF10000), ((0xF100, 0x1000), 0xF00000)].into(),
             iommu_tokens: [((0xFF0, 0xF0), (0x40000, 0x4)), ((0xFF1, 0xF1), (0x50000, 0x5))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
-        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidPhysReg(0xF10000, 0x1000)));
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidRegToken(0xF10000, 0xF00000)));
     }
 
     #[test]
@@ -1642,7 +1713,8 @@
             mmio_tokens: [((0x9, 0xFF), 0x12F00000)].into(),
             iommu_tokens: BTreeMap::new(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
         assert_eq!(device_info, Err(DeviceAssignmentError::InvalidIommus));
     }
@@ -1658,7 +1730,8 @@
             mmio_tokens: [((0x10000, 0x1000), 0xF00000), ((0x20000, 0xFF), 0xF10000)].into(),
             iommu_tokens: [((0xFF, 0xF), (0x40000, 0x4))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
         assert_eq!(device_info, Err(DeviceAssignmentError::DuplicatedPvIommuIds));
     }
@@ -1674,7 +1747,8 @@
             mmio_tokens: [((0x10000, 0x1000), 0xF00000), ((0x20000, 0xFF), 0xF10000)].into(),
             iommu_tokens: [((0xFF, 0xF), (0x40000, 0x4))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
         assert_eq!(device_info, Err(DeviceAssignmentError::UnsupportedIommusDuplication));
     }
@@ -1690,7 +1764,8 @@
             mmio_tokens: [((0xF000, 0x1000), 0xF00000), ((0xF100, 0x1000), 0xF10000)].into(),
             iommu_tokens: [((0xFF0, 0xF0), (0x40000, 0x4)), ((0xFF1, 0xF1), (0x40000, 0x4))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
         assert_eq!(device_info, Err(DeviceAssignmentError::InvalidIommus));
     }
@@ -1706,9 +1781,10 @@
             mmio_tokens: [((0x7fee0000, 0x1000), 0xF00000)].into(),
             iommu_tokens: [((0xFF, 0xF), (0x40000, 0x4))].into(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor);
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE);
 
-        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x7fee0000, 0x1000)));
+        assert_eq!(device_info, Err(DeviceAssignmentError::InvalidReg(0x7fee0000)));
     }
 
     #[test]
@@ -1741,7 +1817,9 @@
             iommu_tokens: Default::default(),
         };
 
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1771,7 +1849,9 @@
             mmio_tokens: [((0xFF000, 0x1), 0xF000), ((0xFF100, 0x1), 0xF100)].into(),
             iommu_tokens: Default::default(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
@@ -1802,7 +1882,9 @@
             mmio_tokens: [((0xFF200, 0x1), 0xF200)].into(),
             iommu_tokens: Default::default(),
         };
-        let device_info = DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor).unwrap().unwrap();
+        const HYP_GRANULE: usize = SIZE_4KB;
+        let device_info =
+            DeviceAssignmentInfo::parse(fdt, vm_dtbo, &hypervisor, HYP_GRANULE).unwrap().unwrap();
         device_info.filter(vm_dtbo).unwrap();
 
         // SAFETY: Damaged VM DTBO wouldn't be used after this unsafe block.
diff --git a/guest/pvmfw/src/fdt.rs b/guest/pvmfw/src/fdt.rs
index 0d934a6..f667d60 100644
--- a/guest/pvmfw/src/fdt.rs
+++ b/guest/pvmfw/src/fdt.rs
@@ -1148,7 +1148,18 @@
     let device_assignment = match vm_dtbo {
         Some(vm_dtbo) => {
             if let Some(hypervisor) = hyp::get_device_assigner() {
-                DeviceAssignmentInfo::parse(fdt, vm_dtbo, hypervisor).map_err(|e| {
+                // TODO(ptosi): Cache the (single?) granule once, in vmbase.
+                let granule = hyp::get_mem_sharer()
+                    .ok_or_else(|| {
+                        error!("No MEM_SHARE found during device assignment validation");
+                        RebootReason::InternalError
+                    })?
+                    .granule()
+                    .map_err(|e| {
+                        error!("Failed to get granule for device assignment validation: {e}");
+                        RebootReason::InternalError
+                    })?;
+                DeviceAssignmentInfo::parse(fdt, vm_dtbo, hypervisor, granule).map_err(|e| {
                     error!("Failed to parse device assignment from DT and VM DTBO: {e}");
                     RebootReason::InvalidFdt
                 })?
diff --git a/libs/service-virtualization/src/com/android/system/virtualmachine/VirtualizationSystemService.java b/libs/service-virtualization/src/com/android/system/virtualmachine/VirtualizationSystemService.java
index 998389b..afa286c 100644
--- a/libs/service-virtualization/src/com/android/system/virtualmachine/VirtualizationSystemService.java
+++ b/libs/service-virtualization/src/com/android/system/virtualmachine/VirtualizationSystemService.java
@@ -120,7 +120,12 @@
 
         @Override
         public void onReceive(Context context, Intent intent) {
-            switch (intent.getAction()) {
+            final String action = intent.getAction();
+            if (action == null) {
+                return;
+            }
+
+            switch (action) {
                 case Intent.ACTION_USER_REMOVED:
                     onUserRemoved(intent);
                     break;
@@ -128,7 +133,7 @@
                     onPackageRemoved(intent);
                     break;
                 default:
-                    Log.e(TAG, "received unexpected intent: " + intent.getAction());
+                    Log.e(TAG, "received unexpected intent: " + intent);
                     break;
             }
         }
diff --git a/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherService.java b/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherService.java
index 3731854..5cd7b92 100644
--- a/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherService.java
+++ b/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherService.java
@@ -33,6 +33,7 @@
 
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
@@ -62,6 +63,11 @@
 
     @Override
     public int onStartCommand(Intent intent, int flags, int startId) {
+        if (Objects.equals(
+                intent.getAction(), VmLauncherServices.ACTION_STOP_VM_LAUNCHER_SERVICE)) {
+            stopSelf();
+            return START_NOT_STICKY;
+        }
         if (mVirtualMachine != null) {
             Log.d(TAG, "VM instance is already started");
             return START_NOT_STICKY;
diff --git a/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherServices.java b/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherServices.java
index 2fa0b32..6eca2b3 100644
--- a/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherServices.java
+++ b/libs/vm_launcher_lib/java/com/android/virtualization/vmlauncher/VmLauncherServices.java
@@ -36,6 +36,8 @@
     private static final String ACTION_START_VM_LAUNCHER_SERVICE =
             "android.virtualization.START_VM_LAUNCHER_SERVICE";
 
+    public static final String ACTION_STOP_VM_LAUNCHER_SERVICE =
+            "android.virtualization.STOP_VM_LAUNCHER_SERVICE";
     private static final int RESULT_START = 0;
     private static final int RESULT_STOP = 1;
     private static final int RESULT_ERROR = 2;
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 974a58c..c5e0171 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
@@ -23,24 +23,21 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
-import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
 import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
-import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.CommandStatus;
-import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
+import com.android.tradefed.util.SearchArtifactUtil;
 
 import org.json.JSONArray;
 import org.json.JSONObject;
 
 import java.io.File;
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -178,28 +175,12 @@
 
     public File findTestFile(String name) {
         String moduleName = getInvocationContext().getConfigurationDescriptor().getModuleName();
-        IBuildInfo buildInfo = getBuild();
-        CompatibilityBuildHelper helper = new CompatibilityBuildHelper(buildInfo);
-
-        // 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 {
-            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) {
+        File testFile = SearchArtifactUtil.searchFile(name, false);
+        if (testFile == null) {
             throw new AssertionError(
-                    "Failed to find test file " + name + " for module " + moduleName, e);
+                    "Failed to find test file " + name + " for module " + moduleName);
         }
-        throw new AssertionError("Failed to find test file " + name + " for module " + moduleName);
+        return testFile;
     }
 
     public String getPathForPackage(String packageName) throws DeviceNotAvailableException {
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index ffcf338..ee2da09 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -1050,9 +1050,28 @@
     }
 
     @Test
-    @Parameters(method = "params")
-    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
-    public void testMicrodroidRamUsage(boolean protectedVm, String gki) throws Exception {
+    public void testMicrodroidRamUsage_protectedVm_true_gki_null() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ true, /* gki= */ "null");
+    }
+
+    @Test
+    public void testMicrodroidRamUsage_protectedVm_false_gki_null() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ false, /* gki= */ "null");
+    }
+
+    @Test
+    public void testMicrodroidRamUsage_protectedVm_true_gki_android15() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ true, /* gki= */ "android15");
+    }
+
+    @Test
+    public void testMicrodroidRamUsage_protectedVm_false_gki_android15() throws Exception {
+        checkMicrodroidRamUsage(/* protectedVm= */ false, /* gki= */ "android15");
+    }
+
+    // TODO(b/209036125): Upgrade this function to a parameterized test once metrics can be
+    // collected with tradefed parameterizer.
+    void checkMicrodroidRamUsage(boolean protectedVm, String gki) throws Exception {
         // Preconditions
         assumeKernelSupported(gki);
         assumeVmTypeSupported(protectedVm);