Merge changes I4e2c8bd2,I6b466c16,I8bb9aa13,I5af9d91a into main

* changes:
  MockHypervisor: Get MMIO tokens by address only
  pvmfw: Validate size of assigned MMIO regions
  pvmfw: Warn for misaligned assigned MMIO regions
  pvmfw: Refactor DeviceAssignmentError for clarity
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
                 })?