refactor(uffd): update UffdPfHandler

Rename `UffdPfHandler` -> `UffdHandler`
Replace direct access to `uffd` field with
separate method.
Rearrange couple functions/struct positions in
the file.

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

Author: ShadowCurse

src/firecracker/examples/uffd/malicious_handler.rs

@@ -9,7 +9,7 @@ mod uffd_utils;
 use std::fs::File;
 use std::os::unix::net::UnixListener;
 
-use uffd_utils::{Runtime, UffdPfHandler};
+use uffd_utils::{Runtime, UffdHandler};
 
 fn main() {
     let mut args = std::env::args();
@@ -23,10 +23,9 @@ fn main() {
     let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
     let mut runtime = Runtime::new(stream, file);
-    runtime.run(|uffd_handler: &mut UffdPfHandler| {
+    runtime.run(|uffd_handler: &mut UffdHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
-            .uffd
             .read_event()
             .expect("Failed to read uffd_msg")
             .expect("uffd_msg not ready");

src/firecracker/examples/uffd/uffd_utils.rs

@@ -11,7 +11,7 @@ use std::os::unix::net::UnixStream;
 use std::ptr;
 
 use serde::Deserialize;
-use userfaultfd::Uffd;
+use userfaultfd::{Error, Event, Uffd};
 use utils::get_page_size;
 use utils::sock_ctrl_msg::ScmSocket;
 
@@ -33,28 +33,28 @@ pub struct GuestRegionUffdMapping {
     pub offset: u64,
 }
 
+#[derive(Debug, Clone)]
+pub enum MemPageState {
+    Uninitialized,
+    FromFile,
+    Removed,
+    Anonymous,
+}
+
 #[derive(Debug)]
 struct MemRegion {
     mapping: GuestRegionUffdMapping,
     page_states: HashMap<u64, MemPageState>,
 }
 
 #[derive(Debug)]
-pub struct UffdPfHandler {
+pub struct UffdHandler {
     mem_regions: Vec<MemRegion>,
     backing_buffer: *const u8,
-    pub uffd: Uffd,
-}
-
-#[derive(Debug, Clone)]
-pub enum MemPageState {
-    Uninitialized,
-    FromFile,
-    Removed,
-    Anonymous,
+    uffd: Uffd,
 }
 
-impl UffdPfHandler {
+impl UffdHandler {
     pub fn from_unix_stream(stream: &UnixStream, backing_buffer: *const u8, size: usize) -> Self {
         let mut message_buf = vec![0u8; 1024];
         let (bytes_read, file) = stream
@@ -83,6 +83,10 @@ impl UffdPfHandler {
         }
     }
 
+    pub fn read_event(&mut self) -> Result<Option<Event>, Error> {
+        self.uffd.read_event()
+    }
+
     pub fn update_mem_state_mappings(&mut self, start: u64, end: u64, state: &MemPageState) {
         for region in self.mem_regions.iter_mut() {
             for (key, value) in region.page_states.iter_mut() {
@@ -93,36 +97,6 @@ impl UffdPfHandler {
         }
     }
 
-    fn populate_from_file(&self, region: &MemRegion, dst: u64, len: usize) -> (u64, u64) {
-        let offset = dst - region.mapping.base_host_virt_addr;
-        let src = self.backing_buffer as u64 + region.mapping.offset + offset;
-
-        let ret = unsafe {
-            self.uffd
-                .copy(src as *const _, dst as *mut _, len, true)
-                .expect("Uffd copy failed")
-        };
-
-        // Make sure the UFFD copied some bytes.
-        assert!(ret > 0);
-
-        (dst, dst + len as u64)
-    }
-
-    fn zero_out(&mut self, addr: u64) -> (u64, u64) {
-        let page_size = get_page_size().unwrap();
-
-        let ret = unsafe {
-            self.uffd
-                .zeropage(addr as *mut _, page_size, true)
-                .expect("Uffd zeropage failed")
-        };
-        // Make sure the UFFD zeroed out some bytes.
-        assert!(ret > 0);
-
-        (addr, addr + page_size as u64)
-    }
-
     pub fn serve_pf(&mut self, addr: *mut u8, len: usize) {
         let page_size = get_page_size().unwrap();
 
@@ -160,6 +134,36 @@ impl UffdPfHandler {
             addr
         );
     }
+
+    fn populate_from_file(&self, region: &MemRegion, dst: u64, len: usize) -> (u64, u64) {
+        let offset = dst - region.mapping.base_host_virt_addr;
+        let src = self.backing_buffer as u64 + region.mapping.offset + offset;
+
+        let ret = unsafe {
+            self.uffd
+                .copy(src as *const _, dst as *mut _, len, true)
+                .expect("Uffd copy failed")
+        };
+
+        // Make sure the UFFD copied some bytes.
+        assert!(ret > 0);
+
+        (dst, dst + len as u64)
+    }
+
+    fn zero_out(&mut self, addr: u64) -> (u64, u64) {
+        let page_size = get_page_size().unwrap();
+
+        let ret = unsafe {
+            self.uffd
+                .zeropage(addr as *mut _, page_size, true)
+                .expect("Uffd zeropage failed")
+        };
+        // Make sure the UFFD zeroed out some bytes.
+        assert!(ret > 0);
+
+        (addr, addr + page_size as u64)
+    }
 }
 
 #[derive(Debug)]
@@ -168,7 +172,7 @@ pub struct Runtime {
     backing_file: File,
     backing_memory: *mut u8,
     backing_memory_size: usize,
-    uffds: HashMap<i32, UffdPfHandler>,
+    uffds: HashMap<i32, UffdHandler>,
 }
 
 impl Runtime {
@@ -207,7 +211,7 @@ impl Runtime {
     /// 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)) {
+    pub fn run(&mut self, pf_event_dispatch: impl Fn(&mut UffdHandler)) {
         let mut pollfds = vec![];
 
         // Poll the stream for incoming uffds
@@ -240,7 +244,7 @@ impl Runtime {
                     nready -= 1;
                     if pollfds[i].fd == self.stream.as_raw_fd() {
                         // Handle new uffd from stream
-                        let handler = UffdPfHandler::from_unix_stream(
+                        let handler = UffdHandler::from_unix_stream(
                             &self.stream,
                             self.backing_memory,
                             self.backing_memory_size,

src/firecracker/examples/uffd/valid_handler.rs

@@ -10,7 +10,7 @@ mod uffd_utils;
 use std::fs::File;
 use std::os::unix::net::UnixListener;
 
-use uffd_utils::{MemPageState, Runtime, UffdPfHandler};
+use uffd_utils::{MemPageState, Runtime, UffdHandler};
 use utils::get_page_size;
 
 fn main() {
@@ -30,10 +30,9 @@ fn main() {
     let len = get_page_size().unwrap();
 
     let mut runtime = Runtime::new(stream, file);
-    runtime.run(|uffd_handler: &mut UffdPfHandler| {
+    runtime.run(|uffd_handler: &mut UffdHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
-            .uffd
             .read_event()
             .expect("Failed to read uffd_msg")
             .expect("uffd_msg not ready");