fix(apf): resolve non exitless errors for apf

There were a couple of error on the non exitless path for APF. We saw
these with the memory hotplug which flodded the APF buffer.

Signed-off-by: Jack Thomson <jackabt@amazon.com>

Author: JackThomson2

resources/seccomp/x86_64-unknown-linux-musl.json

@@ -508,6 +508,18 @@
             {
                 "syscall": "mprotect",
                 "comment": "Used by memory hotplug to protect access to underlying host memory"
+            },
+            {
+                "syscall": "ioctl",
+                "args": [
+                    {
+                        "index": 1,
+                        "type": "dword",
+                        "op": "eq",
+                        "val": 1075883734,
+                        "comment": "KVM_ASYNC_PF (accept/ready/sync_complete)"
+                    }
+                ]
             }
         ]
     },
@@ -580,8 +592,8 @@
                 "comment": "sigaltstack is used by Rust stdlib to remove alternative signal stack during thread teardown."
             },
             {
-              "syscall": "getrandom",
-              "comment": "getrandom is used by `HttpServer` to reinialize `HashMap` after moving to the API thread"
+                "syscall": "getrandom",
+                "comment": "getrandom is used by `HttpServer` to reinialize `HashMap` after moving to the API thread"
             },
             {
                 "syscall": "accept4",
@@ -1242,6 +1254,22 @@
                         "comment": "KVM_IRQFD"
                     }
                 ]
+            },
+            {
+                "syscall": "sendto",
+                "comment": "APF fallback: vCPU sends fault request to handler via socket"
+            },
+            {
+                "syscall": "ioctl",
+                "args": [
+                    {
+                        "index": 1,
+                        "type": "dword",
+                        "op": "eq",
+                        "val": 1075883734,
+                        "comment": "KVM_ASYNC_PF (accept/ready/sync_complete)"
+                    }
+                ]
             }
         ]
     }

src/firecracker/examples/uffd/on_demand_handler.rs

@@ -32,6 +32,9 @@ fn main() {
             .accept()
             .expect("Cannot listen on APF UDS socket");
         apf_stream
+            .set_nonblocking(true)
+            .expect("Cannot set APF stream non-blocking");
+        apf_stream
     } else {
         let (apf_stream, _) = UnixStream::pair().expect("Cannot create APF socket pair");
         apf_stream

src/firecracker/examples/uffd/uffd_utils.rs

@@ -231,8 +231,8 @@ pub struct FaultRequest {
     pub offset: u64,
     /// Flags
     pub flags: u64,
-    /// Async PF token
-    pub token: Option<u32>,
+    /// Async PF GPA (set for APF fallback faults, None for sync faults)
+    pub gpa: Option<u64>,
 }
 
 impl FaultRequest {
@@ -242,7 +242,7 @@ impl FaultRequest {
             offset: self.offset,
             len,
             flags: self.flags,
-            token: self.token,
+            gpa: self.gpa,
             zero: false,
         }
     }
@@ -259,8 +259,8 @@ pub struct FaultReply {
     pub len: u64,
     /// Flags, must be copied from `FaultRequest`, otherwise 0
     pub flags: u64,
-    /// Async PF token, must be copied from `FaultRequest`, otherwise None
-    pub token: Option<u32>,
+    /// Async PF GPA, must be copied from `FaultRequest`, otherwise None
+    pub gpa: Option<u64>,
     /// Whether the populated pages are zero pages
     pub zero: bool,
 }
@@ -423,6 +423,17 @@ impl UffdHandler {
         }
     }
 
+    /// Convert a guest physical address to an offset in the backing memory file.
+    #[inline]
+    pub fn gpa_to_offset(&self, gpa: u64) -> Option<usize> {
+        for region in &self.mem_regions {
+            if region.gpa_start <= gpa && gpa < region.gpa_start + region.size as u64 {
+                return Some((gpa - region.gpa_start + region.offset) as usize);
+            }
+        }
+        None
+    }
+
     pub fn read_event(&mut self) -> Result<Option<Event>, Error> {
         self.uffd.read_event()
     }
@@ -733,6 +744,13 @@ impl Runtime {
         self.stream.write_all(&encoded).unwrap();
     }
 
+    pub fn send_apf_fault_reply(&mut self, fault_reply: FaultReply) {
+        let encoded = bitcode::encode(&fault_reply);
+        let size = (encoded.len() as u32).to_le_bytes();
+        self.apf_stream.write_all(&size).unwrap();
+        self.apf_stream.write_all(&encoded).unwrap();
+    }
+
     pub fn construct_handler(
         stream: &UnixStream,
         backing_memory: *mut u8,
@@ -857,6 +875,11 @@ impl Runtime {
 
         let mut uffd_msg_iter =
             UffdMsgIterBitcode::new(self.stream.try_clone().expect("Failed to clone stream"));
+        let mut apf_msg_iter = UffdMsgIterBitcode::new(
+            self.apf_stream
+                .try_clone()
+                .expect("Failed to clone APF stream"),
+        );
 
         loop {
             let pollfd_ptr = pollfds.as_mut_ptr();
@@ -879,52 +902,40 @@ impl Runtime {
                     if fd.fd == stream_fd {
                         for fault_request in uffd_msg_iter.by_ref() {
                             let page_size = self.handler.page_size;
-
+                            let offset = fault_request
+                                .gpa
+                                .and_then(|gpa| self.handler.gpa_to_offset(gpa))
+                                .unwrap_or(fault_request.offset as usize);
                             assert!(
-                                (fault_request.offset as usize) < self.handler.size(),
+                                offset < self.handler.size(),
                                 "received bogus offset from firecracker"
                             );
-
-                            pf_vcpu_event_dispatch(
-                                &mut self.handler,
-                                fault_request.offset as usize,
-                            );
-
+                            pf_vcpu_event_dispatch(&mut self.handler, offset);
                             self.send_fault_reply(fault_request.into_reply(page_size as u64));
                         }
                     } else if fd.fd == apf_stream_fd {
-                        // APF fallback path: fault requests over socket
-                        for fault_request in uffd_msg_iter.by_ref() {
+                        // APF fallback path: read from APF socket, reply on APF socket
+                        for fault_request in apf_msg_iter.by_ref() {
                             let page_size = self.handler.page_size;
-
+                            let offset = fault_request
+                                .gpa
+                                .and_then(|gpa| self.handler.gpa_to_offset(gpa))
+                                .unwrap_or(fault_request.offset as usize);
                             assert!(
-                                (fault_request.offset as usize) < self.handler.size(),
-                                "received bogus offset from firecracker"
-                            );
-
-                            pf_vcpu_event_dispatch(
-                                &mut self.handler,
-                                fault_request.offset as usize,
+                                offset < self.handler.size(),
+                                "received bogus offset from APF handler"
                             );
-
-                            self.send_fault_reply(fault_request.into_reply(page_size as u64));
+                            pf_vcpu_event_dispatch(&mut self.handler, offset);
+                            self.send_apf_fault_reply(fault_request.into_reply(page_size as u64));
                         }
                     } else if let Some(ctx) = self.exitless_vcpus.get_mut(&fd.fd) {
                         // Exitless APF: drain notify ring and resolve pages
                         ctx.drain_eventfd();
                         while let Some(entry) = ctx.notify_ring().pop() {
-                            let gpa = entry.gpa;
-                            // Find the region and offset for this GPA
-                            for region in &self.handler.mem_regions {
-                                if region.gpa_start <= gpa
-                                    && gpa < region.gpa_start + region.size as u64
-                                {
-                                    let offset = (gpa - region.gpa_start + region.offset) as usize;
-                                    pf_vcpu_event_dispatch(&mut self.handler, offset);
-                                    break;
-                                }
+                            if let Some(offset) = self.handler.gpa_to_offset(entry.gpa) {
+                                pf_vcpu_event_dispatch(&mut self.handler, offset);
                             }
-                            ctx.signal_ready(gpa);
+                            ctx.signal_ready(entry.gpa);
                         }
                     } else {
                         // Handle one of uffd page faults
@@ -977,6 +988,9 @@ mod tests {
             let (apf_stream, _) = apf_listener
                 .accept()
                 .expect("Cannot listen on APF UDS socket");
+            apf_stream
+                .set_nonblocking(true)
+                .expect("Cannot set APF stream non-blocking");
             // Update runtime with actual runtime
             let runtime = uninit_runtime.write(Runtime::new(stream, file, apf_stream));
             runtime.run(|_: &mut UffdHandler| {}, |_: &mut UffdHandler, _: usize| {});
@@ -1024,7 +1038,7 @@ mod tests {
             vcpu: 0,
             offset: 0xDEAD_0000, // way beyond 0x1000 handler size
             flags: 0,
-            token: None,
+            gpa: None,
         };
         let encoded = bitcode::encode(&bogus_request);
         let size = (encoded.len() as u32).to_le_bytes();

src/vmm/src/builder.rs

@@ -430,6 +430,7 @@ pub fn build_microvm_for_boot(
         _guest_memfd: None,
         _eventfd: None,
         apf_stream: None,
+        apf_stream_fd: None,
         apf_supported: false,
         exitless_apf: Vec::new(),
         exitless_apf_setup_done: false,
@@ -769,7 +770,9 @@ pub fn build_microvm_from_snapshot(
         use vmm_sys_util::sock_ctrl_msg::ScmSocket;
         if let Some(uff_socket) = uffd_socket {
             let broker = UffdMessageBroker::new(uff_socket);
-            let apf_exitless = apf_supported && std::env::var("FC_APF_NO_EXITLESS").is_err();
+            let apf_exitless = apf_supported
+                && std::env::var("FC_APF_NO_EXITLESS").is_err()
+                && !std::path::Path::new("/apf_no_exitless").exists();
             let (contexts, done) = if apf_exitless {
                 let mut contexts = Vec::new();
                 let mut handler_fds = Vec::new();
@@ -885,6 +888,12 @@ pub fn build_microvm_from_snapshot(
         device_manager,
         _guest_memfd: None,
         _eventfd: None,
+        apf_stream_fd: apf_stream.as_ref().map(|s| {
+            s.lock()
+                .expect("APF stream lock poisoned")
+                .stream()
+                .as_raw_fd()
+        }),
         apf_stream,
         apf_supported,
         exitless_apf: exitless_apf_contexts,
@@ -1246,6 +1255,7 @@ pub(crate) mod tests {
             _guest_memfd: None,
             _eventfd: None,
             apf_stream: None,
+            apf_stream_fd: None,
             apf_supported: false,
             exitless_apf: Vec::new(),
             exitless_apf_setup_done: false,

src/vmm/src/lib.rs

@@ -135,7 +135,6 @@ 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;
@@ -170,9 +169,7 @@ use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
 use crate::vmm_config::mmds::MmdsConfig;
 use crate::vmm_config::net::NetworkInterfaceConfig;
 use crate::vmm_config::vsock::VsockDeviceConfig;
-use crate::vstate::memory::{
-    GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryRegion,
-};
+use crate::vstate::memory::{GuestMemory, GuestMemoryMmap, GuestMemoryRegion};
 use crate::vstate::memory::{KVM_APF_OP_READY, KVM_ASYNC_PF, KvmAPFReq};
 use crate::vstate::vcpu::VcpuState;
 pub use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuEvent, VcpuHandle, VcpuResponse};
@@ -355,6 +352,8 @@ pub struct Vmm {
     pub(crate) _eventfd: Option<i32>,
     /// APF stream shared with vCPUs
     apf_stream: Option<vstate::vcpu::SharedApfStream>,
+    /// Cached raw fd for APF stream (avoids mutex lock per event loop iteration)
+    apf_stream_fd: Option<RawFd>,
     /// Whether the kernel supports KVM_CAP_ASYNC_PF_USERFAULT
     apf_supported: bool,
     /// Exitless APF contexts (one per vCPU)
@@ -873,18 +872,13 @@ impl Vmm {
     }
 
     fn process_vcpu_userfault(&mut self, vcpu: u32, userfault_data: UserfaultData) {
-        // APF requests are now sent directly by vCPUs, so we only handle sync faults here
-        let offset = self
-            .vm
-            .guest_memory()
-            .gpa_to_offset(GuestAddress(userfault_data.gpa))
-            .expect("Failed to convert GPA to offset");
-
+        // Sync path: vCPU blocked on condvar, send fault to handler for resolution.
+        // Always send GPA — the handler converts GPA→offset via its mem_regions.
         let fault_request = FaultRequest {
             vcpu,
-            offset,
+            offset: 0,
             flags: userfault_data.flags,
-            gpa: None,
+            gpa: Some(userfault_data.gpa),
         };
 
         self.uffd_socket
@@ -1051,15 +1045,11 @@ impl MutEventSubscriber for Vmm {
             self.process_uffd_socket();
         }
 
-        if let Some(apf_stream) = &self.apf_stream {
-            let fd = apf_stream
-                .lock()
-                .expect("APF stream lock poisoned")
-                .stream()
-                .as_raw_fd();
-            if source == fd && event_set == EventSet::IN {
-                self.process_apf_socket();
-            }
+        if let Some(fd) = self.apf_stream_fd
+            && source == fd
+            && event_set == EventSet::IN
+        {
+            self.process_apf_socket();
         }
     }