Merge "Update ARM_SMCCC_KVM_FUNC_DEV_REQ_MMIO ABI" into main
diff --git a/guest/pvmfw/src/device_assignment.rs b/guest/pvmfw/src/device_assignment.rs
index da9462b..9b55cff 100644
--- a/guest/pvmfw/src/device_assignment.rs
+++ b/guest/pvmfw/src/device_assignment.rs
@@ -784,14 +784,16 @@
             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 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)
-            })?;
-            if token != expected_token {
-                return Err(DeviceAssignmentError::InvalidRegToken(token, expected_token));
+            for offset in (0..reg.size).step_by(granule) {
+                let expected_token = phys_reg.addr + offset;
+                // If this call returns successfully, hyp has mapped the MMIO granule.
+                let token = hypervisor.get_phys_mmio_token(reg.addr + offset).map_err(|e| {
+                    error!("Hypervisor error while requesting MMIO token: {e}");
+                    DeviceAssignmentError::InvalidReg(reg.addr)
+                })?;
+                if token != expected_token {
+                    return Err(DeviceAssignmentError::InvalidRegToken(token, expected_token));
+                }
             }
         }
 
@@ -1143,7 +1145,7 @@
 #[cfg(test)]
 trait DeviceAssigningHypervisor {
     /// Returns MMIO token.
-    fn get_phys_mmio_token(&self, base_ipa: u64, size: u64) -> MockHypervisorResult<u64>;
+    fn get_phys_mmio_token(&self, base_ipa: u64) -> MockHypervisorResult<u64>;
 
     /// Returns DMA token as a tuple of (phys_iommu_id, phys_sid).
     fn get_phys_iommu_token(&self, pviommu_id: u64, vsid: u64) -> MockHypervisorResult<(u64, u64)>;
@@ -1206,7 +1208,7 @@
     }
 
     impl DeviceAssigningHypervisor for MockHypervisor {
-        fn get_phys_mmio_token(&self, base_ipa: u64, _size: u64) -> MockHypervisorResult<u64> {
+        fn get_phys_mmio_token(&self, base_ipa: u64) -> MockHypervisorResult<u64> {
             let token = self.get_mmio_token(base_ipa);
 
             Ok(*token.ok_or(MockHypervisorError::FailedGetPhysMmioToken)?)
diff --git a/libs/libvmbase/src/hyp/hypervisor/common.rs b/libs/libvmbase/src/hyp/hypervisor/common.rs
index de0fe12..8f0e4dc 100644
--- a/libs/libvmbase/src/hyp/hypervisor/common.rs
+++ b/libs/libvmbase/src/hyp/hypervisor/common.rs
@@ -69,7 +69,7 @@
 /// Device assigning hypervisor
 pub trait DeviceAssigningHypervisor {
     /// Returns MMIO token.
-    fn get_phys_mmio_token(&self, base_ipa: u64, size: u64) -> Result<u64>;
+    fn get_phys_mmio_token(&self, base_ipa: u64) -> Result<u64>;
 
     /// Returns DMA token as a tuple of (phys_iommu_id, phys_sid).
     fn get_phys_iommu_token(&self, pviommu_id: u64, vsid: u64) -> Result<(u64, u64)>;
diff --git a/libs/libvmbase/src/hyp/hypervisor/kvm.rs b/libs/libvmbase/src/hyp/hypervisor/kvm.rs
index e496f09..7ed829e 100644
--- a/libs/libvmbase/src/hyp/hypervisor/kvm.rs
+++ b/libs/libvmbase/src/hyp/hypervisor/kvm.rs
@@ -173,10 +173,9 @@
 }
 
 impl DeviceAssigningHypervisor for ProtectedKvmHypervisor {
-    fn get_phys_mmio_token(&self, base_ipa: u64, size: u64) -> Result<u64> {
+    fn get_phys_mmio_token(&self, base_ipa: u64) -> Result<u64> {
         let mut args = [0u64; 17];
         args[0] = base_ipa;
-        args[1] = size;
 
         let ret = checked_hvc64_expect_results(VENDOR_HYP_KVM_DEV_REQ_MMIO_FUNC_ID, args)?;
         Ok(ret[0])