From 282d1c989d9744b094e8e1c1beb225480f4d60f1 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sboeuf@meta.com>
Date: Tue, 23 Jun 2026 08:55:52 -0700
Subject: [PATCH] vmm: Add SocketUffdMemorySource implementation

Add the socket-backed UffdMemorySource that resolves each fault by
sending a Command::PageFault request to the peer over a dedicated fault
connection.

Signed-off-by: Sebastien Boeuf <sboeuf@meta.com>
Assisted-by: Claude:claude-opus-4-7
---
 vmm/src/uffd.rs | 111 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 107 insertions(+), 4 deletions(-)

diff --git a/vmm/src/uffd.rs b/vmm/src/uffd.rs
index 6e241c397..ace8d34ed 100644
--- a/vmm/src/uffd.rs
+++ b/vmm/src/uffd.rs
@@ -14,11 +14,14 @@
 //! passthrough and shared-memory-backed guest RAM.
 
 use std::fs::{File, OpenOptions};
-use std::io::{self, Error};
-use std::mem;
-use std::os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd};
+use std::io::{self, Error, Read};
+use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd};
 use std::os::unix::fs::FileExt;
+use std::{fmt, mem};
 
+use vm_migration::protocol::{MemoryRange, Request, Response, Status};
+
+use crate::migration::transport::SocketStream;
 use crate::userfaultfd;
 
 #[repr(C)]
@@ -181,7 +184,7 @@ struct UffdioRange {
 }
 
 /// A guest memory range registered with userfaultfd, plus where its bytes
-/// live in the snapshot file.
+/// live for the data source.
 pub(crate) struct UffdRange {
     pub host_addr: u64,
     pub length: u64,
@@ -274,6 +277,106 @@ impl UffdMemorySource for FileUffdMemorySource {
     }
 }
 
+/// Memory source that provides pages content over a socket.
+// Wired into the receive-migration path in the next commit.
+#[allow(dead_code)]
+pub(crate) struct SocketUffdMemorySource {
+    stream: SocketStream,
+    shared_backing: bool,
+    buf: Vec<u8>,
+}
+
+#[allow(dead_code)]
+impl SocketUffdMemorySource {
+    pub fn new(stream: SocketStream, shared_backing: bool) -> Self {
+        Self {
+            stream,
+            shared_backing,
+            buf: Vec::new(),
+        }
+    }
+
+    /// Returns the length of the inline page payload the peer is about to send
+    /// (0 when the page was written directly into shared memory).
+    fn request_page(&mut self, gpa: u64, len: u64) -> Result<u64, io::Error> {
+        Request::page_fault()
+            .write_to(&mut self.stream)
+            .map_err(io_other)?;
+        MemoryRange { gpa, length: len }
+            .write_to(&mut self.stream)
+            .map_err(io_other)?;
+
+        let resp = Response::read_from(&mut self.stream).map_err(io_other)?;
+        match resp.status() {
+            Status::Ok => Ok(resp.length()),
+            s => Err(io::Error::other(format!(
+                "peer returned {s:?} for PageFault at gpa={gpa:#x} len={len}",
+            ))),
+        }
+    }
+}
+
+impl UffdMemorySource for SocketUffdMemorySource {
+    fn resolve(
+        &mut self,
+        uffd_fd: BorrowedFd<'_>,
+        range: &UffdRange,
+        page_idx: u64,
+    ) -> Result<FaultResolution, io::Error> {
+        let page_size = range.page_size;
+        let page_addr = range.page_addr(page_idx);
+        let page_gpa = range.page_source_offset(page_idx);
+
+        let resp_len = self.request_page(page_gpa, page_size)?;
+
+        if self.shared_backing {
+            if resp_len != 0 {
+                return Err(io::Error::other(format!(
+                    "shared-backing PageFault response carried {resp_len} unexpected bytes",
+                )));
+            }
+            match wake(uffd_fd, page_addr, page_size) {
+                Ok(()) => Ok(FaultResolution::Served),
+                Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => Ok(FaultResolution::Retry),
+                Err(e) => Err(e),
+            }
+        } else {
+            if resp_len != page_size {
+                return Err(io::Error::other(format!(
+                    "inline PageFault response length {resp_len} != page size {page_size}",
+                )));
+            }
+            let len = page_size as usize;
+            if self.buf.len() < len {
+                self.buf.resize(len, 0);
+            }
+            self.stream.read_exact(&mut self.buf[..len])?;
+            match copy(uffd_fd, page_addr, self.buf.as_ptr(), page_size) {
+                Ok(()) => Ok(FaultResolution::Served),
+                Err(e) if e.raw_os_error() == Some(libc::EEXIST) => {
+                    if let Err(e) = wake(uffd_fd, page_addr, page_size) {
+                        log::warn!("UFFDIO_WAKE failed at {page_addr:#x}: {e}");
+                    }
+                    Ok(FaultResolution::Served)
+                }
+                Err(e) if e.raw_os_error() == Some(libc::EAGAIN) => Ok(FaultResolution::Retry),
+                Err(e) => Err(e),
+            }
+        }
+    }
+}
+
+impl Drop for SocketUffdMemorySource {
+    fn drop(&mut self) {
+        // SAFETY: the fd is valid for the duration of this borrow.
+        unsafe { libc::shutdown(self.stream.as_fd().as_raw_fd(), libc::SHUT_RDWR) };
+    }
+}
+
+fn io_other<E: fmt::Display>(e: E) -> io::Error {
+    io::Error::other(e.to_string())
+}
+
 /// Wake threads waiting on a fault in the given range without copying data.
 ///
 /// Needed after UFFDIO_COPY returns EEXIST: the page was already resolved