|
| 1 | +//! Supervisor-side userfaultfd page server. After the restore-stub registers its |
| 2 | +//! anonymous checkpoint regions with a userfaultfd and hands the fd to the |
| 3 | +//! supervisor (via pidfd_getfd), this loop resolves each missing-page fault by |
| 4 | +//! UFFDIO_COPY-ing the page from the in-memory image. Only the anonymous working |
| 5 | +//! set flows through here; file-backed regions are kernel-paged. |
| 6 | +//! |
| 7 | +//! The userfaultfd MUST be created with O_NONBLOCK: poll()/epoll on a blocking |
| 8 | +//! userfaultfd always reports POLLERR (never POLLIN), so a polling pager would |
| 9 | +//! spin without ever reading a fault. The stub creates it O_CLOEXEC|O_NONBLOCK |
| 10 | +//! with a plain -> UFFD_USER_MODE_ONLY fallback (hosts with |
| 11 | +//! vm.unprivileged_userfaultfd=0 reject the plain form). USER_MODE_ONLY serves |
| 12 | +//! only user-mode faults; full-fidelity paging of real programs (kernel-mode |
| 13 | +//! faults from syscalls touching unfilled pages) requires |
| 14 | +//! vm.unprivileged_userfaultfd=1 on the node. |
| 15 | +
|
| 16 | +use std::io; |
| 17 | +use std::sync::atomic::{AtomicBool, Ordering}; |
| 18 | + |
| 19 | +// userfaultfd ioctl request codes (x86_64 _IOWR(0xAA, n, struct)). |
| 20 | +const UFFDIO_API: libc::c_ulong = 0xc018_aa3f; |
| 21 | +const UFFDIO_REGISTER: libc::c_ulong = 0xc020_aa00; |
| 22 | +const UFFDIO_COPY: libc::c_ulong = 0xc028_aa03; |
| 23 | +const UFFD_API: u64 = 0xAA; |
| 24 | +const UFFDIO_REGISTER_MODE_MISSING: u64 = 1; |
| 25 | +const UFFD_EVENT_PAGEFAULT: u8 = 0x12; |
| 26 | + |
| 27 | +#[repr(C)] |
| 28 | +struct UffdioApi { api: u64, features: u64, ioctls: u64 } |
| 29 | +#[repr(C)] |
| 30 | +struct UffdioRange { start: u64, len: u64 } |
| 31 | +#[repr(C)] |
| 32 | +struct UffdioRegister { range: UffdioRange, mode: u64, ioctls: u64 } |
| 33 | +#[repr(C)] |
| 34 | +struct UffdioCopy { dst: u64, src: u64, len: u64, mode: u64, copy: i64 } |
| 35 | + |
| 36 | +// struct uffd_msg is 32 bytes; we only need the fault address at offset 16. |
| 37 | +#[repr(C)] |
| 38 | +struct UffdMsg { event: u8, _pad: [u8; 7], arg: [u64; 3] } |
| 39 | + |
| 40 | +/// In-memory source for the anon working set: `(start, end, bytes)` runs. |
| 41 | +#[derive(Clone)] |
| 42 | +pub(crate) struct PageImage { |
| 43 | + pub regions: Vec<(u64, u64, Vec<u8>)>, |
| 44 | +} |
| 45 | + |
| 46 | +impl PageImage { |
| 47 | + /// The `page`-aligned slice of length `page` covering `addr`, or `None`. |
| 48 | + pub(crate) fn page_at(&self, addr: u64, page: usize) -> Option<&[u8]> { |
| 49 | + let base = addr & !((page as u64) - 1); |
| 50 | + for (start, end, bytes) in &self.regions { |
| 51 | + if base >= *start && base + page as u64 <= *end { |
| 52 | + let off = (base - *start) as usize; |
| 53 | + return Some(&bytes[off..off + page]); |
| 54 | + } |
| 55 | + } |
| 56 | + None |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +/// Enable the uffd API and register `[start, start+len)` in missing mode. |
| 61 | +/// Used by the pager unit test to stand in for the stub. |
| 62 | +pub(crate) fn register_api_and_range(uffd: i32, start: u64, len: u64) -> io::Result<()> { |
| 63 | + let mut api = UffdioApi { api: UFFD_API, features: 0, ioctls: 0 }; |
| 64 | + if unsafe { libc::ioctl(uffd, UFFDIO_API, &mut api) } < 0 { |
| 65 | + return Err(io::Error::last_os_error()); |
| 66 | + } |
| 67 | + let mut reg = UffdioRegister { |
| 68 | + range: UffdioRange { start, len }, |
| 69 | + mode: UFFDIO_REGISTER_MODE_MISSING, |
| 70 | + ioctls: 0, |
| 71 | + }; |
| 72 | + if unsafe { libc::ioctl(uffd, UFFDIO_REGISTER, &mut reg) } < 0 { |
| 73 | + return Err(io::Error::last_os_error()); |
| 74 | + } |
| 75 | + Ok(()) |
| 76 | +} |
| 77 | + |
| 78 | +fn copy_page(uffd: i32, image: &PageImage, addr: u64, page: usize) -> io::Result<()> { |
| 79 | + let base = addr & !((page as u64) - 1); |
| 80 | + let src = match image.page_at(addr, page) { |
| 81 | + Some(s) => s, |
| 82 | + None => { |
| 83 | + // Nothing to serve: zero-fill so the faulting thread makes progress |
| 84 | + // rather than hanging. (Restore images always cover their faults.) |
| 85 | + let zero = vec![0u8; page]; |
| 86 | + let mut c = UffdioCopy { |
| 87 | + dst: base, src: zero.as_ptr() as u64, len: page as u64, mode: 0, copy: 0, |
| 88 | + }; |
| 89 | + if unsafe { libc::ioctl(uffd, UFFDIO_COPY, &mut c) } < 0 { |
| 90 | + return Err(io::Error::last_os_error()); |
| 91 | + } |
| 92 | + return Ok(()); |
| 93 | + } |
| 94 | + }; |
| 95 | + let mut c = UffdioCopy { |
| 96 | + dst: base, src: src.as_ptr() as u64, len: page as u64, mode: 0, copy: 0, |
| 97 | + }; |
| 98 | + if unsafe { libc::ioctl(uffd, UFFDIO_COPY, &mut c) } < 0 { |
| 99 | + return Err(io::Error::last_os_error()); |
| 100 | + } |
| 101 | + Ok(()) |
| 102 | +} |
| 103 | + |
| 104 | +fn poll_once(uffd: i32, timeout_ms: i32) -> io::Result<bool> { |
| 105 | + let mut pfd = libc::pollfd { fd: uffd, events: libc::POLLIN, revents: 0 }; |
| 106 | + let n = unsafe { libc::poll(&mut pfd, 1, timeout_ms) }; |
| 107 | + if n < 0 { |
| 108 | + let e = io::Error::last_os_error(); |
| 109 | + if e.kind() == io::ErrorKind::Interrupted { return Ok(false); } |
| 110 | + return Err(e); |
| 111 | + } |
| 112 | + Ok(n > 0 && (pfd.revents & libc::POLLIN) != 0) |
| 113 | +} |
| 114 | + |
| 115 | +fn handle_ready(uffd: i32, image: &PageImage) -> io::Result<bool> { |
| 116 | + let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as usize; |
| 117 | + let mut msg = UffdMsg { event: 0, _pad: [0; 7], arg: [0; 3] }; |
| 118 | + let n = unsafe { |
| 119 | + libc::read(uffd, &mut msg as *mut _ as *mut libc::c_void, |
| 120 | + std::mem::size_of::<UffdMsg>()) |
| 121 | + }; |
| 122 | + if n == 0 { return Ok(false); } // EOF: all handles closed |
| 123 | + if n < 0 { |
| 124 | + let e = io::Error::last_os_error(); |
| 125 | + if e.kind() == io::ErrorKind::WouldBlock { return Ok(true); } |
| 126 | + return Err(e); |
| 127 | + } |
| 128 | + if msg.event == UFFD_EVENT_PAGEFAULT { |
| 129 | + // struct uffd_msg: 8-byte header (event+reserved), then the pagefault |
| 130 | + // arm { u64 flags; u64 address; ... }. With the 8-byte header the fault |
| 131 | + // address is the SECOND u64 of `arg` (byte offset 16), not the third. |
| 132 | + let addr = msg.arg[1]; |
| 133 | + copy_page(uffd, image, addr, page)?; |
| 134 | + } |
| 135 | + Ok(true) |
| 136 | +} |
| 137 | + |
| 138 | +/// Poll/serve until the uffd reports all handles closed (EOF). |
| 139 | +pub(crate) fn serve(uffd: i32, image: &PageImage) -> io::Result<()> { |
| 140 | + loop { |
| 141 | + if poll_once(uffd, -1)? { |
| 142 | + if !handle_ready(uffd, image)? { return Ok(()); } |
| 143 | + } |
| 144 | + } |
| 145 | +} |
| 146 | + |
| 147 | +/// Poll/serve until `stop` is set. Uses a short poll timeout so it observes the |
| 148 | +/// flag promptly even with no faults arriving. |
| 149 | +pub(crate) fn serve_until(uffd: i32, image: &PageImage, stop: &AtomicBool) { |
| 150 | + while !stop.load(Ordering::SeqCst) { |
| 151 | + match poll_once(uffd, 20) { |
| 152 | + Ok(true) => { let _ = handle_ready(uffd, image); } |
| 153 | + Ok(false) => {} |
| 154 | + Err(_) => break, |
| 155 | + } |
| 156 | + } |
| 157 | +} |
| 158 | + |
| 159 | +#[cfg(test)] |
| 160 | +mod tests { |
| 161 | + use super::*; |
| 162 | + |
| 163 | + #[test] |
| 164 | + fn page_at_returns_aligned_slice_within_run() { |
| 165 | + let img = PageImage { |
| 166 | + regions: vec![(0x1000, 0x3000, { |
| 167 | + let mut v = vec![0u8; 0x2000]; |
| 168 | + // second page all 0xAB |
| 169 | + for b in &mut v[0x1000..0x2000] { *b = 0xAB; } |
| 170 | + v |
| 171 | + })], |
| 172 | + }; |
| 173 | + let page = 0x1000usize; |
| 174 | + // An address in the second page returns that page, all 0xAB. |
| 175 | + let s = img.page_at(0x2500, page).expect("covered"); |
| 176 | + assert_eq!(s.len(), page); |
| 177 | + assert!(s.iter().all(|&b| b == 0xAB)); |
| 178 | + // Out of range. |
| 179 | + assert!(img.page_at(0x9000, page).is_none()); |
| 180 | + } |
| 181 | + |
| 182 | + #[test] |
| 183 | + fn serve_copies_faulted_page_into_a_registered_region() { |
| 184 | + use std::sync::atomic::{AtomicBool, Ordering}; |
| 185 | + use std::sync::Arc; |
| 186 | + |
| 187 | + let page = 4096usize; |
| 188 | + // Map an anon region we will register with uffd in THIS process. |
| 189 | + let len = page; |
| 190 | + let addr = unsafe { |
| 191 | + libc::mmap(std::ptr::null_mut(), len, |
| 192 | + libc::PROT_READ | libc::PROT_WRITE, |
| 193 | + libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, -1, 0) |
| 194 | + }; |
| 195 | + assert_ne!(addr, libc::MAP_FAILED); |
| 196 | + let start = addr as u64; |
| 197 | + |
| 198 | + // O_NONBLOCK is mandatory: poll()/epoll on a *blocking* userfaultfd always |
| 199 | + // returns POLLERR (never POLLIN), so the pager would spin forever. And |
| 200 | + // vm.unprivileged_userfaultfd=0 rejects plain userfaultfd(); UFFD_USER_MODE_ONLY |
| 201 | + // (0x1) is permitted unprivileged and handles user-mode faults, which is all |
| 202 | + // this test triggers. |
| 203 | + const UFFD_USER_MODE_ONLY: libc::c_int = 1; |
| 204 | + let flags = libc::O_CLOEXEC | libc::O_NONBLOCK; |
| 205 | + let uffd = { |
| 206 | + let plain = unsafe { libc::syscall(libc::SYS_userfaultfd, flags) } as i32; |
| 207 | + if plain >= 0 { |
| 208 | + plain |
| 209 | + } else { |
| 210 | + (unsafe { |
| 211 | + libc::syscall(libc::SYS_userfaultfd, flags | UFFD_USER_MODE_ONLY) |
| 212 | + }) as i32 |
| 213 | + } |
| 214 | + }; |
| 215 | + assert!(uffd >= 0, "userfaultfd (tried plain then USER_MODE_ONLY)"); |
| 216 | + register_api_and_range(uffd, start, len as u64).expect("register"); |
| 217 | + |
| 218 | + let img = PageImage { regions: vec![(start, start + len as u64, vec![0x5Au8; len])] }; |
| 219 | + let stop = Arc::new(AtomicBool::new(false)); |
| 220 | + |
| 221 | + // Pager on another thread. |
| 222 | + let stop2 = stop.clone(); |
| 223 | + let uffd_copy = uffd; |
| 224 | + let img_copy = img.clone(); |
| 225 | + let h = std::thread::spawn(move || serve_until(uffd_copy, &img_copy, &stop2)); |
| 226 | + |
| 227 | + // Touch the page -> fault -> pager fills with 0x5A. |
| 228 | + let byte = unsafe { std::ptr::read_volatile(addr as *const u8) }; |
| 229 | + assert_eq!(byte, 0x5A, "faulted page must be filled by the pager"); |
| 230 | + |
| 231 | + stop.store(true, Ordering::SeqCst); |
| 232 | + // Nudge the poll loop so it observes `stop` (write a dummy fault by |
| 233 | + // touching another registered page is unnecessary; serve_until uses a |
| 234 | + // short poll timeout). |
| 235 | + let _ = h.join(); |
| 236 | + unsafe { libc::munmap(addr, len); } |
| 237 | + } |
| 238 | +} |
0 commit comments