feat(vmm): configure kvm userfault if secret free is enabled

This is needed to instruct the kernel to exit to userspace when a vCPU
fault occurs and the corresponding bit in the userfault bitmap is set.

The userfault bitmap is allocated in a memfd by Firecracker and sent to
the UFFD handler.

This also sends 3 fds to the UFFD handler in the handshake:
 - UFFD (original)
 - guest_memfd: for the handler to be able to populate guest memory
 - userfault bitmap memfd: for the handler to be able to disable exits
   to userspace for the pages that have already been populated

Signed-off-by: Nikita Kalyazin <kalyazin@amazon.com>

Author: kalyazin

src/vmm/src/builder.rs

@@ -4,8 +4,9 @@
 //! Enables pre-boot setup, instantiation and booting of a Firecracker VMM.
 
 use std::fmt::Debug;
+use std::fs::File;
 use std::io;
-use std::os::fd::AsFd;
+use std::os::fd::{AsFd, AsRawFd};
 use std::os::unix::fs::MetadataExt;
 #[cfg(feature = "gdb")]
 use std::sync::mpsc;
@@ -14,7 +15,6 @@ use std::sync::{Arc, Mutex};
 use event_manager::SubscriberOps;
 use kvm_ioctls::Cap;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-use userfaultfd::Uffd;
 use utils::time::TimestampUs;
 use vm_allocator::AllocPolicy;
 use vm_memory::GuestAddress;
@@ -29,6 +29,7 @@ use crate::cpu_config::templates::{
 };
 #[cfg(target_arch = "x86_64")]
 use crate::device_manager;
+use crate::device_manager::acpi::ACPIDeviceError;
 use crate::device_manager::pci_mngr::PciManagerError;
 use crate::device_manager::{
     AttachDeviceError, DeviceManager, DeviceManagerCreateError, DevicePersistError,
@@ -45,16 +46,20 @@ use crate::devices::virtio::vsock::{Vsock, VsockUnixBackend};
 use crate::gdb;
 use crate::initrd::{InitrdConfig, InitrdError};
 use crate::logger::debug;
-use crate::persist::{MicrovmState, MicrovmStateError};
+use crate::persist::{
+    GuestMemoryFromFileError, GuestMemoryFromUffdError, MicrovmState, MicrovmStateError,
+    guest_memory_from_file, guest_memory_from_uffd,
+};
 use crate::resources::VmResources;
 use crate::seccomp::BpfThreadMap;
 use crate::snapshot::Persist;
 use crate::utils::{mib_to_bytes, u64_to_usize};
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::MachineConfigError;
 use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
+use crate::vmm_config::snapshot::{LoadSnapshotParams, MemBackendType};
 use crate::vstate::kvm::{Kvm, KvmError};
-use crate::vstate::memory::{GuestRegionMmap, MaybeBounce};
+use crate::vstate::memory::{GuestMemoryState, MaybeBounce, bitmap_size, create_memfd};
 #[cfg(target_arch = "aarch64")]
 use crate::vstate::resources::ResourceAllocator;
 use crate::vstate::vcpu::VcpuError;
@@ -382,6 +387,7 @@ pub fn build_microvm_for_boot(
         kvm,
         vm,
         uffd: None,
+        uffd_socket: None,
         vcpus_handles: Vec::new(),
         vcpus_exit_evt,
         device_manager,
@@ -455,6 +461,17 @@ pub fn build_and_boot_microvm(
     Ok(vmm)
 }
 
+/// Sub-Error type for [`build_microvm_from_snapshot`] to contain either
+/// [`GuestMemoryFromFileError`] or [`GuestMemoryFromUffdError`] within
+/// [`BuildMicrovmFromSnapshotError`].
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
+pub enum BuildMicrovmFromSnapshotErrorGuestMemoryError {
+    /// Error creating guest memory from file: {0}
+    File(#[from] GuestMemoryFromFileError),
+    /// Error creating guest memory from uffd: {0}
+    Uffd(#[from] GuestMemoryFromUffdError),
+}
+
 /// Error type for [`build_microvm_from_snapshot`].
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum BuildMicrovmFromSnapshotError {
@@ -490,6 +507,55 @@ pub enum BuildMicrovmFromSnapshotError {
     SeccompFiltersInternal(#[from] crate::seccomp::InstallationError),
     /// Failed to restore devices: {0}
     RestoreDevices(#[from] DevicePersistError),
+    /// Failed to restore ACPI device manager: {0}
+    ACPIDeviManager(#[from] ACPIDeviceError),
+    /// VMGenID update failed: {0}
+    VMGenIDUpdate(std::io::Error),
+    /// Internal error while restoring microVM: {0}
+    Internal(#[from] VmmError),
+    /// Failed to load guest memory: {0}
+    GuestMemory(#[from] BuildMicrovmFromSnapshotErrorGuestMemoryError),
+    /// Userfault bitmap memfd error: {0}
+    UserfaultBitmapMemfd(#[from] crate::vstate::memory::MemoryError),
+}
+
+fn memfd_to_slice(memfd: &mut Option<File>) -> Option<&mut [u8]> {
+    if let Some(bitmap_file) = memfd {
+        let len = u64_to_usize(
+            bitmap_file
+                .metadata()
+                .expect("Failed to get metadata")
+                .len(),
+        );
+
+        // SAFETY: the arguments to mmap cannot cause any memory unsafety in the rust sense
+        let bitmap_addr = unsafe {
+            libc::mmap(
+                std::ptr::null_mut(),
+                len,
+                libc::PROT_WRITE,
+                libc::MAP_SHARED,
+                bitmap_file.as_raw_fd(),
+                0,
+            )
+        };
+
+        if bitmap_addr == libc::MAP_FAILED {
+            panic!(
+                "Failed to mmap userfault bitmap file: {}",
+                std::io::Error::last_os_error()
+            );
+        }
+
+        // SAFETY: `bitmap_addr` is a valid memory address returned by `mmap`.
+        Some(unsafe { std::slice::from_raw_parts_mut(bitmap_addr.cast(), len) })
+    } else {
+        None
+    }
+}
+
+fn memory_size_from_mem_state(mem_state: &GuestMemoryState) -> usize {
+    mem_state.regions.iter().map(|region| region.size).sum()
 }
 
 /// Builds and starts a microVM based on the provided MicrovmState.
@@ -501,26 +567,108 @@ pub fn build_microvm_from_snapshot(
     instance_info: &InstanceInfo,
     event_manager: &mut EventManager,
     microvm_state: MicrovmState,
-    guest_memory: Vec<GuestRegionMmap>,
-    uffd: Option<Uffd>,
     seccomp_filters: &BpfThreadMap,
+    params: &LoadSnapshotParams,
     vm_resources: &mut VmResources,
 ) -> Result<Arc<Mutex<Vmm>>, BuildMicrovmFromSnapshotError> {
+    // TODO: take it from kvm-bindings when userfault support is merged upstream
+    const KVM_CAP_USERFAULT: u32 = 246;
+
     // Build Vmm.
     debug!("event_start: build microvm from snapshot");
 
-    let kvm = Kvm::new(microvm_state.kvm_state.kvm_cap_modifiers.clone())
-        .map_err(StartMicrovmError::Kvm)?;
+    let secret_free = vm_resources.machine_config.secret_free;
+    let mut kvm_capabilities = microvm_state.kvm_state.kvm_cap_modifiers.clone();
+    if secret_free {
+        kvm_capabilities.push(KvmCapability::Add(Cap::GuestMemfd as u32));
+        kvm_capabilities.push(KvmCapability::Add(KVM_CAP_GUEST_MEMFD_MMAP));
+        kvm_capabilities.push(KvmCapability::Add(KVM_CAP_GUEST_MEMFD_NO_DIRECT_MAP));
+        kvm_capabilities.push(KvmCapability::Add(KVM_CAP_USERFAULT));
+    }
+
+    let kvm = Kvm::new(kvm_capabilities).map_err(StartMicrovmError::Kvm)?;
+
     // Set up Kvm Vm and register memory regions.
     // Build custom CPU config if a custom template is provided.
-    let mut vm = Vm::new(&kvm, false).map_err(StartMicrovmError::Vm)?;
+    let mut vm = Vm::new(&kvm, secret_free).map_err(StartMicrovmError::Vm)?;
 
     let (mut vcpus, vcpus_exit_evt) = vm
         .create_vcpus(vm_resources.machine_config.vcpu_count)
         .map_err(StartMicrovmError::Vm)?;
 
-    vm.restore_memory_regions(guest_memory, &microvm_state.vm_state.memory, None)
-        .map_err(StartMicrovmError::Vm)?;
+    let guest_memfd = match secret_free {
+        true => Some(
+            vm.create_guest_memfd(
+                memory_size_from_mem_state(&microvm_state.vm_state.memory),
+                GUEST_MEMFD_FLAG_MMAP | GUEST_MEMFD_FLAG_NO_DIRECT_MAP,
+            )
+            .map_err(VmmError::Vm)?,
+        ),
+        false => None,
+    };
+
+    let mut userfault_bitmap_memfd = if secret_free {
+        let bitmap_size =
+            bitmap_size(memory_size_from_mem_state(&microvm_state.vm_state.memory) as u64);
+        let bitmap_file = create_memfd(bitmap_size, None)?;
+
+        Some(bitmap_file.into_file())
+    } else {
+        None
+    };
+
+    let mem_backend_path = &params.mem_backend.backend_path;
+    let mem_state = &microvm_state.vm_state.memory;
+    let track_dirty_pages = params.track_dirty_pages;
+
+    let (guest_memory, uffd, uffd_socket) = match params.mem_backend.backend_type {
+        MemBackendType::File => {
+            if vm_resources.machine_config.huge_pages.is_hugetlbfs() {
+                return Err(BuildMicrovmFromSnapshotErrorGuestMemoryError::File(
+                    GuestMemoryFromFileError::HugetlbfsSnapshot,
+                )
+                .into());
+            }
+            (
+                guest_memory_from_file(mem_backend_path, mem_state, track_dirty_pages)
+                    .map_err(BuildMicrovmFromSnapshotErrorGuestMemoryError::File)?,
+                None,
+                None,
+            )
+        }
+        MemBackendType::Uffd => {
+            if vm_resources.machine_config.huge_pages.is_hugetlbfs() && guest_memfd.is_some() {
+                return Err(BuildMicrovmFromSnapshotErrorGuestMemoryError::Uffd(
+                    GuestMemoryFromUffdError::HugetlbfsSnapshot,
+                )
+                .into());
+            }
+            guest_memory_from_uffd(
+                mem_backend_path,
+                mem_state,
+                track_dirty_pages,
+                vm_resources.machine_config.huge_pages,
+                guest_memfd,
+                userfault_bitmap_memfd.as_ref(),
+            )
+            .map_err(BuildMicrovmFromSnapshotErrorGuestMemoryError::Uffd)?
+        }
+    };
+
+    let mut userfault_bitmap_slice = memfd_to_slice(&mut userfault_bitmap_memfd);
+    if let Some(ref mut slice) = userfault_bitmap_slice {
+        // Set all bits so a fault on any page will cause a VM exit
+        slice.fill(0xffu8);
+    }
+
+    let userfault_bitmap: Option<u64> = userfault_bitmap_slice.map(|s| s.as_ptr() as u64);
+
+    vm.restore_memory_regions(
+        guest_memory,
+        &microvm_state.vm_state.memory,
+        userfault_bitmap,
+    )
+    .map_err(StartMicrovmError::Vm)?;
 
     #[cfg(target_arch = "x86_64")]
     {
@@ -582,6 +730,7 @@ pub fn build_microvm_from_snapshot(
         kvm,
         vm,
         uffd,
+        uffd_socket,
         vcpus_handles: Vec::new(),
         vcpus_exit_evt,
         device_manager,
@@ -921,6 +1070,7 @@ pub(crate) mod tests {
             kvm,
             vm: Arc::new(vm),
             uffd: None,
+            uffd_socket: None,
             vcpus_handles: Vec::new(),
             vcpus_exit_evt,
             device_manager: default_device_manager(),

src/vmm/src/lib.rs

@@ -119,6 +119,7 @@ pub mod initrd;
 use std::collections::HashMap;
 use std::io;
 use std::os::unix::io::AsRawFd;
+use std::os::unix::net::UnixStream;
 use std::sync::mpsc::RecvTimeoutError;
 use std::sync::{Arc, Barrier, Mutex};
 use std::time::Duration;
@@ -307,6 +308,8 @@ pub struct Vmm {
     // Save UFFD in order to keep it open in the Firecracker process, as well.
     #[allow(unused)]
     uffd: Option<Uffd>,
+    // Used for userfault communication with the UFFD handler when secret freedom is enabled
+    uffd_socket: Option<UnixStream>,
     /// Handles to the vcpu threads with vcpu_fds inside them.
     pub vcpus_handles: Vec<VcpuHandle>,
     // Used by Vcpus and devices to initiate teardown; Vmm should never write here.