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

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 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 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 <[email protected]>

Author: kalyazin

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);
+}

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);
+            }
+        },
     );
 }