test(uffd_utils): add handling for FaultRequest in secret freedom

kalyazin · kalyazin · commit 01e381aee75e · 2025-06-26T13:37:09.000+01:00
There are two ways a UFFD handler receives a fault notification if
Secret Fredom is enabled (which is inferred from 3 fds sent by
Firecracker instead of 1):
 - a VMM- or KVM-triggered fault is delivered via a minor UFFD fault
   event.  The handler is supposed to respond to it via memcpying the
   content of the page (if the page hasn't already been populated)
   followed by a UFFDIO_CONTINUE call.
 - a vCPU-triggered fault is delievered via a FaultRequest message on
   the UDS socket.  The handler is supposed to reply with a pwrite64
   call on the guest_memfd to populate the page followed by a FaultReply
   message on the UDS socket.

In both cases, the handler also needs to clear the bit in the userfault
bitmap at the corresponding offset in order to stop further fault
notifications for the same page.

UFFD handlers use the userfault bitmap for two purposes:
 - communicate to the kernel whether a fault at the corresponding
   guest_memfd offset will cause a VM exit
 - keep track of pages that have already been populated in order to
   avoid overwriting the content of the page that is already
   initialised.

Signed-off-by: Nikita Kalyazin &lt;kalyazin@amazon.com&gt;
diff --git a/src/firecracker/examples/uffd/fault_all_handler.rs b/src/firecracker/examples/uffd/fault_all_handler.rs
@@ -8,11 +8,14 @@
 mod uffd_utils;
 
 use std::fs::File;
+use std::os::fd::AsRawFd;
 use std::os::unix::net::UnixListener;
 
 use uffd_utils::{Runtime, UffdHandler};
 use utils::time::{ClockType, get_time_us};
 
+use crate::uffd_utils::uffd_continue;
+
 fn main() {
     let mut args = std::env::args();
     let uffd_sock_path = args.nth(1).expect("No socket path given");
@@ -37,19 +40,69 @@ fn main() {
                 .expect("Failed to read uffd_msg")
                 .expect("uffd_msg not ready");
 
-            match event {
-                userfaultfd::Event::Pagefault { .. } => {
-                    let start = get_time_us(ClockType::Monotonic);
-                    for region in uffd_handler.mem_regions.clone() {
-                        uffd_handler.serve_pf(region.base_host_virt_addr as _, region.size);
-                    }
-                    let end = get_time_us(ClockType::Monotonic);
+            if let userfaultfd::Event::Pagefault { addr, .. } = event {
+                let bit =
+                    uffd_handler.addr_to_offset(addr.cast()) as usize / uffd_handler.page_size;
+
+                // If Secret Free, we know if this is the first fault based on the userfault
+                // bitmap state. Otherwise, we assume that we will ever only receive a single fault
+                // event via UFFD.
+                let are_we_faulted_yet = uffd_handler
+                    .userfault_bitmap
+                    .as_mut()
+                    .map_or(false, |bitmap| !bitmap.is_bit_set(bit));
 
-                    println!("Finished Faulting All: {}us", end - start);
+                if are_we_faulted_yet {
+                    // TODO: we currently ignore the result as we may attempt to
+                    // populate the page that is already present as we may receive
+                    // multiple minor fault events per page.
+                    let _ = uffd_continue(
+                        uffd_handler.uffd.as_raw_fd(),
+                        addr as _,
+                        uffd_handler.page_size as u64,
+                    )
+                    .inspect_err(|err| println!("Error during uffdio_continue: {:?}", err));
+                } else {
+                    fault_all(uffd_handler, addr);
                 }
-                _ => panic!("Unexpected event on userfaultfd"),
             }
         },
         |_uffd_handler: &mut UffdHandler, _offset: usize| {},
     );
 }
+
+fn fault_all(uffd_handler: &mut UffdHandler, fault_addr: *mut libc::c_void) {
+    let start = get_time_us(ClockType::Monotonic);
+    for region in uffd_handler.mem_regions.clone() {
+        match uffd_handler.guest_memfd {
+            None => {
+                uffd_handler.serve_pf(region.base_host_virt_addr as _, region.size);
+            }
+            Some(_) => {
+                let written = uffd_handler.populate_via_write(region.offset as usize, region.size);
+
+                // This code is written under the assumption that the first fault triggered by
+                // Firecracker is either due to an MSR write (on x86) or due to device restoration
+                // reading from guest memory to check the virtio queues are sane (on
+                // ARM). This will be reported via a UFFD minor fault which needs to
+                // be handled via memcpy. Importantly, we get to the UFFD handler
+                // with the actual guest_memfd page already faulted in, meaning pwrite will stop
+                // once it gets to the offset of that page (e.g. written < region.size above).
+                // Thus, to fault in everything, we now need to skip this one page, write the
+                // remaining region, and then deal with the "gap" via uffd_handler.serve_pf().
+
+                if written < region.size - uffd_handler.page_size {
+                    let r = uffd_handler.populate_via_write(
+                        region.offset as usize + written + uffd_handler.page_size,
+                        region.size - written - uffd_handler.page_size,
+                    );
+                    assert_eq!(written + r, region.size - uffd_handler.page_size);
+                }
+            }
+        }
+    }
+    uffd_handler.serve_pf(fault_addr.cast(), uffd_handler.page_size);
+    let end = get_time_us(ClockType::Monotonic);
+
+    println!("Finished Faulting All: {}us", end - start);
+}
diff --git a/src/firecracker/examples/uffd/on_demand_handler.rs b/src/firecracker/examples/uffd/on_demand_handler.rs
@@ -8,10 +8,13 @@
 mod uffd_utils;
 
 use std::fs::File;
+use std::os::fd::AsRawFd;
 use std::os::unix::net::UnixListener;
 
 use uffd_utils::{Runtime, UffdHandler};
 
+use crate::uffd_utils::uffd_continue;
+
 fn main() {
     let mut args = std::env::args();
     let uffd_sock_path = args.nth(1).expect("No socket path given");
@@ -90,8 +93,36 @@ fn main() {
                     // event (if the balloon device is enabled).
                     match event {
                         userfaultfd::Event::Pagefault { addr, .. } => {
-                            if !uffd_handler.serve_pf(addr.cast(), uffd_handler.page_size) {
-                                deferred_events.push(event);
+                            let bit = uffd_handler.addr_to_offset(addr.cast()) as usize
+                                / uffd_handler.page_size;
+
+                            if uffd_handler.userfault_bitmap.is_some() {
+                                if uffd_handler
+                                    .userfault_bitmap
+                                    .as_mut()
+                                    .unwrap()
+                                    .is_bit_set(bit)
+                                {
+                                    if !uffd_handler.serve_pf(addr.cast(), uffd_handler.page_size) {
+                                        deferred_events.push(event);
+                                    }
+                                } else {
+                                    // TODO: we currently ignore the result as we may attempt to
+                                    // populate the page that is already present as we may receive
+                                    // multiple minor fault events per page.
+                                    let _ = uffd_continue(
+                                        uffd_handler.uffd.as_raw_fd(),
+                                        addr as _,
+                                        uffd_handler.page_size as u64,
+                                    )
+                                    .inspect_err(|err| {
+                                        println!("uffdio_continue error: {:?}", err)
+                                    });
+                                }
+                            } else {
+                                if !uffd_handler.serve_pf(addr.cast(), uffd_handler.page_size) {
+                                    deferred_events.push(event);
+                                }
                             }
                         }
                         userfaultfd::Event::Remove { start, end } => {
@@ -111,6 +142,17 @@ fn main() {
                 }
             }
         },
-        |_uffd_handler: &mut UffdHandler, _offset: usize| {},
+        |uffd_handler: &mut UffdHandler, offset: usize| {
+            let bytes_written = uffd_handler.populate_via_write(offset, uffd_handler.page_size);
+
+            if bytes_written == 0 {
+                println!(
+                    "got a vcpu fault for an already populated page at offset {}",
+                    offset
+                );
+            } else {
+                assert_eq!(bytes_written, uffd_handler.page_size);
+            }
+        },
     );
 }
diff --git a/src/firecracker/examples/uffd/uffd_utils.rs b/src/firecracker/examples/uffd/uffd_utils.rs
@@ -21,10 +21,47 @@ use std::time::Duration;
 use serde::{Deserialize, Serialize};
 use serde_json::{Deserializer, StreamDeserializer};
 use userfaultfd::{Error, Event, Uffd};
+use vmm_sys_util::ioctl::ioctl_with_mut_ref;
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
+use vmm_sys_util::{ioctl_ioc_nr, ioctl_iowr_nr};
 
 use crate::uffd_utils::userfault_bitmap::UserfaultBitmap;
 
+// TODO: remove when UFFDIO_CONTINUE for guest_memfd is available in the crate
+#[repr(C)]
+struct uffdio_continue {
+    range: uffdio_range,
+    mode: u64,
+    mapped: u64,
+}
+
+ioctl_iowr_nr!(UFFDIO_CONTINUE, 0xAA, 0x7, uffdio_continue);
+
+#[repr(C)]
+struct uffdio_range {
+    start: u64,
+    len: u64,
+}
+
+pub fn uffd_continue(uffd: RawFd, fault_addr: u64, len: u64) -> std::io::Result<()> {
+    let mut cont = uffdio_continue {
+        range: uffdio_range {
+            start: fault_addr,
+            len,
+        },
+        mode: 0, // Normal continuation mode
+        mapped: 0,
+    };
+
+    let ret = unsafe { ioctl_with_mut_ref(&uffd, UFFDIO_CONTINUE(), &mut cont) };
+
+    if ret == -1 {
+        return Err(std::io::Error::last_os_error());
+    }
+
+    Ok(())
+}
+
 // This is the same with the one used in src/vmm.
 /// This describes the mapping between Firecracker base virtual address and offset in the
 /// buffer or file backend for a guest memory region. It is used to tell an external
@@ -117,7 +154,7 @@ pub struct UffdHandler {
     pub mem_regions: Vec<GuestRegionUffdMapping>,
     pub page_size: usize,
     backing_buffer: *const u8,
-    uffd: Uffd,
+    pub uffd: Uffd,
     removed_pages: HashSet<u64>,
     pub guest_memfd: Option<File>,
     pub guest_memfd_addr: Option<*mut u8>,
@@ -261,6 +298,20 @@ impl UffdHandler {
         }
     }
 
+    pub fn addr_to_offset(&self, addr: *mut u8) -> u64 {
+        let addr = addr as u64;
+        for region in &self.mem_regions {
+            if region.contains(addr) {
+                return addr - region.base_host_virt_addr + region.offset as u64;
+            }
+        }
+
+        panic!(
+            "Could not find addr: {:#x} within guest region mappings.",
+            addr
+        );
+    }
+
     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;
@@ -273,7 +324,7 @@ impl UffdHandler {
         } else {
             for region in self.mem_regions.iter() {
                 if region.contains(fault_page_addr) {
-                    return self.populate_from_file(region, fault_page_addr, len);
+                    return self.populate_from_file(&region.clone(), fault_page_addr, len);
                 }
             }
         }
@@ -288,12 +339,61 @@ impl UffdHandler {
         self.mem_regions.iter().map(|r| r.size).sum()
     }
 
-    fn populate_from_file(&self, region: &GuestRegionUffdMapping, dst: u64, len: usize) -> bool {
-        let offset = dst - region.base_host_virt_addr;
-        let src = self.backing_buffer as u64 + region.offset + offset;
+    pub fn populate_via_write(&mut self, offset: usize, len: usize) -> usize {
+        // man 2 write:
+        //
+        //    On Linux, write() (and similar system calls) will transfer at most
+        //    0x7ffff000 (2,147,479,552) bytes, returning the number of bytes
+        //    actually transferred.  (This is true on both 32-bit and 64-bit
+        //    systems.)
+        const MAX_WRITE_LEN: usize = 2_147_479_552;
+
+        assert!(
+            offset.checked_add(len).unwrap() <= self.size(),
+            "{} + {} >= {}",
+            offset,
+            len,
+            self.size()
+        );
+
+        let mut total_written = 0;
+
+        while total_written < len {
+            let src = unsafe { self.backing_buffer.add(offset + total_written) };
+            let len_to_write = (len - total_written).min(MAX_WRITE_LEN);
+            let bytes_written = unsafe {
+                libc::pwrite64(
+                    self.guest_memfd.as_ref().unwrap().as_raw_fd(),
+                    src.cast(),
+                    len_to_write,
+                    (offset + total_written) as libc::off64_t,
+                )
+            };
+
+            let bytes_written = match bytes_written {
+                -1 if vmm_sys_util::errno::Error::last().errno() == libc::ENOSPC => 0,
+                written @ 0.. => written as usize,
+                _ => panic!("{:?}", std::io::Error::last_os_error()),
+            };
+
+            self.userfault_bitmap
+                .as_mut()
+                .unwrap()
+                .reset_addr_range(offset + total_written, bytes_written);
+
+            total_written += bytes_written;
+
+            if bytes_written != len_to_write {
+                break;
+            }
+        }
+
+        total_written
+    }
 
+    fn populate_via_uffdio_copy(&self, src: *const u8, dst: u64, len: usize) -> bool {
         unsafe {
-            match self.uffd.copy(src as *const _, dst as *mut _, len, true) {
+            match self.uffd.copy(src.cast(), dst as *mut _, len, true) {
                 // Make sure the UFFD copied some bytes.
                 Ok(value) => assert!(value > 0),
                 // Catch EAGAIN errors, which occur when a `remove` event lands in the UFFD
@@ -318,6 +418,42 @@ impl UffdHandler {
         true
     }
 
+    fn populate_via_memcpy(&mut self, src: *const u8, dst: u64, offset: usize, len: usize) -> bool {
+        let dst_memcpy = unsafe {
+            self.guest_memfd_addr
+                .expect("no guest_memfd addr")
+                .add(offset)
+        };
+
+        unsafe {
+            std::ptr::copy_nonoverlapping(src, dst_memcpy, len);
+        }
+
+        self.userfault_bitmap
+            .as_mut()
+            .unwrap()
+            .reset_addr_range(offset, len);
+
+        uffd_continue(self.uffd.as_raw_fd(), dst, len as u64).expect("uffd_continue");
+
+        true
+    }
+
+    fn populate_from_file(
+        &mut self,
+        region: &GuestRegionUffdMapping,
+        dst: u64,
+        len: usize,
+    ) -> bool {
+        let offset = (region.offset + dst - region.base_host_virt_addr) as usize;
+        let src = unsafe { self.backing_buffer.add(offset) };
+
+        match self.guest_memfd {
+            Some(_) => self.populate_via_memcpy(src, dst, offset, len),
+            None => self.populate_via_uffdio_copy(src, dst, len),
+        }
+    }
+
     fn zero_out(&mut self, addr: u64) {
         let ret = unsafe {
             self.uffd
@@ -619,7 +755,7 @@ mod tests {
             let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
             // Update runtime with actual runtime
             let runtime = uninit_runtime.write(Runtime::new(stream, file));
-            runtime.run(|_: &mut UffdHandler| {});
+            runtime.run(|_: &mut UffdHandler| {}, |_: &mut UffdHandler, _: usize| {});
         });
 
         // wait for runtime thread to initialize itself
