[WIP] fix(vmm): Do not store UFFD handle in VMM

src/vmm/src/builder.rs

@@ -198,7 +198,6 @@ fn create_vmm_and_vcpus(
         shutdown_exit_code: None,
         kvm,
         vm,
-        uffd: None,
         uffd_socket: None,
         vcpus_handles: Vec::new(),
         vcpus_exit_evt,
@@ -582,7 +581,7 @@ pub fn build_microvm_from_snapshot(
     let mem_state = &microvm_state.vm_state.memory;
     let track_dirty_pages = params.track_dirty_pages;
 
-    let (guest_memory, uffd, socket) = match params.mem_backend.backend_type {
+    let (guest_memory, socket) = match params.mem_backend.backend_type {
         MemBackendType::File => {
             if vm_resources.machine_config.huge_pages.is_hugetlbfs() {
                 return Err(BuildMicrovmFromSnapshotErrorGuestMemoryError::File(
@@ -594,7 +593,6 @@ pub fn build_microvm_from_snapshot(
                 guest_memory_from_file(mem_backend_path, mem_state, track_dirty_pages)
                     .map_err(BuildMicrovmFromSnapshotErrorGuestMemoryError::File)?,
                 None,
-                None,
             )
         }
         MemBackendType::Uffd => {
@@ -626,7 +624,6 @@ pub fn build_microvm_from_snapshot(
         .register_memory_regions(guest_memory, userfault_bitmap)
         .map_err(VmmError::Vm)
         .map_err(StartMicrovmError::Internal)?;
-    vmm.uffd = uffd;
     vmm.uffd_socket = socket;
 
     #[cfg(target_arch = "x86_64")]
@@ -683,7 +680,7 @@ pub fn build_microvm_from_snapshot(
         resource_allocator: &mut vmm.resource_allocator,
         vm_resources,
         instance_id: &instance_info.id,
-        restored_from_file: vmm.uffd.is_none(),
+        restored_from_file: vmm.uffd_socket.is_none(),
     };
 
     vmm.mmio_device_manager =
@@ -1092,7 +1089,6 @@ pub(crate) mod tests {
             shutdown_exit_code: None,
             kvm,
             vm,
-            uffd: None,
             uffd_socket: None,
             vcpus_handles: Vec::new(),
             vcpus_exit_evt,

src/vmm/src/lib.rs

@@ -128,7 +128,6 @@ use device_manager::resources::ResourceAllocator;
 use devices::acpi::vmgenid::VmGenIdError;
 use event_manager::{EventManager as BaseEventManager, EventOps, Events, MutEventSubscriber};
 use seccomp::BpfProgram;
-use userfaultfd::Uffd;
 use vm_memory::GuestAddress;
 use vmm_sys_util::epoll::EventSet;
 use vmm_sys_util::eventfd::EventFd;
@@ -314,8 +313,6 @@ pub struct Vmm {
     kvm: Kvm,
     /// VM object
     pub vm: Vm,
-    // Save UFFD in order to keep it open in the Firecracker process, as well.
-    uffd: Option<Uffd>,
     // Used for userfault communication with the UFFD handler when secret freedom is enabled
     uffd_socket: Option<UnixStream>,
     vcpus_handles: Vec<VcpuHandle>,

src/vmm/src/persist.rs

@@ -14,7 +14,7 @@ use std::sync::{Arc, Mutex};
 
 use semver::Version;
 use serde::{Deserialize, Serialize};
-use userfaultfd::{FeatureFlags, RegisterMode, Uffd, UffdBuilder};
+use userfaultfd::{FeatureFlags, RegisterMode, UffdBuilder};
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
 #[cfg(target_arch = "aarch64")]
@@ -515,9 +515,6 @@ pub enum GuestMemoryFromUffdError {
 // TODO remove these when the UFFD crate supports minor faults for guest_memfd
 const UFFDIO_REGISTER_MODE_MINOR: u64 = 1 << 2;
 
-type GuestMemoryResult =
-    Result<(Vec<GuestRegionMmap>, Option<Uffd>, Option<UnixStream>), GuestMemoryFromUffdError>;
-
 /// Creates guest memory using a UDS socket provided by a UFFD handler.
 pub fn guest_memory_from_uffd(
     mem_uds_path: &Path,
@@ -526,7 +523,7 @@ pub fn guest_memory_from_uffd(
     huge_pages: HugePageConfig,
     guest_memfd: Option<File>,
     userfault_bitmap_memfd: Option<&File>,
-) -> GuestMemoryResult {
+) -> Result<(Vec<GuestRegionMmap>, Option<UnixStream>), GuestMemoryFromUffdError> {
     let guest_memfd_fd = guest_memfd.as_ref().map(|f| f.as_raw_fd());
     let (guest_memory, backend_mappings) =
         create_guest_memory(mem_state, track_dirty_pages, huge_pages, guest_memfd)?;
@@ -566,7 +563,7 @@ pub fn guest_memory_from_uffd(
 
     let socket = send_uffd_handshake(mem_uds_path, &backend_mappings, fds)?;
 
-    Ok((guest_memory, Some(uffd), Some(socket)))
+    Ok((guest_memory, Some(socket)))
 }
 
 fn create_guest_memory(

tests/integration_tests/functional/test_uffd.py

@@ -4,6 +4,7 @@
 
 import os
 import re
+import time
 
 import pytest
 import requests
@@ -142,3 +143,36 @@ def test_malicious_handler(uvm_plain, snapshot):
             assert False, "Firecracker should freeze"
     except (TimeoutError, requests.exceptions.ReadTimeout):
         vm.uffd_handler.mark_killed()
+
+
+def test_fault_all_handler_exit(uvm_plain, snapshot):
+    """
+    Test that the VM is functional if the fault-all handler exits
+    after prepopulating the guest memory.
+    """
+    vm = uvm_plain
+    vm.memory_monitor = None
+    vm.spawn()
+    vm.restore_from_snapshot(snapshot, resume=True, uffd_handler_name="fault_all")
+
+    # Verify if the restored guest works.
+    vm.ssh.check_output("true")
+
+    # Kill the UFFD handler.
+    vm.uffd_handler.kill()
+
+    # Give UFFD time to unregister all guest memory.
+    time.sleep(1)
+
+    # Verify if the restored guest works after the handler exited.
+    #
+    # It is empirically known that invoking `ps` first time after snapshot restore
+    # will likely trigger an access to a guest memory page via userspace mappings
+    # either due to an MMIO instruction lookup (x86 only) or due to
+    # Firecracker accessing the guest memory.
+    #
+    # On Secret Free VMs, we do not preinstall userspace mappings when prepopulating
+    # guest memory in the fault-all handler. If we fail to unregister all guest memory
+    # with UFFD on the handler exit, accessing the userspace mapping will trigger
+    # a UFFD notification that will never be handled.
+    vm.ssh.check_output("ps")