feat(uffd): new Runtime

The new `Runtime` struct can contain multiple
uffd objects received from the UDS socket during
its runtime.

This is done by polling the socket and uffd
file descriptors. When the socket event is
polled, a new uffd object is retrieved. When
the uffd fd is polled, a page fault for
corresponding uffd is handled.

The behaviour for page faults can be configured
with a parameter passed to the `Runtime::run`.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

src/firecracker/examples/uffd/malicious_handler.rs

@@ -6,37 +6,33 @@
 
 mod uffd_utils;
 
-use std::os::unix::io::AsRawFd;
+use std::fs::File;
+use std::os::unix::net::UnixListener;
 
-use uffd_utils::create_pf_handler;
+use uffd_utils::{Runtime, UffdPfHandler};
 
 fn main() {
-    let uffd_handler = create_pf_handler();
+    let mut args = std::env::args();
+    let uffd_sock_path = args.nth(1).expect("No socket path given");
+    let mem_file_path = args.next().expect("No memory file given");
 
-    let mut pollfd = libc::pollfd {
-        fd: uffd_handler.uffd.as_raw_fd(),
-        events: libc::POLLIN,
-        revents: 0,
-    };
+    let file = File::open(mem_file_path).expect("Cannot open memfile");
 
-    // Loop, handling incoming events on the userfaultfd file descriptor.
-    loop {
-        let nready = unsafe { libc::poll(&mut pollfd, 1, -1) };
+    // Get Uffd from UDS. We'll use the uffd to handle PFs for Firecracker.
+    let listener = UnixListener::bind(uffd_sock_path).expect("Cannot bind to socket path");
+    let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
-        if nready == -1 {
-            panic!("Could not poll for events!")
-        }
+    let mut runtime = Runtime::new(stream, file);
+    runtime.run(|uffd_handler: &mut UffdPfHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
             .uffd
             .read_event()
             .expect("Failed to read uffd_msg")
             .expect("uffd_msg not ready");
 
-        // We expect to receive either a Page Fault or Removed
-        // event (if the balloon device is enabled).
         if let userfaultfd::Event::Pagefault { .. } = event {
             panic!("Fear me! I am the malicious page fault handler.")
         }
-    }
+    });
 }

src/firecracker/examples/uffd/uffd_utils.rs

@@ -7,7 +7,7 @@
 use std::collections::HashMap;
 use std::fs::File;
 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
-use std::os::unix::net::{UnixListener, UnixStream};
+use std::os::unix::net::UnixStream;
 use std::ptr;
 
 use serde::Deserialize;
@@ -162,6 +162,110 @@ impl UffdPfHandler {
     }
 }
 
+#[derive(Debug)]
+pub struct Runtime {
+    stream: UnixStream,
+    backing_file: File,
+    backing_memory: *mut u8,
+    backing_memory_size: usize,
+    uffds: HashMap<i32, UffdPfHandler>,
+}
+
+impl Runtime {
+    pub fn new(stream: UnixStream, backing_file: File) -> Self {
+        let file_meta = backing_file
+            .metadata()
+            .expect("can not get backing file metadata");
+        let backing_memory_size = file_meta.len() as usize;
+        // # Safety:
+        // File size and fd are valid
+        let ret = unsafe {
+            libc::mmap(
+                ptr::null_mut(),
+                backing_memory_size,
+                libc::PROT_READ,
+                libc::MAP_PRIVATE,
+                backing_file.as_raw_fd(),
+                0,
+            )
+        };
+        if ret == libc::MAP_FAILED {
+            panic!("mmap on backing file failed");
+        }
+
+        Self {
+            stream,
+            backing_file,
+            backing_memory: ret.cast(),
+            backing_memory_size,
+            uffds: HashMap::default(),
+        }
+    }
+
+    /// Polls the `UnixStream` and UFFD fds in a loop.
+    /// When stream is polled, new uffd is retrieved.
+    /// When uffd is polled, page fault is handled by
+    /// calling `pf_event_dispatch` with corresponding
+    /// uffd object passed in.
+    pub fn run(&mut self, pf_event_dispatch: impl Fn(&mut UffdPfHandler)) {
+        let mut pollfds = vec![];
+
+        // Poll the stream for incoming uffds
+        pollfds.push(libc::pollfd {
+            fd: self.stream.as_raw_fd(),
+            events: libc::POLLIN,
+            revents: 0,
+        });
+
+        // We can skip polling on stream fd if
+        // the connection is closed.
+        let mut skip_stream: usize = 0;
+        loop {
+            let pollfd_ptr = pollfds[skip_stream..].as_mut_ptr();
+            let pollfd_size = pollfds[skip_stream..].len() as u64;
+
+            // # Safety:
+            // Pollfds vector is valid
+            let mut nready = unsafe { libc::poll(pollfd_ptr, pollfd_size, -1) };
+
+            if nready == -1 {
+                panic!("Could not poll for events!")
+            }
+
+            for i in skip_stream..pollfds.len() {
+                if nready == 0 {
+                    break;
+                }
+                if pollfds[i].revents & libc::POLLIN != 0 {
+                    nready -= 1;
+                    if pollfds[i].fd == self.stream.as_raw_fd() {
+                        // Handle new uffd from stream
+                        let handler = UffdPfHandler::from_unix_stream(
+                            &self.stream,
+                            self.backing_memory,
+                            self.backing_memory_size,
+                        );
+                        pollfds.push(libc::pollfd {
+                            fd: handler.uffd.as_raw_fd(),
+                            events: libc::POLLIN,
+                            revents: 0,
+                        });
+                        self.uffds.insert(handler.uffd.as_raw_fd(), handler);
+
+                        // If connection is closed, we can skip the socket from being polled.
+                        if pollfds[i].revents & (libc::POLLRDHUP | libc::POLLHUP) != 0 {
+                            skip_stream = 1;
+                        }
+                    } else {
+                        // Handle one of uffd page faults
+                        pf_event_dispatch(self.uffds.get_mut(&pollfds[i].fd).unwrap());
+                    }
+                }
+            }
+        }
+    }
+}
+
 fn create_mem_regions(mappings: &Vec<GuestRegionUffdMapping>) -> Vec<MemRegion> {
     let page_size = get_page_size().unwrap();
     let mut mem_regions: Vec<MemRegion> = Vec::with_capacity(mappings.len());
@@ -184,34 +288,3 @@ fn create_mem_regions(mappings: &Vec<GuestRegionUffdMapping>) -> Vec<MemRegion>
 
     mem_regions
 }
-
-pub fn create_pf_handler() -> UffdPfHandler {
-    let uffd_sock_path = std::env::args().nth(1).expect("No socket path given");
-    let mem_file_path = std::env::args().nth(2).expect("No memory file given");
-
-    let file = File::open(mem_file_path).expect("Cannot open memfile");
-    let size = file.metadata().unwrap().len() as usize;
-
-    // mmap a memory area used to bring in the faulting regions.
-    let ret = unsafe {
-        libc::mmap(
-            ptr::null_mut(),
-            size,
-            libc::PROT_READ,
-            libc::MAP_PRIVATE,
-            file.as_raw_fd(),
-            0,
-        )
-    };
-    if ret == libc::MAP_FAILED {
-        panic!("mmap failed");
-    }
-    let memfile_buffer = ret as *const u8;
-
-    // Get Uffd from UDS. We'll use the uffd to handle PFs for Firecracker.
-    let listener = UnixListener::bind(uffd_sock_path).expect("Cannot bind to socket path");
-
-    let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
-
-    UffdPfHandler::from_unix_stream(&stream, memfile_buffer, size)
-}

src/firecracker/examples/uffd/valid_handler.rs

@@ -7,42 +7,30 @@
 
 mod uffd_utils;
 
-use std::os::unix::io::AsRawFd;
+use std::fs::File;
+use std::os::unix::net::UnixListener;
 
-use uffd_utils::{create_pf_handler, MemPageState};
+use uffd_utils::{MemPageState, Runtime, UffdPfHandler};
 use utils::get_page_size;
 
 fn main() {
-    let mut uffd_handler = create_pf_handler();
+    let mut args = std::env::args();
+    let uffd_sock_path = args.nth(1).expect("No socket path given");
+    let mem_file_path = args.next().expect("No memory file given");
+
+    let file = File::open(mem_file_path).expect("Cannot open memfile");
+
+    // Get Uffd from UDS. We'll use the uffd to handle PFs for Firecracker.
+    let listener = UnixListener::bind(uffd_sock_path).expect("Cannot bind to socket path");
+    let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
     // Populate a single page from backing memory file.
     // This is just an example, probably, with the worst-case latency scenario,
     // of how memory can be loaded in guest RAM.
     let len = get_page_size().unwrap();
 
-    let mut pollfd = libc::pollfd {
-        fd: uffd_handler.uffd.as_raw_fd(),
-        events: libc::POLLIN,
-        revents: 0,
-    };
-    // Loop, handling incoming events on the userfaultfd file descriptor.
-    loop {
-        // See what poll() tells us about the userfaultfd.
-        let nready = unsafe { libc::poll(&mut pollfd, 1, -1) };
-
-        if nready == -1 {
-            panic!("Could not poll for events!")
-        }
-
-        let revents = pollfd.revents;
-
-        println!(
-            "poll() returns: nready = {}; POLLIN = {}; POLLERR = {}",
-            nready,
-            revents & libc::POLLIN,
-            revents & libc::POLLERR,
-        );
-
+    let mut runtime = Runtime::new(stream, file);
+    runtime.run(|uffd_handler: &mut UffdPfHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
             .uffd
@@ -53,15 +41,13 @@ fn main() {
         // We expect to receive either a Page Fault or Removed
         // event (if the balloon device is enabled).
         match event {
-            userfaultfd::Event::Pagefault { addr, .. } => {
-                uffd_handler.serve_pf(addr as *mut u8, len)
-            }
+            userfaultfd::Event::Pagefault { addr, .. } => uffd_handler.serve_pf(addr.cast(), len),
             userfaultfd::Event::Remove { start, end } => uffd_handler.update_mem_state_mappings(
-                start as *mut u8 as u64,
-                end as *mut u8 as u64,
+                start as u64,
+                end as u64,
                 &MemPageState::Removed,
             ),
             _ => panic!("Unexpected event on userfaultfd"),
         }
-    }
+    });
 }