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 a message to the VMM thread via the userfault
   channel
 - The VMM thread forwards the message to the UFFD handler via the UDS
   socket
 - The UFFD hnadler populates the page, clears the corresponding bit in
   the userfault bitmap and sends a reply to Firecracker
 - The VMM thread receives the reply and forwards it to the vCPU via the
   userfault channel
 - The vCPU resumes execution

Signed-off-by: Nikita Kalyazin <[email protected]>

Author: kalyazin

src/vmm/src/lib.rs

@@ -115,7 +115,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;
@@ -128,6 +129,7 @@ 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;
 use vmm_sys_util::terminal::Terminal;
@@ -147,15 +149,17 @@ use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::{TYPE_BALLOON, TYPE_BLOCK, TYPE_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::snapshot::Persist;
 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::UserfaultChannel;
 pub use crate::vstate::vm::Vm;
+use crate::vstate::vm::{UserfaultChannel, UserfaultChannelError, UserfaultData};
 
 /// Shorthand type for the EventManager flavour used by Firecracker.
 pub type EventManager = BaseEventManager<Arc<Mutex<dyn MutEventSubscriber>>>;
@@ -803,6 +807,168 @@ impl Vmm {
         self.shutdown_exit_code = Some(exit_code);
     }
 
+    fn active_event_in_userfault_channel(
+        &self,
+        source: RawFd,
+        event_set: EventSet,
+    ) -> Option<usize> {
+        if let Some(userfault_channels) = &self.userfault_channels {
+            userfault_channels.iter().position(|channel| {
+                let receiver = &channel.receiver;
+                source == receiver.as_raw_fd() && event_set == EventSet::IN
+            })
+        } else {
+            None
+        }
+    }
+
+    fn process_userfault_channels(&mut self, vcpu: usize) {
+        loop {
+            match self
+                .userfault_channels
+                .as_mut()
+                .expect("Userfault channels must be set")[vcpu]
+                .recv()
+            {
+                Ok(userfault_data) => {
+                    let offset = self
+                        .vm
+                        .guest_memory()
+                        .gpa_to_offset(GuestAddress(userfault_data.gpa))
+                        .unwrap();
+
+                    let fault_request = FaultRequest {
+                        vcpu: vcpu.try_into().expect("Invalid vCPU index"),
+                        offset,
+                        flags: userfault_data.flags,
+                        token: None,
+                    };
+                    let fault_request_json = serde_json::to_string(&fault_request).unwrap();
+
+                    let written = self
+                        .uffd_socket
+                        .as_ref()
+                        .unwrap()
+                        .write(fault_request_json.as_bytes())
+                        .unwrap();
+
+                    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
+                        );
+                    }
+                }
+                Err(ref e) => match e {
+                    UserfaultChannelError::IO(io_e) if io_e.kind() == io::ErrorKind::WouldBlock => {
+                        break;
+                    }
+                    _ => panic!("Error receiving userfault data: {}", e),
+                },
+            }
+        }
+    }
+
+    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;
+        let mut exit_loop = false;
+
+        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 exit_loop {
+                            break;
+                        }
+                    }
+                    Err(e) => panic!("Read error: {}", e),
+                }
+
+                exit_loop = false;
+            }
+
+            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 gpa = self
+                            .vm
+                            .common
+                            .guest_memory
+                            .offset_to_gpa(fault_reply.offset)
+                            .expect("Failed to convert offset to GPA");
+
+                        let userfaultfd_data = UserfaultData {
+                            flags: fault_reply.flags,
+                            gpa: gpa.0,
+                            size: fault_reply.len,
+                        };
+
+                        let vcpu = fault_reply.vcpu.expect("vCPU must be set");
+
+                        self.userfault_channels
+                            .as_mut()
+                            .expect("userfault_channels are not set")
+                            .get_mut(vcpu as usize)
+                            .expect("Invalid vcpu index")
+                            .send(userfaultfd_data)
+                            .expect("Failed to send userfault data");
+
+                        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;
+            }
+
+            // We consumed all data in the buffer, but the socket may have remaining unread data so
+            // we attempt to read from it and exit the loop only if we confirm that nothing is in
+            // there.
+            if current_pos == 0 {
+                exit_loop = true;
+            }
+        }
+    }
+
     /// Gets a reference to kvm-ioctls Vm
     #[cfg(feature = "gdb")]
     pub fn vm(&self) -> &Vm {
@@ -909,14 +1075,34 @@ impl MutEventSubscriber for Vmm {
                 FcExitCode::Ok
             };
             self.stop(exit_code);
-        } else {
-            error!("Spurious EventManager event for handler: Vmm");
+        }
+
+        if let Some(vcpu) = self.active_event_in_userfault_channel(source, event_set) {
+            self.process_userfault_channels(vcpu);
+        }
+
+        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() {
+            if let Err(err) = ops.add(Events::new(uffd_socket, EventSet::IN)) {
+                panic!("Failed to register UFFD socket: {}", err);
+            }
+        }
+
+        if let Some(userfault_channels) = self.userfault_channels.as_ref() {
+            for channel in userfault_channels {
+                if let Err(err) = ops.add(Events::new(&channel.receiver, EventSet::IN)) {
+                    panic!("Failed to register userfault events: {}", err);
+                }
+            }
+        }
     }
 }

src/vmm/src/persist.rs

@@ -652,6 +652,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

src/vmm/src/vstate/vcpu.rs

@@ -31,11 +31,15 @@ use crate::logger::{IncMetric, METRICS};
 use crate::seccomp::{BpfProgram, BpfProgramRef};
 use crate::utils::signal::{Killable, register_signal_handler, sigrtmin};
 use crate::utils::sm::StateMachine;
-use crate::vstate::vm::{UserfaultChannel, Vm};
+use crate::vstate::vm::{UserfaultChannel, UserfaultData, Vm};
 
 /// Signal number (SIGRTMIN) used to kick Vcpus.
 pub const VCPU_RTSIG_OFFSET: i32 = 0;
 
+// TODO: remove when KVM userfault support is merged upstream.
+/// VM exit due to a userfault.
+const KVM_MEMORY_EXIT_FLAG_USERFAULT: u64 = 1 << 4;
+
 /// Errors associated with the wrappers over KVM ioctls.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum VcpuError {
@@ -312,6 +316,7 @@ impl Vcpu {
                 // - the other vCPUs won't ever exit out of `KVM_RUN`, but they won't consume CPU.
                 // So we pause vCPU0 and send a signal to the emulation thread to stop the VMM.
                 Ok(VcpuEmulation::Stopped) => return self.exit(FcExitCode::Ok),
+                Ok(VcpuEmulation::Userfault(_)) => unreachable!(),
                 // If the emulation requests a pause lets do this
                 #[cfg(feature = "gdb")]
                 Ok(VcpuEmulation::Paused) => {
@@ -495,6 +500,26 @@ impl Vcpu {
         StateMachine::finish()
     }
 
+    fn handle_userfault(
+        &mut self,
+        userfaultfd_data: UserfaultData,
+    ) -> Result<VcpuEmulation, VcpuError> {
+        let userfault_channel = self
+            .userfault_channel
+            .as_mut()
+            .expect("userfault channel not set");
+
+        userfault_channel
+            .send(userfaultfd_data)
+            .expect("Failed to send userfault data");
+
+        let _ = userfault_channel
+            .recv()
+            .expect("Failed to receive userfault response");
+
+        Ok(VcpuEmulation::Handled)
+    }
+
     /// Runs the vCPU in KVM context and handles the kvm exit reason.
     ///
     /// Returns error or enum specifying whether emulation was handled or interrupted.
@@ -505,7 +530,7 @@ impl Vcpu {
             return Ok(VcpuEmulation::Interrupted);
         }
 
-        match self.kvm_vcpu.fd.run() {
+        let result = match self.kvm_vcpu.fd.run() {
             Err(ref err) if err.errno() == libc::EINTR => {
                 self.kvm_vcpu.fd.set_kvm_immediate_exit(0);
                 // Notify that this KVM_RUN was interrupted.
@@ -522,7 +547,14 @@ impl Vcpu {
                 Ok(VcpuEmulation::Paused)
             }
             emulation_result => handle_kvm_exit(&mut self.kvm_vcpu.peripherals, emulation_result),
-        }
+        };
+
+        let userfault_data = match result {
+            Ok(VcpuEmulation::Userfault(userfault_data)) => userfault_data,
+            _ => return result,
+        };
+
+        self.handle_userfault(userfault_data)
     }
 }
 
@@ -600,6 +632,16 @@ fn handle_kvm_exit(
                     )))
                 }
             },
+            VcpuExit::MemoryFault { flags, gpa, size } => {
+                if flags & KVM_MEMORY_EXIT_FLAG_USERFAULT == 0 {
+                    Err(VcpuError::UnhandledKvmExit(format!(
+                        "flags {:x} gpa {:x} size {:x}",
+                        flags, gpa, size
+                    )))
+                } else {
+                    Ok(VcpuEmulation::Userfault(UserfaultData { flags, gpa, size }))
+                }
+            }
             arch_specific_reason => {
                 // run specific architecture emulation.
                 peripherals.run_arch_emulation(arch_specific_reason)
@@ -761,6 +803,8 @@ pub enum VcpuEmulation {
     Interrupted,
     /// Stopped.
     Stopped,
+    /// Userfault
+    Userfault(UserfaultData),
     /// Pause request
     #[cfg(feature = "gdb")]
     Paused,