fix(examples/uffd): unregister range on UFFD Remove event

If the handler receives a UFFD remove event, it currently stores the PFN
and will reply with a zero page whenever it receives a pagefault event
for that page.
This works well with 4k pages, but zeropage is not supported on
hugepages. In order to support hugepages, let's just unregister from
UFFD whenever we get a remove event. By doing so, the handler won't
receive a notification for the removed page, and the VM will get a new
zero page from the kernel.

Signed-off-by: Riccardo Mancini <[email protected]>

Author: Manciukic

src/firecracker/examples/uffd/on_demand_handler.rs

@@ -87,7 +87,7 @@ fn main() {
                         }
                     }
                     userfaultfd::Event::Remove { start, end } => {
-                        uffd_handler.mark_range_removed(start as u64, end as u64)
+                        uffd_handler.unregister_range(start, end)
                     }
                     _ => panic!("Unexpected event on userfaultfd"),
                 }

src/firecracker/examples/uffd/uffd_utils.rs

@@ -9,7 +9,7 @@
     dead_code
 )]
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::ffi::c_void;
 use std::fs::File;
 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
@@ -54,7 +54,6 @@ pub struct UffdHandler {
     pub page_size: usize,
     backing_buffer: *const u8,
     uffd: Uffd,
-    removed_pages: HashSet<u64>,
 }
 
 impl UffdHandler {
@@ -125,32 +124,25 @@ impl UffdHandler {
             page_size,
             backing_buffer,
             uffd,
-            removed_pages: HashSet::new(),
         }
     }
 
     pub fn read_event(&mut self) -> Result<Option<Event>, Error> {
         self.uffd.read_event()
     }
 
-    pub fn mark_range_removed(&mut self, start: u64, end: u64) {
-        let pfn_start = start / self.page_size as u64;
-        let pfn_end = end / self.page_size as u64;
-
-        for pfn in pfn_start..pfn_end {
-            self.removed_pages.insert(pfn);
-        }
+    pub fn unregister_range(&mut self, start: *mut c_void, end: *mut c_void) {
+        // SAFETY: start and end are valid and provided by UFFD
+        let len = unsafe { end.offset_from_unsigned(start) };
+        self.uffd
+            .unregister(start, len)
+            .expect("range should be valid");
     }
 
     pub fn serve_pf(&mut self, addr: *mut u8, len: usize) -> bool {
         // Find the start of the page that the current faulting address belongs to.
         let dst = (addr as usize & !(self.page_size - 1)) as *mut libc::c_void;
         let fault_page_addr = dst as u64;
-        let fault_pfn = fault_page_addr / self.page_size as u64;
-
-        if self.removed_pages.contains(&fault_pfn) {
-            return self.zero_out(fault_page_addr);
-        }
 
         for region in self.mem_regions.iter() {
             if region.contains(fault_page_addr) {
@@ -193,14 +185,6 @@ impl UffdHandler {
 
         true
     }
-
-    fn zero_out(&mut self, addr: u64) -> bool {
-        match unsafe { self.uffd.zeropage(addr as *mut _, self.page_size, true) } {
-            Ok(_) => true,
-            Err(Error::ZeropageFailed(error)) if error as i32 == libc::EAGAIN => false,
-            r => panic!("Unexpected zeropage result: {:?}", r),
-        }
-    }
 }
 
 #[derive(Debug)]