feat(vmm): extend register_memory_regions with userfault bitmap

kalyazin · kalyazin · commit 8a7300e05114 · 2025-06-26T13:37:09.000+01:00
If configured, userfault bitmap is registered with KVM and controls
whether KVM will exit to userspace on a fault of the corresponding page.

We are going to allocate the bitmap in a memfd in Firecracker, set bits
for all pages to request notifications for vCPU faults and send
it to the UFFD handler to delegate clearing the bits as pages get
populated.

Since the KVM userfault patches are still in review,
set_user_memory_region2 is not aware of the userfault flag and the
userfault bitmap address in its input structure.  Define it in
Firecracker code temporarily.

Signed-off-by: Nikita Kalyazin &lt;kalyazin@amazon.com&gt;
diff --git a/src/vmm/src/builder.rs b/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;
diff --git a/src/vmm/src/device_manager/mmio.rs b/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();
diff --git a/src/vmm/src/vstate/vm.rs b/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,61 @@ impl Vm {
     pub fn register_memory_regions(
         &mut self,
         regions: Vec<GuestRegionMmap>,
+        mut userfault_bitmap: Option<&mut [u8]>,
     ) -> Result<(), VmError> {
         for region in regions {
-            self.register_memory_region(region)?
+            let bitmap_slice = if let Some(remaining) = userfault_bitmap {
+                let region_len = u64_to_usize(region.len());
+                // Firecracker does not allow sub-MB granularity when allocating guest memory
+                assert_eq!(region_len % (host_page_size() * u8::BITS as usize), 0);
+                let bitmap_len = region_len / host_page_size() / (u8::BITS as usize);
+                let (head, tail) = remaining.split_at_mut(bitmap_len);
+                userfault_bitmap = Some(tail);
+                Some(head)
+            } else {
+                None
+            };
+            self.register_memory_region(region, bitmap_slice)?
         }
-
         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
+        );
+
+        #[allow(clippy::undocumented_unsafe_blocks)]
+        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_bitmap: Option<&mut [u8]>,
+    ) -> 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 +282,14 @@ impl Vm {
             (0, 0)
         };
 
+        let userfault_bitmap = match userfault_bitmap {
+            Some(addr) => {
+                flags |= KVM_MEM_USERFAULT;
+                addr.as_ptr() as u64
+            }
+            None => 0,
+        };
+
         let memory_region = kvm_userspace_memory_region2 {
             slot: next_slot,
             guest_phys_addr: region.start_addr().raw_value(),
@@ -226,24 +298,22 @@ 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!
             assert_eq!(memory_region.guest_memfd, 0);
             assert_eq!(memory_region.guest_memfd_offset, 0);
+            assert_eq!(memory_region.userfault_bitmap, 0);
             assert_eq!(memory_region.flags & KVM_MEM_GUEST_MEMFD, 0);
+            assert_eq!(memory_region.flags & KVM_MEM_USERFAULT, 0);
 
             // SAFETY: We are passing a valid memory region and operate on a valid KVM FD.
             unsafe {
@@ -417,7 +487,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 +517,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 +559,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!(
