feat(vmm): implement secret-free fault handling protocol

It contains two parts:
 - external: between the VMM thread and the UFFD handler
 - internal: between vCPUs and the VMM thread

An outline of the workflow:
 - When a vCPU fault occurs, vCPU exits to userspace
 - The vCPU thread sends sends the exit syndrome in the vCPU to VMM
   channel and writes to the eventfd
 - The VMM thread forwards the syndrome to the UFFD handler via the UDS
   socket
 - The UFFD handler populates the page, clears the corresponding bit in
   the userfault bitmap and sends a reply to Firecracker
 - The VMM thread receives the reply and updates a vCPU condvar to
   notify the vCPU that the fault has been resolved
 - The vCPU resumes execution

Note that as a result of this change, an ability to exit the VM
gracefully is lost (at least on x86).  In the existing implementation,
the VMM thread initiated an exit if an event was read from the eventfd,
but no VcpuResponse::Exited responses were read for unknown reason.
Since the exit_evt eventfd is now also used by vCPUs to notify the VMM
thread of the VM exits caused by pagefaults, this situation (an eventfd
event, but response in the channel) can occur also because we have read
all VcpuResponse::Userfault in response to the previous eventfd event.

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

Author: kalyazin

src/vmm/src/builder.rs

@@ -189,7 +189,8 @@ pub fn build_microvm_for_boot(
     // Set up Kvm Vm and register memory regions.
     // Build custom CPU config if a custom template is provided.
     let mut vm = Vm::new(&kvm, secret_free)?;
-    let (mut vcpus, vcpus_exit_evt) = vm.create_vcpus(vm_resources.machine_config.vcpu_count)?;
+    let (mut vcpus, vcpus_exit_evt) =
+        vm.create_vcpus(vm_resources.machine_config.vcpu_count, secret_free)?;
 
     let guest_memfd = match secret_free {
         true => Some(Arc::new(
@@ -593,7 +594,7 @@ pub fn build_microvm_from_snapshot(
     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)
+        .create_vcpus(vm_resources.machine_config.vcpu_count, secret_free)
         .map_err(StartMicrovmError::Vm)?;
 
     let guest_memfd = match secret_free {
@@ -1062,7 +1063,7 @@ pub(crate) mod tests {
     pub(crate) fn default_vmm() -> Vmm {
         let (kvm, mut vm) = setup_vm_with_memory(mib_to_bytes(128));
 
-        let (_, vcpus_exit_evt) = vm.create_vcpus(1).unwrap();
+        let (_, vcpus_exit_evt) = vm.create_vcpus(1, false).unwrap();
 
         Vmm {
             instance_info: InstanceInfo::default(),

src/vmm/src/lib.rs

@@ -117,7 +117,8 @@ pub mod vstate;
 pub mod initrd;
 
 use std::collections::HashMap;
-use std::io;
+use std::io::{self, Read, Write};
+use std::os::fd::RawFd;
 use std::os::unix::io::AsRawFd;
 use std::os::unix::net::UnixStream;
 use std::sync::mpsc::RecvTimeoutError;
@@ -129,6 +130,7 @@ use event_manager::{EventManager as BaseEventManager, EventOps, Events, MutEvent
 use seccomp::BpfProgram;
 use snapshot::Persist;
 use userfaultfd::Uffd;
+use vm_memory::GuestAddress;
 use vmm_sys_util::epoll::EventSet;
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::terminal::Terminal;
@@ -145,12 +147,15 @@ use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::mem::{VIRTIO_MEM_DEV_ID, VirtioMem, VirtioMemError, VirtioMemStatus};
 use crate::devices::virtio::net::Net;
 use crate::logger::{METRICS, MetricsError, error, info, warn};
-use crate::persist::{MicrovmState, MicrovmStateError, VmInfo};
+use crate::persist::{FaultReply, FaultRequest, MicrovmState, MicrovmStateError, VmInfo};
 use crate::rate_limiter::BucketUpdate;
 use crate::vmm_config::instance_info::{InstanceInfo, VmState};
-use crate::vstate::memory::{GuestMemory, GuestMemoryMmap, GuestMemoryRegion};
+use crate::vstate::memory::{
+    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryRegion,
+};
 use crate::vstate::vcpu::VcpuState;
 pub use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuEvent, VcpuHandle, VcpuResponse};
+use crate::vstate::vm::UserfaultData;
 pub use crate::vstate::vm::Vm;
 
 /// Shorthand type for the EventManager flavour used by Firecracker.
@@ -688,6 +693,111 @@ impl Vmm {
         self.shutdown_exit_code = Some(exit_code);
     }
 
+    fn process_vcpu_userfault(&mut self, vcpu: u32, userfault_data: UserfaultData) {
+        let offset = self
+            .vm
+            .guest_memory()
+            .gpa_to_offset(GuestAddress(userfault_data.gpa))
+            .expect("Failed to convert GPA to offset");
+
+        let fault_request = FaultRequest {
+            vcpu,
+            offset,
+            flags: userfault_data.flags,
+            token: None,
+        };
+        let fault_request_json =
+            serde_json::to_string(&fault_request).expect("Failed to serialize fault request");
+
+        let written = self
+            .uffd_socket
+            .as_ref()
+            .expect("Uffd socket is not set")
+            .write(fault_request_json.as_bytes())
+            .expect("Failed to write to uffd socket");
+
+        if written != fault_request_json.len() {
+            panic!(
+                "Failed to write the entire fault request to the uffd socket: expected {}, \
+                 written {}",
+                fault_request_json.len(),
+                written
+            );
+        }
+    }
+
+    fn active_event_in_uffd_socket(&self, source: RawFd, event_set: EventSet) -> bool {
+        if let Some(uffd_socket) = &self.uffd_socket {
+            uffd_socket.as_raw_fd() == source && event_set == EventSet::IN
+        } else {
+            false
+        }
+    }
+
+    fn process_uffd_socket(&mut self) {
+        const BUFFER_SIZE: usize = 4096;
+
+        let stream = self.uffd_socket.as_mut().expect("Uffd socket is not set");
+
+        let mut buffer = [0u8; BUFFER_SIZE];
+        let mut current_pos = 0;
+
+        loop {
+            if current_pos < BUFFER_SIZE {
+                match stream.read(&mut buffer[current_pos..]) {
+                    Ok(0) => break,
+                    Ok(n) => current_pos += n,
+                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        if current_pos == 0 {
+                            break;
+                        }
+                    }
+                    Err(e) => panic!("Read error: {}", e),
+                }
+            }
+
+            let mut parser = serde_json::Deserializer::from_slice(&buffer[..current_pos])
+                .into_iter::<FaultReply>();
+            let mut total_consumed = 0;
+            let mut needs_more = false;
+
+            while let Some(result) = parser.next() {
+                match result {
+                    Ok(fault_reply) => {
+                        let vcpu = fault_reply.vcpu.expect("vCPU must be set");
+
+                        self.vcpus_handles
+                            .get(vcpu as usize)
+                            .expect("Invalid vcpu index")
+                            .send_userfault_resolved();
+
+                        total_consumed = parser.byte_offset();
+                    }
+                    Err(e) if e.is_eof() => {
+                        needs_more = true;
+                        break;
+                    }
+                    Err(e) => {
+                        println!(
+                            "Buffer content: {:?}",
+                            std::str::from_utf8(&buffer[..current_pos])
+                        );
+                        panic!("Invalid JSON: {}", e);
+                    }
+                }
+            }
+
+            if total_consumed > 0 {
+                buffer.copy_within(total_consumed..current_pos, 0);
+                current_pos -= total_consumed;
+            }
+
+            if needs_more {
+                continue;
+            }
+        }
+    }
+
     /// Gets a reference to kvm-ioctls Vm
     #[cfg(feature = "gdb")]
     pub fn vm(&self) -> &Vm {
@@ -765,38 +875,55 @@ impl MutEventSubscriber for Vmm {
         let event_set = event.event_set();
 
         if source == self.vcpus_exit_evt.as_raw_fd() && event_set == EventSet::IN {
-            // Exit event handling should never do anything more than call 'self.stop()'.
             let _ = self.vcpus_exit_evt.read();
 
-            let exit_code = 'exit_code: {
-                // Query each vcpu for their exit_code.
-                for handle in &self.vcpus_handles {
-                    // Drain all vcpu responses that are pending from this vcpu until we find an
-                    // exit status.
-                    for response in handle.response_receiver().try_iter() {
-                        if let VcpuResponse::Exited(status) = response {
-                            // It could be that some vcpus exited successfully while others
-                            // errored out. Thus make sure that error exits from one vcpu always
-                            // takes precedence over "ok" exits
+            let mut pending_userfaults = Vec::with_capacity(self.vcpus_handles.len());
+            let mut should_exit = false;
+            let mut final_exit_code = FcExitCode::Ok;
+
+            // First pass: collect all responses and determine exit status
+            for (handle, index) in self.vcpus_handles.iter().zip(0u32..) {
+                for response in handle.response_receiver().try_iter() {
+                    match response {
+                        VcpuResponse::Exited(status) => {
+                            should_exit = true;
                             if status != FcExitCode::Ok {
-                                break 'exit_code status;
+                                final_exit_code = status;
                             }
                         }
+                        VcpuResponse::Userfault(userfault_data) => {
+                            pending_userfaults.push((index, userfault_data));
+                        }
+                        _ => panic!("Unexpected response from vcpu: {:?}", response),
                     }
                 }
+            }
 
-                // No CPUs exited with error status code, report "Ok"
-                FcExitCode::Ok
-            };
-            self.stop(exit_code);
-        } else {
-            error!("Spurious EventManager event for handler: Vmm");
+            // Process any pending userfaults
+            for (index, userfault_data) in pending_userfaults {
+                self.process_vcpu_userfault(index, userfault_data);
+            }
+
+            // Stop if we received an exit event
+            if should_exit {
+                self.stop(final_exit_code);
+            }
+        }
+
+        if self.active_event_in_uffd_socket(source, event_set) {
+            self.process_uffd_socket();
         }
     }
 
     fn init(&mut self, ops: &mut EventOps) {
         if let Err(err) = ops.add(Events::new(&self.vcpus_exit_evt, EventSet::IN)) {
             error!("Failed to register vmm exit event: {}", err);
         }
+
+        if let Some(uffd_socket) = self.uffd_socket.as_ref()
+            && let Err(err) = ops.add(Events::new(uffd_socket, EventSet::IN))
+        {
+            panic!("Failed to register UFFD socket: {}", err);
+        }
     }
 }

src/vmm/src/persist.rs

@@ -616,6 +616,10 @@ fn send_uffd_handshake(
     let backend_mappings = serde_json::to_string(backend_mappings).unwrap();
 
     let socket = UnixStream::connect(mem_uds_path)?;
+    socket
+        .set_nonblocking(true)
+        .expect("Cannot set non-blocking");
+
     socket.send_with_fds(
         &[backend_mappings.as_bytes()],
         // In the happy case we can close the fd since the other process has it open and is