Keep track of kvm memslot index inside GuestMemoryMmap

Use a special KvmRegion type, which allows us to manage
kvm_userspace_memory_region2 inside of GuestRegionCollection.

Signed-off-by: Patrick Roy <roypat@amazon.co.uk>

Author: roypat

Cargo.lock

@@ -117,7 +117,9 @@ pub fn configure_system_for_boot(
     )?;
 
     let fdt_address = GuestAddress(get_fdt_addr(vmm.vm.guest_memory()));
-    vmm.vm.guest_memory().write_slice(fdt.as_slice(), fdt_address)?;
+    vmm.vm
+        .guest_memory()
+        .write_slice(fdt.as_slice(), fdt_address)?;
 
     Ok(())
 }

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

@@ -378,7 +378,7 @@ mod tests {
     use crate::devices::virtio::mmio::VIRTIO_MMIO_INT_CONFIG;
     use crate::test_utils::create_tmp_socket;
     use crate::vstate::memory;
-    use crate::vstate::memory::GuestAddress;
+    use crate::vstate::memory::{GuestAddress, KvmRegion};
 
     #[test]
     fn test_from_config() {
@@ -780,7 +780,11 @@ mod tests {
         file.set_len(region_size as u64).unwrap();
         let regions = vec![(GuestAddress(0x0), region_size)];
         let guest_memory = GuestMemoryMmap::from_regions(
-            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
+                .unwrap()
+                .into_iter()
+                .map(|region| KvmRegion::from_mmap_region(region, 0))
+                .collect(),
         )
         .unwrap();
 

src/vmm/src/devices/virtio/block/vhost_user/device.rs

@@ -188,14 +188,15 @@ pub mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
     use std::os::unix::ffi::OsStrExt;
 
+    use vm_memory::GuestMemoryRegion;
     use vmm_sys_util::tempfile::TempFile;
 
     use super::*;
     use crate::devices::virtio::block::virtio::device::FileEngineType;
     use crate::utils::u64_to_usize;
     use crate::vmm_config::machine_config::HugePageConfig;
     use crate::vstate::memory;
-    use crate::vstate::memory::{Bitmap, Bytes, GuestMemory};
+    use crate::vstate::memory::{Bitmap, Bytes, GuestMemory, KvmRegion};
 
     const FILE_LEN: u32 = 1024;
     // 2 pages of memory should be enough to test read/write ops and also dirty tracking.
@@ -237,7 +238,10 @@ pub mod tests {
                 true,
                 HugePageConfig::None,
             )
-            .unwrap(),
+            .unwrap()
+            .into_iter()
+            .map(|region| KvmRegion::from_mmap_region(region, 0))
+            .collect(),
         )
         .unwrap()
     }

src/vmm/src/devices/virtio/block/virtio/io/mod.rs

@@ -380,7 +380,7 @@ impl<T: VhostUserHandleBackend> VhostUserHandleImpl<T> {
             let vhost_user_net_reg = VhostUserMemoryRegionInfo {
                 guest_phys_addr: region.start_addr().raw_value(),
                 memory_size: region.len(),
-                userspace_addr: region.as_ptr() as u64,
+                userspace_addr: region.inner().userspace_addr,
                 mmap_offset,
                 mmap_handle,
             };
@@ -467,7 +467,7 @@ mod tests {
     use super::*;
     use crate::test_utils::create_tmp_socket;
     use crate::vstate::memory;
-    use crate::vstate::memory::GuestAddress;
+    use crate::vstate::memory::{GuestAddress, KvmRegion};
 
     #[test]
     fn test_new() {
@@ -765,7 +765,11 @@ mod tests {
         ];
 
         let guest_memory = GuestMemoryMmap::from_regions(
-            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
+                .unwrap()
+                .into_iter()
+                .map(|region| KvmRegion::from_mmap_region(region, 0))
+                .collect(),
         )
         .unwrap();
 
@@ -777,7 +781,7 @@ mod tests {
             .map(|region| VhostUserMemoryRegionInfo {
                 guest_phys_addr: region.start_addr().raw_value(),
                 memory_size: region.len(),
-                userspace_addr: region.as_ptr() as u64,
+                userspace_addr: region.inner().userspace_addr,
                 mmap_offset: region.file_offset().unwrap().start(),
                 mmap_handle: region.file_offset().unwrap().file().as_raw_fd(),
             })
@@ -882,7 +886,11 @@ mod tests {
         let regions = vec![(GuestAddress(0x0), region_size)];
 
         let guest_memory = GuestMemoryMmap::from_regions(
-            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false).unwrap(),
+            memory::create(regions.into_iter(), libc::MAP_PRIVATE, Some(file), false)
+                .unwrap()
+                .into_iter()
+                .map(|region| KvmRegion::from_mmap_region(region, 0))
+                .collect(),
         )
         .unwrap();
 

src/vmm/src/devices/virtio/vhost_user.rs

@@ -63,7 +63,9 @@ const PTE_ADDRESS_MASK: u64 = !0b111u64;
 /// Read a u64 value from a guest memory address
 fn read_address(vmm: &Vmm, address: u64) -> Result<u64, GdbTargetError> {
     let mut buf = [0; 8];
-    vmm.vm.guest_memory().read(&mut buf, GuestAddress(address))?;
+    vmm.vm
+        .guest_memory()
+        .read(&mut buf, GuestAddress(address))?;
 
     Ok(u64::from_le_bytes(buf))
 }

src/vmm/src/gdb/arch/aarch64.rs

@@ -399,7 +399,8 @@ impl MultiThreadBase for FirecrackerTarget {
                 GUEST_PAGE_SIZE - (u64_to_usize(gpa) & (GUEST_PAGE_SIZE - 1)),
             );
 
-            vmm.vm.guest_memory()
+            vmm.vm
+                .guest_memory()
                 .read(&mut data[..read_len], GuestAddress(gpa as u64))
                 .map_err(|e| {
                     error!("Error reading memory {e:?} gpa is {gpa}");
@@ -433,7 +434,8 @@ impl MultiThreadBase for FirecrackerTarget {
                 GUEST_PAGE_SIZE - (u64_to_usize(gpa) & (GUEST_PAGE_SIZE - 1)),
             );
 
-            vmm.vm.guest_memory()
+            vmm.vm
+                .guest_memory()
                 .write(&data[..write_len], GuestAddress(gpa))
                 .map_err(|e| {
                     error!("Error {e:?} writing memory at {gpa:#X}");

src/vmm/src/gdb/target.rs

@@ -5,7 +5,7 @@
 
 use std::sync::{Arc, Mutex};
 
-use vm_memory::GuestAddress;
+use vm_memory::{GuestAddress, GuestRegionCollection};
 use vmm_sys_util::tempdir::TempDir;
 
 use crate::builder::build_microvm_for_boot;
@@ -16,7 +16,7 @@ use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::HugePageConfig;
 use crate::vstate::memory;
-use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap};
+use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap, KvmRegion};
 use crate::{EventManager, Vmm};
 
 pub mod mock_resources;
@@ -43,9 +43,12 @@ pub fn single_region_mem_at_raw(at: u64, size: usize) -> Vec<GuestRegionMmap> {
 
 /// Creates a [`GuestMemoryMmap`] with multiple regions and without dirty page tracking.
 pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
-    GuestMemoryMmap::from_regions(
+    GuestRegionCollection::from_regions(
         memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
-            .expect("Cannot initialize memory"),
+            .expect("Cannot initialize memory")
+            .into_iter()
+            .map(|region| KvmRegion::from_mmap_region(region, 0))
+            .collect(),
     )
     .unwrap()
 }