switch to using kvm_userspace_region2

Fall back to kvm_user_memory_region in case the 2 version of the struct
isnt supported.

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

Author: roypat

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

@@ -240,7 +240,7 @@ pub mod tests {
             )
             .unwrap()
             .into_iter()
-            .map(|region| KvmRegion::from_mmap_region(region, 0))
+            .map(|region| KvmRegion::from_mmap_region(region, 0, None))
             .collect(),
         )
         .unwrap()

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

@@ -482,7 +482,7 @@ pub(crate) mod tests {
             )
             .unwrap()
             .into_iter()
-            .map(|region| KvmRegion::from_mmap_region(region, 0))
+            .map(|region| KvmRegion::from_mmap_region(region, 0, None))
             .collect(),
         )
         .unwrap()

src/vmm/src/test_utils/mod.rs

@@ -48,7 +48,7 @@ pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
         memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
             .expect("Cannot initialize memory")
             .into_iter()
-            .map(|region| KvmRegion::from_mmap_region(region, 0))
+            .map(|region| KvmRegion::from_mmap_region(region, 0, None))
             .collect(),
     )
     .unwrap()

src/vmm/src/vstate/memory.rs

@@ -8,10 +8,10 @@
 use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::mem::ManuallyDrop;
-use std::os::fd::AsFd;
+use std::os::fd::{AsFd, AsRawFd};
 use std::sync::Arc;
 
-use kvm_bindings::{KVM_MEM_LOG_DIRTY_PAGES, kvm_userspace_memory_region};
+use kvm_bindings::{KVM_MEM_LOG_DIRTY_PAGES, kvm_userspace_memory_region2};
 use serde::{Deserialize, Serialize};
 pub use vm_memory::bitmap::{AtomicBitmap, BS, Bitmap, BitmapSlice};
 pub use vm_memory::mmap::MmapRegionBuilder;
@@ -115,7 +115,7 @@ impl<S: Seek> Seek for Bounce<S> {
 /// A memory region, described in terms of `kvm_userspace_memory_region`
 #[derive(Debug)]
 pub struct KvmRegion {
-    region: kvm_userspace_memory_region,
+    region: kvm_userspace_memory_region2,
     bitmap: Option<AtomicBitmap>,
     file_offset: Option<FileOffset>,
 }
@@ -129,7 +129,7 @@ impl KvmRegion {
     /// `kvm_region.userspace_addr as *mut u8` is valid for reads and writes of length
     /// `kvm_region.memory_size`.
     pub unsafe fn new(
-        region: kvm_userspace_memory_region,
+        region: kvm_userspace_memory_region2,
         bitmap: Option<AtomicBitmap>,
         file_offset: Option<FileOffset>,
     ) -> Self {
@@ -140,34 +140,46 @@ impl KvmRegion {
         }
     }
 
-    pub(crate) fn from_mmap_region(region: GuestRegionMmap, slot: u32) -> Self {
+    pub(crate) fn from_mmap_region(
+        region: GuestRegionMmap,
+        slot: u32,
+        guest_memfd: Option<&FileOffset>,
+    ) -> Self {
         let region = ManuallyDrop::new(region);
         let flags = if region.bitmap().is_some() {
             KVM_MEM_LOG_DIRTY_PAGES
         } else {
             0
         };
 
+        #[allow(clippy::cast_sign_loss)]
+        let (guest_memfd, guest_memfd_offset) = guest_memfd
+            .map(|fo| (fo.file().as_raw_fd() as u32, fo.start()))
+            .unwrap_or((0, 0));
+
         // SAFETY: `GuestRegionMmap` is essentially a fat pointer, and ensures that
         // region.as_ptr() is valid for reads and writes of length region.len(),
         // and by placing our region into a `ManuallyDrop` we ensure that its `Drop`
         // impl won't run and free the memory away from underneath us.
         unsafe {
             Self::new(
-                kvm_userspace_memory_region {
+                kvm_userspace_memory_region2 {
                     slot,
                     flags,
                     guest_phys_addr: region.start_addr().0,
                     memory_size: region.len(),
                     userspace_addr: region.as_ptr() as u64,
+                    guest_memfd,
+                    guest_memfd_offset,
+                    ..Default::default()
                 },
                 region.bitmap().clone(),
                 region.file_offset().cloned(),
             )
         }
     }
 
-    pub(crate) fn inner(&self) -> &kvm_userspace_memory_region {
+    pub(crate) fn inner(&self) -> &kvm_userspace_memory_region2 {
         &self.region
     }
 }
@@ -546,7 +558,7 @@ mod tests {
             regions
                 .into_iter()
                 .zip(0u32..) // assign dummy slots
-                .map(|(region, slot)| KvmRegion::from_mmap_region(region, slot))
+                .map(|(region, slot)| KvmRegion::from_mmap_region(region, slot, None))
                 .collect(),
         )
         .unwrap()

src/vmm/src/vstate/vm.rs

@@ -12,7 +12,7 @@ use std::os::fd::FromRawFd;
 use std::path::Path;
 use std::sync::Arc;
 
-use kvm_bindings::kvm_create_guest_memfd;
+use kvm_bindings::{kvm_create_guest_memfd, kvm_userspace_memory_region};
 use kvm_ioctls::{Cap, VmFd};
 use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
 use vmm_sys_util::eventfd::EventFd;
@@ -186,15 +186,35 @@ impl Vm {
             return Err(VmError::NotEnoughMemorySlots);
         }
 
-        Ok(KvmRegion::from_mmap_region(region, next_slot))
+        Ok(KvmRegion::from_mmap_region(region, next_slot, None))
     }
 
     fn register_kvm_region(&mut self, region: &KvmRegion) -> Result<(), VmError> {
-        // SAFETY: Safe because the fd is a valid KVM file descriptor.
-        unsafe {
-            self.fd()
-                .set_user_memory_region(*region.inner())
-                .map_err(VmError::SetUserMemoryRegion)?;
+        if self.fd().check_extension(Cap::UserMemory2) {
+            // SAFETY: Safe because the fd is a valid KVM file descriptor.
+            unsafe {
+                self.fd()
+                    .set_user_memory_region2(*region.inner())
+                    .map_err(VmError::SetUserMemoryRegion)?;
+            }
+        } else {
+            // Something is seriously wrong if we manage to set these fields on a host that doesn't
+            // even allow creation of guest_memfds!
+            assert_eq!(region.inner().guest_memfd, 0);
+            assert_eq!(region.inner().guest_memfd_offset, 0);
+
+            // SAFETY: We are passing a valid memory region and operate on a valid KVM FD.
+            unsafe {
+                self.fd()
+                    .set_user_memory_region(kvm_userspace_memory_region {
+                        slot: region.inner().slot,
+                        flags: region.inner().flags,
+                        guest_phys_addr: region.inner().guest_phys_addr,
+                        memory_size: region.inner().memory_size,
+                        userspace_addr: region.inner().userspace_addr,
+                    })
+                    .map_err(VmError::SetUserMemoryRegion)?;
+            }
         }
 
         Ok(())