resources: add address_from_pci_offset function

Refactor current code and add tests.

BUG=chromium:924405
TEST=compile and test

Change-Id: I9476f3a4ffd8ae85fc95d6889ada6b056613bbfa
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2216447
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>
Commit-Queue: Gurchetan Singh <gurchetansingh@chromium.org>
Tested-by: Gurchetan Singh <gurchetansingh@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Auto-Submit: Gurchetan Singh <gurchetansingh@chromium.org>

3 files changed, 96 insertions, 40 deletions

diff --git a/resources/src/address_allocator.rs b/resources/src/address_allocator.rs
index 88f652a..dce904f 100644
--- a/resources/src/address_allocator.rs
+++ b/resources/src/address_allocator.rs
@@ -215,6 +215,27 @@ impl AddressAllocator {
 
         Ok(())
     }
+
+    /// Returns an address from associated PCI `alloc` given an allocation offset and size.
+    pub fn address_from_pci_offset(&self, alloc: Alloc, offset: u64, size: u64) -> Result<u64> {
+        match alloc {
+            Alloc::PciBar { .. } => (),
+            _ => return Err(Error::InvalidAlloc(alloc)),
+        };
+
+        match self.allocs.get(&alloc) {
+            Some((start_addr, length, _)) => {
+                let address = start_addr.checked_add(offset).ok_or(Error::OutOfBounds)?;
+                let range = *start_addr..*start_addr + *length;
+                let end = address.checked_add(size).ok_or(Error::OutOfBounds)?;
+                match (range.contains(&address), range.contains(&end)) {
+                    (true, true) => Ok(address),
+                    _ => return Err(Error::OutOfBounds),
+                }
+            }
+            None => return Err(Error::InvalidAlloc(alloc)),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -422,4 +443,65 @@ mod tests {
             Ok(0x1000)
         );
     }
+
+    #[test]
+    fn allocate_and_verify_pci_offset() {
+        let mut pool = AddressAllocator::new(0x1000, 0x10000, None).unwrap();
+        let pci_bar0 = Alloc::PciBar {
+            bus: 1,
+            dev: 2,
+            func: 0,
+            bar: 0,
+        };
+        let pci_bar1 = Alloc::PciBar {
+            bus: 1,
+            dev: 2,
+            func: 0,
+            bar: 1,
+        };
+        let pci_bar2 = Alloc::PciBar {
+            bus: 1,
+            dev: 2,
+            func: 0,
+            bar: 2,
+        };
+        let anon = Alloc::Anon(1);
+
+        assert_eq!(
+            pool.allocate(0x800, pci_bar0, String::from("bar0")),
+            Ok(0x1000)
+        );
+        assert_eq!(
+            pool.allocate(0x800, pci_bar1, String::from("bar1")),
+            Ok(0x1800)
+        );
+        assert_eq!(pool.allocate(0x800, anon, String::from("anon")), Ok(0x2000));
+
+        assert_eq!(
+            pool.address_from_pci_offset(pci_bar0, 0x600, 0x100),
+            Ok(0x1600)
+        );
+        assert_eq!(
+            pool.address_from_pci_offset(pci_bar1, 0x600, 0x100),
+            Ok(0x1E00)
+        );
+        assert_eq!(
+            pool.address_from_pci_offset(pci_bar0, 0x7FE, 0x001),
+            Ok(0x17FE)
+        );
+        assert_eq!(
+            pool.address_from_pci_offset(pci_bar0, 0x7FF, 0x001),
+            Err(Error::OutOfBounds)
+        );
+
+        assert_eq!(
+            pool.address_from_pci_offset(pci_bar2, 0x7FF, 0x001),
+            Err(Error::InvalidAlloc(pci_bar2))
+        );
+
+        assert_eq!(
+            pool.address_from_pci_offset(anon, 0x600, 0x100),
+            Err(Error::InvalidAlloc(anon))
+        );
+    }
 }
diff --git a/resources/src/lib.rs b/resources/src/lib.rs
index 4da5d22..95e4b98 100644
--- a/resources/src/lib.rs
+++ b/resources/src/lib.rs
@@ -46,10 +46,12 @@ pub enum Error {
     BadAlignment,
     CreateGpuAllocator(GpuAllocatorError),
     ExistingAlloc(Alloc),
+    InvalidAlloc(Alloc),
     MissingHighMMIOAddresses,
     MissingLowMMIOAddresses,
     NoIoAllocator,
     OutOfSpace,
+    OutOfBounds,
     PoolOverflow { base: u64, size: u64 },
     PoolSizeZero,
     RegionOverlap { base: u64, size: u64 },
@@ -66,10 +68,12 @@ impl Display for Error {
             BadAlignment => write!(f, "Pool alignment must be a power of 2"),
             CreateGpuAllocator(e) => write!(f, "Failed to create GPU allocator: {:?}", e),
             ExistingAlloc(tag) => write!(f, "Alloc already exists: {:?}", tag),
+            InvalidAlloc(tag) => write!(f, "Invalid Alloc: {:?}", tag),
             MissingHighMMIOAddresses => write!(f, "High MMIO address range not specified"),
             MissingLowMMIOAddresses => write!(f, "Low MMIO address range not specified"),
             NoIoAllocator => write!(f, "No IO address range specified"),
             OutOfSpace => write!(f, "Out of space"),
+            OutOfBounds => write!(f, "Out of bounds"),
             PoolOverflow { base, size } => write!(f, "base={} + size={} overflows", base, size),
             PoolSizeZero => write!(f, "Pool cannot have size of 0"),
             RegionOverlap { base, size } => {
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index f7514c2..5336f3b 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -535,7 +535,7 @@ fn register_memory(
     allocator: &mut SystemAllocator,
     fd: &dyn AsRawFd,
     size: usize,
-    allocation: Option<(Alloc, u64)>,
+    pci_allocation: Option<(Alloc, u64)>,
 ) -> Result<(u64, u32)> {
     let mmap = match MemoryMapping::from_fd(fd, size) {
         Ok(v) => v,
@@ -543,48 +543,18 @@ fn register_memory(
         _ => return Err(SysError::new(EINVAL)),
     };
 
-    let addr = match allocation {
-        Some((
-            Alloc::PciBar {
-                bus,
-                dev,
-                func,
-                bar,
-            },
-            offset,
-        )) => {
-            match allocator
-                .mmio_allocator(MmioType::High)
-                .get(&Alloc::PciBar {
-                    bus,
-                    dev,
-                    func,
-                    bar,
-                }) {
-                Some((start_addr, length, _)) => {
-                    let address = *start_addr + offset;
-                    let range = *start_addr..*start_addr + *length;
-                    let end = address + (size as u64);
-                    match (range.contains(&address), range.contains(&end)) {
-                        (true, true) => address,
-                        _ => return Err(SysError::new(EINVAL)),
-                    }
-                }
-                None => return Err(SysError::new(EINVAL)),
-            }
-        }
+    let addr = match pci_allocation {
+        Some(pci_allocation) => allocator
+            .mmio_allocator(MmioType::High)
+            .address_from_pci_offset(pci_allocation.0, pci_allocation.1, size as u64)
+            .map_err(|_e| SysError::new(EINVAL))?,
         None => {
             let alloc = allocator.get_anon_alloc();
-            match allocator.mmio_allocator(MmioType::High).allocate(
-                size as u64,
-                alloc,
-                "vmcontrol_register_memory".to_string(),
-            ) {
-                Ok(a) => a,
-                _ => return Err(SysError::new(EINVAL)),
-            }
+            allocator
+                .mmio_allocator(MmioType::High)
+                .allocate(size as u64, alloc, "vmcontrol_register_memory".to_string())
+                .map_err(|_e| SysError::new(EINVAL))?
         }
-        _ => return Err(SysError::new(EINVAL)),
     };
 
     let slot = vm.add_memory_region(GuestAddress(addr), Box::new(mmap), false, false)?;