add offset argument to arch_regions

This allows having regions that start somewhere other than guest
physical address 0. Useful in case not all regions are allocated at
once, and each batch needs to be placed after all already allocated
regions.

Since the resulting code has some intricacies (and I actually got it
wrong the first time around), write a kani harness to validate that
we're computing the regions correctly.

On ARM, add a warning in case we hit some questionable logic around
truncation of memory sizes if they exceed architectural bounds (although
if anyone is hitting this, they're already doing something wrong anyway,
because it'd require a VM with more than 1022GiB of memory).

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

src/vmm/src/arch/aarch64/mod.rs

@@ -32,7 +32,7 @@ use crate::utils::{align_up, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
 use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
 use crate::vstate::vcpu::KvmVcpuError;
-use crate::{Vcpu, VcpuConfig, Vmm};
+use crate::{Vcpu, VcpuConfig, Vmm, logger};
 
 /// Errors thrown while configuring aarch64 system.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -58,9 +58,35 @@ pub const MMIO_MEM_SIZE: u64 = layout::DRAM_MEM_START - layout::MAPPED_IO_START;
 
 /// Returns a Vec of the valid memory addresses for aarch64.
 /// See [`layout`](layout) module for a drawing of the specific memory model for this platform.
-pub fn arch_memory_regions(size: usize) -> Vec<(GuestAddress, usize)> {
-    let dram_size = min(size, layout::DRAM_MEM_MAX_SIZE);
-    vec![(GuestAddress(layout::DRAM_MEM_START), dram_size)]
+///
+/// The `offset` parameter specified the offset from [`layout::DRAM_MEM_START`].
+pub fn arch_memory_regions(offset: usize, size: usize) -> Vec<(GuestAddress, usize)> {
+    assert!(size > 0, "Attempt to allocate guest memory of length 0");
+    assert!(
+        offset.checked_add(size).is_some(),
+        "Attempt to allocate guest memory such that the address space would wrap around"
+    );
+    assert!(
+        offset < layout::DRAM_MEM_MAX_SIZE,
+        "offset outside allowed DRAM range"
+    );
+
+    let dram_size = min(size, layout::DRAM_MEM_MAX_SIZE - offset);
+
+    if dram_size != size {
+        logger::warn!(
+            "Requested offset/memory size {}/{} exceeds architectural maximum (1022GiB). Size has \
+             been truncated to {}",
+            offset,
+            size,
+            dram_size
+        );
+    }
+
+    vec![(
+        GuestAddress(layout::DRAM_MEM_START + offset as u64),
+        dram_size,
+    )]
 }
 
 /// Configures the system for booting Linux.
@@ -184,22 +210,61 @@ pub fn load_kernel(
     })
 }
 
+#[cfg(kani)]
+mod verification {
+    use vm_memory::GuestAddress;
+
+    use crate::arch::aarch64::layout;
+    use crate::arch::arch_memory_regions;
+
+    #[kani::proof]
+    #[kani::unwind(3)]
+    fn verify_arch_memory_regions() {
+        let offset: u64 = kani::any::<u64>();
+        let len: u64 = kani::any::<u64>();
+
+        kani::assume(len > 0);
+        kani::assume(offset.checked_add(len).is_some());
+        kani::assume(offset < layout::DRAM_MEM_MAX_SIZE as u64);
+
+        let regions = arch_memory_regions(offset as usize, len as usize);
+
+        // No MMIO gap on ARM
+        assert_eq!(regions.len(), 1);
+
+        let (GuestAddress(start), actual_len) = regions[0];
+        let actual_len = actual_len as u64;
+
+        assert_eq!(start, layout::DRAM_MEM_START + offset);
+        assert!(actual_len <= layout::DRAM_MEM_MAX_SIZE as u64);
+        assert!(actual_len <= len);
+
+        if actual_len < len {
+            assert_eq!(
+                start + actual_len,
+                layout::DRAM_MEM_START + layout::DRAM_MEM_MAX_SIZE as u64
+            );
+            assert!(offset + len >= layout::DRAM_MEM_MAX_SIZE as u64);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::test_utils::arch_mem;
 
     #[test]
     fn test_regions_lt_1024gb() {
-        let regions = arch_memory_regions(1usize << 29);
+        let regions = arch_memory_regions(0, 1usize << 29);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(1usize << 29, regions[0].1);
     }
 
     #[test]
     fn test_regions_gt_1024gb() {
-        let regions = arch_memory_regions(1usize << 41);
+        let regions = arch_memory_regions(0, 1usize << 41);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(super::layout::DRAM_MEM_MAX_SIZE, regions[0].1);

src/vmm/src/arch/x86_64/mod.rs

@@ -110,17 +110,33 @@ pub const MMIO_MEM_SIZE: u64 = MEM_32BIT_GAP_SIZE;
 /// These should be used to configure the GuestMemoryMmap structure for the platform.
 /// For x86_64 all addresses are valid from the start of the kernel except a
 /// carve out at the end of 32bit address space.
-pub fn arch_memory_regions(size: usize) -> Vec<(GuestAddress, usize)> {
+pub fn arch_memory_regions(offset: usize, size: usize) -> Vec<(GuestAddress, usize)> {
+    // If we get here with size == 0 something has seriously gone wrong. Firecracker should never
+    // try to allocate guest memory of size 0
+    assert!(size > 0, "Attempt to allocate guest memory of length 0");
+    assert!(
+        offset.checked_add(size).is_some(),
+        "Attempt to allocate guest memory such that the address space would wrap around"
+    );
+
     // It's safe to cast MMIO_MEM_START to usize because it fits in a u32 variable
     // (It points to an address in the 32 bit space).
-    match size.checked_sub(usize::try_from(MMIO_MEM_START).unwrap()) {
+    match (size + offset).checked_sub(u64_to_usize(MMIO_MEM_START)) {
         // case1: guest memory fits before the gap
-        None | Some(0) => vec![(GuestAddress(0), size)],
-        // case2: guest memory extends beyond the gap
-        Some(remaining) => vec![
-            (GuestAddress(0), usize::try_from(MMIO_MEM_START).unwrap()),
+        None | Some(0) => vec![(GuestAddress(offset as u64), size)],
+        // case2: starts before the gap, but doesn't completely fit
+        Some(remaining) if (offset as u64) < MMIO_MEM_START => vec![
+            (
+                GuestAddress(offset as u64),
+                u64_to_usize(MMIO_MEM_START) - offset,
+            ),
             (GuestAddress(FIRST_ADDR_PAST_32BITS), remaining),
         ],
+        // case3: guest memory start after the gap
+        Some(_) => vec![(
+            GuestAddress(FIRST_ADDR_PAST_32BITS.max(offset as u64)),
+            size,
+        )],
     }
 }
 
@@ -456,6 +472,56 @@ pub fn load_kernel(
     })
 }
 
+#[cfg(kani)]
+mod verification {
+    use crate::arch::x86_64::FIRST_ADDR_PAST_32BITS;
+    use crate::arch::{MMIO_MEM_START, arch_memory_regions};
+
+    #[kani::proof]
+    #[kani::unwind(3)]
+    fn verify_arch_memory_regions() {
+        let offset: u64 = kani::any::<u64>();
+        let len: u64 = kani::any::<u64>();
+
+        kani::assume(len > 0);
+        kani::assume(offset.checked_add(len).is_some());
+
+        let regions = arch_memory_regions(offset as usize, len as usize);
+
+        // There's only one MMIO gap, so we can get either 1 or 2 regions
+        assert!(regions.len() <= 2);
+        assert!(regions.len() >= 1);
+
+        // The total length of all regions is what we requested
+        assert_eq!(
+            regions.iter().map(|&(_, len)| len).sum::<usize>(),
+            len as usize
+        );
+
+        // No region overlaps the MMIO gap
+        assert!(
+            regions
+                .iter()
+                .all(|&(start, len)| start.0 >= FIRST_ADDR_PAST_32BITS
+                    || start.0 + len as u64 <= MMIO_MEM_START)
+        );
+
+        // All regions start after our specified offset
+        assert!(regions.iter().all(|&(start, _)| start.0 >= offset as u64));
+
+        // All regions have non-zero length
+        assert!(regions.iter().all(|&(_, len)| len > 0));
+
+        // If there's two regions, they perfectly snuggle up to the MMIO gap
+        if regions.len() == 2 {
+            kani::cover!();
+
+            assert_eq!(regions[0].0.0 + regions[0].1 as u64, MMIO_MEM_START);
+            assert_eq!(regions[1].0.0, FIRST_ADDR_PAST_32BITS);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use linux_loader::loader::bootparam::boot_e820_entry;
@@ -466,18 +532,41 @@ mod tests {
 
     #[test]
     fn regions_lt_4gb() {
-        let regions = arch_memory_regions(1usize << 29);
+        let regions = arch_memory_regions(0, 1usize << 29);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(0), regions[0].0);
         assert_eq!(1usize << 29, regions[0].1);
+
+        let regions = arch_memory_regions(1 << 28, 1 << 29);
+        assert_eq!(1, regions.len());
+        assert_eq!(regions[0], (GuestAddress(1 << 28), 1 << 29));
     }
 
     #[test]
     fn regions_gt_4gb() {
-        let regions = arch_memory_regions((1usize << 32) + 0x8000);
+        const MEMORY_SIZE: usize = (1 << 32) + 0x8000;
+
+        let regions = arch_memory_regions(0, MEMORY_SIZE);
         assert_eq!(2, regions.len());
         assert_eq!(GuestAddress(0), regions[0].0);
         assert_eq!(GuestAddress(1u64 << 32), regions[1].0);
+
+        let regions = arch_memory_regions(1 << 31, MEMORY_SIZE);
+        assert_eq!(2, regions.len());
+        assert_eq!(
+            regions[0],
+            (
+                GuestAddress(1 << 31),
+                u64_to_usize(MMIO_MEM_START) - (1 << 31)
+            )
+        );
+        assert_eq!(
+            regions[1],
+            (
+                GuestAddress(FIRST_ADDR_PAST_32BITS),
+                MEMORY_SIZE - regions[0].1
+            )
+        )
     }
 
     #[test]

src/vmm/src/resources.rs

@@ -458,7 +458,7 @@ impl VmResources {
         // a single way of backing guest memory for vhost-user and non-vhost-user cases,
         // that would not be worth the effort.
         let regions =
-            crate::arch::arch_memory_regions(mib_to_bytes(self.machine_config.mem_size_mib));
+            crate::arch::arch_memory_regions(0, mib_to_bytes(self.machine_config.mem_size_mib));
         if vhost_user_device_used {
             memory::memfd_backed(
                 regions.as_ref(),

src/vmm/src/test_utils/mod.rs

@@ -58,11 +58,11 @@ pub fn multi_region_mem_raw(regions: &[(GuestAddress, usize)]) -> Vec<GuestRegio
 /// Creates a [`GuestMemoryMmap`] of the given size with the contained regions laid out in
 /// accordance with the requirements of the architecture on which the tests are being run.
 pub fn arch_mem(mem_size_bytes: usize) -> GuestMemoryMmap {
-    multi_region_mem(&crate::arch::arch_memory_regions(mem_size_bytes))
+    multi_region_mem(&crate::arch::arch_memory_regions(0, mem_size_bytes))
 }
 
 pub fn arch_mem_raw(mem_size_bytes: usize) -> Vec<GuestRegionMmap> {
-    multi_region_mem_raw(&crate::arch::arch_memory_regions(mem_size_bytes))
+    multi_region_mem_raw(&crate::arch::arch_memory_regions(0, mem_size_bytes))
 }
 
 pub fn create_vmm(