feat(vmm): extend register_memory_regions with userfault bitmap

If configured, userfault bitmap is registered with KVM and controls
whether KVM will exit to userspace on a fault of the corresponding page.

Signed-off-by: Nikita Kalyazin <[email protected]>

Author: kalyazin

src/vmm/src/builder.rs

@@ -260,7 +260,7 @@ pub fn build_microvm_for_boot(
         .map_err(StartMicrovmError::GuestMemory)?;
 
     vmm.vm
-        .register_memory_regions(guest_memory)
+        .register_memory_regions(guest_memory, None)
         .map_err(VmmError::Vm)?;
 
     #[cfg(target_arch = "x86_64")]
@@ -487,7 +487,7 @@ pub fn build_microvm_from_snapshot(
     .map_err(StartMicrovmError::Internal)?;
 
     vmm.vm
-        .register_memory_regions(guest_memory)
+        .register_memory_regions(guest_memory, None)
         .map_err(VmmError::Vm)
         .map_err(StartMicrovmError::Internal)?;
     vmm.uffd = uffd;

src/vmm/src/device_manager/mmio.rs

@@ -660,7 +660,7 @@ mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm, false).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem, None).unwrap();
         let mut device_manager = MMIODeviceManager::new();
         let mut resource_allocator = ResourceAllocator::new().unwrap();
 
@@ -691,7 +691,7 @@ mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm, false).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem, None).unwrap();
         let mut device_manager = MMIODeviceManager::new();
         let mut resource_allocator = ResourceAllocator::new().unwrap();
 
@@ -747,7 +747,7 @@ mod tests {
         let guest_mem = multi_region_mem_raw(&[(start_addr1, 0x1000), (start_addr2, 0x1000)]);
         let kvm = Kvm::new(vec![]).expect("Cannot create Kvm");
         let mut vm = Vm::new(&kvm, false).unwrap();
-        vm.register_memory_regions(guest_mem).unwrap();
+        vm.register_memory_regions(guest_mem, None).unwrap();
 
         #[cfg(target_arch = "x86_64")]
         vm.setup_irqchip().unwrap();

src/vmm/src/vstate/vm.rs

@@ -13,12 +13,13 @@ use std::path::Path;
 use std::sync::Arc;
 
 use kvm_bindings::{
-    KVM_MEM_GUEST_MEMFD, KVM_MEM_LOG_DIRTY_PAGES, KVM_MEMORY_ATTRIBUTE_PRIVATE,
+    KVM_MEM_GUEST_MEMFD, KVM_MEM_LOG_DIRTY_PAGES, KVM_MEMORY_ATTRIBUTE_PRIVATE, KVMIO,
     kvm_create_guest_memfd, kvm_memory_attributes, kvm_userspace_memory_region,
-    kvm_userspace_memory_region2,
 };
 use kvm_ioctls::{Cap, VmFd};
 use vmm_sys_util::eventfd::EventFd;
+use vmm_sys_util::ioctl::ioctl_with_ref;
+use vmm_sys_util::{ioctl_ioc_nr, ioctl_iow_nr};
 
 pub use crate::arch::{ArchVm as Vm, ArchVmError, VmState};
 use crate::arch::{VM_TYPE_FOR_SECRET_FREEDOM, host_page_size};
@@ -73,6 +74,24 @@ pub enum VmError {
     SetMemoryAttributes(kvm_ioctls::Error),
 }
 
+// Upstream `kvm_userspace_memory_region2` definition does not include `userfault_bitmap` field yet.
+// TODO: revert to `kvm_userspace_memory_region2` from kvm-bindings
+#[allow(non_camel_case_types)]
+#[repr(C)]
+#[derive(Debug, Default, Copy, Clone, PartialEq)]
+struct kvm_userspace_memory_region2 {
+    slot: u32,
+    flags: u32,
+    guest_phys_addr: u64,
+    memory_size: u64,
+    userspace_addr: u64,
+    guest_memfd_offset: u64,
+    guest_memfd: u32,
+    pad1: u32,
+    userfault_bitmap: u64,
+    pad2: [u64; 13],
+}
+
 /// Contains Vm functions that are usable across CPU architectures
 impl Vm {
     /// Create a KVM VM
@@ -181,16 +200,75 @@ impl Vm {
     pub fn register_memory_regions(
         &mut self,
         regions: Vec<GuestRegionMmap>,
+        userfault_bitmap_memfd: Option<&File>,
     ) -> Result<(), VmError> {
+        let addr = match userfault_bitmap_memfd {
+            Some(file) => {
+                let addr = unsafe {
+                    libc::mmap(
+                        std::ptr::null_mut(),
+                        file.metadata().unwrap().len() as usize,
+                        libc::PROT_WRITE,
+                        libc::MAP_SHARED,
+                        file.as_raw_fd(),
+                        0,
+                    )
+                };
+
+                if addr == libc::MAP_FAILED {
+                    panic!(
+                        "Failed to mmap userfault bitmap file: {}",
+                        std::io::Error::last_os_error()
+                    );
+                }
+
+                Some(addr as u64)
+            }
+            None => None,
+        };
+
         for region in regions {
-            self.register_memory_region(region)?
+            self.register_memory_region(region, addr)?
         }
 
         Ok(())
     }
 
+    // TODO: remove when userfault support is merged upstream
+    fn set_user_memory_region2(
+        &self,
+        user_memory_region2: kvm_userspace_memory_region2,
+    ) -> Result<(), VmError> {
+        ioctl_iow_nr!(
+            KVM_SET_USER_MEMORY_REGION2,
+            KVMIO,
+            0x49,
+            kvm_userspace_memory_region2
+        );
+
+        let ret = unsafe {
+            ioctl_with_ref(
+                self.fd(),
+                KVM_SET_USER_MEMORY_REGION2(),
+                &user_memory_region2,
+            )
+        };
+        if ret == 0 {
+            Ok(())
+        } else {
+            Err(VmError::SetUserMemoryRegion(kvm_ioctls::Error::last()))
+        }
+    }
+
     /// Register a new memory region to this [`Vm`].
-    pub fn register_memory_region(&mut self, region: GuestRegionMmap) -> Result<(), VmError> {
+    pub fn register_memory_region(
+        &mut self,
+        region: GuestRegionMmap,
+        userfault_addr: Option<u64>,
+    ) -> Result<(), VmError> {
+        // TODO: take it from kvm-bindings when merged upstream
+        const KVM_MEM_USERFAULT: u32 = 1 << 3;
+
         let next_slot = self
             .guest_memory()
             .num_regions()
@@ -218,6 +296,18 @@ impl Vm {
             (0, 0)
         };
 
+        let userfault_bitmap = match userfault_addr {
+            Some(addr) => {
+                flags |= KVM_MEM_USERFAULT;
+
+                let file_offset_start = region.file_offset().unwrap().start();
+                let pages_offset = file_offset_start / (host_page_size() as u64);
+                let bytes_offset = pages_offset / (u8::BITS as u64);
+                addr + bytes_offset
+            }
+            None => 0,
+        };
+
         let memory_region = kvm_userspace_memory_region2 {
             slot: next_slot,
             guest_phys_addr: region.start_addr().raw_value(),
@@ -226,18 +316,14 @@ impl Vm {
             flags,
             guest_memfd,
             guest_memfd_offset,
+            userfault_bitmap,
             ..Default::default()
         };
 
         let new_guest_memory = self.common.guest_memory.insert_region(Arc::new(region))?;
 
         if self.fd().check_extension(Cap::UserMemory2) {
-            // SAFETY: We are passing a valid memory region and operate on a valid KVM FD.
-            unsafe {
-                self.fd()
-                    .set_user_memory_region2(memory_region)
-                    .map_err(VmError::SetUserMemoryRegion)?;
-            }
+            self.set_user_memory_region2(memory_region)?;
         } else {
             // Something is seriously wrong if we manage to set these fields on a host that doesn't
             // even allow creation of guest_memfds!
@@ -417,7 +503,7 @@ pub(crate) mod tests {
     pub(crate) fn setup_vm_with_memory(mem_size: usize) -> (Kvm, Vm) {
         let (kvm, mut vm) = setup_vm();
         let gm = single_region_mem_raw(mem_size);
-        vm.register_memory_regions(gm).unwrap();
+        vm.register_memory_regions(gm, None).unwrap();
         (kvm, vm)
     }
 
@@ -447,14 +533,14 @@ pub(crate) mod tests {
         // Trying to set a memory region with a size that is not a multiple of GUEST_PAGE_SIZE
         // will result in error.
         let gm = single_region_mem_raw(0x10);
-        let res = vm.register_memory_regions(gm);
+        let res = vm.register_memory_regions(gm, None);
         assert_eq!(
             res.unwrap_err().to_string(),
             "Cannot set the memory regions: Invalid argument (os error 22)"
         );
 
         let gm = single_region_mem_raw(0x1000);
-        let res = vm.register_memory_regions(gm);
+        let res = vm.register_memory_regions(gm, None);
         res.unwrap();
     }
 
@@ -489,7 +575,7 @@ pub(crate) mod tests {
 
             let region = GuestRegionMmap::new(region, GuestAddress(i as u64 * 0x1000)).unwrap();
 
-            let res = vm.register_memory_region(region);
+            let res = vm.register_memory_region(region, None);
 
             if i >= max_nr_regions {
                 assert!(