From dc1aaced93d7fe962dfd513aec5ddc24ae0f3e8c Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sat, 8 Aug 2026 11:31:23 +0800 Subject: [PATCH 1/7] feat(fspy): add fspy_alloc, a lock-free async-signal-safe allocator for the preload library MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The preload library interposes libc functions that POSIX declares async-signal-safe (open, stat, execve, ...), so its Rust allocations must not take libc malloc's locks: a signal handler, or the child of fork() in a multithreaded process, would deadlock on locks held by suspended or vanished threads. fspy_alloc routes every allocation in the preload cdylib through a lock-free size-class pool: power-of-two classes carve blocks from 1 MiB slabs, freed blocks recycle through per-class Treiber free lists made ABA-resistant by a 40-bit generation tag, and larger or over-aligned requests map directly. All memory comes from anonymous mappings issued as raw syscalls via rustix's linux_raw backend — on Linux the allocator relies on nothing from libc, discovering even the page size with an mprotect probe (no getauxval, no /proc, no minimum kernel version). macOS goes through the thin libSystem stubs, its only syscall interface. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 8 + Cargo.toml | 2 + crates/fspy_alloc/Cargo.toml | 14 + crates/fspy_alloc/src/class.rs | 38 ++ crates/fspy_alloc/src/lib.rs | 124 +++++ crates/fspy_alloc/src/mapping.rs | 62 +++ crates/fspy_alloc/src/mmap.rs | 211 +++++++ crates/fspy_alloc/src/pool.rs | 705 ++++++++++++++++++++++++ crates/fspy_alloc/src/slab.rs | 203 +++++++ crates/fspy_alloc/src/sys.rs | 25 + crates/fspy_alloc/tests/global_alloc.rs | 84 +++ crates/fspy_preload_unix/Cargo.toml | 1 + crates/fspy_preload_unix/src/lib.rs | 10 + 13 files changed, 1487 insertions(+) create mode 100644 crates/fspy_alloc/Cargo.toml create mode 100644 crates/fspy_alloc/src/class.rs create mode 100644 crates/fspy_alloc/src/lib.rs create mode 100644 crates/fspy_alloc/src/mapping.rs create mode 100644 crates/fspy_alloc/src/mmap.rs create mode 100644 crates/fspy_alloc/src/pool.rs create mode 100644 crates/fspy_alloc/src/slab.rs create mode 100644 crates/fspy_alloc/src/sys.rs create mode 100644 crates/fspy_alloc/tests/global_alloc.rs diff --git a/Cargo.lock b/Cargo.lock index 00ee74ab3..6182d52d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1243,6 +1243,13 @@ dependencies = [ "winsafe 0.0.27", ] +[[package]] +name = "fspy_alloc" +version = "0.0.0" +dependencies = [ + "rustix", +] + [[package]] name = "fspy_benchmark" version = "0.0.0" @@ -1299,6 +1306,7 @@ dependencies = [ "artifact_profile", "bstr", "ctor", + "fspy_alloc", "fspy_shared", "fspy_shared_unix", "libc", diff --git a/Cargo.toml b/Cargo.toml index ee6abee66..4385383c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,6 +73,7 @@ materialized_artifact = { path = "crates/materialized_artifact" } materialized_artifact_build = { path = "crates/materialized_artifact_build" } flate2 = "1.0.35" fspy = { path = "crates/fspy" } +fspy_alloc = { path = "crates/fspy_alloc" } fspy_benchmark_launcher = { path = "crates/fspy_benchmark_launcher", artifact = "bin" } fspy_benchmark_target = { path = "crates/fspy_benchmark_target", artifact = "bin" } fspy_detours_sys = { path = "crates/fspy_detours_sys" } @@ -120,6 +121,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" +rustix = { version = "1", default-features = false, features = ["mm", "param"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml new file mode 100644 index 000000000..809d0f7e8 --- /dev/null +++ b/crates/fspy_alloc/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "fspy_alloc" +edition = "2024" +license.workspace = true +publish = false + +[lib] +doctest = false + +[target.'cfg(unix)'.dependencies] +rustix = { workspace = true } + +[lints] +workspace = true diff --git a/crates/fspy_alloc/src/class.rs b/crates/fspy_alloc/src/class.rs new file mode 100644 index 000000000..df8d3797a --- /dev/null +++ b/crates/fspy_alloc/src/class.rs @@ -0,0 +1,38 @@ +//! Size-class policy: which layouts the pool serves, and at what block size. + +use core::alloc::Layout; + +const MIN_CLASS_SHIFT: u32 = 4; +const MAX_CLASS_SHIFT: u32 = 16; +pub const CLASS_COUNT: usize = (MAX_CLASS_SHIFT - MIN_CLASS_SHIFT) as usize + 1; +const MIN_BLOCK_SIZE: usize = 1 << MIN_CLASS_SHIFT; +pub const MAX_BLOCK_SIZE: usize = 1 << MAX_CLASS_SHIFT; +/// Block areas start at this alignment within a slab, making it the largest +/// alignment the pool can serve; stricter layouts map directly. +pub const MAX_POOL_ALIGN: usize = 4096; + +pub const fn block_size(class: usize) -> usize { + 1 << (MIN_CLASS_SHIFT as usize + class) +} + +/// Returns the size class for `layout`, or `None` if the request must be +/// mapped directly (too large or over-aligned). +pub const fn class_of(layout: Layout) -> Option { + if layout.align() > MAX_POOL_ALIGN { + return None; + } + let mut size = layout.size(); + // A block of `size >= align` at a `min(block size, 4 KiB)` boundary is + // aligned to `align` (both are powers of two and `align <= 4 KiB`). + if size < layout.align() { + size = layout.align(); + } + if size < MIN_BLOCK_SIZE { + size = MIN_BLOCK_SIZE; + } + if size > MAX_BLOCK_SIZE { + return None; + } + let shift = size.next_power_of_two().trailing_zeros(); + Some((shift - MIN_CLASS_SHIFT) as usize) +} diff --git a/crates/fspy_alloc/src/lib.rs b/crates/fspy_alloc/src/lib.rs new file mode 100644 index 000000000..90b2524cf --- /dev/null +++ b/crates/fspy_alloc/src/lib.rs @@ -0,0 +1,124 @@ +//! Lock-free, async-signal-safe global allocator for the fspy preload library. +//! +//! The preload library interposes libc functions that POSIX declares +//! async-signal-safe (`open`, `stat`, `execve`, ...). Programs may call these +//! from signal handlers, and — more commonly — from the child of `fork()` in a +//! multithreaded process, where only async-signal-safe calls are permitted: +//! the libc allocator's locks may be held forever by threads that no longer +//! exist after the fork. Routing the preload's Rust allocations through this +//! allocator keeps them safe in both contexts: +//! +//! - **No locks.** Every state transition is a lock-free compare-and-swap +//! loop: an attempt only retries because another running thread completed +//! its operation, so nothing ever waits on state that a thread which +//! vanished at `fork()` — or sits suspended under a signal handler — would +//! have to release. (Lock-free, not wait-free: an individual operation has +//! no fixed retry bound under active contention.) +//! - **No thread-locals.** TLS first-touch allocates through libc malloc on +//! some platforms (macOS thread-local variables), which would reintroduce +//! the hazard this crate exists to remove. +//! - **mmap-backed.** Memory comes straight from the kernel. On Linux the +//! allocator relies on nothing from libc: mapping syscalls are issued +//! directly (rustix's raw backend) and even the page size is discovered by +//! probing with raw syscalls. On macOS, which has no stable raw-syscall +//! ABI, calls go through the thin libSystem stubs. libc malloc is never +//! called anywhere. +//! +//! Design: power-of-two size classes (16 B ..= 64 KiB) carve blocks out of +//! 1 MiB slabs; freed blocks recycle through a per-class Treiber free list +//! made ABA-safe by a generation tag. Requests larger than the biggest class +//! (or over-aligned beyond 4 KiB) map and unmap directly. See the `pool` +//! module for the details. +//! +//! Because the allocator is a `const`-initialized static with no lazy setup, +//! it works from the very first allocation in the process — even before the +//! preload library's constructor runs. + +#![cfg_attr(not(test), no_std)] + +// Compile as an empty crate on non-unix targets: the allocator backs the unix +// preload library. A Windows backend can be added alongside `sys::Mmap` if +// the Windows preload ever needs one. + +#[cfg(unix)] +mod class; +#[cfg(unix)] +mod mapping; +#[cfg(unix)] +mod mmap; +#[cfg(unix)] +mod pool; +#[cfg(unix)] +mod slab; +#[cfg(unix)] +mod sys; + +#[cfg(unix)] +use core::{ + alloc::{GlobalAlloc, Layout}, + ptr::{self, NonNull}, +}; + +#[cfg(unix)] +use crate::{mmap::Mmap, pool::Pool}; + +/// A lock-free, async-signal-safe, fork-safe [`GlobalAlloc`] implementation. +/// +/// Intended to be installed as the `#[global_allocator]` of the fspy preload +/// library. All memory comes from anonymous mappings; libc malloc is never +/// called, no locks are taken, and no thread-local state is used. +/// +/// Capacity is bounded by design: each size class can hold at most 256 slabs +/// of 1 MiB (roughly 200 MiB per class). Requests beyond that — far outside +/// anything the preload library does — fail like any other out-of-memory +/// condition (`alloc` returns null). +#[cfg(unix)] +pub struct FspyAlloc { + pool: Pool, +} + +#[cfg(unix)] +impl FspyAlloc { + /// Creates the allocator. `const` so it can back a `static` with no + /// runtime initialization. + #[must_use] + pub const fn new() -> Self { + Self { pool: Pool::new() } + } +} + +#[cfg(unix)] +impl Default for FspyAlloc { + fn default() -> Self { + Self::new() + } +} + +// SAFETY: `Pool` hands out blocks that are non-null, at least `layout.size()` +// bytes large, aligned to at least `layout.align()`, and exclusively owned +// until returned via `dealloc`. Allocation failure is reported as null, and +// none of the methods unwind. +#[cfg(unix)] +unsafe impl GlobalAlloc for FspyAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + self.pool.alloc(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let Some(ptr) = NonNull::new(ptr) else { return }; + // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this + // allocator for this `layout`. + unsafe { self.pool.dealloc(ptr, layout) } + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + self.pool.alloc_zeroed(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let Some(ptr) = NonNull::new(ptr) else { return ptr::null_mut() }; + // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this + // allocator for this `layout`, and `new_size` is non-zero. + unsafe { self.pool.realloc(ptr, layout, new_size) }.map_or(ptr::null_mut(), NonNull::as_ptr) + } +} diff --git a/crates/fspy_alloc/src/mapping.rs b/crates/fspy_alloc/src/mapping.rs new file mode 100644 index 000000000..d3f53e726 --- /dev/null +++ b/crates/fspy_alloc/src/mapping.rs @@ -0,0 +1,62 @@ +//! Owned memory regions obtained from a [`Sys`] provider. + +use core::{marker::PhantomData, mem, ptr::NonNull}; + +use crate::sys::Sys; + +/// An owned region obtained from `S`, released on drop. +/// +/// This is the only place that calls [`Sys::unmap`]: pool code either lets a +/// `Mapping` drop (probe scratch, install races, freed large allocations) or +/// deliberately leaks it with [`Mapping::into_raw`] (published slabs, live +/// large allocations). Reconstructing ownership from a raw pointer via +/// [`Mapping::from_raw`] is the single unsafe step. +pub struct Mapping { + ptr: NonNull, + size: usize, + align: usize, + sys: PhantomData S>, +} + +impl Mapping { + /// Maps `size` bytes of zero-initialized memory aligned to `align` + /// (a power of two). Returns `None` when memory is exhausted. + pub fn new(size: usize, align: usize) -> Option { + let ptr = S::map(size, align)?; + Some(Self { ptr, size, align, sys: PhantomData }) + } + + /// Reclaims ownership of a mapping previously released with + /// [`Mapping::into_raw`]. + /// + /// # Safety + /// + /// `ptr` must have come from `Mapping::::into_raw` (or `Sys::map`) + /// with exactly this `size` and `align`, the region must not be in use, + /// and ownership must not be reclaimed twice. + pub unsafe fn from_raw(ptr: NonNull, size: usize, align: usize) -> Self { + Self { ptr, size, align, sys: PhantomData } + } + + /// The mapped region's base address. + pub const fn ptr(&self) -> NonNull { + self.ptr + } + + /// Releases ownership without unmapping; the region lives until (unless) + /// [`Mapping::from_raw`] reclaims it. + pub const fn into_raw(self) -> NonNull { + let ptr = self.ptr; + mem::forget(self); + ptr + } +} + +impl Drop for Mapping { + fn drop(&mut self) { + // SAFETY: this type owns the mapping (constructed from `Sys::map` + // directly or via the `from_raw` contract), and after drop nothing + // can use it. + unsafe { S::unmap(self.ptr, self.size, self.align) } + } +} diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/fspy_alloc/src/mmap.rs new file mode 100644 index 000000000..053356b56 --- /dev/null +++ b/crates/fspy_alloc/src/mmap.rs @@ -0,0 +1,211 @@ +//! Kernel-backed [`Sys`] provider: anonymous mappings via direct syscalls. +//! +//! On Linux the provider relies on nothing from libc — rustix issues raw +//! syscalls, and even the page size is discovered with raw syscalls (see +//! [`page_size`]). macOS has no stable raw-syscall ABI, so calls go through +//! the thin libSystem stubs there. + +use core::{ + ptr::{self, NonNull}, + sync::atomic::{AtomicUsize, Ordering}, +}; + +use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; + +use crate::sys::Sys; + +/// Returns the kernel page size, or `None` when it cannot be determined. +/// +/// On Linux the value is discovered with raw syscalls only — no libc, no +/// `/proc`, no minimum kernel version (see [`query_page_size`]). On other +/// unix platforms (macOS, where every syscall goes through libSystem by +/// platform contract anyway) it comes from `sysconf` via rustix. Either +/// way the result is a process constant, validated as a power of two and +/// cached, so `map` and `unmap` can never disagree on rounding; an +/// undeterminable page size fails the allocation rather than guessing. +fn page_size() -> Option { + static CACHE: AtomicUsize = AtomicUsize::new(0); + let cached = CACHE.load(Ordering::Relaxed); + if cached != 0 { + return Some(cached); + } + let page = query_page_size()?; + if !page.is_power_of_two() { + return None; + } + // First store wins; every query returns the same value, so the cache + // only avoids repeated probing. + let _ = CACHE.compare_exchange(0, page, Ordering::Relaxed, Ordering::Relaxed); + Some(page) +} + +/// Discovers the page size by probing, using nothing but raw syscalls. +/// +/// `mprotect` fails with `EINVAL` unless its address is a multiple of the +/// page size, so re-protecting a scratch mapping at increasing +/// power-of-two offsets identifies the page size as the first offset the +/// kernel accepts (for a power-of-two page size P, the first power of two +/// P divides is P itself). A handful of syscalls, once per process: +/// async-signal-safe, fork-safe, and independent of libc, `/proc` +/// availability, and kernel version. +/// +/// Why not `rustix::param::page_size()` here (it's fine on macOS)? On +/// Linux the allocator must not rely on libc, which rules out rustix's +/// `use-libc-auxv` (`getauxval`) configuration — and rustix's libc-free +/// fallback is unusable *inside* a global allocator: +/// +/// - Its lazy init tries `prctl(PR_GET_AUXV)` (kernel 6.4+ only) and +/// otherwise reads `/proc/self/auxv`; with rustix's `alloc` feature +/// enabled, that read path heap-allocates (`Vec`) — through *this* +/// allocator, whose `map` is the caller waiting on the page size — +/// recursing unboundedly. And we cannot pin `alloc` off: Cargo feature +/// unification lets any other rustix user in the build graph enable it +/// for our copy. +/// - It panics on read errors, truncated auxv, or both sources being +/// unavailable, where this allocator requires failure to surface as +/// `None` (panic formatting itself allocates, re-entering the same +/// uninitialized path). +/// +/// The probe has neither problem: no allocation, no panic, no minimum +/// kernel, and its worst case is `None`. +#[cfg(target_os = "linux")] +fn query_page_size() -> Option { + use rustix::mm::{MprotectFlags, mprotect}; + + /// Smallest page size of any Linux configuration; the probe starts + /// here. + const MIN_PROBE_PAGE: usize = 4096; + /// Generous upper bound for the probe; no supported configuration + /// uses larger pages. + const MAX_PROBE_PAGE: usize = 1 << 20; + + // Large enough that every probe below stays inside the mapping even + // after the kernel rounds the one-byte length up to a full page. + let scratch = map_anonymous(2 * MAX_PROBE_PAGE)?; + let mut page = None; + let mut offset = MIN_PROBE_PAGE; + while offset <= MAX_PROBE_PAGE { + // SAFETY: `scratch + offset` (plus the page-rounded single byte) + // lies within the scratch mapping, which we exclusively own; the + // protection flags match the mapping's existing ones. + let accepted = unsafe { + mprotect( + scratch.as_ptr().add(offset).cast(), + 1, + MprotectFlags::READ | MprotectFlags::WRITE, + ) + } + .is_ok(); + if accepted { + page = Some(offset); + break; + } + offset *= 2; + } + // SAFETY: releasing the scratch mapping created above. + let _ = unsafe { munmap(scratch.as_ptr().cast(), 2 * MAX_PROBE_PAGE) }; + page +} + +#[cfg(not(target_os = "linux"))] +#[expect( + clippy::unnecessary_wraps, + reason = "must match the signature of the fallible Linux probe variant" +)] +fn query_page_size() -> Option { + Some(rustix::param::page_size()) +} + +const fn round_up(value: usize, align: usize) -> usize { + (value + align - 1) & !(align - 1) +} + +fn map_anonymous(len: usize) -> Option> { + // SAFETY: a fresh anonymous private mapping at no particular address + // has no memory-safety preconditions. + let ptr = unsafe { + mmap_anonymous(ptr::null_mut(), len, ProtFlags::READ | ProtFlags::WRITE, MapFlags::PRIVATE) + } + .ok()?; + NonNull::new(ptr.cast::()) +} + +/// Kernel-backed provider: anonymous mappings obtained through direct +/// syscalls, alignment achieved by over-mapping and trimming. +pub struct Mmap; + +impl Sys for Mmap { + fn map(size: usize, align: usize) -> Option> { + debug_assert!(align.is_power_of_two()); + let page = page_size()?; + let size = round_up(size.max(1), page); + + if align <= page { + // Mapping results are aligned to the (verified real) page + // size. + return map_anonymous(size); + } + + // Over-map by `align`, then unmap the misaligned head and the + // leftover tail. All cut points are page-aligned: `raw` and + // `aligned` are page-aligned, and `size`/`align` are multiples of + // the page size. + let raw = map_anonymous(size.checked_add(align)?)?; + let raw_addr = raw.as_ptr().addr(); + let aligned_addr = round_up(raw_addr, align); + let head = aligned_addr - raw_addr; + let tail = align - head; + if head > 0 { + // SAFETY: `[raw_addr, raw_addr + head)` lies within the fresh + // mapping and `raw_addr` is page-aligned. Failure is + // impossible for a region we own; if it happened anyway the + // pages would merely stay mapped. + let _ = unsafe { munmap(raw.as_ptr().cast(), head) }; + } + if tail > 0 { + let tail_start = raw.as_ptr().with_addr(aligned_addr + size); + // SAFETY: `[aligned_addr + size, raw_addr + size + align)` + // lies within the fresh mapping and its start is page-aligned + // (see above). Failure is impossible for a region we own. + let _ = unsafe { munmap(tail_start.cast(), tail) }; + } + // `aligned_addr` lies inside a successful mapping and so can + // never be zero, but checked construction costs nothing here. + core::num::NonZero::new(aligned_addr).map(|addr| raw.with_addr(addr)) + } + + unsafe fn unmap(ptr: NonNull, size: usize, _align: usize) { + // A successful `map` proved the page size, so this cannot fail + // for a live mapping; if it somehow did, leaking the region is + // the only safe response. + let Some(page) = page_size() else { return }; + // Whether or not `map` trimmed for alignment, the retained region + // is exactly `[ptr, ptr + round_up(size, page))`. + let size = round_up(size.max(1), page); + // SAFETY: caller contract — `ptr`/`size` describe a live mapping + // returned by `map`. Failure is impossible for a region we own. + let _ = unsafe { munmap(ptr.as_ptr().cast(), size) }; + } +} + +#[cfg(all(test, not(miri)))] +mod tests { + use super::Mmap; + use crate::sys::Sys; + + #[test] + fn real_mappings_are_aligned_zeroed_and_writable() { + for (size, align) in [(1, 1), (4096, 4096), (100, 1 << 20), (5 << 20, 4096)] { + let ptr = Mmap::map(size, align).unwrap(); + assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); + for i in 0..size { + // SAFETY: fresh exclusive mapping of at least `size` bytes. + assert_eq!(unsafe { ptr.as_ptr().add(i).read() }, 0); + } + // SAFETY: fresh exclusive mapping of at least `size` bytes. + unsafe { ptr.as_ptr().write_bytes(0x5A, size) }; + // SAFETY: mapped above with the same size and alignment. + unsafe { Mmap::unmap(ptr, size, align) }; + } + } +} diff --git a/crates/fspy_alloc/src/pool.rs b/crates/fspy_alloc/src/pool.rs new file mode 100644 index 000000000..6bb4f82f1 --- /dev/null +++ b/crates/fspy_alloc/src/pool.rs @@ -0,0 +1,705 @@ +//! Lock-free size-class pool. +//! +//! Layouts round up to a power-of-two class (see the `class` module) whose +//! blocks are carved from slabs (see the `slab` module for the memory +//! layout). +//! +//! # Concurrency +//! +//! Each class has a Treiber-stack free list plus a bump cursor over the +//! newest ("active") slab: +//! +//! - **alloc** pops the free list, or carves the next block off the active +//! slab, installing a fresh slab when the active one is exhausted. +//! - **dealloc** pushes the block back onto its class's free list. +//! +//! The 64-bit list head packs a 40-bit generation tag next to the 24-bit +//! block reference; the tag advances on every successful push and pop, which +//! makes the classic Treiber-stack ABA failure require a thread to stall +//! between its head load and CAS while other threads perform an exact +//! multiple of 2^40 head mutations. That is a probabilistic defense, not a +//! formal impossibility — see [`HEAD_TAG_SHIFT`]. Slab installation races +//! are resolved with a compare-and-swap on the slab table slot; the loser +//! simply returns its (never-published) mapping. +//! +//! The pool is **lock-free, not wait-free**: a CAS loop can retry +//! indefinitely under contention, but a retry only ever happens because +//! another thread completed an operation, so system-wide progress never +//! stalls — and, the property fork- and signal-safety actually require, no +//! operation ever waits on state that only a suspended or vanished thread +//! could release. Blocks are aligned to `min(block size, 4 KiB)`; requests +//! with stricter alignment (or size beyond the largest class) bypass the +//! pool and map directly. + +use core::{ + alloc::Layout, + marker::PhantomData, + ptr, + ptr::NonNull, + sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, Ordering}, +}; + +use crate::{ + class::{CLASS_COUNT, class_of}, + mapping::Mapping, + slab::{BLOCKS_PER_SLAB, SLAB_SIZE, Slab}, + sys::Sys, +}; + +/// Bounded by the 8 bits reserved for slab indices in a packed block +/// reference. Caps each class at ~200 MiB. +const MAX_SLABS_PER_CLASS: usize = 256; + +/// Sentinel for "no block" in the 24-bit packed-reference field of a +/// free-list head. Never collides with a real reference: block indices stay +/// below `0xFFFF` (asserted with the slab geometry below). +const NO_BLOCK: u32 = 0x00FF_FFFF; +/// Sentinel for "no slab installed yet" in [`ClassState::active`]. +const NO_SLAB: u32 = u32::MAX; + +const _: () = { + assert!(MAX_SLABS_PER_CLASS <= 1 << 8, "slab index must fit in 8 bits"); + let mut class = 0; + while class < CLASS_COUNT { + assert!( + BLOCKS_PER_SLAB[class] <= 0xFFFF, + "block indices must fit in 16 bits and stay below the NO_BLOCK sentinel" + ); + class += 1; + } +}; + +/// Packs a block's location into 24 bits: slab index in bits 16..24, block +/// index in bits 0..16. +#[expect( + clippy::cast_possible_truncation, + reason = "callers pass indices bounded by MAX_SLABS_PER_CLASS and BLOCKS_PER_SLAB" +)] +const fn pack_ref(slab_idx: usize, block_idx: usize) -> u32 { + ((slab_idx as u32) << 16) | (block_idx as u32) +} + +const fn unpack_ref(packed: u32) -> (usize, usize) { + (((packed >> 16) & 0xFF) as usize, (packed & 0xFFFF) as usize) +} + +/// A free-list head is `[generation tag : 40 | packed block reference : 24]`. +/// The tag advances on every successful push and pop, so a stale +/// compare-and-swap can only succeed if its thread stalls between head load +/// and CAS while others perform an exact multiple of 2^40 head mutations — +/// not a formal impossibility, but hours of maximum-rate churn inside one +/// stalled instruction window. +const HEAD_TAG_SHIFT: u32 = 24; +const HEAD_TAG_MASK: u64 = (1 << 40) - 1; +const HEAD_REF_MASK: u64 = (1 << HEAD_TAG_SHIFT) - 1; + +/// Splits a free-list head into `(generation tag, packed block reference)`. +const fn head_parts(head: u64) -> (u64, u32) { + // The mask keeps the reference within 24 bits, so the cast is lossless. + (head >> HEAD_TAG_SHIFT, (head & HEAD_REF_MASK) as u32) +} + +const fn head_from_parts(tag: u64, block_ref: u32) -> u64 { + ((tag & HEAD_TAG_MASK) << HEAD_TAG_SHIFT) | block_ref as u64 +} + +/// Aligned to its own cache-line region so hot heads of different classes +/// don't false-share. +#[repr(align(128))] +struct ClassState { + /// Treiber free-list head: `[generation tag : 40 | packed ref : 24]`. + head: AtomicU64, + /// Index of the slab currently being carved, or [`NO_SLAB`] before the + /// first slab is installed. + active: AtomicU32, + /// Base addresses of installed slabs. Written once (null → mapping) and + /// never cleared. + slabs: [AtomicPtr; MAX_SLABS_PER_CLASS], +} + +impl ClassState { + const fn new() -> Self { + Self { + head: AtomicU64::new(head_from_parts(0, NO_BLOCK)), + active: AtomicU32::new(NO_SLAB), + slabs: [const { AtomicPtr::new(ptr::null_mut()) }; MAX_SLABS_PER_CLASS], + } + } +} + +/// A block handed out by [`Pool::alloc_block`], remembering whether it came +/// off the free list (may contain stale data) or was carved fresh off a slab +/// (still kernel-zeroed). +enum ClassBlock { + Recycled(NonNull), + Carved(NonNull), +} + +enum Carve { + Block(NonNull), + /// A new slab was (or concurrently got) installed; retry the allocation. + Retry, + /// Out of slab table entries or out of memory. + Exhausted, +} + +/// The allocator core, generic over its [`Sys`] memory provider. +pub struct Pool { + classes: [ClassState; CLASS_COUNT], + sys: PhantomData S>, +} + +impl Pool { + #[expect( + clippy::large_stack_arrays, + reason = "the class table (~30 KiB) is only ever materialized into a const-initialized static, never built on a runtime stack" + )] + pub const fn new() -> Self { + Self { classes: [const { ClassState::new() }; CLASS_COUNT], sys: PhantomData } + } + + /// Allocates memory for `layout`. Returns `None` on exhaustion. + pub fn alloc(&self, layout: Layout) -> Option> { + match class_of(layout) { + Some(class) => match self.alloc_block(class)? { + ClassBlock::Recycled(ptr) | ClassBlock::Carved(ptr) => Some(ptr), + }, + // Deliberately leaked until `dealloc` reclaims ownership. + None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), + } + } + + /// Like [`Pool::alloc`], but the returned memory is zeroed. + pub fn alloc_zeroed(&self, layout: Layout) -> Option> { + match class_of(layout) { + Some(class) => match self.alloc_block(class)? { + // Freshly carved blocks are still kernel-zeroed. + ClassBlock::Carved(ptr) => Some(ptr), + ClassBlock::Recycled(ptr) => { + // SAFETY: the block is freshly allocated, exclusively + // ours, and at least `layout.size()` bytes. + unsafe { ptr.as_ptr().write_bytes(0, layout.size()) }; + Some(ptr) + } + }, + // Fresh mappings are zeroed by the kernel; deliberately leaked + // until `dealloc` reclaims ownership. + None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), + } + } + + /// Releases memory obtained from this pool. + /// + /// # Safety + /// + /// `ptr` must have been returned by this pool for exactly this `layout` + /// and must not be used afterwards. + pub unsafe fn dealloc(&self, ptr: NonNull, layout: Layout) { + match class_of(layout) { + // SAFETY: caller contract — a `Some(class)` layout was served + // from the pool, so `ptr` is a live block of this class. + Some(class) => unsafe { self.push_free(ptr, class) }, + // SAFETY: caller contract — a `None` layout was served by the + // mapping path with these parameters and is no longer in use, so + // ownership can be reclaimed (and the region dropped). + None => drop(unsafe { Mapping::::from_raw(ptr, layout.size(), layout.align()) }), + } + } + + /// Grows or shrinks an allocation, preserving contents up to the smaller + /// of the old and new sizes. Returns `None` on exhaustion (the original + /// allocation stays valid). + /// + /// # Safety + /// + /// `ptr` must have been returned by this pool for exactly `layout`, and + /// `new_size` must be non-zero. + pub unsafe fn realloc( + &self, + ptr: NonNull, + layout: Layout, + new_size: usize, + ) -> Option> { + let new_layout = Layout::from_size_align(new_size, layout.align()).ok()?; + let class = class_of(layout); + if class.is_some() && class == class_of(new_layout) { + // Same size class: the existing block already fits. + return Some(ptr); + } + let new_ptr = self.alloc(new_layout)?; + // SAFETY: `new_ptr` is a fresh exclusive allocation of at least + // `new_size` bytes; `ptr` is valid for `layout.size()` bytes (caller + // contract); distinct allocations never overlap. + unsafe { + new_ptr.as_ptr().copy_from_nonoverlapping(ptr.as_ptr(), layout.size().min(new_size)); + } + // SAFETY: caller contract — `ptr` came from this pool with `layout`. + unsafe { self.dealloc(ptr, layout) }; + Some(new_ptr) + } + + fn alloc_block(&self, class: usize) -> Option { + loop { + if let Some(ptr) = self.pop_free(class) { + return Some(ClassBlock::Recycled(ptr)); + } + match self.carve(class) { + Carve::Block(ptr) => return Some(ClassBlock::Carved(ptr)), + Carve::Retry => {} + Carve::Exhausted => return None, + } + } + } + + /// Returns the published slab of `class` at `slab_idx`, if installed. + fn published_slab(&self, class: usize, slab_idx: usize) -> Option { + let base = NonNull::new(self.classes[class].slabs[slab_idx].load(Ordering::Acquire))?; + // SAFETY: non-null slab-table entries are only published (with + // `Release`, paired with the `Acquire` above) after + // `Slab::init_header` ran on a fresh SLAB_SIZE-aligned mapping, and + // are never cleared or unmapped while the pool is in use. + Some(unsafe { Slab::from_published(base, class) }) + } + + /// Pops a block off the class's free list. + fn pop_free(&self, class: usize) -> Option> { + let state = &self.classes[class]; + loop { + let head = state.head.load(Ordering::Acquire); + let (tag, block_ref) = head_parts(head); + if block_ref == NO_BLOCK { + return None; + } + let (slab_idx, block_idx) = unpack_ref(block_ref); + // A listed block was carved from its slab, so the slab is always + // published; `None` here is unreachable in practice. + let slab = self.published_slab(class, slab_idx)?; + // Read the successor link before the CAS. If another thread pops + // this block first, the tag comparison below fails and the value + // read here is discarded; since links live in the atomic side + // table, the racing read itself is well-defined. + let next = slab.link(block_idx).load(Ordering::Relaxed); + let new_head = head_from_parts(tag.wrapping_add(1), next); + if state + .head + .compare_exchange_weak(head, new_head, Ordering::AcqRel, Ordering::Relaxed) + .is_ok() + { + return Some(slab.block(block_idx)); + } + } + } + + /// Pushes a block onto the class's free list. + /// + /// # Safety + /// + /// `ptr` must be a block of `class` previously returned by this pool and + /// no longer in use. + unsafe fn push_free(&self, ptr: NonNull, class: usize) { + // SAFETY: caller contract — `ptr` is a live block of `class` from + // this pool. (The impossible `None` would merely leak the block.) + let Some(slab) = (unsafe { Slab::of_block(ptr, class) }) else { return }; + let block_idx = slab.block_index(ptr); + let packed = pack_ref(slab.slab_idx(), block_idx); + let link = slab.link(block_idx); + let state = &self.classes[class]; + loop { + let head = state.head.load(Ordering::Relaxed); + let (tag, head_ref) = head_parts(head); + link.store(head_ref, Ordering::Relaxed); + let new_head = head_from_parts(tag.wrapping_add(1), packed); + // `Release` publishes the link store above to the eventual popper. + if state + .head + .compare_exchange_weak(head, new_head, Ordering::Release, Ordering::Relaxed) + .is_ok() + { + return; + } + } + } + + /// Carves the next block off the class's active slab, installing a new + /// slab if the active one is exhausted (or none exists yet). + fn carve(&self, class: usize) -> Carve { + let state = &self.classes[class]; + let active = state.active.load(Ordering::Acquire); + // `active` is only published after its slab pointer, so a live + // `active` always resolves to a published slab. + if active != NO_SLAB + && let Some(slab) = self.published_slab(class, active as usize) + { + let carved_idx = slab.carved().fetch_add(1, Ordering::Relaxed); + if let Some(block_idx) = carved_to_block_idx(carved_idx, class) { + return Carve::Block(slab.block(block_idx)); + } + // Active slab exhausted; fall through to install the next one. + } + let next_idx = if active == NO_SLAB { 0 } else { active as usize + 1 }; + if next_idx >= MAX_SLABS_PER_CLASS { + return Carve::Exhausted; + } + self.install_slab(class, next_idx); + if self.published_slab(class, next_idx).is_none() { + // Our mapping failed and no other thread succeeded either. + return Carve::Exhausted; + } + #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] + let next_active = next_idx as u32; + // Advance `active`; losing the race just means another thread already + // advanced it. Either way the retry re-reads it. + let _ = + state.active.compare_exchange(active, next_active, Ordering::AcqRel, Ordering::Relaxed); + Carve::Retry + } + + /// Maps and publishes the slab at `slab_idx`, unless another thread beats + /// us to it (or the mapping fails, leaving the slot null). + fn install_slab(&self, class: usize, slab_idx: usize) { + let state = &self.classes[class]; + if !state.slabs[slab_idx].load(Ordering::Acquire).is_null() { + return; + } + let Some(mapping) = Mapping::::new(SLAB_SIZE, SLAB_SIZE) else { return }; + #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] + let idx = slab_idx as u32; + // SAFETY: `mapping` is a fresh, exclusive, zero-initialized, + // SLAB_SIZE-byte and SLAB_SIZE-aligned mapping. + unsafe { Slab::init_header(mapping.ptr(), idx) }; + if state.slabs[slab_idx] + .compare_exchange( + ptr::null_mut(), + mapping.ptr().as_ptr(), + Ordering::Release, + Ordering::Relaxed, + ) + .is_ok() + { + // Published: the slab now lives for the rest of the process. + let _ = mapping.into_raw(); + } + // Otherwise another thread installed this slot first; our + // never-published mapping is released when `mapping` drops. + } + + /// Tears down all slab mappings. Test-only: the global allocator lives in + /// a static and never releases its slabs, but tests (and Miri's leak + /// checker) want a clean shutdown. + #[cfg(test)] + fn unmap_all_slabs(&mut self) { + for state in &mut self.classes { + *state.head.get_mut() = head_from_parts(0, NO_BLOCK); + *state.active.get_mut() = NO_SLAB; + for slot in &mut state.slabs { + let slab = core::mem::replace(slot.get_mut(), ptr::null_mut()); + if let Some(base) = NonNull::new(slab) { + // SAFETY: `base` was leaked into the table by + // `install_slab` with these parameters; `&mut self` + // guarantees no concurrent (or future) use of its blocks. + drop(unsafe { Mapping::::from_raw(base, SLAB_SIZE, SLAB_SIZE) }); + } + } + } + } +} + +/// Converts a raw carve-counter value into a block index, or `None` if the +/// slab is exhausted. +fn carved_to_block_idx(carved: u64, class: usize) -> Option { + let idx = usize::try_from(carved).ok()?; + (idx < BLOCKS_PER_SLAB[class]).then_some(idx) +} + +#[cfg(test)] +mod tests { + use std::{boxed::Box, sync::mpsc, thread, vec::Vec}; + + use super::*; + use crate::class::{MAX_BLOCK_SIZE, MAX_POOL_ALIGN, block_size}; + + /// Host-allocator-backed [`Sys`] so the pool core runs under Miri and on + /// any platform. + struct TestSys; + + impl Sys for TestSys { + fn map(size: usize, align: usize) -> Option> { + let layout = Layout::from_size_align(size.max(1), align).ok()?; + // SAFETY: `layout` has non-zero size. + NonNull::new(unsafe { std::alloc::alloc_zeroed(layout) }) + } + + unsafe fn unmap(ptr: NonNull, size: usize, align: usize) { + let layout = Layout::from_size_align(size.max(1), align).unwrap(); + // SAFETY: caller contract — `ptr` was returned by `map`, which + // used exactly this layout. + unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) } + } + } + + fn with_pool(test: impl FnOnce(&Pool)) { + let mut pool = Box::new(Pool::::new()); + test(&pool); + pool.unmap_all_slabs(); + } + + fn layout(size: usize, align: usize) -> Layout { + Layout::from_size_align(size, align).unwrap() + } + + #[test] + fn head_tag_wraps_within_its_field() { + let head = head_from_parts(HEAD_TAG_MASK, 42); + assert_eq!(head_parts(head), (HEAD_TAG_MASK, 42)); + // Advancing the maximal tag must wrap to zero without touching the + // reference bits. + let wrapped = head_from_parts(HEAD_TAG_MASK.wrapping_add(1), 42); + assert_eq!(head_parts(wrapped), (0, 42)); + } + + #[test] + fn packed_refs_cannot_collide_with_the_sentinel() { + for &blocks in &BLOCKS_PER_SLAB { + let max_ref = pack_ref(MAX_SLABS_PER_CLASS - 1, blocks - 1); + assert_ne!(max_ref, NO_BLOCK); + assert!(u64::from(max_ref) <= HEAD_REF_MASK); + } + } + + #[test] + fn classify_boundaries() { + assert_eq!(class_of(layout(1, 1)), Some(0)); + assert_eq!(class_of(layout(16, 1)), Some(0)); + assert_eq!(class_of(layout(17, 1)), Some(1)); + assert_eq!(class_of(layout(MAX_BLOCK_SIZE, 8)), Some(CLASS_COUNT - 1)); + assert_eq!(class_of(layout(MAX_BLOCK_SIZE + 1, 8)), None); + // Alignment can raise the class. + assert_eq!(class_of(layout(8, 1024)), class_of(layout(1024, 8))); + // Over-aligned layouts bypass the pool. + assert_eq!(class_of(layout(16, MAX_POOL_ALIGN * 2)), None); + } + + #[test] + fn round_trip_every_class() { + with_pool(|pool| { + for class in 0..CLASS_COUNT { + let l = layout(block_size(class), 8); + let ptr = pool.alloc(l).unwrap(); + // SAFETY: fresh exclusive allocation of `block_size` bytes. + unsafe { ptr.as_ptr().write_bytes(0xAB, l.size()) }; + // SAFETY: reading back the block we just wrote. + let last = unsafe { ptr.as_ptr().add(l.size() - 1).read() }; + assert_eq!(last, 0xAB); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn block_alignment() { + with_pool(|pool| { + for align in [8, 64, 1024, MAX_POOL_ALIGN] { + let l = layout(24, align); + let ptr = pool.alloc(l).unwrap(); + assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn free_list_recycles_lifo() { + with_pool(|pool| { + let l = layout(100, 8); + let first = pool.alloc(l).unwrap(); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(first, l) }; + let second = pool.alloc(l).unwrap(); + assert_eq!(first, second); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(second, l) }; + }); + } + + #[test] + fn carves_across_multiple_slabs() { + with_pool(|pool| { + // The largest class has the fewest blocks per slab, so a couple + // dozen live allocations force several slab installations. + let l = layout(MAX_BLOCK_SIZE, 8); + let count = BLOCKS_PER_SLAB[CLASS_COUNT - 1] * 3 + 1; + let blocks: Vec> = (0..count).map(|_| pool.alloc(l).unwrap()).collect(); + for (i, ptr) in blocks.iter().enumerate() { + assert!(blocks[..i].iter().all(|other| other != ptr), "duplicate block"); + // SAFETY: live exclusive allocation of `MAX_BLOCK_SIZE` bytes. + unsafe { ptr.as_ptr().write_bytes(0x5A, l.size()) }; + } + for ptr in blocks { + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn large_allocations_bypass_pool() { + with_pool(|pool| { + for l in [layout(MAX_BLOCK_SIZE + 1, 8), layout(5 << 20, 8), layout(64, 8192)] { + let ptr = pool.alloc(l).unwrap(); + assert_eq!(ptr.as_ptr().addr() % l.align(), 0); + // SAFETY: fresh exclusive allocation of `l.size()` bytes. + unsafe { ptr.as_ptr().write_bytes(0xCD, l.size()) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn realloc_within_class_keeps_block() { + with_pool(|pool| { + let l = layout(100, 8); + let ptr = pool.alloc(l).unwrap(); + // SAFETY: realloc contract — `ptr` allocated with `l`. + let grown = unsafe { pool.realloc(ptr, l, 120) }.unwrap(); + assert_eq!(ptr, grown, "same class must realloc in place"); + // SAFETY: allocated above; 120 rounds to the same class as 100. + unsafe { pool.dealloc(grown, layout(120, 8)) }; + }); + } + + #[test] + fn realloc_across_classes_preserves_contents() { + with_pool(|pool| { + let l = layout(64, 8); + let ptr = pool.alloc(l).unwrap(); + for i in 0..64u8 { + // SAFETY: live exclusive allocation of 64 bytes. + unsafe { ptr.as_ptr().add(usize::from(i)).write(i) }; + } + // SAFETY: realloc contract — `ptr` allocated with `l`. + let grown = unsafe { pool.realloc(ptr, l, 4096) }.unwrap(); + for i in 0..64u8 { + // SAFETY: `grown` is live for 4096 bytes. + let got = unsafe { grown.as_ptr().add(usize::from(i)).read() }; + assert_eq!(got, i); + } + // Shrink across classes, and from the large path back into a class. + // SAFETY: `grown` allocated with the 4096 layout above. + let shrunk = unsafe { pool.realloc(grown, layout(4096, 8), 8) }.unwrap(); + // SAFETY: `shrunk` is live for 8 bytes. + assert_eq!(unsafe { shrunk.as_ptr().read() }, 0); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(shrunk, layout(8, 8)) }; + }); + } + + #[test] + fn alloc_zeroed_scrubs_recycled_blocks() { + with_pool(|pool| { + let l = layout(256, 8); + let dirty = pool.alloc(l).unwrap(); + // SAFETY: live exclusive allocation of 256 bytes. + unsafe { dirty.as_ptr().write_bytes(0xFF, l.size()) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(dirty, l) }; + + let zeroed = pool.alloc_zeroed(l).unwrap(); + assert_eq!( + zeroed, dirty, + "must recycle the dirty block for this test to be meaningful" + ); + for i in 0..l.size() { + // SAFETY: live exclusive allocation of 256 bytes. + assert_eq!(unsafe { zeroed.as_ptr().add(i).read() }, 0); + } + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(zeroed, l) }; + }); + } + + /// Cheap deterministic PRNG so the stress test needs no dependencies. + fn lcg(state: &mut u64) -> u64 { + *state = + state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407); + *state >> 33 + } + + /// An owned allocation in flight between threads. + struct SendBlock(NonNull, Layout); + // SAFETY: SendBlock represents exclusive ownership of the block, which is + // transferred wholesale to the receiving thread. + unsafe impl Send for SendBlock {} + + #[test] + fn concurrent_stress_with_cross_thread_frees() { + const THREADS: usize = if cfg!(miri) { 3 } else { 8 }; + const OPS: usize = if cfg!(miri) { 60 } else { 20_000 }; + + with_pool(|pool| { + thread::scope(|scope| { + // Each thread frees blocks allocated by its neighbor, so + // pushes and pops of the same list run on different threads. + // Bounded channels keep the number of in-flight blocks well + // under the pool's per-class capacity regardless of thread + // scheduling. + let (senders, receivers): (Vec<_>, Vec<_>) = + (0..THREADS).map(|_| mpsc::sync_channel::(64)).unzip(); + let mut senders_rotated: Vec<_> = senders.into_iter().map(Some).collect(); + senders_rotated.rotate_left(1); + + for (thread_idx, (receiver, sender)) in + receivers.into_iter().zip(&mut senders_rotated).enumerate() + { + let sender = sender.take().unwrap(); + let pool = &*pool; + scope.spawn(move || { + let mut rng = 0x9E37_79B9_7F4A_7C15_u64 ^ thread_idx as u64; + for _ in 0..OPS { + let size = 1 + usize::try_from(lcg(&mut rng)).unwrap() + % (3 * MAX_BLOCK_SIZE / 2); + let align = 1 << (lcg(&mut rng) % 7); + let l = layout(size, align); + let ptr = pool.alloc(l).unwrap(); + let marker = u8::try_from(lcg(&mut rng) & 0xFF).unwrap(); + // SAFETY: fresh exclusive allocation of `size` bytes. + unsafe { + ptr.as_ptr().write(marker); + ptr.as_ptr().add(size - 1).write(marker); + } + // SAFETY: we still exclusively own the block. + let (first, last) = + unsafe { (ptr.as_ptr().read(), ptr.as_ptr().add(size - 1).read()) }; + assert_eq!((first, last), (marker, marker), "block corrupted"); + if let Err(returned) = sender.try_send(SendBlock(ptr, l)) { + // Neighbor's queue is full (or it finished); + // free locally instead of blocking, which in + // a ring of senders could deadlock. + let (mpsc::TrySendError::Full(SendBlock(ptr, l)) + | mpsc::TrySendError::Disconnected(SendBlock(ptr, l))) = returned; + // SAFETY: allocated above with layout `l`. + unsafe { pool.dealloc(ptr, l) }; + } + // Drain what our own producer has sent so far, + // interleaving cross-thread pops with pushes. + while let Ok(SendBlock(ptr, l)) = receiver.try_recv() { + // SAFETY: the neighbor allocated `ptr` with + // `l` and transferred ownership over the + // channel. + unsafe { pool.dealloc(ptr, l) }; + } + } + drop(sender); + for SendBlock(ptr, l) in receiver { + // SAFETY: the neighbor allocated `ptr` with `l` + // and transferred ownership over the channel. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + }); + }); + } +} diff --git a/crates/fspy_alloc/src/slab.rs b/crates/fspy_alloc/src/slab.rs new file mode 100644 index 000000000..3aca476dc --- /dev/null +++ b/crates/fspy_alloc/src/slab.rs @@ -0,0 +1,203 @@ +//! Slab memory layout and the [`Slab`] view type. +//! +//! Each size class carves fixed-size blocks out of [`SLAB_SIZE`]-byte, +//! `SLAB_SIZE`-aligned slabs. A slab looks like this: +//! +//! ```text +//! | SlabHeader | links: [AtomicU32; blocks] | pad to 4 KiB | block 0 | block 1 | ... | +//! ``` +//! +//! Free blocks are chained through the `links` side table — never through +//! block memory itself — so block payloads are only ever plain data. That +//! keeps every concurrent access well-defined under the memory model (no +//! mixed atomic/non-atomic reads of the same bytes) and Miri-clean. + +use core::{ + num::NonZero, + ptr::NonNull, + sync::atomic::{AtomicU32, AtomicU64}, +}; + +use crate::class::{CLASS_COUNT, MAX_POOL_ALIGN, block_size}; + +/// Slab size and alignment. Masking a block address with `!(SLAB_SIZE - 1)` +/// recovers its slab base. +pub const SLAB_SIZE: usize = 1 << 20; + +const fn round_up(value: usize, align: usize) -> usize { + (value + align - 1) & !(align - 1) +} + +/// Byte offset of the block area inside a slab holding `blocks` blocks. +const fn block_area_offset(blocks: usize) -> usize { + round_up(size_of::() + blocks * size_of::(), MAX_POOL_ALIGN) +} + +const fn compute_blocks_per_slab(class: usize) -> usize { + let size = block_size(class); + // Upper bound ignoring metadata, then shrink until header, link table, + // padding, and blocks all fit in one slab. + let mut blocks = SLAB_SIZE / size; + while block_area_offset(blocks) + blocks * size > SLAB_SIZE { + blocks -= 1; + } + blocks +} + +/// Blocks per slab, for each size class. +pub const BLOCKS_PER_SLAB: [usize; CLASS_COUNT] = { + let mut table = [0usize; CLASS_COUNT]; + let mut class = 0; + while class < CLASS_COUNT { + table[class] = compute_blocks_per_slab(class); + class += 1; + } + table +}; + +const _: () = { + assert!(size_of::() == 16); + assert!(MAX_POOL_ALIGN >= align_of::()); + let mut class = 0; + while class < CLASS_COUNT { + assert!(BLOCKS_PER_SLAB[class] > 0); + assert!( + block_area_offset(BLOCKS_PER_SLAB[class]) + BLOCKS_PER_SLAB[class] * block_size(class) + <= SLAB_SIZE + ); + class += 1; + } +}; + +/// Lives at the base of every slab mapping, ahead of the link table. +#[repr(C)] +struct SlabHeader { + /// This slab's index in its class's slab table. + slab_idx: u32, + _reserved: u32, + /// Number of blocks ever carved off this slab (monotonic; values at or + /// beyond the class's blocks-per-slab mean the slab is exhausted). 64-bit + /// so over-counting by racing threads can never wrap it around. + carved: AtomicU64, +} + +/// A view of one live slab mapping of a particular class. +/// +/// # Invariant +/// +/// `base` points to a [`SLAB_SIZE`]-byte, `SLAB_SIZE`-aligned mapping laid +/// out for `class` (initialized header, link table, block area) that is +/// never unmapped while the pool is in use. All raw-pointer arithmetic on +/// slab memory lives in this type's methods; the unsafe constructors are the +/// only places the invariant is asserted, and every accessor then relies on +/// it for bounds and liveness. +#[derive(Clone, Copy)] +pub struct Slab { + base: NonNull, + class: usize, +} + +impl Slab { + /// Wraps a slab pointer loaded from a class's slab table. + /// + /// # Safety + /// + /// `base` must be a non-null entry of the class's slab table. Such + /// entries are only published after [`Slab::init_header`] ran on a + /// suitable mapping, and are never cleared or unmapped while the pool is + /// in use, so the type invariant holds. + pub const unsafe fn from_published(base: NonNull, class: usize) -> Self { + Self { base, class } + } + + /// Recovers the slab containing a live pool block: slabs are + /// `SLAB_SIZE`-aligned, so masking the block address's low bits yields + /// the slab base (`with_addr` keeps the block pointer's provenance, which + /// covers the whole slab mapping it was carved from). Returns `None` only + /// for an address whose masked base would be null, which no real block + /// can produce. + /// + /// # Safety + /// + /// `block` must be a block of `class` previously handed out by this pool + /// and thus carved from a published slab of this class. + pub unsafe fn of_block(block: NonNull, class: usize) -> Option { + let base_addr = NonZero::new(block.as_ptr().addr() & !(SLAB_SIZE - 1))?; + Some(Self { base: block.with_addr(base_addr), class }) + } + + /// Initializes the header of a fresh slab mapping, making it publishable + /// into a slab table. + /// + /// # Safety + /// + /// `base` must point to a fresh, exclusive, zero-initialized, + /// `SLAB_SIZE`-byte and `SLAB_SIZE`-aligned mapping. The zero fill + /// doubles as the initial state of the carve counter and the link table. + pub unsafe fn init_header(base: NonNull, slab_idx: u32) { + // SAFETY: caller contract — the fresh exclusive mapping is aligned + // (SLAB_SIZE-aligned, far beyond SlabHeader's needs) and large + // enough for the header. + unsafe { + (*base.cast::().as_ptr()).slab_idx = slab_idx; + } + } + + const fn header(&self) -> &SlabHeader { + // SAFETY: type invariant — the mapping is live, SLAB_SIZE-aligned + // (far beyond SlabHeader's needs), and its header was initialized + // before publication; the only non-atomic field is never written + // again afterwards. + unsafe { self.base.cast::().as_ref() } + } + + /// This slab's index in its class's slab table. + pub const fn slab_idx(&self) -> usize { + self.header().slab_idx as usize + } + + /// The monotonic carve counter. + pub const fn carved(&self) -> &AtomicU64 { + &self.header().carved + } + + const fn blocks(&self) -> usize { + BLOCKS_PER_SLAB[self.class] + } + + /// The free-list link slot of `block_idx`. + pub fn link(&self, block_idx: usize) -> &AtomicU32 { + debug_assert!(block_idx < self.blocks()); + #[expect( + clippy::cast_ptr_alignment, + reason = "the link table starts at offset 16 of a SLAB_SIZE-aligned slab, so entries are 4-byte aligned" + )] + // SAFETY: type invariant plus the bound above put the slot within + // the slab's link table; slots are only ever accessed atomically, + // and the mapping outlives any borrow. + unsafe { + AtomicU32::from_ptr( + self.base.as_ptr().add(size_of::()).cast::().add(block_idx), + ) + } + } + + /// The address of block `block_idx`. + pub fn block(&self, block_idx: usize) -> NonNull { + debug_assert!(block_idx < self.blocks()); + let offset = block_area_offset(self.blocks()) + block_idx * block_size(self.class); + // SAFETY: type invariant plus the bound above keep `offset` within + // the SLAB_SIZE mapping. + unsafe { self.base.add(offset) } + } + + /// The index of a block previously returned by [`Slab::block`]. + pub fn block_index(&self, block: NonNull) -> usize { + let offset = + block.as_ptr().addr() - self.base.as_ptr().addr() - block_area_offset(self.blocks()); + debug_assert_eq!(offset % block_size(self.class), 0); + let block_idx = offset / block_size(self.class); + debug_assert!(block_idx < self.blocks()); + block_idx + } +} diff --git a/crates/fspy_alloc/src/sys.rs b/crates/fspy_alloc/src/sys.rs new file mode 100644 index 000000000..4fecdc085 --- /dev/null +++ b/crates/fspy_alloc/src/sys.rs @@ -0,0 +1,25 @@ +//! The backing-memory provider abstraction. +//! +//! The pool is generic over [`Sys`] so tests (and Miri) can substitute a +//! provider backed by the host allocator, while the real allocator obtains +//! anonymous mappings from the kernel (see the `mmap` module). + +use core::ptr::NonNull; + +/// Provides zero-initialized, aligned memory regions. +/// +/// Implementations must themselves be async-signal-safe and fork-safe: no +/// locks, no thread-locals, no libc malloc. +pub trait Sys { + /// Maps `size` bytes of zero-initialized memory aligned to `align` + /// (a power of two). Returns `None` when memory is exhausted. + fn map(size: usize, align: usize) -> Option>; + + /// Releases a region previously returned by [`Sys::map`]. + /// + /// # Safety + /// + /// `ptr` must have been returned by `Sys::map(size, align)` with the same + /// `size` and `align`, and must not be accessed afterwards. + unsafe fn unmap(ptr: NonNull, size: usize, align: usize); +} diff --git a/crates/fspy_alloc/tests/global_alloc.rs b/crates/fspy_alloc/tests/global_alloc.rs new file mode 100644 index 000000000..ab0670ae5 --- /dev/null +++ b/crates/fspy_alloc/tests/global_alloc.rs @@ -0,0 +1,84 @@ +//! Installs [`FspyAlloc`] as this test binary's global allocator, so every +//! allocation — including the test harness's own — exercises the allocator +//! end-to-end over real memory mappings. +#![cfg(all(unix, not(miri)))] // Miri covers the pool core via its mockable backend instead. + +use std::{collections::BTreeMap, sync::mpsc, thread}; + +use fspy_alloc::FspyAlloc; + +#[global_allocator] +static GLOBAL: FspyAlloc = FspyAlloc::new(); + +#[test] +fn collections_round_trip() { + let mut map = BTreeMap::new(); + for i in 0..1000_u32 { + let len = usize::try_from(i % 300).unwrap(); + map.insert(i, vec![0xAB_u8; len]); + } + assert_eq!(map.len(), 1000); + for (i, bytes) in &map { + assert_eq!(bytes.len(), usize::try_from(i % 300).unwrap()); + assert!(bytes.iter().all(|byte| *byte == 0xAB)); + } +} + +#[test] +fn vec_growth_reallocs_preserve_contents() { + let mut bytes = Vec::new(); + for i in 0..1_000_000_usize { + bytes.push(u8::try_from(i % 251).unwrap()); + } + for (i, byte) in bytes.iter().enumerate() { + assert_eq!(usize::from(*byte), i % 251); + } +} + +#[test] +fn large_and_zeroed_allocations() { + // Direct-mapped (beyond the largest size class), via the zeroing path. + let large = vec![0_u8; 5 * 1024 * 1024]; + assert!(large.iter().all(|byte| *byte == 0)); + // And through the plain path. + let boxed: Box<[u8]> = vec![7_u8; 300 * 1024].into_boxed_slice(); + assert!(boxed.iter().all(|byte| *byte == 7)); +} + +#[test] +fn many_small_allocations_span_slabs() { + // More live 16-byte-class blocks than one slab holds, forcing the pool + // through several slab installations (and the holding Vec through the + // direct-mapped path as it grows). + let boxes: Vec> = (0..120_000).map(|_| Box::new([0xEE_u8; 8])).collect(); + assert!(boxes.iter().all(|block| block.iter().all(|byte| *byte == 0xEE))); +} + +#[test] +fn threaded_producers_and_consumers() { + let threads = 8; + let per_thread = 5_000_usize; + let (sender, receiver) = mpsc::channel::>(); + thread::scope(|scope| { + for t in 0..threads { + let sender = sender.clone(); + scope.spawn(move || { + for i in 0..per_thread { + // Vary sizes across classes; blocks are freed (and often + // allocated) on the consumer thread below. + let len = 1 + (i * 37 + t * 101) % 5000; + sender.send(vec![u8::try_from(t % 251).unwrap(); len]).unwrap(); + } + }); + } + drop(sender); + let mut message_count = 0_usize; + for bytes in receiver { + message_count += 1; + assert!(!bytes.is_empty()); + let first = bytes[0]; + assert!(bytes.iter().all(|byte| *byte == first)); + } + assert_eq!(message_count, threads * per_thread); + }); +} diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 4b89bbd85..70a554563 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -12,6 +12,7 @@ anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } ctor = { workspace = true } +fspy_alloc = { workspace = true } fspy_shared = { workspace = true } fspy_shared_unix = { workspace = true } libc = { workspace = true } diff --git a/crates/fspy_preload_unix/src/lib.rs b/crates/fspy_preload_unix/src/lib.rs index 6c4e10b3e..a9e667970 100644 --- a/crates/fspy_preload_unix/src/lib.rs +++ b/crates/fspy_preload_unix/src/lib.rs @@ -1,6 +1,16 @@ // Compile as an empty crate on non-unix targets and on musl (where seccomp // alone handles access tracking). +/// This library interposes libc functions that POSIX declares +/// async-signal-safe (`open`, `stat`, `execve`, ...), so its own code may run +/// inside signal handlers and in the child of `fork()` in a multithreaded +/// process — contexts where taking the libc allocator's locks can deadlock. +/// Route every Rust allocation in this cdylib through fspy's lock-free, +/// mmap-backed allocator instead. +#[cfg(all(unix, not(target_env = "musl")))] +#[global_allocator] +static GLOBAL_ALLOCATOR: fspy_alloc::FspyAlloc = fspy_alloc::FspyAlloc::new(); + #[cfg(all(unix, not(target_env = "musl")))] mod client; #[cfg(all(unix, not(target_env = "musl")))] From 87be7e0147587f05424affcd50ce2dbf776a4302 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 08:05:46 +0800 Subject: [PATCH 2/7] feat(fspy): replace the pool allocator with per-call arenas over a chunk pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework fspy_alloc from a size-class global allocator into three layers: MmapAllocator (stateless, every allocation is a fresh kernel mapping), ChunkPool (a lock-free cache of 64 KiB chunks), and arena(), the only public entry, which hands each intercepted call its own bump arena (bump_scope::Bump) drawing chunks from the process-wide pool. Use the arena for the first preload call site: joining a directory and a relative path when resolving fd-relative opens. Add an access-relative benchmark suite so this lane — working-directory resolution plus path joining — is measured; the existing absolute-path suite never enters it. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 12 + Cargo.toml | 4 +- crates/fspy_alloc/Cargo.toml | 7 + crates/fspy_alloc/src/class.rs | 38 - crates/fspy_alloc/src/lib.rs | 223 +++-- crates/fspy_alloc/src/mapping.rs | 62 -- crates/fspy_alloc/src/mmap.rs | 310 +++--- crates/fspy_alloc/src/pool.rs | 914 ++++++------------ crates/fspy_alloc/src/slab.rs | 203 ---- crates/fspy_alloc/src/sys.rs | 25 - crates/fspy_alloc/tests/global_alloc.rs | 84 -- crates/fspy_preload_unix/Cargo.toml | 1 + .../fspy_preload_unix/src/client/convert.rs | 22 +- crates/fspy_preload_unix/src/lib.rs | 10 - 14 files changed, 580 insertions(+), 1335 deletions(-) delete mode 100644 crates/fspy_alloc/src/class.rs delete mode 100644 crates/fspy_alloc/src/mapping.rs delete mode 100644 crates/fspy_alloc/src/slab.rs delete mode 100644 crates/fspy_alloc/src/sys.rs delete mode 100644 crates/fspy_alloc/tests/global_alloc.rs diff --git a/Cargo.lock b/Cargo.lock index 6182d52d4..b783495c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -349,6 +349,15 @@ dependencies = [ "serde", ] +[[package]] +name = "bump-scope" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4747cadbe8ffeff2a58cb86bcb35f9ec490f2329e12a88a696bccfd3dc97bf4d" +dependencies = [ + "allocator-api2", +] + [[package]] name = "bumpalo" version = "3.19.1" @@ -1247,6 +1256,8 @@ dependencies = [ name = "fspy_alloc" version = "0.0.0" dependencies = [ + "allocator-api2", + "bump-scope", "rustix", ] @@ -1302,6 +1313,7 @@ dependencies = [ name = "fspy_preload_unix" version = "0.0.0" dependencies = [ + "allocator-api2", "anyhow", "artifact_profile", "bstr", diff --git a/Cargo.toml b/Cargo.toml index 4385383c4..f1c766099 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ multiple_crate_versions = "allow" future_not_send = "allow" [workspace.dependencies] +allocator-api2 = { version = "0.2", default-features = false } artifact_profile = { path = "crates/artifact_profile" } anstream = "1.0.0" anyhow = "1.0.103" @@ -51,6 +52,7 @@ bindgen = "0.72.1" bitflags = "2.10.0" brush-parser = "0.4.0" bstr = { version = "1.12.0", default-features = false, features = ["alloc", "std"] } +bump-scope = { version = "2", default-features = false, features = ["allocator-api2-02"] } bumpalo = { version = "3.17.0", features = ["collections"] } bytemuck = { version = "1.23.0", features = ["extern_crate_alloc", "must_cast"] } cc = "1.2.39" @@ -121,7 +123,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" -rustix = { version = "1", default-features = false, features = ["mm", "param"] } +rustix = { version = "1", default-features = false, features = ["mm", "param", "use-libc-auxv"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml index 809d0f7e8..edcb89eaa 100644 --- a/crates/fspy_alloc/Cargo.toml +++ b/crates/fspy_alloc/Cargo.toml @@ -8,7 +8,14 @@ publish = false doctest = false [target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true } +bump-scope = { workspace = true } rustix = { workspace = true } +[target.'cfg(unix)'.dev-dependencies] +# The `alloc` feature provides `Global`, letting tests run the pool against +# the host allocator (and thus under Miri). +allocator-api2 = { workspace = true, features = ["alloc"] } + [lints] workspace = true diff --git a/crates/fspy_alloc/src/class.rs b/crates/fspy_alloc/src/class.rs deleted file mode 100644 index df8d3797a..000000000 --- a/crates/fspy_alloc/src/class.rs +++ /dev/null @@ -1,38 +0,0 @@ -//! Size-class policy: which layouts the pool serves, and at what block size. - -use core::alloc::Layout; - -const MIN_CLASS_SHIFT: u32 = 4; -const MAX_CLASS_SHIFT: u32 = 16; -pub const CLASS_COUNT: usize = (MAX_CLASS_SHIFT - MIN_CLASS_SHIFT) as usize + 1; -const MIN_BLOCK_SIZE: usize = 1 << MIN_CLASS_SHIFT; -pub const MAX_BLOCK_SIZE: usize = 1 << MAX_CLASS_SHIFT; -/// Block areas start at this alignment within a slab, making it the largest -/// alignment the pool can serve; stricter layouts map directly. -pub const MAX_POOL_ALIGN: usize = 4096; - -pub const fn block_size(class: usize) -> usize { - 1 << (MIN_CLASS_SHIFT as usize + class) -} - -/// Returns the size class for `layout`, or `None` if the request must be -/// mapped directly (too large or over-aligned). -pub const fn class_of(layout: Layout) -> Option { - if layout.align() > MAX_POOL_ALIGN { - return None; - } - let mut size = layout.size(); - // A block of `size >= align` at a `min(block size, 4 KiB)` boundary is - // aligned to `align` (both are powers of two and `align <= 4 KiB`). - if size < layout.align() { - size = layout.align(); - } - if size < MIN_BLOCK_SIZE { - size = MIN_BLOCK_SIZE; - } - if size > MAX_BLOCK_SIZE { - return None; - } - let shift = size.next_power_of_two().trailing_zeros(); - Some((shift - MIN_CLASS_SHIFT) as usize) -} diff --git a/crates/fspy_alloc/src/lib.rs b/crates/fspy_alloc/src/lib.rs index 90b2524cf..24b97020e 100644 --- a/crates/fspy_alloc/src/lib.rs +++ b/crates/fspy_alloc/src/lib.rs @@ -1,124 +1,147 @@ -//! Lock-free, async-signal-safe global allocator for the fspy preload library. +//! Async-signal-safe allocation for the fspy preload library. //! //! The preload library interposes libc functions that POSIX declares -//! async-signal-safe (`open`, `stat`, `execve`, ...). Programs may call these -//! from signal handlers, and — more commonly — from the child of `fork()` in a -//! multithreaded process, where only async-signal-safe calls are permitted: -//! the libc allocator's locks may be held forever by threads that no longer -//! exist after the fork. Routing the preload's Rust allocations through this -//! allocator keeps them safe in both contexts: +//! async-signal-safe (`open`, `stat`, `execve`, ...). Its code therefore runs +//! inside signal handlers and in the child of `fork()` in a multithreaded +//! process — contexts where taking libc malloc's locks can deadlock on a +//! suspended or vanished lock holder, so the preload's own allocations must +//! never go through libc malloc. //! -//! - **No locks.** Every state transition is a lock-free compare-and-swap -//! loop: an attempt only retries because another running thread completed -//! its operation, so nothing ever waits on state that a thread which -//! vanished at `fork()` — or sits suspended under a signal handler — would -//! have to release. (Lock-free, not wait-free: an individual operation has -//! no fixed retry bound under active contention.) -//! - **No thread-locals.** TLS first-touch allocates through libc malloc on -//! some platforms (macOS thread-local variables), which would reintroduce -//! the hazard this crate exists to remove. -//! - **mmap-backed.** Memory comes straight from the kernel. On Linux the -//! allocator relies on nothing from libc: mapping syscalls are issued -//! directly (rustix's raw backend) and even the page size is discovered by -//! probing with raw syscalls. On macOS, which has no stable raw-syscall -//! ABI, calls go through the thin libSystem stubs. libc malloc is never -//! called anywhere. -//! -//! Design: power-of-two size classes (16 B ..= 64 KiB) carve blocks out of -//! 1 MiB slabs; freed blocks recycle through a per-class Treiber free list -//! made ABA-safe by a generation tag. Requests larger than the biggest class -//! (or over-aligned beyond 4 KiB) map and unmap directly. See the `pool` -//! module for the details. -//! -//! Because the allocator is a `const`-initialized static with no lazy setup, -//! it works from the very first allocation in the process — even before the -//! preload library's constructor runs. +//! This crate stacks three layers, and exposes only the top one, [`arena`]. +//! `MmapAllocator` is the bottom: a stateless allocator where every +//! allocation is a fresh anonymous mapping from the kernel. `ChunkPool` +//! sits on top of it and caches fixed-size chunks, so that frequent short +//! tracing calls can reuse memory instead of paying two syscalls per call. +//! [`arena`] creates one `bump_scope::Bump` per intercepted call, drawing +//! its chunks from the process-wide pool and returning them on drop. +// Compile as an empty crate on non-unix targets: the allocator backs the +// unix preload library. +#![cfg(unix)] #![cfg_attr(not(test), no_std)] -// Compile as an empty crate on non-unix targets: the allocator backs the unix -// preload library. A Windows backend can be added alongside `sys::Mmap` if -// the Windows preload ever needs one. - -#[cfg(unix)] -mod class; -#[cfg(unix)] -mod mapping; -#[cfg(unix)] mod mmap; -#[cfg(unix)] mod pool; -#[cfg(unix)] -mod slab; -#[cfg(unix)] -mod sys; -#[cfg(unix)] -use core::{ - alloc::{GlobalAlloc, Layout}, - ptr::{self, NonNull}, +use allocator_api2::alloc::Allocator; +use bump_scope::{ + Bump, + alloc::compat::AllocatorApi2V02Compat, + settings::{BumpAllocatorSettings, BumpSettings}, }; +use mmap::MmapAllocator; +use pool::ChunkPool; -#[cfg(unix)] -use crate::{mmap::Mmap, pool::Pool}; +/// Every cached chunk is 64 KiB: a whole multiple of the page size on all +/// supported targets, and big enough that most intercepted calls fit their +/// allocations into a single chunk. +const CHUNK_SIZE: usize = 64 * 1024; +/// The alignment chunks are allocated with. Must be at least the alignment +/// `bump_scope::Bump` uses for its chunk requests — 16 (see +/// [`MmapAllocator`]'s docs for the links). bump-scope does not export that +/// constant, so the `bump_chunk_requests_fit_the_pool_gates` test pins the +/// fit instead: it fails if a bump-scope upgrade ever requests chunks the +/// pool would refuse. +const CHUNK_ALIGN: usize = 16; +/// At most this many chunks stay cached, capping retained memory at +/// `SLOTS * CHUNK_SIZE` = 4 MiB. +const SLOTS: usize = 64; -/// A lock-free, async-signal-safe, fork-safe [`GlobalAlloc`] implementation. -/// -/// Intended to be installed as the `#[global_allocator]` of the fspy preload -/// library. All memory comes from anonymous mappings; libc malloc is never -/// called, no locks are taken, and no thread-local state is used. -/// -/// Capacity is bounded by design: each size class can hold at most 256 slabs -/// of 1 MiB (roughly 200 MiB per class). Requests beyond that — far outside -/// anything the preload library does — fail like any other out-of-memory -/// condition (`alloc` returns null). -#[cfg(unix)] -pub struct FspyAlloc { - pool: Pool, -} +/// The process-wide chunk pool. `const`-initialized, so it works from the +/// first allocation on — even before any constructor has run. +static CHUNK_POOL: ChunkPool = ChunkPool::new(); -#[cfg(unix)] -impl FspyAlloc { - /// Creates the allocator. `const` so it can back a `static` with no - /// runtime initialization. - #[must_use] - pub const fn new() -> Self { - Self { pool: Pool::new() } +/// `Bump::unallocated` requires its base allocator to implement `Default` +/// (an arena without chunks has nowhere to store an allocator value, so it +/// conjures one on first use). Point defaulted references at the +/// process-wide pool. As an allocator, `&ChunkPool` already works through +/// allocator-api2's blanket `impl Allocator for &A`. +impl Default for &'static ChunkPool { + fn default() -> Self { + &CHUNK_POOL } } -#[cfg(unix)] -impl Default for FspyAlloc { - fn default() -> Self { - Self::new() - } +/// Creates a fresh bump arena for one intercepted call, backed by the +/// process-wide chunk pool. +/// +/// Creating the arena allocates nothing; the first allocation grabs a whole +/// chunk — usually a recycled one, so most calls touch no syscalls at all. +/// Deallocation only takes back the most recent allocation (bump-arena +/// semantics); everything is freed at once when the arena is dropped, and +/// its chunks go back to the pool. +/// +/// The arena itself is single-owner — use one per call, do not share it +/// across threads. Creating one is safe anywhere, any time: the pool +/// underneath works in signal handlers and in the child of `fork()` (see +/// `ChunkPool` and `MmapAllocator` in this crate's source for why). +#[must_use] +pub fn arena() -> impl Allocator { + // The default `Bump` settings, except the arena starts life without a + // chunk, so creating one allocates nothing. + type Settings = ::WithGuaranteedAllocated; + Bump::< + AllocatorApi2V02Compat<&'static ChunkPool>, + Settings, + >::unallocated() } -// SAFETY: `Pool` hands out blocks that are non-null, at least `layout.size()` -// bytes large, aligned to at least `layout.align()`, and exclusively owned -// until returned via `dealloc`. Allocation failure is reported as null, and -// none of the methods unwind. -#[cfg(unix)] -unsafe impl GlobalAlloc for FspyAlloc { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - self.pool.alloc(layout).map_or(ptr::null_mut(), NonNull::as_ptr) - } +#[cfg(test)] +mod tests { + use core::alloc::Layout; - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - let Some(ptr) = NonNull::new(ptr) else { return }; - // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this - // allocator for this `layout`. - unsafe { self.pool.dealloc(ptr, layout) } - } + use allocator_api2::alloc::Global; - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - self.pool.alloc_zeroed(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + use super::*; + + /// Pins the fit between `Bump`'s chunk requests and the pool's gates + /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]). + /// bump-scope keeps those request parameters private, so this test is + /// the enforcement: it fails if an upgrade ever changes them. + #[test] + fn bump_chunk_requests_fit_the_pool_gates() { + let pool = ChunkPool::::new_in(Global); + // `try_new_in` requests the first chunk right away; if that request + // were over-aligned or oversized, the pool would refuse and this + // would be an error. + let bump: Bump>> = + Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); + let block = bump.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); + let block_addr = block.cast::().as_ptr().addr(); + drop(bump); + + // The chunk went back into the pool's cache (proving it was served + // as a chunk, not passed through): the next chunk-sized request + // returns the block's surroundings. + let recycled = pool.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); + assert_eq!(recycled.len(), CHUNK_SIZE); + let base = recycled.cast::().as_ptr().addr(); + assert!((base..base + CHUNK_SIZE).contains(&block_addr)); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(recycled.cast(), Layout::from_size_align(100, 8).unwrap()) }; } - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - let Some(ptr) = NonNull::new(ptr) else { return ptr::null_mut() }; - // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this - // allocator for this `layout`, and `new_size` is non-zero. - unsafe { self.pool.realloc(ptr, layout, new_size) }.map_or(ptr::null_mut(), NonNull::as_ptr) + #[test] + #[cfg(not(miri))] + fn arena_allocates_and_returns_chunks_to_the_pool() { + let layout = Layout::from_size_align(100, 8).unwrap(); + + let first_arena = arena(); + let first = first_arena.allocate(layout).unwrap(); + assert!(first.len() >= 100); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { first.cast::().as_ptr().write_bytes(0x5A, 100) }; + let second = first_arena.allocate(layout).unwrap(); + assert_ne!(first.cast::().as_ptr().addr(), second.cast::().as_ptr().addr()); + let first_addr = first.cast::().as_ptr().addr(); + // Everything dies at once; the chunk goes back to the pool. + drop(first_arena); + + // No other test touches the process-wide pool, so a new arena draws + // the same chunk back and its first allocation lands at the same + // address. + let second_arena = arena(); + let again = second_arena.allocate(layout).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), first_addr); } } diff --git a/crates/fspy_alloc/src/mapping.rs b/crates/fspy_alloc/src/mapping.rs deleted file mode 100644 index d3f53e726..000000000 --- a/crates/fspy_alloc/src/mapping.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! Owned memory regions obtained from a [`Sys`] provider. - -use core::{marker::PhantomData, mem, ptr::NonNull}; - -use crate::sys::Sys; - -/// An owned region obtained from `S`, released on drop. -/// -/// This is the only place that calls [`Sys::unmap`]: pool code either lets a -/// `Mapping` drop (probe scratch, install races, freed large allocations) or -/// deliberately leaks it with [`Mapping::into_raw`] (published slabs, live -/// large allocations). Reconstructing ownership from a raw pointer via -/// [`Mapping::from_raw`] is the single unsafe step. -pub struct Mapping { - ptr: NonNull, - size: usize, - align: usize, - sys: PhantomData S>, -} - -impl Mapping { - /// Maps `size` bytes of zero-initialized memory aligned to `align` - /// (a power of two). Returns `None` when memory is exhausted. - pub fn new(size: usize, align: usize) -> Option { - let ptr = S::map(size, align)?; - Some(Self { ptr, size, align, sys: PhantomData }) - } - - /// Reclaims ownership of a mapping previously released with - /// [`Mapping::into_raw`]. - /// - /// # Safety - /// - /// `ptr` must have come from `Mapping::::into_raw` (or `Sys::map`) - /// with exactly this `size` and `align`, the region must not be in use, - /// and ownership must not be reclaimed twice. - pub unsafe fn from_raw(ptr: NonNull, size: usize, align: usize) -> Self { - Self { ptr, size, align, sys: PhantomData } - } - - /// The mapped region's base address. - pub const fn ptr(&self) -> NonNull { - self.ptr - } - - /// Releases ownership without unmapping; the region lives until (unless) - /// [`Mapping::from_raw`] reclaims it. - pub const fn into_raw(self) -> NonNull { - let ptr = self.ptr; - mem::forget(self); - ptr - } -} - -impl Drop for Mapping { - fn drop(&mut self) { - // SAFETY: this type owns the mapping (constructed from `Sys::map` - // directly or via the `from_raw` contract), and after drop nothing - // can use it. - unsafe { S::unmap(self.ptr, self.size, self.align) } - } -} diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/fspy_alloc/src/mmap.rs index 053356b56..3cc09d876 100644 --- a/crates/fspy_alloc/src/mmap.rs +++ b/crates/fspy_alloc/src/mmap.rs @@ -1,211 +1,161 @@ -//! Kernel-backed [`Sys`] provider: anonymous mappings via direct syscalls. -//! -//! On Linux the provider relies on nothing from libc — rustix issues raw -//! syscalls, and even the page size is discovered with raw syscalls (see -//! [`page_size`]). macOS has no stable raw-syscall ABI, so calls go through -//! the thin libSystem stubs there. +//! Page-granularity allocator backed by anonymous memory mappings. use core::{ + alloc::Layout, ptr::{self, NonNull}, - sync::atomic::{AtomicUsize, Ordering}, }; +use allocator_api2::alloc::{AllocError, Allocator}; use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; -use crate::sys::Sys; - -/// Returns the kernel page size, or `None` when it cannot be determined. +/// A stateless allocator: every allocation is a fresh anonymous mapping and +/// every deallocation an `munmap`. /// -/// On Linux the value is discovered with raw syscalls only — no libc, no -/// `/proc`, no minimum kernel version (see [`query_page_size`]). On other -/// unix platforms (macOS, where every syscall goes through libSystem by -/// platform contract anyway) it comes from `sysconf` via rustix. Either -/// way the result is a process constant, validated as a power of two and -/// cached, so `map` and `unmap` can never disagree on rounding; an -/// undeterminable page size fails the allocation rather than guessing. -fn page_size() -> Option { - static CACHE: AtomicUsize = AtomicUsize::new(0); - let cached = CACHE.load(Ordering::Relaxed); - if cached != 0 { - return Some(cached); - } - let page = query_page_size()?; - if !page.is_power_of_two() { - return None; - } - // First store wins; every query returns the same value, so the cache - // only avoids repeated probing. - let _ = CACHE.compare_exchange(0, page, Ordering::Relaxed, Ordering::Relaxed); - Some(page) -} - -/// Discovers the page size by probing, using nothing but raw syscalls. +/// # Why it is safe in signal handlers and forked children /// -/// `mprotect` fails with `EINVAL` unless its address is a multiple of the -/// page size, so re-protecting a scratch mapping at increasing -/// power-of-two offsets identifies the page size as the first offset the -/// kernel accepts (for a power-of-two page size P, the first power of two -/// P divides is P itself). A handful of syscalls, once per process: -/// async-signal-safe, fork-safe, and independent of libc, `/proc` -/// availability, and kernel version. +/// It holds no state at all — no locks, no free lists, no thread-locals; +/// the kernel does all the bookkeeping. A signal or a `fork()` can never +/// catch it holding a lock or a half-written structure, because there is +/// nothing to hold. Mapping syscalls go straight to the kernel on Linux +/// (rustix's raw backend) and through the thin libSystem wrappers on macOS. +/// The page size comes from static process data (`getauxval` on Linux, +/// `sysconf` on macOS) — no locks or allocation there either. /// -/// Why not `rustix::param::page_size()` here (it's fine on macOS)? On -/// Linux the allocator must not rely on libc, which rules out rustix's -/// `use-libc-auxv` (`getauxval`) configuration — and rustix's libc-free -/// fallback is unusable *inside* a global allocator: +/// # What it accepts /// -/// - Its lazy init tries `prctl(PR_GET_AUXV)` (kernel 6.4+ only) and -/// otherwise reads `/proc/self/auxv`; with rustix's `alloc` feature -/// enabled, that read path heap-allocates (`Vec`) — through *this* -/// allocator, whose `map` is the caller waiting on the page size — -/// recursing unboundedly. And we cannot pin `alloc` off: Cargo feature -/// unification lets any other rustix user in the build graph enable it -/// for our copy. -/// - It panics on read errors, truncated auxv, or both sources being -/// unavailable, where this allocator requires failure to surface as -/// `None` (panic formatting itself allocates, re-entering the same -/// uninitialized path). +/// The only intended caller is `bump_scope::Bump`, and `Bump` only ever +/// asks its base allocator for chunks ([single call site][site]). Chunks +/// are never zero-sized (a chunk always contains its own header), and +/// their alignment is [`max(MIN_CHUNK_ALIGN, header alignment)`][chunk-align], +/// [where `MIN_CHUNK_ALIGN` is 16][mca] — so 16 bytes in practice. /// -/// The probe has neither problem: no allocation, no panic, no minimum -/// kernel, and its worst case is `None`. -#[cfg(target_os = "linux")] -fn query_page_size() -> Option { - use rustix::mm::{MprotectFlags, mprotect}; - - /// Smallest page size of any Linux configuration; the probe starts - /// here. - const MIN_PROBE_PAGE: usize = 4096; - /// Generous upper bound for the probe; no supported configuration - /// uses larger pages. - const MAX_PROBE_PAGE: usize = 1 << 20; +/// This allocator accepts more than that, because mapped memory gives the +/// extra range away for free: any non-zero size, and any alignment up to +/// the page size (mappings are always page-aligned). It refuses only two +/// kinds of request, which would each need extra code and never happen: +/// zero-sized layouts (they would need fake dangling blocks) and alignment +/// above the page size (it would need mapping extra space and trimming the +/// misaligned edges). The [`Allocator`] contract allows refusing any +/// request; refused requests get [`AllocError`]. +/// +/// # Cost +/// +/// Every allocation takes whole pages (4 KiB at least, 16 KiB on Apple +/// silicon) and one syscall, and so does every deallocation. This +/// allocator is meant to sit below a chunk pool and bump arenas, not to +/// serve small allocations directly. +/// +/// `allocate` returns the whole page-rounded block, and `bump_scope` [uses +/// the full returned length][fit], so none of the page is wasted. +/// +/// [mca]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/chunk/size_config.rs.html#8 +/// [chunk-align]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/chunk/size_config.rs.html#56-58 +/// [site]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/raw_bump.rs.html#865 +/// [fit]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/raw_bump.rs.html#873-884 +#[derive(Clone, Copy, Debug, Default)] +pub struct MmapAllocator; + +pub fn page_size() -> usize { + // getauxval on Linux, sysconf on macOS: reads of static process data — + // no locks, no allocation, valid from the first instruction on. + rustix::param::page_size() +} - // Large enough that every probe below stays inside the mapping even - // after the kernel rounds the one-byte length up to a full page. - let scratch = map_anonymous(2 * MAX_PROBE_PAGE)?; - let mut page = None; - let mut offset = MIN_PROBE_PAGE; - while offset <= MAX_PROBE_PAGE { - // SAFETY: `scratch + offset` (plus the page-rounded single byte) - // lies within the scratch mapping, which we exclusively own; the - // protection flags match the mapping's existing ones. - let accepted = unsafe { - mprotect( - scratch.as_ptr().add(offset).cast(), - 1, - MprotectFlags::READ | MprotectFlags::WRITE, - ) +// SAFETY: returned blocks are non-null, page-aligned (at least +// `layout.align()` for every served layout), at least `layout.size()` bytes +// large (the returned length reports the exact mapped size), stay valid +// until deallocated, and distinct allocations never overlap. +unsafe impl Allocator for MmapAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + let page = page_size(); + // Outside the served request profile (see the type docs): refuse + // rather than carry an over-alignment or dangling-block code path. + if layout.size() == 0 || layout.align() > page { + return Err(AllocError); } - .is_ok(); - if accepted { - page = Some(offset); - break; + // `checked_next_multiple_of` is total: `None` on overflow (already + // impossible — `Layout` caps sizes at `isize::MAX`) or a zero page + // size, with no power-of-two assumption to uphold. + let size = layout.size().checked_next_multiple_of(page).ok_or(AllocError)?; + // SAFETY: a fresh anonymous private mapping at no particular + // address has no memory-safety preconditions. + let ptr = unsafe { + mmap_anonymous( + ptr::null_mut(), + size, + ProtFlags::READ | ProtFlags::WRITE, + MapFlags::PRIVATE, + ) } - offset *= 2; + .map_err(|_| AllocError)?; + // Mapping results are page-aligned, which covers every served + // layout. + NonNull::new(ptr.cast::()) + .map(|ptr| NonNull::slice_from_raw_parts(ptr, size)) + .ok_or(AllocError) } - // SAFETY: releasing the scratch mapping created above. - let _ = unsafe { munmap(scratch.as_ptr().cast(), 2 * MAX_PROBE_PAGE) }; - page -} - -#[cfg(not(target_os = "linux"))] -#[expect( - clippy::unnecessary_wraps, - reason = "must match the signature of the fallible Linux probe variant" -)] -fn query_page_size() -> Option { - Some(rustix::param::page_size()) -} -const fn round_up(value: usize, align: usize) -> usize { - (value + align - 1) & !(align - 1) -} - -fn map_anonymous(len: usize) -> Option> { - // SAFETY: a fresh anonymous private mapping at no particular address - // has no memory-safety preconditions. - let ptr = unsafe { - mmap_anonymous(ptr::null_mut(), len, ProtFlags::READ | ProtFlags::WRITE, MapFlags::PRIVATE) - } - .ok()?; - NonNull::new(ptr.cast::()) -} - -/// Kernel-backed provider: anonymous mappings obtained through direct -/// syscalls, alignment achieved by over-mapping and trimming. -pub struct Mmap; - -impl Sys for Mmap { - fn map(size: usize, align: usize) -> Option> { - debug_assert!(align.is_power_of_two()); - let page = page_size()?; - let size = round_up(size.max(1), page); - - if align <= page { - // Mapping results are aligned to the (verified real) page - // size. - return map_anonymous(size); - } - - // Over-map by `align`, then unmap the misaligned head and the - // leftover tail. All cut points are page-aligned: `raw` and - // `aligned` are page-aligned, and `size`/`align` are multiples of - // the page size. - let raw = map_anonymous(size.checked_add(align)?)?; - let raw_addr = raw.as_ptr().addr(); - let aligned_addr = round_up(raw_addr, align); - let head = aligned_addr - raw_addr; - let tail = align - head; - if head > 0 { - // SAFETY: `[raw_addr, raw_addr + head)` lies within the fresh - // mapping and `raw_addr` is page-aligned. Failure is - // impossible for a region we own; if it happened anyway the - // pages would merely stay mapped. - let _ = unsafe { munmap(raw.as_ptr().cast(), head) }; - } - if tail > 0 { - let tail_start = raw.as_ptr().with_addr(aligned_addr + size); - // SAFETY: `[aligned_addr + size, raw_addr + size + align)` - // lies within the fresh mapping and its start is page-aligned - // (see above). Failure is impossible for a region we own. - let _ = unsafe { munmap(tail_start.cast(), tail) }; - } - // `aligned_addr` lies inside a successful mapping and so can - // never be zero, but checked construction costs nothing here. - core::num::NonZero::new(aligned_addr).map(|addr| raw.with_addr(addr)) + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + // Fresh anonymous mappings are already zero-filled by the kernel. + self.allocate(layout) } - unsafe fn unmap(ptr: NonNull, size: usize, _align: usize) { - // A successful `map` proved the page size, so this cannot fail - // for a live mapping; if it somehow did, leaking the region is - // the only safe response. - let Some(page) = page_size() else { return }; - // Whether or not `map` trimmed for alignment, the retained region - // is exactly `[ptr, ptr + round_up(size, page))`. - let size = round_up(size.max(1), page); - // SAFETY: caller contract — `ptr`/`size` describe a live mapping - // returned by `map`. Failure is impossible for a region we own. + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // The contract lets callers pass any size that fits the block, i.e. + // anything in `[requested, mapped]`; every such size rounds up to + // the mapped length. Zero-sized blocks are never allocated, so no + // dangling pointers can arrive here. Rounding cannot fail for a + // layout that fits a block we mapped; leaking the region is the + // safe response if it somehow did. + let Some(size) = layout.size().checked_next_multiple_of(page_size()) else { return }; + // SAFETY: caller contract — `ptr` was returned by `allocate` with a + // fitting layout and the block is no longer in use. Failure is + // impossible for a region we own. let _ = unsafe { munmap(ptr.as_ptr().cast(), size) }; } } #[cfg(all(test, not(miri)))] mod tests { - use super::Mmap; - use crate::sys::Sys; + use super::*; #[test] - fn real_mappings_are_aligned_zeroed_and_writable() { - for (size, align) in [(1, 1), (4096, 4096), (100, 1 << 20), (5 << 20, 4096)] { - let ptr = Mmap::map(size, align).unwrap(); - assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); - for i in 0..size { - // SAFETY: fresh exclusive mapping of at least `size` bytes. - assert_eq!(unsafe { ptr.as_ptr().add(i).read() }, 0); + fn blocks_are_aligned_zeroed_and_page_rounded() { + for (size, align) in [(1, 1), (100, 64), (4096, 4096), (5 << 20, 8)] { + let layout = Layout::from_size_align(size, align).unwrap(); + let block = MmapAllocator.allocate(layout).unwrap(); + assert!(block.len() >= size, "size {size}"); + assert_eq!(block.len() % page_size(), 0); + assert_eq!(block.cast::().as_ptr().addr() % align, 0, "align {align}"); + for i in 0..block.len() { + // SAFETY: fresh exclusive block of `block.len()` bytes. + assert_eq!(unsafe { block.cast::().as_ptr().add(i).read() }, 0); } - // SAFETY: fresh exclusive mapping of at least `size` bytes. - unsafe { ptr.as_ptr().write_bytes(0x5A, size) }; - // SAFETY: mapped above with the same size and alignment. - unsafe { Mmap::unmap(ptr, size, align) }; + // SAFETY: fresh exclusive block of at least `size` bytes. + unsafe { block.cast::().as_ptr().write_bytes(0x5A, size) }; + // SAFETY: allocated above; the layout fits the block. + unsafe { MmapAllocator.deallocate(block.cast(), layout) }; } } + + #[test] + fn deallocate_accepts_any_fitting_size() { + let requested = Layout::from_size_align(100, 8).unwrap(); + let block = MmapAllocator.allocate(requested).unwrap(); + // Deallocate with the *returned* size instead of the requested one — + // both are within the fit range the contract allows. + let fitting = Layout::from_size_align(block.len(), 8).unwrap(); + // SAFETY: allocated above; `fitting` is within the block's fit range. + unsafe { MmapAllocator.deallocate(block.cast(), fitting) }; + } + + #[test] + fn out_of_profile_requests_are_refused() { + // Zero-sized and over-page-aligned layouts are outside the served + // request profile and must fail cleanly, not misbehave. + let zero = Layout::from_size_align(0, 16).unwrap(); + assert!(MmapAllocator.allocate(zero).is_err()); + let over_aligned = Layout::from_size_align(64, 1 << 24).unwrap(); + assert!(MmapAllocator.allocate(over_aligned).is_err()); + } } diff --git a/crates/fspy_alloc/src/pool.rs b/crates/fspy_alloc/src/pool.rs index 6bb4f82f1..f0e3e4112 100644 --- a/crates/fspy_alloc/src/pool.rs +++ b/crates/fspy_alloc/src/pool.rs @@ -1,705 +1,363 @@ -//! Lock-free size-class pool. -//! -//! Layouts round up to a power-of-two class (see the `class` module) whose -//! blocks are carved from slabs (see the `slab` module for the memory -//! layout). -//! -//! # Concurrency -//! -//! Each class has a Treiber-stack free list plus a bump cursor over the -//! newest ("active") slab: -//! -//! - **alloc** pops the free list, or carves the next block off the active -//! slab, installing a fresh slab when the active one is exhausted. -//! - **dealloc** pushes the block back onto its class's free list. -//! -//! The 64-bit list head packs a 40-bit generation tag next to the 24-bit -//! block reference; the tag advances on every successful push and pop, which -//! makes the classic Treiber-stack ABA failure require a thread to stall -//! between its head load and CAS while other threads perform an exact -//! multiple of 2^40 head mutations. That is a probabilistic defense, not a -//! formal impossibility — see [`HEAD_TAG_SHIFT`]. Slab installation races -//! are resolved with a compare-and-swap on the slab table slot; the loser -//! simply returns its (never-published) mapping. -//! -//! The pool is **lock-free, not wait-free**: a CAS loop can retry -//! indefinitely under contention, but a retry only ever happens because -//! another thread completed an operation, so system-wide progress never -//! stalls — and, the property fork- and signal-safety actually require, no -//! operation ever waits on state that only a suspended or vanished thread -//! could release. Blocks are aligned to `min(block size, 4 KiB)`; requests -//! with stricter alignment (or size beyond the largest class) bypass the -//! pool and map directly. +//! A small cache of memory chunks between the bump arenas and the kernel. use core::{ alloc::Layout, - marker::PhantomData, - ptr, - ptr::NonNull, - sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, Ordering}, + ptr::{self, NonNull}, + sync::atomic::{AtomicPtr, Ordering}, }; -use crate::{ - class::{CLASS_COUNT, class_of}, - mapping::Mapping, - slab::{BLOCKS_PER_SLAB, SLAB_SIZE, Slab}, - sys::Sys, -}; - -/// Bounded by the 8 bits reserved for slab indices in a packed block -/// reference. Caps each class at ~200 MiB. -const MAX_SLABS_PER_CLASS: usize = 256; - -/// Sentinel for "no block" in the 24-bit packed-reference field of a -/// free-list head. Never collides with a real reference: block indices stay -/// below `0xFFFF` (asserted with the slab geometry below). -const NO_BLOCK: u32 = 0x00FF_FFFF; -/// Sentinel for "no slab installed yet" in [`ClassState::active`]. -const NO_SLAB: u32 = u32::MAX; - -const _: () = { - assert!(MAX_SLABS_PER_CLASS <= 1 << 8, "slab index must fit in 8 bits"); - let mut class = 0; - while class < CLASS_COUNT { - assert!( - BLOCKS_PER_SLAB[class] <= 0xFFFF, - "block indices must fit in 16 bits and stay below the NO_BLOCK sentinel" - ); - class += 1; - } -}; - -/// Packs a block's location into 24 bits: slab index in bits 16..24, block -/// index in bits 0..16. -#[expect( - clippy::cast_possible_truncation, - reason = "callers pass indices bounded by MAX_SLABS_PER_CLASS and BLOCKS_PER_SLAB" -)] -const fn pack_ref(slab_idx: usize, block_idx: usize) -> u32 { - ((slab_idx as u32) << 16) | (block_idx as u32) -} - -const fn unpack_ref(packed: u32) -> (usize, usize) { - (((packed >> 16) & 0xFF) as usize, (packed & 0xFFFF) as usize) -} - -/// A free-list head is `[generation tag : 40 | packed block reference : 24]`. -/// The tag advances on every successful push and pop, so a stale -/// compare-and-swap can only succeed if its thread stalls between head load -/// and CAS while others perform an exact multiple of 2^40 head mutations — -/// not a formal impossibility, but hours of maximum-rate churn inside one -/// stalled instruction window. -const HEAD_TAG_SHIFT: u32 = 24; -const HEAD_TAG_MASK: u64 = (1 << 40) - 1; -const HEAD_REF_MASK: u64 = (1 << HEAD_TAG_SHIFT) - 1; - -/// Splits a free-list head into `(generation tag, packed block reference)`. -const fn head_parts(head: u64) -> (u64, u32) { - // The mask keeps the reference within 24 bits, so the cast is lossless. - (head >> HEAD_TAG_SHIFT, (head & HEAD_REF_MASK) as u32) -} - -const fn head_from_parts(tag: u64, block_ref: u32) -> u64 { - ((tag & HEAD_TAG_MASK) << HEAD_TAG_SHIFT) | block_ref as u64 -} - -/// Aligned to its own cache-line region so hot heads of different classes -/// don't false-share. -#[repr(align(128))] -struct ClassState { - /// Treiber free-list head: `[generation tag : 40 | packed ref : 24]`. - head: AtomicU64, - /// Index of the slab currently being carved, or [`NO_SLAB`] before the - /// first slab is installed. - active: AtomicU32, - /// Base addresses of installed slabs. Written once (null → mapping) and - /// never cleared. - slabs: [AtomicPtr; MAX_SLABS_PER_CLASS], -} - -impl ClassState { - const fn new() -> Self { - Self { - head: AtomicU64::new(head_from_parts(0, NO_BLOCK)), - active: AtomicU32::new(NO_SLAB), - slabs: [const { AtomicPtr::new(ptr::null_mut()) }; MAX_SLABS_PER_CLASS], - } - } -} - -/// A block handed out by [`Pool::alloc_block`], remembering whether it came -/// off the free list (may contain stale data) or was carved fresh off a slab -/// (still kernel-zeroed). -enum ClassBlock { - Recycled(NonNull), - Carved(NonNull), -} - -enum Carve { - Block(NonNull), - /// A new slab was (or concurrently got) installed; retry the allocation. - Retry, - /// Out of slab table entries or out of memory. - Exhausted, +use allocator_api2::alloc::{AllocError, Allocator}; + +use crate::MmapAllocator; + +/// A fixed-size cache of memory chunks. +/// +/// Sits between the bump arenas and the underlying allocator `A` — +/// [`MmapAllocator`] in real use, any allocator in tests — so the chunks +/// that a finished arena gives back reach the next arena without going back +/// to the kernel. +/// +/// The parameters are chosen by the layer above (the arena layer): every +/// cached chunk is exactly `CHUNK_SIZE` bytes and allocated with +/// `CHUNK_ALIGN`, which is also the strictest alignment the pool accepts +/// (`bump_scope::Bump` requests 16 — see [`MmapAllocator`]'s docs for the +/// links). The cache holds at most `SLOTS` chunks, so at most +/// `SLOTS * CHUNK_SIZE` bytes stay retained. +/// +/// Requests of up to `CHUNK_SIZE` bytes are served with a whole chunk — +/// a cached one, or a fresh one when the cache is empty. Bigger requests go +/// straight to the underlying allocator. Deallocation mirrors that: chunks +/// go back into the cache (or to the underlying allocator when it is full), +/// bigger blocks go to the underlying allocator directly. The size passed +/// to `deallocate` tells the two apart exactly: for a block served as a +/// chunk, every size the caller may legally pass is at most `CHUNK_SIZE`; +/// for a pass-through block, every legal size is bigger. +/// +/// Zero-sized requests and alignments above `CHUNK_ALIGN` are refused — +/// its only intended caller, `bump_scope::Bump`, never sends either. +/// +/// # Why it is safe in signal handlers and forked children +/// +/// The only state of the pool itself is a fixed array of atomic pointers, +/// one slot per cached chunk. Taking a chunk swaps a slot to null; +/// returning one swaps a null slot back. Every operation visits each slot +/// at most once and never waits — there are no locks to leave locked and no +/// lists to leave half-linked, so a thread that disappears mid-operation +/// (`fork`, a signal) can strand at most the one chunk it was holding, +/// never the pool. (The underlying allocator must give the same guarantee; +/// [`MmapAllocator`] does.) +/// +/// The pool is `const`-constructible, so it can live in a `static` and +/// work before anything else has run. +pub struct ChunkPool< + A: Allocator, + const CHUNK_SIZE: usize, + const CHUNK_ALIGN: usize, + const SLOTS: usize, +> { + slots: [AtomicPtr; SLOTS], + allocator: A, } -/// The allocator core, generic over its [`Sys`] memory provider. -pub struct Pool { - classes: [ClassState; CLASS_COUNT], - sys: PhantomData S>, -} - -impl Pool { - #[expect( - clippy::large_stack_arrays, - reason = "the class table (~30 KiB) is only ever materialized into a const-initialized static, never built on a runtime stack" - )] +impl + ChunkPool +{ + #[must_use] pub const fn new() -> Self { - Self { classes: [const { ClassState::new() }; CLASS_COUNT], sys: PhantomData } + Self::new_in(MmapAllocator) } +} - /// Allocates memory for `layout`. Returns `None` on exhaustion. - pub fn alloc(&self, layout: Layout) -> Option> { - match class_of(layout) { - Some(class) => match self.alloc_block(class)? { - ClassBlock::Recycled(ptr) | ClassBlock::Carved(ptr) => Some(ptr), - }, - // Deliberately leaked until `dealloc` reclaims ownership. - None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), +impl + ChunkPool +{ + pub const fn new_in(allocator: A) -> Self { + const { + assert!(CHUNK_SIZE > 0, "CHUNK_SIZE must not be zero"); + assert!(CHUNK_ALIGN.is_power_of_two(), "CHUNK_ALIGN must be a power of two"); } + Self { slots: [const { AtomicPtr::new(ptr::null_mut()) }; SLOTS], allocator } } - /// Like [`Pool::alloc`], but the returned memory is zeroed. - pub fn alloc_zeroed(&self, layout: Layout) -> Option> { - match class_of(layout) { - Some(class) => match self.alloc_block(class)? { - // Freshly carved blocks are still kernel-zeroed. - ClassBlock::Carved(ptr) => Some(ptr), - ClassBlock::Recycled(ptr) => { - // SAFETY: the block is freshly allocated, exclusively - // ours, and at least `layout.size()` bytes. - unsafe { ptr.as_ptr().write_bytes(0, layout.size()) }; - Some(ptr) - } - }, - // Fresh mappings are zeroed by the kernel; deliberately leaked - // until `dealloc` reclaims ownership. - None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), - } + fn chunk_layout() -> Result { + Layout::from_size_align(CHUNK_SIZE, CHUNK_ALIGN).map_err(|_| AllocError) } +} - /// Releases memory obtained from this pool. - /// - /// # Safety - /// - /// `ptr` must have been returned by this pool for exactly this `layout` - /// and must not be used afterwards. - pub unsafe fn dealloc(&self, ptr: NonNull, layout: Layout) { - match class_of(layout) { - // SAFETY: caller contract — a `Some(class)` layout was served - // from the pool, so `ptr` is a live block of this class. - Some(class) => unsafe { self.push_free(ptr, class) }, - // SAFETY: caller contract — a `None` layout was served by the - // mapping path with these parameters and is no longer in use, so - // ownership can be reclaimed (and the region dropped). - None => drop(unsafe { Mapping::::from_raw(ptr, layout.size(), layout.align()) }), - } +impl + Default for ChunkPool +{ + fn default() -> Self { + Self::new_in(A::default()) } +} - /// Grows or shrinks an allocation, preserving contents up to the smaller - /// of the old and new sizes. Returns `None` on exhaustion (the original - /// allocation stays valid). - /// - /// # Safety - /// - /// `ptr` must have been returned by this pool for exactly `layout`, and - /// `new_size` must be non-zero. - pub unsafe fn realloc( - &self, - ptr: NonNull, - layout: Layout, - new_size: usize, - ) -> Option> { - let new_layout = Layout::from_size_align(new_size, layout.align()).ok()?; - let class = class_of(layout); - if class.is_some() && class == class_of(new_layout) { - // Same size class: the existing block already fits. - return Some(ptr); +// SAFETY: returned blocks are non-null, aligned to at least the accepted +// `layout.align()` (chunks are allocated with CHUNK_ALIGN, and stricter +// requests are refused), and at least `layout.size()` bytes large (chunks +// are exactly CHUNK_SIZE bytes and only serve requests up to that size; +// everything else comes from the underlying allocator with its own +// guarantees). A chunk leaves its slot before it is handed out and returns +// only on deallocation, so every block has exactly one owner and stays +// valid until deallocated. +unsafe impl + Allocator for ChunkPool +{ + fn allocate(&self, layout: Layout) -> Result, AllocError> { + if layout.size() == 0 || layout.align() > CHUNK_ALIGN { + return Err(AllocError); } - let new_ptr = self.alloc(new_layout)?; - // SAFETY: `new_ptr` is a fresh exclusive allocation of at least - // `new_size` bytes; `ptr` is valid for `layout.size()` bytes (caller - // contract); distinct allocations never overlap. - unsafe { - new_ptr.as_ptr().copy_from_nonoverlapping(ptr.as_ptr(), layout.size().min(new_size)); + if layout.size() > CHUNK_SIZE { + return self.allocator.allocate(layout); } - // SAFETY: caller contract — `ptr` came from this pool with `layout`. - unsafe { self.dealloc(ptr, layout) }; - Some(new_ptr) - } - - fn alloc_block(&self, class: usize) -> Option { - loop { - if let Some(ptr) = self.pop_free(class) { - return Some(ClassBlock::Recycled(ptr)); + for slot in &self.slots { + if slot.load(Ordering::Relaxed).is_null() { + continue; } - match self.carve(class) { - Carve::Block(ptr) => return Some(ClassBlock::Carved(ptr)), - Carve::Retry => {} - Carve::Exhausted => return None, + // Acquire pairs with the Release in `deallocate`; the swap makes + // this thread the chunk's only owner. A racing taker may empty + // the slot between the load and the swap — then the swap returns + // null and the scan moves on. + let chunk = slot.swap(ptr::null_mut(), Ordering::Acquire); + if let Some(chunk) = NonNull::new(chunk) { + return Ok(NonNull::slice_from_raw_parts(chunk, CHUNK_SIZE)); } } + // Cache empty: get a fresh chunk. + self.allocator.allocate(Self::chunk_layout()?) } - /// Returns the published slab of `class` at `slab_idx`, if installed. - fn published_slab(&self, class: usize, slab_idx: usize) -> Option { - let base = NonNull::new(self.classes[class].slabs[slab_idx].load(Ordering::Acquire))?; - // SAFETY: non-null slab-table entries are only published (with - // `Release`, paired with the `Acquire` above) after - // `Slab::init_header` ran on a fresh SLAB_SIZE-aligned mapping, and - // are never cleared or unmapped while the pool is in use. - Some(unsafe { Slab::from_published(base, class) }) - } - - /// Pops a block off the class's free list. - fn pop_free(&self, class: usize) -> Option> { - let state = &self.classes[class]; - loop { - let head = state.head.load(Ordering::Acquire); - let (tag, block_ref) = head_parts(head); - if block_ref == NO_BLOCK { - return None; - } - let (slab_idx, block_idx) = unpack_ref(block_ref); - // A listed block was carved from its slab, so the slab is always - // published; `None` here is unreachable in practice. - let slab = self.published_slab(class, slab_idx)?; - // Read the successor link before the CAS. If another thread pops - // this block first, the tag comparison below fails and the value - // read here is discarded; since links live in the atomic side - // table, the racing read itself is well-defined. - let next = slab.link(block_idx).load(Ordering::Relaxed); - let new_head = head_from_parts(tag.wrapping_add(1), next); - if state - .head - .compare_exchange_weak(head, new_head, Ordering::AcqRel, Ordering::Relaxed) - .is_ok() - { - return Some(slab.block(block_idx)); - } + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + if layout.size() > CHUNK_SIZE { + // A pass-through block: only requests bigger than CHUNK_SIZE get + // one, and all their legal deallocation sizes are bigger too. + // SAFETY: caller contract — the block came from `allocate` + // (which forwarded to the underlying allocator) with a fitting + // layout. + return unsafe { self.allocator.deallocate(ptr, layout) }; } - } - - /// Pushes a block onto the class's free list. - /// - /// # Safety - /// - /// `ptr` must be a block of `class` previously returned by this pool and - /// no longer in use. - unsafe fn push_free(&self, ptr: NonNull, class: usize) { - // SAFETY: caller contract — `ptr` is a live block of `class` from - // this pool. (The impossible `None` would merely leak the block.) - let Some(slab) = (unsafe { Slab::of_block(ptr, class) }) else { return }; - let block_idx = slab.block_index(ptr); - let packed = pack_ref(slab.slab_idx(), block_idx); - let link = slab.link(block_idx); - let state = &self.classes[class]; - loop { - let head = state.head.load(Ordering::Relaxed); - let (tag, head_ref) = head_parts(head); - link.store(head_ref, Ordering::Relaxed); - let new_head = head_from_parts(tag.wrapping_add(1), packed); - // `Release` publishes the link store above to the eventual popper. - if state - .head - .compare_exchange_weak(head, new_head, Ordering::Release, Ordering::Relaxed) + // A chunk: put it back into an empty slot, or give it back to the + // underlying allocator if the cache is full. Release pairs with the + // Acquire in `allocate`. A failed exchange means another thread just + // filled the slot; move on. + for slot in &self.slots { + if !slot.load(Ordering::Relaxed).is_null() { + continue; + } + if slot + .compare_exchange( + ptr::null_mut(), + ptr.as_ptr(), + Ordering::Release, + Ordering::Relaxed, + ) .is_ok() { return; } } + let Ok(chunk_layout) = Self::chunk_layout() else { return }; + // SAFETY: every block in this branch is a CHUNK_SIZE chunk that came + // from the underlying allocator with `chunk_layout` (caller contract + // plus the size routing above), and it no longer has an owner. + unsafe { self.allocator.deallocate(ptr, chunk_layout) }; } +} - /// Carves the next block off the class's active slab, installing a new - /// slab if the active one is exhausted (or none exists yet). - fn carve(&self, class: usize) -> Carve { - let state = &self.classes[class]; - let active = state.active.load(Ordering::Acquire); - // `active` is only published after its slab pointer, so a live - // `active` always resolves to a published slab. - if active != NO_SLAB - && let Some(slab) = self.published_slab(class, active as usize) - { - let carved_idx = slab.carved().fetch_add(1, Ordering::Relaxed); - if let Some(block_idx) = carved_to_block_idx(carved_idx, class) { - return Carve::Block(slab.block(block_idx)); - } - // Active slab exhausted; fall through to install the next one. - } - let next_idx = if active == NO_SLAB { 0 } else { active as usize + 1 }; - if next_idx >= MAX_SLABS_PER_CLASS { - return Carve::Exhausted; - } - self.install_slab(class, next_idx); - if self.published_slab(class, next_idx).is_none() { - // Our mapping failed and no other thread succeeded either. - return Carve::Exhausted; - } - #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] - let next_active = next_idx as u32; - // Advance `active`; losing the race just means another thread already - // advanced it. Either way the retry re-reads it. - let _ = - state.active.compare_exchange(active, next_active, Ordering::AcqRel, Ordering::Relaxed); - Carve::Retry - } - - /// Maps and publishes the slab at `slab_idx`, unless another thread beats - /// us to it (or the mapping fails, leaving the slot null). - fn install_slab(&self, class: usize, slab_idx: usize) { - let state = &self.classes[class]; - if !state.slabs[slab_idx].load(Ordering::Acquire).is_null() { - return; - } - let Some(mapping) = Mapping::::new(SLAB_SIZE, SLAB_SIZE) else { return }; - #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] - let idx = slab_idx as u32; - // SAFETY: `mapping` is a fresh, exclusive, zero-initialized, - // SLAB_SIZE-byte and SLAB_SIZE-aligned mapping. - unsafe { Slab::init_header(mapping.ptr(), idx) }; - if state.slabs[slab_idx] - .compare_exchange( - ptr::null_mut(), - mapping.ptr().as_ptr(), - Ordering::Release, - Ordering::Relaxed, - ) - .is_ok() - { - // Published: the slab now lives for the rest of the process. - let _ = mapping.into_raw(); - } - // Otherwise another thread installed this slot first; our - // never-published mapping is released when `mapping` drops. - } - - /// Tears down all slab mappings. Test-only: the global allocator lives in - /// a static and never releases its slabs, but tests (and Miri's leak - /// checker) want a clean shutdown. - #[cfg(test)] - fn unmap_all_slabs(&mut self) { - for state in &mut self.classes { - *state.head.get_mut() = head_from_parts(0, NO_BLOCK); - *state.active.get_mut() = NO_SLAB; - for slot in &mut state.slabs { - let slab = core::mem::replace(slot.get_mut(), ptr::null_mut()); - if let Some(base) = NonNull::new(slab) { - // SAFETY: `base` was leaked into the table by - // `install_slab` with these parameters; `&mut self` - // guarantees no concurrent (or future) use of its blocks. - drop(unsafe { Mapping::::from_raw(base, SLAB_SIZE, SLAB_SIZE) }); - } +impl Drop + for ChunkPool +{ + fn drop(&mut self) { + // Statics never drop; this runs for pools created in tests. + for slot in &mut self.slots { + let chunk = core::mem::replace(slot.get_mut(), ptr::null_mut()); + if let Some(chunk) = NonNull::new(chunk) { + let Ok(chunk_layout) = Self::chunk_layout() else { return }; + // SAFETY: cached chunks came from the underlying allocator + // with `chunk_layout`, and `&mut self` means no owner exists. + unsafe { self.allocator.deallocate(chunk, chunk_layout) }; } } } } -/// Converts a raw carve-counter value into a block index, or `None` if the -/// slab is exhausted. -fn carved_to_block_idx(carved: u64, class: usize) -> Option { - let idx = usize::try_from(carved).ok()?; - (idx < BLOCKS_PER_SLAB[class]).then_some(idx) -} - #[cfg(test)] mod tests { - use std::{boxed::Box, sync::mpsc, thread, vec::Vec}; - - use super::*; - use crate::class::{MAX_BLOCK_SIZE, MAX_POOL_ALIGN, block_size}; - - /// Host-allocator-backed [`Sys`] so the pool core runs under Miri and on - /// any platform. - struct TestSys; - - impl Sys for TestSys { - fn map(size: usize, align: usize) -> Option> { - let layout = Layout::from_size_align(size.max(1), align).ok()?; - // SAFETY: `layout` has non-zero size. - NonNull::new(unsafe { std::alloc::alloc_zeroed(layout) }) - } + use std::{thread, vec::Vec}; - unsafe fn unmap(ptr: NonNull, size: usize, align: usize) { - let layout = Layout::from_size_align(size.max(1), align).unwrap(); - // SAFETY: caller contract — `ptr` was returned by `map`, which - // used exactly this layout. - unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) } - } - } + use allocator_api2::alloc::Global; - fn with_pool(test: impl FnOnce(&Pool)) { - let mut pool = Box::new(Pool::::new()); - test(&pool); - pool.unmap_all_slabs(); - } + use super::*; - fn layout(size: usize, align: usize) -> Layout { - Layout::from_size_align(size, align).unwrap() - } + const CHUNK_SIZE: usize = 64 * 1024; + const CHUNK_ALIGN: usize = 16; - #[test] - fn head_tag_wraps_within_its_field() { - let head = head_from_parts(HEAD_TAG_MASK, 42); - assert_eq!(head_parts(head), (HEAD_TAG_MASK, 42)); - // Advancing the maximal tag must wrap to zero without touching the - // reference bits. - let wrapped = head_from_parts(HEAD_TAG_MASK.wrapping_add(1), 42); - assert_eq!(head_parts(wrapped), (0, 42)); + fn pool() -> ChunkPool { + ChunkPool::new_in(Global) } - #[test] - fn packed_refs_cannot_collide_with_the_sentinel() { - for &blocks in &BLOCKS_PER_SLAB { - let max_ref = pack_ref(MAX_SLABS_PER_CLASS - 1, blocks - 1); - assert_ne!(max_ref, NO_BLOCK); - assert!(u64::from(max_ref) <= HEAD_REF_MASK); - } + fn layout(size: usize) -> Layout { + Layout::from_size_align(size, 8).unwrap() } #[test] - fn classify_boundaries() { - assert_eq!(class_of(layout(1, 1)), Some(0)); - assert_eq!(class_of(layout(16, 1)), Some(0)); - assert_eq!(class_of(layout(17, 1)), Some(1)); - assert_eq!(class_of(layout(MAX_BLOCK_SIZE, 8)), Some(CLASS_COUNT - 1)); - assert_eq!(class_of(layout(MAX_BLOCK_SIZE + 1, 8)), None); - // Alignment can raise the class. - assert_eq!(class_of(layout(8, 1024)), class_of(layout(1024, 8))); - // Over-aligned layouts bypass the pool. - assert_eq!(class_of(layout(16, MAX_POOL_ALIGN * 2)), None); + fn small_requests_get_whole_recycled_chunks() { + let pool = pool(); + let first = pool.allocate(layout(100)).unwrap(); + assert_eq!(first.len(), CHUNK_SIZE); + let first_addr = first.cast::().as_ptr().addr(); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { first.cast::().as_ptr().write_bytes(0xAB, 100) }; + // SAFETY: allocated above; the full returned size is a legal size to + // pass back. + unsafe { pool.deallocate(first.cast(), layout(CHUNK_SIZE)) }; + + // The chunk is cached now; the next small request must get it back. + let second = pool.allocate(layout(5000)).unwrap(); + assert_eq!(second.cast::().as_ptr().addr(), first_addr); + assert_eq!(second.len(), CHUNK_SIZE); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(second.cast(), layout(5000)) }; } #[test] - fn round_trip_every_class() { - with_pool(|pool| { - for class in 0..CLASS_COUNT { - let l = layout(block_size(class), 8); - let ptr = pool.alloc(l).unwrap(); - // SAFETY: fresh exclusive allocation of `block_size` bytes. - unsafe { ptr.as_ptr().write_bytes(0xAB, l.size()) }; - // SAFETY: reading back the block we just wrote. - let last = unsafe { ptr.as_ptr().add(l.size() - 1).read() }; - assert_eq!(last, 0xAB); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn deallocating_with_the_requested_size_also_recycles() { + let pool = pool(); + let block = pool.allocate(layout(100)).unwrap(); + let addr = block.cast::().as_ptr().addr(); + // SAFETY: allocated above; the requested size is a legal size to + // pass back, and must still be recognized as a chunk. + unsafe { pool.deallocate(block.cast(), layout(100)) }; + let again = pool.allocate(layout(100)).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), addr); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(again.cast(), layout(100)) }; } #[test] - fn block_alignment() { - with_pool(|pool| { - for align in [8, 64, 1024, MAX_POOL_ALIGN] { - let l = layout(24, align); - let ptr = pool.alloc(l).unwrap(); - assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn big_requests_bypass_the_cache() { + let pool = pool(); + let big = pool.allocate(layout(CHUNK_SIZE + 1)).unwrap(); + assert!(big.len() > CHUNK_SIZE); + // SAFETY: fresh exclusive block of at least CHUNK_SIZE + 1 + // bytes. + unsafe { big.cast::().as_ptr().write_bytes(0x5A, CHUNK_SIZE + 1) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(big.cast(), layout(CHUNK_SIZE + 1)) }; + + // A small request afterwards gets a chunk, not the big block. + let small = pool.allocate(layout(100)).unwrap(); + assert_eq!(small.len(), CHUNK_SIZE); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(small.cast(), layout(100)) }; } #[test] - fn free_list_recycles_lifo() { - with_pool(|pool| { - let l = layout(100, 8); - let first = pool.alloc(l).unwrap(); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(first, l) }; - let second = pool.alloc(l).unwrap(); - assert_eq!(first, second); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(second, l) }; - }); + fn refuses_zero_size_and_over_chunk_alignment() { + let pool = pool(); + let zero = Layout::from_size_align(0, 16).unwrap(); + assert!(pool.allocate(zero).is_err()); + let over_aligned = Layout::from_size_align(64, CHUNK_ALIGN * 2).unwrap(); + assert!(pool.allocate(over_aligned).is_err()); } #[test] - fn carves_across_multiple_slabs() { - with_pool(|pool| { - // The largest class has the fewest blocks per slab, so a couple - // dozen live allocations force several slab installations. - let l = layout(MAX_BLOCK_SIZE, 8); - let count = BLOCKS_PER_SLAB[CLASS_COUNT - 1] * 3 + 1; - let blocks: Vec> = (0..count).map(|_| pool.alloc(l).unwrap()).collect(); - for (i, ptr) in blocks.iter().enumerate() { - assert!(blocks[..i].iter().all(|other| other != ptr), "duplicate block"); - // SAFETY: live exclusive allocation of `MAX_BLOCK_SIZE` bytes. - unsafe { ptr.as_ptr().write_bytes(0x5A, l.size()) }; - } - for ptr in blocks { - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn allocate_zeroed_scrubs_recycled_chunks() { + let pool = pool(); + let block = pool.allocate(layout(256)).unwrap(); + // SAFETY: fresh exclusive block of at least 256 bytes. + unsafe { block.cast::().as_ptr().write_bytes(0xFF, 256) }; + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(256)) }; + + // Recycled chunks are dirty, so the default `allocate_zeroed` (which + // we deliberately do not override here) must scrub them. + let zeroed = pool.allocate_zeroed(layout(256)).unwrap(); + for i in 0..256 { + // SAFETY: fresh exclusive block of at least 256 bytes. + assert_eq!(unsafe { zeroed.cast::().as_ptr().add(i).read() }, 0); + } + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(zeroed.cast(), layout(256)) }; } #[test] - fn large_allocations_bypass_pool() { - with_pool(|pool| { - for l in [layout(MAX_BLOCK_SIZE + 1, 8), layout(5 << 20, 8), layout(64, 8192)] { - let ptr = pool.alloc(l).unwrap(); - assert_eq!(ptr.as_ptr().addr() % l.align(), 0); - // SAFETY: fresh exclusive allocation of `l.size()` bytes. - unsafe { ptr.as_ptr().write_bytes(0xCD, l.size()) }; - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; + fn overflowing_the_cache_is_fine() { + // A tiny 4-slot pool so the full-cache path is genuinely exercised. + let pool = ChunkPool::::new_in(Global); + for _ in 0..2 { + let blocks: Vec<_> = (0..8).map(|_| pool.allocate(layout(1000)).unwrap()).collect(); + for block in blocks { + // SAFETY: allocated above with a fitting layout. The first 4 + // chunks fill the cache; the rest go back to the underlying + // allocator. + unsafe { pool.deallocate(block.cast(), layout(1000)) }; } - }); - } - - #[test] - fn realloc_within_class_keeps_block() { - with_pool(|pool| { - let l = layout(100, 8); - let ptr = pool.alloc(l).unwrap(); - // SAFETY: realloc contract — `ptr` allocated with `l`. - let grown = unsafe { pool.realloc(ptr, l, 120) }.unwrap(); - assert_eq!(ptr, grown, "same class must realloc in place"); - // SAFETY: allocated above; 120 rounds to the same class as 100. - unsafe { pool.dealloc(grown, layout(120, 8)) }; - }); + } } #[test] - fn realloc_across_classes_preserves_contents() { - with_pool(|pool| { - let l = layout(64, 8); - let ptr = pool.alloc(l).unwrap(); - for i in 0..64u8 { - // SAFETY: live exclusive allocation of 64 bytes. - unsafe { ptr.as_ptr().add(usize::from(i)).write(i) }; - } - // SAFETY: realloc contract — `ptr` allocated with `l`. - let grown = unsafe { pool.realloc(ptr, l, 4096) }.unwrap(); - for i in 0..64u8 { - // SAFETY: `grown` is live for 4096 bytes. - let got = unsafe { grown.as_ptr().add(usize::from(i)).read() }; - assert_eq!(got, i); - } - // Shrink across classes, and from the large path back into a class. - // SAFETY: `grown` allocated with the 4096 layout above. - let shrunk = unsafe { pool.realloc(grown, layout(4096, 8), 8) }.unwrap(); - // SAFETY: `shrunk` is live for 8 bytes. - assert_eq!(unsafe { shrunk.as_ptr().read() }, 0); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(shrunk, layout(8, 8)) }; - }); + fn custom_chunk_size_and_alignment_work() { + let pool = ChunkPool::::new_in(Global); + let block = pool.allocate(layout(10)).unwrap(); + assert_eq!(block.len(), 1024); + assert_eq!(block.cast::().as_ptr().addr() % 64, 0); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(10)) }; + + // Bigger than this pool's chunk size: passes through. + let big = pool.allocate(layout(2000)).unwrap(); + assert!(big.len() >= 2000); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(big.cast(), layout(2000)) }; } #[test] - fn alloc_zeroed_scrubs_recycled_blocks() { - with_pool(|pool| { - let l = layout(256, 8); - let dirty = pool.alloc(l).unwrap(); - // SAFETY: live exclusive allocation of 256 bytes. - unsafe { dirty.as_ptr().write_bytes(0xFF, l.size()) }; - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(dirty, l) }; - - let zeroed = pool.alloc_zeroed(l).unwrap(); - assert_eq!( - zeroed, dirty, - "must recycle the dirty block for this test to be meaningful" - ); - for i in 0..l.size() { - // SAFETY: live exclusive allocation of 256 bytes. - assert_eq!(unsafe { zeroed.as_ptr().add(i).read() }, 0); + fn concurrent_use_smoke() { + const OPS: usize = if cfg!(miri) { 40 } else { 500 }; + + let pool = pool(); + thread::scope(|scope| { + for t in 0..4_usize { + let pool = &pool; + scope.spawn(move || { + for i in 0..OPS { + let size = 1 + (i * 37 + t * 101) % (2 * CHUNK_SIZE); + let l = Layout::from_size_align(size, 8).unwrap(); + let block = pool.allocate(l).unwrap(); + assert!(block.len() >= size); + let marker = u8::try_from((t * 50 + i) % 251).unwrap(); + // SAFETY: exclusive fresh block of at least `size` + // bytes, allocated above with layout `l`. + unsafe { + block.cast::().as_ptr().write(marker); + block.cast::().as_ptr().add(size - 1).write(marker); + assert_eq!(block.cast::().as_ptr().read(), marker); + assert_eq!(block.cast::().as_ptr().add(size - 1).read(), marker); + pool.deallocate(block.cast(), l); + } + } + }); } - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(zeroed, l) }; }); } - /// Cheap deterministic PRNG so the stress test needs no dependencies. - fn lcg(state: &mut u64) -> u64 { - *state = - state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407); - *state >> 33 - } - - /// An owned allocation in flight between threads. - struct SendBlock(NonNull, Layout); - // SAFETY: SendBlock represents exclusive ownership of the block, which is - // transferred wholesale to the receiving thread. - unsafe impl Send for SendBlock {} - + /// The one test against the real kernel-backed default; everything else + /// runs against `Global` so Miri can check it. #[test] - fn concurrent_stress_with_cross_thread_frees() { - const THREADS: usize = if cfg!(miri) { 3 } else { 8 }; - const OPS: usize = if cfg!(miri) { 60 } else { 20_000 }; - - with_pool(|pool| { - thread::scope(|scope| { - // Each thread frees blocks allocated by its neighbor, so - // pushes and pops of the same list run on different threads. - // Bounded channels keep the number of in-flight blocks well - // under the pool's per-class capacity regardless of thread - // scheduling. - let (senders, receivers): (Vec<_>, Vec<_>) = - (0..THREADS).map(|_| mpsc::sync_channel::(64)).unzip(); - let mut senders_rotated: Vec<_> = senders.into_iter().map(Some).collect(); - senders_rotated.rotate_left(1); - - for (thread_idx, (receiver, sender)) in - receivers.into_iter().zip(&mut senders_rotated).enumerate() - { - let sender = sender.take().unwrap(); - let pool = &*pool; - scope.spawn(move || { - let mut rng = 0x9E37_79B9_7F4A_7C15_u64 ^ thread_idx as u64; - for _ in 0..OPS { - let size = 1 + usize::try_from(lcg(&mut rng)).unwrap() - % (3 * MAX_BLOCK_SIZE / 2); - let align = 1 << (lcg(&mut rng) % 7); - let l = layout(size, align); - let ptr = pool.alloc(l).unwrap(); - let marker = u8::try_from(lcg(&mut rng) & 0xFF).unwrap(); - // SAFETY: fresh exclusive allocation of `size` bytes. - unsafe { - ptr.as_ptr().write(marker); - ptr.as_ptr().add(size - 1).write(marker); - } - // SAFETY: we still exclusively own the block. - let (first, last) = - unsafe { (ptr.as_ptr().read(), ptr.as_ptr().add(size - 1).read()) }; - assert_eq!((first, last), (marker, marker), "block corrupted"); - if let Err(returned) = sender.try_send(SendBlock(ptr, l)) { - // Neighbor's queue is full (or it finished); - // free locally instead of blocking, which in - // a ring of senders could deadlock. - let (mpsc::TrySendError::Full(SendBlock(ptr, l)) - | mpsc::TrySendError::Disconnected(SendBlock(ptr, l))) = returned; - // SAFETY: allocated above with layout `l`. - unsafe { pool.dealloc(ptr, l) }; - } - // Drain what our own producer has sent so far, - // interleaving cross-thread pops with pushes. - while let Ok(SendBlock(ptr, l)) = receiver.try_recv() { - // SAFETY: the neighbor allocated `ptr` with - // `l` and transferred ownership over the - // channel. - unsafe { pool.dealloc(ptr, l) }; - } - } - drop(sender); - for SendBlock(ptr, l) in receiver { - // SAFETY: the neighbor allocated `ptr` with `l` - // and transferred ownership over the channel. - unsafe { pool.dealloc(ptr, l) }; - } - }); - } - }); - }); + #[cfg(not(miri))] + fn mmap_backed_pool_smoke() { + let pool = ChunkPool::::new(); + let block = pool.allocate(layout(100)).unwrap(); + assert_eq!(block.len(), CHUNK_SIZE); + let addr = block.cast::().as_ptr().addr(); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { block.cast::().as_ptr().write_bytes(0xCD, 100) }; + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(100)) }; + let again = pool.allocate(layout(200)).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), addr); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(again.cast(), layout(200)) }; } } diff --git a/crates/fspy_alloc/src/slab.rs b/crates/fspy_alloc/src/slab.rs deleted file mode 100644 index 3aca476dc..000000000 --- a/crates/fspy_alloc/src/slab.rs +++ /dev/null @@ -1,203 +0,0 @@ -//! Slab memory layout and the [`Slab`] view type. -//! -//! Each size class carves fixed-size blocks out of [`SLAB_SIZE`]-byte, -//! `SLAB_SIZE`-aligned slabs. A slab looks like this: -//! -//! ```text -//! | SlabHeader | links: [AtomicU32; blocks] | pad to 4 KiB | block 0 | block 1 | ... | -//! ``` -//! -//! Free blocks are chained through the `links` side table — never through -//! block memory itself — so block payloads are only ever plain data. That -//! keeps every concurrent access well-defined under the memory model (no -//! mixed atomic/non-atomic reads of the same bytes) and Miri-clean. - -use core::{ - num::NonZero, - ptr::NonNull, - sync::atomic::{AtomicU32, AtomicU64}, -}; - -use crate::class::{CLASS_COUNT, MAX_POOL_ALIGN, block_size}; - -/// Slab size and alignment. Masking a block address with `!(SLAB_SIZE - 1)` -/// recovers its slab base. -pub const SLAB_SIZE: usize = 1 << 20; - -const fn round_up(value: usize, align: usize) -> usize { - (value + align - 1) & !(align - 1) -} - -/// Byte offset of the block area inside a slab holding `blocks` blocks. -const fn block_area_offset(blocks: usize) -> usize { - round_up(size_of::() + blocks * size_of::(), MAX_POOL_ALIGN) -} - -const fn compute_blocks_per_slab(class: usize) -> usize { - let size = block_size(class); - // Upper bound ignoring metadata, then shrink until header, link table, - // padding, and blocks all fit in one slab. - let mut blocks = SLAB_SIZE / size; - while block_area_offset(blocks) + blocks * size > SLAB_SIZE { - blocks -= 1; - } - blocks -} - -/// Blocks per slab, for each size class. -pub const BLOCKS_PER_SLAB: [usize; CLASS_COUNT] = { - let mut table = [0usize; CLASS_COUNT]; - let mut class = 0; - while class < CLASS_COUNT { - table[class] = compute_blocks_per_slab(class); - class += 1; - } - table -}; - -const _: () = { - assert!(size_of::() == 16); - assert!(MAX_POOL_ALIGN >= align_of::()); - let mut class = 0; - while class < CLASS_COUNT { - assert!(BLOCKS_PER_SLAB[class] > 0); - assert!( - block_area_offset(BLOCKS_PER_SLAB[class]) + BLOCKS_PER_SLAB[class] * block_size(class) - <= SLAB_SIZE - ); - class += 1; - } -}; - -/// Lives at the base of every slab mapping, ahead of the link table. -#[repr(C)] -struct SlabHeader { - /// This slab's index in its class's slab table. - slab_idx: u32, - _reserved: u32, - /// Number of blocks ever carved off this slab (monotonic; values at or - /// beyond the class's blocks-per-slab mean the slab is exhausted). 64-bit - /// so over-counting by racing threads can never wrap it around. - carved: AtomicU64, -} - -/// A view of one live slab mapping of a particular class. -/// -/// # Invariant -/// -/// `base` points to a [`SLAB_SIZE`]-byte, `SLAB_SIZE`-aligned mapping laid -/// out for `class` (initialized header, link table, block area) that is -/// never unmapped while the pool is in use. All raw-pointer arithmetic on -/// slab memory lives in this type's methods; the unsafe constructors are the -/// only places the invariant is asserted, and every accessor then relies on -/// it for bounds and liveness. -#[derive(Clone, Copy)] -pub struct Slab { - base: NonNull, - class: usize, -} - -impl Slab { - /// Wraps a slab pointer loaded from a class's slab table. - /// - /// # Safety - /// - /// `base` must be a non-null entry of the class's slab table. Such - /// entries are only published after [`Slab::init_header`] ran on a - /// suitable mapping, and are never cleared or unmapped while the pool is - /// in use, so the type invariant holds. - pub const unsafe fn from_published(base: NonNull, class: usize) -> Self { - Self { base, class } - } - - /// Recovers the slab containing a live pool block: slabs are - /// `SLAB_SIZE`-aligned, so masking the block address's low bits yields - /// the slab base (`with_addr` keeps the block pointer's provenance, which - /// covers the whole slab mapping it was carved from). Returns `None` only - /// for an address whose masked base would be null, which no real block - /// can produce. - /// - /// # Safety - /// - /// `block` must be a block of `class` previously handed out by this pool - /// and thus carved from a published slab of this class. - pub unsafe fn of_block(block: NonNull, class: usize) -> Option { - let base_addr = NonZero::new(block.as_ptr().addr() & !(SLAB_SIZE - 1))?; - Some(Self { base: block.with_addr(base_addr), class }) - } - - /// Initializes the header of a fresh slab mapping, making it publishable - /// into a slab table. - /// - /// # Safety - /// - /// `base` must point to a fresh, exclusive, zero-initialized, - /// `SLAB_SIZE`-byte and `SLAB_SIZE`-aligned mapping. The zero fill - /// doubles as the initial state of the carve counter and the link table. - pub unsafe fn init_header(base: NonNull, slab_idx: u32) { - // SAFETY: caller contract — the fresh exclusive mapping is aligned - // (SLAB_SIZE-aligned, far beyond SlabHeader's needs) and large - // enough for the header. - unsafe { - (*base.cast::().as_ptr()).slab_idx = slab_idx; - } - } - - const fn header(&self) -> &SlabHeader { - // SAFETY: type invariant — the mapping is live, SLAB_SIZE-aligned - // (far beyond SlabHeader's needs), and its header was initialized - // before publication; the only non-atomic field is never written - // again afterwards. - unsafe { self.base.cast::().as_ref() } - } - - /// This slab's index in its class's slab table. - pub const fn slab_idx(&self) -> usize { - self.header().slab_idx as usize - } - - /// The monotonic carve counter. - pub const fn carved(&self) -> &AtomicU64 { - &self.header().carved - } - - const fn blocks(&self) -> usize { - BLOCKS_PER_SLAB[self.class] - } - - /// The free-list link slot of `block_idx`. - pub fn link(&self, block_idx: usize) -> &AtomicU32 { - debug_assert!(block_idx < self.blocks()); - #[expect( - clippy::cast_ptr_alignment, - reason = "the link table starts at offset 16 of a SLAB_SIZE-aligned slab, so entries are 4-byte aligned" - )] - // SAFETY: type invariant plus the bound above put the slot within - // the slab's link table; slots are only ever accessed atomically, - // and the mapping outlives any borrow. - unsafe { - AtomicU32::from_ptr( - self.base.as_ptr().add(size_of::()).cast::().add(block_idx), - ) - } - } - - /// The address of block `block_idx`. - pub fn block(&self, block_idx: usize) -> NonNull { - debug_assert!(block_idx < self.blocks()); - let offset = block_area_offset(self.blocks()) + block_idx * block_size(self.class); - // SAFETY: type invariant plus the bound above keep `offset` within - // the SLAB_SIZE mapping. - unsafe { self.base.add(offset) } - } - - /// The index of a block previously returned by [`Slab::block`]. - pub fn block_index(&self, block: NonNull) -> usize { - let offset = - block.as_ptr().addr() - self.base.as_ptr().addr() - block_area_offset(self.blocks()); - debug_assert_eq!(offset % block_size(self.class), 0); - let block_idx = offset / block_size(self.class); - debug_assert!(block_idx < self.blocks()); - block_idx - } -} diff --git a/crates/fspy_alloc/src/sys.rs b/crates/fspy_alloc/src/sys.rs deleted file mode 100644 index 4fecdc085..000000000 --- a/crates/fspy_alloc/src/sys.rs +++ /dev/null @@ -1,25 +0,0 @@ -//! The backing-memory provider abstraction. -//! -//! The pool is generic over [`Sys`] so tests (and Miri) can substitute a -//! provider backed by the host allocator, while the real allocator obtains -//! anonymous mappings from the kernel (see the `mmap` module). - -use core::ptr::NonNull; - -/// Provides zero-initialized, aligned memory regions. -/// -/// Implementations must themselves be async-signal-safe and fork-safe: no -/// locks, no thread-locals, no libc malloc. -pub trait Sys { - /// Maps `size` bytes of zero-initialized memory aligned to `align` - /// (a power of two). Returns `None` when memory is exhausted. - fn map(size: usize, align: usize) -> Option>; - - /// Releases a region previously returned by [`Sys::map`]. - /// - /// # Safety - /// - /// `ptr` must have been returned by `Sys::map(size, align)` with the same - /// `size` and `align`, and must not be accessed afterwards. - unsafe fn unmap(ptr: NonNull, size: usize, align: usize); -} diff --git a/crates/fspy_alloc/tests/global_alloc.rs b/crates/fspy_alloc/tests/global_alloc.rs deleted file mode 100644 index ab0670ae5..000000000 --- a/crates/fspy_alloc/tests/global_alloc.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! Installs [`FspyAlloc`] as this test binary's global allocator, so every -//! allocation — including the test harness's own — exercises the allocator -//! end-to-end over real memory mappings. -#![cfg(all(unix, not(miri)))] // Miri covers the pool core via its mockable backend instead. - -use std::{collections::BTreeMap, sync::mpsc, thread}; - -use fspy_alloc::FspyAlloc; - -#[global_allocator] -static GLOBAL: FspyAlloc = FspyAlloc::new(); - -#[test] -fn collections_round_trip() { - let mut map = BTreeMap::new(); - for i in 0..1000_u32 { - let len = usize::try_from(i % 300).unwrap(); - map.insert(i, vec![0xAB_u8; len]); - } - assert_eq!(map.len(), 1000); - for (i, bytes) in &map { - assert_eq!(bytes.len(), usize::try_from(i % 300).unwrap()); - assert!(bytes.iter().all(|byte| *byte == 0xAB)); - } -} - -#[test] -fn vec_growth_reallocs_preserve_contents() { - let mut bytes = Vec::new(); - for i in 0..1_000_000_usize { - bytes.push(u8::try_from(i % 251).unwrap()); - } - for (i, byte) in bytes.iter().enumerate() { - assert_eq!(usize::from(*byte), i % 251); - } -} - -#[test] -fn large_and_zeroed_allocations() { - // Direct-mapped (beyond the largest size class), via the zeroing path. - let large = vec![0_u8; 5 * 1024 * 1024]; - assert!(large.iter().all(|byte| *byte == 0)); - // And through the plain path. - let boxed: Box<[u8]> = vec![7_u8; 300 * 1024].into_boxed_slice(); - assert!(boxed.iter().all(|byte| *byte == 7)); -} - -#[test] -fn many_small_allocations_span_slabs() { - // More live 16-byte-class blocks than one slab holds, forcing the pool - // through several slab installations (and the holding Vec through the - // direct-mapped path as it grows). - let boxes: Vec> = (0..120_000).map(|_| Box::new([0xEE_u8; 8])).collect(); - assert!(boxes.iter().all(|block| block.iter().all(|byte| *byte == 0xEE))); -} - -#[test] -fn threaded_producers_and_consumers() { - let threads = 8; - let per_thread = 5_000_usize; - let (sender, receiver) = mpsc::channel::>(); - thread::scope(|scope| { - for t in 0..threads { - let sender = sender.clone(); - scope.spawn(move || { - for i in 0..per_thread { - // Vary sizes across classes; blocks are freed (and often - // allocated) on the consumer thread below. - let len = 1 + (i * 37 + t * 101) % 5000; - sender.send(vec![u8::try_from(t % 251).unwrap(); len]).unwrap(); - } - }); - } - drop(sender); - let mut message_count = 0_usize; - for bytes in receiver { - message_count += 1; - assert!(!bytes.is_empty()); - let first = bytes[0]; - assert!(bytes.iter().all(|byte| *byte == first)); - } - assert_eq!(message_count, threads * per_thread); - }); -} diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 70a554563..7b54a06f5 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -8,6 +8,7 @@ publish = false crate-type = ["cdylib"] [target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true, features = ["alloc"] } anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index c42854b32..be1f2cc80 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "linux")] use std::ffi::CString; use std::{ - ffi::{CStr, OsStr}, + ffi::CStr, os::{fd::RawFd, unix::ffi::OsStrExt as _}, path::PathBuf, }; @@ -71,13 +71,27 @@ impl ToAbsolutePath for PathAt { if pathname.first().copied() == Some(b'/') { f(pathname.into()) } else { - let Some(mut abs_path) = get_fd_path(self.0)? else { + let Some(dir) = get_fd_path(self.0)? else { return f(None); }; + // Join `dir` and the relative `pathname` in a per-call bump + // arena instead of `PathBuf::push` on the global allocator. This + // runs on every fd-relative open/stat, including inside signal + // handlers and fork children, where the global allocator's locks + // are unsafe to take. The joined path only lives for the `f` + // call — nothing escapes the arena. + let arena = fspy_alloc::arena(); + let mut joined = allocator_api2::vec::Vec::new_in(&arena); + joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { - abs_path.push(OsStr::from_bytes(pathname)); + // Mirror `PathBuf::push`: exactly one separator between the + // (absolute, non-empty) dir and the relative pathname. + if joined.last() != Some(&b'/') { + joined.push(b'/'); + } + joined.extend_from_slice(pathname); } - f(Some(abs_path.as_os_str().as_bytes().as_bstr())) + f(Some(joined.as_bstr())) } } } diff --git a/crates/fspy_preload_unix/src/lib.rs b/crates/fspy_preload_unix/src/lib.rs index a9e667970..6c4e10b3e 100644 --- a/crates/fspy_preload_unix/src/lib.rs +++ b/crates/fspy_preload_unix/src/lib.rs @@ -1,16 +1,6 @@ // Compile as an empty crate on non-unix targets and on musl (where seccomp // alone handles access tracking). -/// This library interposes libc functions that POSIX declares -/// async-signal-safe (`open`, `stat`, `execve`, ...), so its own code may run -/// inside signal handlers and in the child of `fork()` in a multithreaded -/// process — contexts where taking the libc allocator's locks can deadlock. -/// Route every Rust allocation in this cdylib through fspy's lock-free, -/// mmap-backed allocator instead. -#[cfg(all(unix, not(target_env = "musl")))] -#[global_allocator] -static GLOBAL_ALLOCATOR: fspy_alloc::FspyAlloc = fspy_alloc::FspyAlloc::new(); - #[cfg(all(unix, not(target_env = "musl")))] mod client; #[cfg(all(unix, not(target_env = "musl")))] From 7f63a9316f62a28d7c4f73cbadf93d28eb00141f Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 11:06:01 +0800 Subject: [PATCH 3/7] feat(sigsafe): move the allocator into a crate of signal-safe syscall wrappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The allocator is the first piece of a broader need: the preload library runs inside intercepted libc calls, so everything it uses must work in signal handlers, in the child of fork() in a multithreaded process, and before libc has finished initializing. sigsafe is where such code now lives. fspy_alloc becomes sigsafe::alloc, and the mmap/munmap/page-size calls it makes move behind sigsafe::mm and sigsafe::param, the first safe syscall wrappers the crate offers on its own. sigsafe promises that on Linux none of its calls go through libc, and enforces the promise at compile time: lib.rs references rustix::runtime, a module that exists only in rustix's raw-syscall backend, so any build configuration that selects rustix's libc backend — including a feature enabled outside this repository, which no build script can see — fails to compile instead of silently keeping libc in the picture. The README explains the purpose, the rules, and the enforcement. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 20 +-- Cargo.toml | 4 +- crates/fspy_alloc/Cargo.toml | 21 --- crates/fspy_preload_unix/Cargo.toml | 2 +- .../fspy_preload_unix/src/client/convert.rs | 2 +- crates/sigsafe/Cargo.toml | 35 +++++ crates/sigsafe/README.md | 60 +++++++ .../src => sigsafe/src/alloc}/mmap.rs | 19 +-- .../src/lib.rs => sigsafe/src/alloc/mod.rs} | 82 ++++++---- .../src => sigsafe/src/alloc}/pool.rs | 2 +- crates/sigsafe/src/lib.rs | 40 +++++ crates/sigsafe/src/mm.rs | 11 ++ crates/sigsafe/src/param.rs | 147 ++++++++++++++++++ 13 files changed, 365 insertions(+), 80 deletions(-) delete mode 100644 crates/fspy_alloc/Cargo.toml create mode 100644 crates/sigsafe/Cargo.toml create mode 100644 crates/sigsafe/README.md rename crates/{fspy_alloc/src => sigsafe/src/alloc}/mmap.rs (92%) rename crates/{fspy_alloc/src/lib.rs => sigsafe/src/alloc/mod.rs} (61%) rename crates/{fspy_alloc/src => sigsafe/src/alloc}/pool.rs (99%) create mode 100644 crates/sigsafe/src/lib.rs create mode 100644 crates/sigsafe/src/mm.rs create mode 100644 crates/sigsafe/src/param.rs diff --git a/Cargo.lock b/Cargo.lock index b783495c7..5b2ecd7dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1252,15 +1252,6 @@ dependencies = [ "winsafe 0.0.27", ] -[[package]] -name = "fspy_alloc" -version = "0.0.0" -dependencies = [ - "allocator-api2", - "bump-scope", - "rustix", -] - [[package]] name = "fspy_benchmark" version = "0.0.0" @@ -1318,11 +1309,11 @@ dependencies = [ "artifact_profile", "bstr", "ctor", - "fspy_alloc", "fspy_shared", "fspy_shared_unix", "libc", "nix 0.31.2", + "sigsafe", "wincode", ] @@ -3448,6 +3439,15 @@ dependencies = [ "libc", ] +[[package]] +name = "sigsafe" +version = "0.0.0" +dependencies = [ + "allocator-api2", + "bump-scope", + "rustix", +] + [[package]] name = "simd-adler32" version = "0.3.8" diff --git a/Cargo.toml b/Cargo.toml index f1c766099..457ba5cc3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,7 +75,6 @@ materialized_artifact = { path = "crates/materialized_artifact" } materialized_artifact_build = { path = "crates/materialized_artifact_build" } flate2 = "1.0.35" fspy = { path = "crates/fspy" } -fspy_alloc = { path = "crates/fspy_alloc" } fspy_benchmark_launcher = { path = "crates/fspy_benchmark_launcher", artifact = "bin" } fspy_benchmark_target = { path = "crates/fspy_benchmark_target", artifact = "bin" } fspy_detours_sys = { path = "crates/fspy_detours_sys" } @@ -123,7 +122,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" -rustix = { version = "1", default-features = false, features = ["mm", "param", "use-libc-auxv"] } +rustix = { version = "1", default-features = false, features = ["mm"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" @@ -131,6 +130,7 @@ serde_json = "1.0.140" serde_norway = "0.9.42" sha2 = "0.11.0" shell-escape = "0.1.5" +sigsafe = { path = "crates/sigsafe" } similar = "3.0.0" smallvec = { version = "2.0.0-alpha.12", features = ["std"] } snapshot_test = { path = "crates/snapshot_test" } diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml deleted file mode 100644 index edcb89eaa..000000000 --- a/crates/fspy_alloc/Cargo.toml +++ /dev/null @@ -1,21 +0,0 @@ -[package] -name = "fspy_alloc" -edition = "2024" -license.workspace = true -publish = false - -[lib] -doctest = false - -[target.'cfg(unix)'.dependencies] -allocator-api2 = { workspace = true } -bump-scope = { workspace = true } -rustix = { workspace = true } - -[target.'cfg(unix)'.dev-dependencies] -# The `alloc` feature provides `Global`, letting tests run the pool against -# the host allocator (and thus under Miri). -allocator-api2 = { workspace = true, features = ["alloc"] } - -[lints] -workspace = true diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 7b54a06f5..89ffaabdc 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -13,11 +13,11 @@ anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } ctor = { workspace = true } -fspy_alloc = { workspace = true } fspy_shared = { workspace = true } fspy_shared_unix = { workspace = true } libc = { workspace = true } nix = { workspace = true, features = ["signal", "fs", "socket", "mman", "time"] } +sigsafe = { workspace = true } [build-dependencies] artifact_profile = { workspace = true } diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index be1f2cc80..70fc82a8d 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -80,7 +80,7 @@ impl ToAbsolutePath for PathAt { // handlers and fork children, where the global allocator's locks // are unsafe to take. The joined path only lives for the `f` // call — nothing escapes the arena. - let arena = fspy_alloc::arena(); + let arena = sigsafe::alloc::arena(); let mut joined = allocator_api2::vec::Vec::new_in(&arena); joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { diff --git a/crates/sigsafe/Cargo.toml b/crates/sigsafe/Cargo.toml new file mode 100644 index 000000000..9bd1af418 --- /dev/null +++ b/crates/sigsafe/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "sigsafe" +edition = "2024" +license.workspace = true +publish = false + +[lib] +doctest = false + +[target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true } +bump-scope = { workspace = true } +rustix = { workspace = true } + +# On Linux the page size is probed from the kernel directly (see param.rs); +# rustix's `param` is only needed where sysconf is the platform interface. +[target.'cfg(all(unix, not(target_os = "linux")))'.dependencies] +rustix = { workspace = true, features = ["param"] } + +# The compile-time backend check in lib.rs needs a `linux_raw`-gated rustix +# item to reference; `runtime` is the module that has one. +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true, features = ["runtime"] } + +# Cross-validates the page-size probe against rustix's auxv-based answer. +[target.'cfg(target_os = "linux")'.dev-dependencies] +rustix = { workspace = true, features = ["param"] } + +[target.'cfg(unix)'.dev-dependencies] +# The `alloc` feature provides `Global`, letting tests run the pool against +# the host allocator (and thus under Miri). +allocator-api2 = { workspace = true, features = ["alloc"] } + +[lints] +workspace = true diff --git a/crates/sigsafe/README.md b/crates/sigsafe/README.md new file mode 100644 index 000000000..083c4adc6 --- /dev/null +++ b/crates/sigsafe/README.md @@ -0,0 +1,60 @@ +# sigsafe + +Unix syscall wrappers that are safe to call where libc is not. + +## Why this crate exists + +The fspy preload library injects itself into traced programs and intercepts +their libc calls (`open`, `stat`, `execve`, ...). POSIX declares those +functions async-signal-safe, so programs are allowed to call them: + +- inside a signal handler, +- in the child of `fork()` of a multithreaded program, +- while the process is still starting up, before libc is fully initialized. + +Most of libc is off limits in those places. `malloc` is the classic trap: a +signal can pause a thread while it holds malloc's lock, and `fork()` copies a +locked lock into a child that has no thread left to unlock it — the next +`malloc` waits forever. Interception code runs exactly there, so anything it +calls must work without libc's machinery, or the traced program can hang. + +## The rules + +Every function in this crate follows three rules: + +1. **Syscalls only.** On Linux, nothing goes through libc — the syscall + instructions are emitted directly (rustix's raw backend). On macOS there + is no stable syscall interface, so calls go through libSystem's wrappers; + for the calls exposed here those are thin stubs with no locks and no + state. +2. **No locks, no hidden state.** Nothing a signal or a `fork()` could catch + locked or half-written. Where shared state is unavoidable it is a fixed + set of atomics, each touched by single complete operations. +3. **No global allocation.** No function touches a heap behind the caller's + back. Code that needs memory gets it from an explicit allocator — + [`alloc`](src/alloc/mod.rs) provides one built on `mmap`. + +## How rule 1 is enforced on Linux + +rustix can be built with a libc backend instead of raw syscalls, and anything +in the dependency graph — including crates outside this repository — can +select it (the `rustix/use-libc` feature, or +`RUSTFLAGS=--cfg=rustix_use_libc`). No build script can detect that reliably, +so [`lib.rs`](src/lib.rs) checks at compile time instead: it references +`rustix::runtime`, a module that exists only in rustix's raw-syscall build. +Selecting the libc backend makes this crate fail to compile, rather than +silently losing the guarantee. + +## What's inside + +Functions whose rustix implementation already meets the rules are re-exposed +as-is; being listed in a module here is what marks a call as allowed, and the +backend check above is what keeps that true. + +- `mm` — anonymous memory mappings: `mmap_anonymous`, `munmap`. +- `param` — `page_size`. +- `alloc` — allocation without malloc: `alloc::arena()` gives one + intercepted call a bump arena that draws 64 KiB chunks from a process-wide + lock-free pool and returns them when the call ends. Taking or returning a + chunk is one atomic swap; when the pool is empty, chunks come straight + from the kernel through `mm`. diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/sigsafe/src/alloc/mmap.rs similarity index 92% rename from crates/fspy_alloc/src/mmap.rs rename to crates/sigsafe/src/alloc/mmap.rs index 3cc09d876..40791d287 100644 --- a/crates/fspy_alloc/src/mmap.rs +++ b/crates/sigsafe/src/alloc/mmap.rs @@ -6,7 +6,11 @@ use core::{ }; use allocator_api2::alloc::{AllocError, Allocator}; -use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; + +use crate::{ + mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}, + param::page_size, +}; /// A stateless allocator: every allocation is a fresh anonymous mapping and /// every deallocation an `munmap`. @@ -16,10 +20,9 @@ use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; /// It holds no state at all — no locks, no free lists, no thread-locals; /// the kernel does all the bookkeeping. A signal or a `fork()` can never /// catch it holding a lock or a half-written structure, because there is -/// nothing to hold. Mapping syscalls go straight to the kernel on Linux -/// (rustix's raw backend) and through the thin libSystem wrappers on macOS. -/// The page size comes from static process data (`getauxval` on Linux, -/// `sysconf` on macOS) — no locks or allocation there either. +/// nothing to hold. The mapping calls and the page-size read come from +/// [`crate::mm`] and [`crate::param`], which carry the same guarantee (see +/// their docs). /// /// # What it accepts /// @@ -55,12 +58,6 @@ use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; #[derive(Clone, Copy, Debug, Default)] pub struct MmapAllocator; -pub fn page_size() -> usize { - // getauxval on Linux, sysconf on macOS: reads of static process data — - // no locks, no allocation, valid from the first instruction on. - rustix::param::page_size() -} - // SAFETY: returned blocks are non-null, page-aligned (at least // `layout.align()` for every served layout), at least `layout.size()` bytes // large (the returned length reports the exact mapped size), stay valid diff --git a/crates/fspy_alloc/src/lib.rs b/crates/sigsafe/src/alloc/mod.rs similarity index 61% rename from crates/fspy_alloc/src/lib.rs rename to crates/sigsafe/src/alloc/mod.rs index 24b97020e..1a815f36b 100644 --- a/crates/fspy_alloc/src/lib.rs +++ b/crates/sigsafe/src/alloc/mod.rs @@ -1,24 +1,15 @@ -//! Async-signal-safe allocation for the fspy preload library. +//! Allocation that never touches libc malloc. //! -//! The preload library interposes libc functions that POSIX declares -//! async-signal-safe (`open`, `stat`, `execve`, ...). Its code therefore runs -//! inside signal handlers and in the child of `fork()` in a multithreaded -//! process — contexts where taking libc malloc's locks can deadlock on a -//! suspended or vanished lock holder, so the preload's own allocations must -//! never go through libc malloc. -//! -//! This crate stacks three layers, and exposes only the top one, [`arena`]. -//! `MmapAllocator` is the bottom: a stateless allocator where every -//! allocation is a fresh anonymous mapping from the kernel. `ChunkPool` -//! sits on top of it and caches fixed-size chunks, so that frequent short -//! tracing calls can reuse memory instead of paying two syscalls per call. -//! [`arena`] creates one `bump_scope::Bump` per intercepted call, drawing -//! its chunks from the process-wide pool and returning them on drop. - -// Compile as an empty crate on non-unix targets: the allocator backs the -// unix preload library. -#![cfg(unix)] -#![cfg_attr(not(test), no_std)] +//! Taking malloc's lock is the classic way for interposed code to deadlock a +//! traced program (see the crate docs), so the preload library allocates +//! through this module instead. It stacks three layers and exposes only the +//! top one, [`arena`]. `MmapAllocator` is the bottom: a stateless allocator +//! where every allocation is a fresh anonymous mapping from [`crate::mm`]. +//! `ChunkPool` sits on top of it and caches fixed-size chunks, so that +//! frequent short tracing calls can reuse memory instead of paying two +//! syscalls per call. [`arena`] creates one `bump_scope::Bump` per +//! intercepted call, drawing its chunks from the process-wide pool and +//! returning them on drop. mod mmap; mod pool; @@ -34,7 +25,8 @@ use pool::ChunkPool; /// Every cached chunk is 64 KiB: a whole multiple of the page size on all /// supported targets, and big enough that most intercepted calls fit their -/// allocations into a single chunk. +/// allocations into a single chunk. [`ArenaSettings`] pins the arenas' own +/// chunk sizing to this same value. const CHUNK_SIZE: usize = 64 * 1024; /// The alignment chunks are allocated with. Must be at least the alignment /// `bump_scope::Bump` uses for its chunk requests — 16 (see @@ -51,6 +43,22 @@ const SLOTS: usize = 64; /// first allocation on — even before any constructor has run. static CHUNK_POOL: ChunkPool = ChunkPool::new(); +/// The `Bump` settings the arenas use — the defaults, with two changes: +/// +/// - `WithGuaranteedAllocated`: an arena starts life without a chunk, +/// so creating one allocates nothing. +/// - `WithMinimumChunkSize`: the arena's first chunk request is +/// sized to the pool's chunks, making the coupling explicit — rather than +/// relying on the pool rounding the default 512-byte first request up to a +/// whole chunk. (Both end up serving the same memory: the pool answers any +/// request up to `CHUNK_SIZE` with a whole chunk, and `Bump` uses the full +/// returned length.) The request comes out slightly *under* `CHUNK_SIZE` — +/// bump-scope deducts an assumed allocator-header overhead from its +/// minimum — which is exactly what keeps a minimum-sized request within +/// the pool's `size <= CHUNK_SIZE` gate; the +/// `bump_chunk_requests_fit_the_pool_gates` test pins that fit. +type ArenaSettings = <::WithGuaranteedAllocated as BumpAllocatorSettings>::WithMinimumChunkSize; + /// `Bump::unallocated` requires its base allocator to implement `Default` /// (an arena without chunks has nowhere to store an allocator value, so it /// conjures one on first use). Point defaulted references at the @@ -77,12 +85,9 @@ impl Default for &'static ChunkPool impl Allocator { - // The default `Bump` settings, except the arena starts life without a - // chunk, so creating one allocates nothing. - type Settings = ::WithGuaranteedAllocated; Bump::< AllocatorApi2V02Compat<&'static ChunkPool>, - Settings, + ArenaSettings, >::unallocated() } @@ -95,17 +100,28 @@ mod tests { use super::*; /// Pins the fit between `Bump`'s chunk requests and the pool's gates - /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]). - /// bump-scope keeps those request parameters private, so this test is - /// the enforcement: it fails if an upgrade ever changes them. + /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]), + /// under the same [`ArenaSettings`] the arenas use — including that a + /// request under `WithMinimumChunkSize` still fits the + /// `size <= CHUNK_SIZE` gate. bump-scope keeps its request parameters + /// private, so this test is the enforcement: it fails if an upgrade + /// ever changes them. + /// [`ArenaSettings`] with `GuaranteedAllocated` flipped back on: without + /// it, `Bump` demands a `Default` base allocator, and a reference to the + /// test's stack-local pool cannot provide one. Chunk request sizing — + /// what the test pins — is unaffected by that flag. + type TestSettings = ::WithGuaranteedAllocated; + #[test] fn bump_chunk_requests_fit_the_pool_gates() { let pool = ChunkPool::::new_in(Global); - // `try_new_in` requests the first chunk right away; if that request - // were over-aligned or oversized, the pool would refuse and this - // would be an error. - let bump: Bump>> = - Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); + // `try_new_in` requests the first chunk right away, at the settings' + // minimum chunk size; if that request were over-aligned or + // oversized, the pool would refuse and this would be an error. + let bump: Bump< + AllocatorApi2V02Compat<&ChunkPool>, + TestSettings, + > = Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); let block = bump.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); let block_addr = block.cast::().as_ptr().addr(); drop(bump); diff --git a/crates/fspy_alloc/src/pool.rs b/crates/sigsafe/src/alloc/pool.rs similarity index 99% rename from crates/fspy_alloc/src/pool.rs rename to crates/sigsafe/src/alloc/pool.rs index f0e3e4112..3dcf31573 100644 --- a/crates/fspy_alloc/src/pool.rs +++ b/crates/sigsafe/src/alloc/pool.rs @@ -8,7 +8,7 @@ use core::{ use allocator_api2::alloc::{AllocError, Allocator}; -use crate::MmapAllocator; +use super::mmap::MmapAllocator; /// A fixed-size cache of memory chunks. /// diff --git a/crates/sigsafe/src/lib.rs b/crates/sigsafe/src/lib.rs new file mode 100644 index 000000000..9ee868345 --- /dev/null +++ b/crates/sigsafe/src/lib.rs @@ -0,0 +1,40 @@ +//! Unix syscall wrappers that are safe to call where libc is not: in signal +//! handlers, in the child of `fork()` in a multithreaded process, and before +//! libc has finished initializing. +//! +//! The fspy preload library interposes libc functions that POSIX declares +//! async-signal-safe (`open`, `stat`, `execve`, ...), so its code runs in all +//! of those places, where libc's own machinery — locks, lazy initialization, +//! malloc — is off limits. Everything this crate exposes follows three rules: +//! syscalls only (never through libc on Linux), no locks and no hidden state, +//! and no global allocation. See README.md for the full approach. + +// Compile as an empty crate on non-unix targets: the crate backs the unix +// preload library. +#![cfg(unix)] +#![cfg_attr(not(test), no_std)] + +pub mod alloc; +pub mod mm; +pub mod param; + +pub use rustix::io::Errno; + +// Compile-time proof that rustix uses its raw-syscall backend (`linux_raw`) +// on Linux — and with it, that no call in this crate goes through libc there. +// `rustix::runtime` is gated on that backend (`#[cfg(linux_raw)]`), so this +// reference fails to resolve, failing the whole build, whenever anything +// selects the libc backend instead: the `rustix/use-libc` feature (which any +// crate in the dependency graph can enable through feature unification, where +// no build script could ever see it), `RUSTFLAGS=--cfg=rustix_use_libc`, or a +// target rustix has no raw backend for. Miri also forces the libc backend, so +// it is exempted: it type-checks rather than ships code. +// +// The module's contents are not covered by rustix's stability promise, so a +// rustix upgrade may break this line. If that happens, re-point it at any +// other `rustix::runtime` item — do not delete it: it is the only enforcement +// of the no-libc rule above. +#[cfg(all(target_os = "linux", not(miri)))] +const _: () = { + let _ = rustix::runtime::exit_group; +}; diff --git a/crates/sigsafe/src/mm.rs b/crates/sigsafe/src/mm.rs new file mode 100644 index 000000000..c7cdecd05 --- /dev/null +++ b/crates/sigsafe/src/mm.rs @@ -0,0 +1,11 @@ +//! Anonymous memory mappings. +//! +//! Re-exposed from rustix as-is: these are single syscalls against the +//! kernel's own address-space bookkeeping — no libc state, no locks, no +//! allocation — so they already meet this crate's rules everywhere it +//! promises to work. What this module adds is the curation (being listed +//! here is what marks them safe for signal handlers, fork children, and +//! pre-libc startup) and the crate-level backend check, which guarantees +//! they cannot silently turn into libc calls on Linux. + +pub use rustix::mm::{MapFlags, MprotectFlags, ProtFlags, mmap_anonymous, mprotect, munmap}; diff --git a/crates/sigsafe/src/param.rs b/crates/sigsafe/src/param.rs new file mode 100644 index 000000000..0e907689a --- /dev/null +++ b/crates/sigsafe/src/param.rs @@ -0,0 +1,147 @@ +//! Process parameters. +//! +//! On Linux, `page_size` asks the kernel itself — this crate stands in for +//! libc, so it cannot lean on libc's startup knowledge. The kernel hands a +//! process its auxiliary vector exactly once, above the stack pointer at +//! entry, and only the code owning the entry point (a libc, or ld.so) sees +//! it; the retroactive interfaces (`prctl(PR_GET_AUXV)`, `/proc/self/auxv`) +//! need a 6.4 kernel or a mounted `/proc`. So instead of the auxiliary +//! vector, the value is learned from syscall behavior that any kernel +//! version guarantees: see [`linux::page_size`]. +//! +//! On other unix platforms (macOS, where every syscall goes through +//! libSystem by platform contract anyway) it is rustix's `sysconf`, a +//! lock-free read of startup data. + +#[cfg(target_os = "linux")] +pub use linux::page_size; +#[cfg(not(target_os = "linux"))] +pub use rustix::param::page_size; + +#[cfg(target_os = "linux")] +mod linux { + use core::{ + ptr, + sync::atomic::{AtomicUsize, Ordering}, + }; + + use crate::mm::{self, MapFlags, MprotectFlags, ProtFlags}; + + /// Smallest page size of any Linux configuration; the probe starts + /// here. + const MIN_PROBE_PAGE: usize = 4096; + /// Generous upper bound for the probe; no supported configuration uses + /// larger pages. + const MAX_PROBE_PAGE: usize = 1 << 20; + + /// Returns the size of a memory page, or 0 if it cannot be determined. + /// + /// Discovered from the kernel with a handful of raw syscalls, once per + /// process, and cached in an atomic: no libc, no auxiliary vector, no + /// `/proc`, no minimum kernel version, no allocation, and no init to + /// call first — safe from the first instruction on, in signal handlers, + /// and in the child of `fork()`. + /// + /// Zero — the kernel refusing a scratch mapping, i.e. out of memory — + /// is not cached, so a later call retries; callers treat it as "fail + /// this operation" (an allocator refuses the allocation). + #[must_use] + pub fn page_size() -> usize { + static CACHE: AtomicUsize = AtomicUsize::new(0); + let cached = CACHE.load(Ordering::Relaxed); + if cached != 0 { + return cached; + } + let page = probe_page_size(); + if page != 0 { + // Every probe returns the same value; the cache only avoids + // repeated probing. + CACHE.store(page, Ordering::Relaxed); + } + page + } + + /// Discovers the page size by probing, using nothing but raw syscalls. + /// + /// `mprotect` fails with `EINVAL` unless its address is a multiple of + /// the page size, and `mmap` returns page-aligned bases — both + /// kernel-guaranteed on every version. So re-protecting a scratch + /// mapping at increasing power-of-two offsets identifies the page size + /// as the first offset the kernel accepts: for a power-of-two page + /// size P, the first power of two P divides is P itself. + /// + /// Why not `rustix::param::page_size()` (fine on macOS, used there)? + /// Its Linux sources are exactly the ones this crate must not assume: + /// with `use-libc-auxv` it is libc's `sysconf`; without, its lazy init + /// needs `prctl(PR_GET_AUXV)` (kernel 6.4+) or `/proc/self/auxv`, it + /// panics when both are unavailable, and with rustix's `alloc` feature + /// enabled — which Cargo feature unification lets any other rustix + /// user in the build graph turn on for our copy — the `/proc` path + /// heap-allocates, which is forbidden in the contexts this crate + /// serves. The probe has none of those modes: no allocation, no panic, + /// no minimum kernel, and its worst case is 0. + #[cold] + fn probe_page_size() -> usize { + // Large enough that every probe below stays inside the mapping + // even after the kernel rounds the one-byte length up to a full + // page. + const SCRATCH_LEN: usize = 2 * MAX_PROBE_PAGE; + + // SAFETY: a fresh anonymous private mapping at no particular + // address has no memory-safety preconditions. + let Ok(base) = (unsafe { + mm::mmap_anonymous( + ptr::null_mut(), + SCRATCH_LEN, + ProtFlags::READ | ProtFlags::WRITE, + MapFlags::PRIVATE, + ) + }) else { + return 0; + }; + let mut page = 0; + let mut offset = MIN_PROBE_PAGE; + while offset <= MAX_PROBE_PAGE { + // SAFETY: `base + offset` (plus the page-rounded single byte) + // lies within the scratch mapping, which we exclusively own; + // the protection flags match the mapping's existing ones. + let accepted = unsafe { + mm::mprotect( + base.cast::().add(offset).cast(), + 1, + MprotectFlags::READ | MprotectFlags::WRITE, + ) + } + .is_ok(); + if accepted { + page = offset; + break; + } + offset *= 2; + } + // SAFETY: releasing the scratch mapping created above; nothing + // uses it afterwards. + let _ = unsafe { mm::munmap(base, SCRATCH_LEN) }; + page + } + + #[cfg(test)] + mod tests { + use super::*; + + /// Cross-validates the probe against rustix's auxv-based answer + /// (a dev-dependency here; CI Linux runners always have `/proc`). + #[test] + fn probe_matches_rustix() { + assert_eq!(page_size(), rustix::param::page_size()); + #[cfg(target_arch = "x86_64")] + assert_eq!(page_size(), 4096); + } + + #[test] + fn probe_is_stable_and_cached() { + assert_eq!(probe_page_size(), page_size()); + assert_eq!(page_size(), page_size()); + } + } +} From 1f8f298844fef549e82d6983094faafaba0733b6 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 12:49:10 +0800 Subject: [PATCH 4/7] fix(fspy): keep execveat's resolved path alive across the exec, in the per-call arena MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The arena's first user moves from the fd-relative path join to execveat, for two reasons. The join was the wrong showcase: the benchmark suite added for it showed the arena costing ~3% on that lane, because the code it replaced is nearly free — getcwd() already returns a buffer with room to append, so PathBuf::push rarely allocates at all. That code comes back unchanged. execveat is the right one. Its resolved absolute path must be copied into fresh NUL-terminated storage, exec runs in the child of fork() in multithreaded programs (posix_spawn does exactly that) where malloc can deadlock on a lock held by a vanished thread — and the previous code had a use-after-free: it built a CString, moved it into a match arm that extracted the raw pointer, dropped it at the arm's end, and passed the dangling pointer to the real exec. The arena copy is borrowed from a binding that provably outlives the exec call, and dropping CString::new also drops its interior-NUL panic path from an interception. execveat currently has no test coverage — which is how the dangling pointer survived. That gap is left for a follow-up. Co-Authored-By: Claude Fable 5 --- .../fspy_preload_unix/src/client/convert.rs | 22 +++------------- .../src/interceptions/spawn/exec/mod.rs | 26 +++++++++++++++---- 2 files changed, 25 insertions(+), 23 deletions(-) diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index 70fc82a8d..c42854b32 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "linux")] use std::ffi::CString; use std::{ - ffi::CStr, + ffi::{CStr, OsStr}, os::{fd::RawFd, unix::ffi::OsStrExt as _}, path::PathBuf, }; @@ -71,27 +71,13 @@ impl ToAbsolutePath for PathAt { if pathname.first().copied() == Some(b'/') { f(pathname.into()) } else { - let Some(dir) = get_fd_path(self.0)? else { + let Some(mut abs_path) = get_fd_path(self.0)? else { return f(None); }; - // Join `dir` and the relative `pathname` in a per-call bump - // arena instead of `PathBuf::push` on the global allocator. This - // runs on every fd-relative open/stat, including inside signal - // handlers and fork children, where the global allocator's locks - // are unsafe to take. The joined path only lives for the `f` - // call — nothing escapes the arena. - let arena = sigsafe::alloc::arena(); - let mut joined = allocator_api2::vec::Vec::new_in(&arena); - joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { - // Mirror `PathBuf::push`: exactly one separator between the - // (absolute, non-empty) dir and the relative pathname. - if joined.last() != Some(&b'/') { - joined.push(b'/'); - } - joined.extend_from_slice(pathname); + abs_path.push(OsStr::from_bytes(pathname)); } - f(Some(joined.as_bstr())) + f(Some(abs_path.as_os_str().as_bytes().as_bstr())) } } } diff --git a/crates/fspy_preload_unix/src/interceptions/spawn/exec/mod.rs b/crates/fspy_preload_unix/src/interceptions/spawn/exec/mod.rs index 182226eea..78824397d 100644 --- a/crates/fspy_preload_unix/src/interceptions/spawn/exec/mod.rs +++ b/crates/fspy_preload_unix/src/interceptions/spawn/exec/mod.rs @@ -1,8 +1,5 @@ mod with_argv; -#[cfg(target_os = "linux")] -use std::ffi::CString; - use fspy_shared_unix::exec::ExecResolveConfig; use libc::{c_char, c_int}; use with_argv::with_argv; @@ -195,13 +192,27 @@ mod linux_only { reason = "suppresses unused warning on *::original" )] let _unused = execveat::original; + // One bump arena for this intercepted call. The NUL-terminated copy + // of the resolved path made below must not go through libc malloc: + // programs exec from the child of `fork()` in multithreaded + // processes (posix_spawn does exactly that), where malloc's lock may + // be held by a thread that no longer exists. + let arena = sigsafe::alloc::arena(); // SAFETY: PathAt wraps a valid dirfd and pathname pointer from the interposed execveat call let abs_path_result = unsafe { PathAt(dirfd, pathname).to_absolute_path(|path| { let Some(path) = path else { return Ok(None); }; - Ok(Some(CString::new(&**path).unwrap())) + // The resolved path plus a NUL terminator, allocated in the + // arena so it stays valid past this callback. Interior NULs + // cannot occur: the bytes come from NUL-terminated C strings + // and fd symlink targets. + let mut abs_path = + allocator_api2::vec::Vec::with_capacity_in(path.len() + 1, &arena); + abs_path.extend_from_slice(path); + abs_path.push(0); + Ok(Some(abs_path)) }) }; let abs_path = match abs_path_result { @@ -209,7 +220,12 @@ mod linux_only { // SAFETY: forwarding the original arguments to the real execveat syscall return unsafe { execveat::original()(dirfd, pathname, argv, envp, flags) }; } - Ok(Some(path)) => path.as_ptr(), + // Borrowed out of `abs_path_result`, which lives until the end + // of the function — past the `handle_exec` below that reads the + // pointer. (The previous version built a `CString`, moved it + // into this match arm, and dropped it here — `handle_exec` then + // read freed memory.) + Ok(Some(ref path)) => path.as_ptr().cast(), Err(errno) => { errno.set(); return -1; From e7cb59c540e2baf8d89663915c0333a4e4210ee8 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 13:40:24 +0800 Subject: [PATCH 5/7] test(fspy): cover execveat and fexecve interception MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These two exec entry points name their program by file descriptor and had no coverage — the use-after-free fixed in the previous commit lived in execveat's resolved-path handling and was only ever found by reading the code. The interposer reports a program's path while resolving it, before the program is opened, so a missing program still produces the access and no executable needs to be staged. execveat passes a dirfd plus a relative pathname — the resolve-and-copy lane that dangled — and asserts the joined absolute path is captured. fexecve names its program only as /proc/self/fd/N (which the client does not report), so the test points it at a shebang script whose interpreter is reported instead, proving the fd was resolved. Both execs fail as intended and the child exits cleanly. Linux-gnu only: the interfaces do not exist on macOS, and musl uses the seccomp backend with an empty preload. execveat goes through the libc symbol so the interposer runs, rather than nix's raw SYS_execveat. Co-Authored-By: Claude Fable 5 --- crates/fspy/tests/exec_fd.rs | 108 +++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 crates/fspy/tests/exec_fd.rs diff --git a/crates/fspy/tests/exec_fd.rs b/crates/fspy/tests/exec_fd.rs new file mode 100644 index 000000000..9c1a75cfb --- /dev/null +++ b/crates/fspy/tests/exec_fd.rs @@ -0,0 +1,108 @@ +//! Tests for the exec entry points that name their program by file +//! descriptor: `execveat` (directory fd plus relative pathname) and +//! `fexecve` (program fd). Under test are the preload library's interposers +//! for them, which had no coverage at all — a use-after-free in `execveat`'s +//! resolved-path handling was only ever found by reading the code. +//! +//! The interposer reports a program's path while *resolving* it, before the +//! program is opened, so a program that does not exist still produces the +//! path access — no real executable needs to be staged. The exec then fails, +//! the child returns from the closure, and the harness sees a clean exit. +//! +//! Scope: these interfaces do not exist on macOS, and on musl the preload is +//! empty (the seccomp backend traces syscalls instead), so the tests cover +//! the gnu preload lane only. `execveat` must be reached through the libc +//! symbol — a raw `SYS_execveat` syscall (what `nix::unistd::execveat` emits) +//! would bypass the interposer under test. +#![cfg(all(target_os = "linux", not(target_env = "musl")))] + +use std::path::Path; + +use fspy::AccessMode; +use test_log::test; + +use crate::test_utils::assert_contains; + +mod test_utils; + +#[test(tokio::test)] +async fn execveat_resolves_and_captures_the_relative_pathname() -> anyhow::Result<()> { + let tmp_dir = tempfile::tempdir()?; + let dir = std::fs::canonicalize(tmp_dir.path())?; + // The resolved path the interposer should report: the directory fd's path + // joined with the relative pathname. The file need not exist — it is + // reported during resolution, before it is opened. + let expected = dir.join("ghost-program"); + + let accesses = track_fn!(dir.to_str().unwrap().to_owned(), |dir: String| { + use std::os::fd::AsRawFd as _; + + let dirfd = nix::fcntl::open( + dir.as_str(), + nix::fcntl::OFlag::O_RDONLY | nix::fcntl::OFlag::O_DIRECTORY, + nix::sys::stat::Mode::empty(), + ) + .expect("failed to open the directory fd"); + + // dirfd + relative pathname: the lane where the interposer resolves + // the directory and copies the joined path — the copy that used to + // dangle. Through the libc symbol so the interposer runs; expected to + // fail with ENOENT, which is fine — the access is already reported. + let args: [*mut libc::c_char; 2] = + [c"ghost-program".as_ptr().cast_mut(), core::ptr::null_mut()]; + let env: [*mut libc::c_char; 1] = [core::ptr::null_mut()]; + // SAFETY: dirfd is a valid directory fd; the argument and environment + // arrays are NULL-terminated arrays of valid NUL-terminated strings, + // alive across the call. execveat does not mutate them. + let ret = unsafe { + libc::execveat( + dirfd.as_raw_fd(), + c"ghost-program".as_ptr(), + args.as_ptr(), + env.as_ptr(), + 0, + ) + }; + // execveat only returns on failure; the child exits cleanly here so + // the harness's success assertion holds. + assert_eq!(ret, -1, "execveat unexpectedly succeeded on a missing program"); + }) + .await?; + + assert_contains(&accesses, &expected, AccessMode::READ); + Ok(()) +} + +#[test(tokio::test)] +async fn fexecve_resolves_through_the_program_fd() -> anyhow::Result<()> { + let tmp_dir = tempfile::tempdir()?; + let dir = std::fs::canonicalize(tmp_dir.path())?; + // A shebang script whose interpreter does not exist. `fexecve` names the + // program only as `/proc/self/fd/N` (which the client deliberately does + // not report), but resolving the script reads its `#!` line and reports + // the interpreter path — a non-`/proc` path that is reported and is the + // observable proof the fexecve interposer ran and resolved the fd. + let interpreter = Path::new("/fspy-exec-fd-missing-interpreter"); + let script = dir.join("script"); + std::fs::write(&script, format!("#!{}\n", interpreter.display()))?; + + let accesses = track_fn!(script.to_str().unwrap().to_owned(), |script: String| { + let program_fd = nix::fcntl::open( + script.as_str(), + nix::fcntl::OFlag::O_RDONLY, + nix::sys::stat::Mode::empty(), + ) + .expect("failed to open the script fd"); + + // nix's fexecve goes through the libc symbol, so the interposer runs. + // It fails because the interpreter is missing; the child exits + // cleanly afterwards. + let err = nix::unistd::fexecve(program_fd, &[c"script"], &[c"X=1"]) + .expect_err("fexecve unexpectedly succeeded"); + assert_eq!(err, nix::Error::ENOENT, "unexpected fexecve error: {err}"); + }) + .await?; + + assert_contains(&accesses, interpreter, AccessMode::READ); + Ok(()) +} From 8ba13559689b264d8dee248c9857403496598116 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 13:57:00 +0800 Subject: [PATCH 6/7] fix(fspy): resolve execveat via dlsym so the exec_fd test links on glibc 2.17 CI links tests against a glibc 2.17 baseline, which predates glibc's execveat wrapper (added in 2.34), so a direct libc::execveat reference failed with 'undefined symbol: execveat'. Resolve the symbol with dlsym(RTLD_DEFAULT) instead: it links against ancient glibc and, under LD_PRELOAD, still resolves to the preload's interposer exactly as a PLT call would. fexecve is unaffected (glibc 2.3.2). Co-Authored-By: Claude Fable 5 --- crates/fspy/tests/exec_fd.rs | 54 +++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 16 deletions(-) diff --git a/crates/fspy/tests/exec_fd.rs b/crates/fspy/tests/exec_fd.rs index 9c1a75cfb..17386baf0 100644 --- a/crates/fspy/tests/exec_fd.rs +++ b/crates/fspy/tests/exec_fd.rs @@ -12,11 +12,16 @@ //! Scope: these interfaces do not exist on macOS, and on musl the preload is //! empty (the seccomp backend traces syscalls instead), so the tests cover //! the gnu preload lane only. `execveat` must be reached through the libc -//! symbol — a raw `SYS_execveat` syscall (what `nix::unistd::execveat` emits) -//! would bypass the interposer under test. +//! *symbol* — a raw `SYS_execveat` syscall (what `nix::unistd::execveat` +//! emits) would bypass the interposer under test — but the symbol is resolved +//! with `dlsym` rather than called directly: the CI links against a glibc +//! 2.17 baseline, which predates glibc's `execveat` wrapper (added in 2.34), +//! so a direct `libc::execveat` reference fails to link. Under `LD_PRELOAD`, +//! `dlsym(RTLD_DEFAULT, "execveat")` resolves to the interposer, exactly as a +//! PLT call would. #![cfg(all(target_os = "linux", not(target_env = "musl")))] -use std::path::Path; +use std::{ffi::c_char, path::Path}; use fspy::AccessMode; use test_log::test; @@ -25,6 +30,28 @@ use crate::test_utils::assert_contains; mod test_utils; +/// The glibc `execveat` prototype: `int execveat(int, const char *, char *const +/// [], char *const [], int)`. +type ExecveatFn = unsafe extern "C" fn( + libc::c_int, + *const c_char, + *const *const c_char, + *const *const c_char, + libc::c_int, +) -> libc::c_int; + +/// Resolves `execveat` from the global symbol scope. Under `LD_PRELOAD` this +/// is the preload's interposer; resolving it dynamically avoids a link-time +/// dependency on a glibc new enough to export the wrapper. +fn resolve_execveat() -> ExecveatFn { + // SAFETY: `dlsym` with a valid NUL-terminated symbol name; the resolved + // pointer has the C prototype transmuted onto it below. + let sym = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c"execveat".as_ptr()) }; + assert!(!sym.is_null(), "execveat symbol not present in the global scope"); + // SAFETY: `sym` points at the C `execveat`, whose ABI matches `ExecveatFn`. + unsafe { std::mem::transmute::<*mut libc::c_void, ExecveatFn>(sym) } +} + #[test(tokio::test)] async fn execveat_resolves_and_captures_the_relative_pathname() -> anyhow::Result<()> { let tmp_dir = tempfile::tempdir()?; @@ -46,22 +73,17 @@ async fn execveat_resolves_and_captures_the_relative_pathname() -> anyhow::Resul // dirfd + relative pathname: the lane where the interposer resolves // the directory and copies the joined path — the copy that used to - // dangle. Through the libc symbol so the interposer runs; expected to - // fail with ENOENT, which is fine — the access is already reported. - let args: [*mut libc::c_char; 2] = - [c"ghost-program".as_ptr().cast_mut(), core::ptr::null_mut()]; - let env: [*mut libc::c_char; 1] = [core::ptr::null_mut()]; + // dangle. Through the (dlsym-resolved) libc symbol so the interposer + // runs; expected to fail with ENOENT, which is fine — the access is + // already reported by then. + let execveat = resolve_execveat(); + let args = [c"ghost-program".as_ptr(), core::ptr::null()]; + let env = [core::ptr::null()]; // SAFETY: dirfd is a valid directory fd; the argument and environment // arrays are NULL-terminated arrays of valid NUL-terminated strings, - // alive across the call. execveat does not mutate them. + // alive across the call. let ret = unsafe { - libc::execveat( - dirfd.as_raw_fd(), - c"ghost-program".as_ptr(), - args.as_ptr(), - env.as_ptr(), - 0, - ) + execveat(dirfd.as_raw_fd(), c"ghost-program".as_ptr(), args.as_ptr(), env.as_ptr(), 0) }; // execveat only returns on failure; the child exits cleanly here so // the harness's success assertion holds. From 0942c2faa021b7e8e9a8306d75ba41fc7946517d Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 14:04:00 +0800 Subject: [PATCH 7/7] test(fspy): assert fexecve carries tracking into the exec'd program MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first version pointed fexecve at a shebang script with a missing interpreter and expected ENOENT, but got EACCES: the script was written 0644, so the kernel refused the exec before the interpreter mattered. Chmod would have fixed the errno, but the assertion was weak anyway — it depended on shebang-resolution internals to surface a reportable path, since fexecve names its program only as /proc/self/fd/N, which the client deliberately drops. Assert the interposer's actual job instead: exec a real /bin/sh through a program fd, with an environment that carries none of fspy's state, and check the shell's own file read is still captured. That only holds if the interposer ran and re-injected the tracking environment, and it exercises the exec success path rather than a failure path. Co-Authored-By: Claude Fable 5 --- crates/fspy/tests/exec_fd.rs | 66 +++++++++++++++++++----------------- 1 file changed, 35 insertions(+), 31 deletions(-) diff --git a/crates/fspy/tests/exec_fd.rs b/crates/fspy/tests/exec_fd.rs index 17386baf0..812305625 100644 --- a/crates/fspy/tests/exec_fd.rs +++ b/crates/fspy/tests/exec_fd.rs @@ -4,10 +4,15 @@ //! for them, which had no coverage at all — a use-after-free in `execveat`'s //! resolved-path handling was only ever found by reading the code. //! -//! The interposer reports a program's path while *resolving* it, before the -//! program is opened, so a program that does not exist still produces the -//! path access — no real executable needs to be staged. The exec then fails, -//! the child returns from the closure, and the harness sees a clean exit. +//! The two are checked differently because they report differently. +//! `execveat` resolves its dirfd and relative pathname into an absolute path +//! and reports that, before opening it — so a program that does not exist +//! still produces the access, and nothing needs to be staged. `fexecve` +//! names its program only as `/proc/self/fd/N`, which the client +//! deliberately drops, so there is no path of its own to assert on; instead +//! it execs a real shell with an environment stripped of everything fspy +//! needs, and the test asserts the shell is still traced — which only holds +//! if the interposer ran and re-injected the tracking environment. //! //! Scope: these interfaces do not exist on macOS, and on musl the preload is //! empty (the seccomp backend traces syscalls instead), so the tests cover @@ -21,7 +26,7 @@ //! PLT call would. #![cfg(all(target_os = "linux", not(target_env = "musl")))] -use std::{ffi::c_char, path::Path}; +use std::ffi::c_char; use fspy::AccessMode; use test_log::test; @@ -96,35 +101,34 @@ async fn execveat_resolves_and_captures_the_relative_pathname() -> anyhow::Resul } #[test(tokio::test)] -async fn fexecve_resolves_through_the_program_fd() -> anyhow::Result<()> { +async fn fexecve_carries_tracking_into_the_new_program() -> anyhow::Result<()> { let tmp_dir = tempfile::tempdir()?; - let dir = std::fs::canonicalize(tmp_dir.path())?; - // A shebang script whose interpreter does not exist. `fexecve` names the - // program only as `/proc/self/fd/N` (which the client deliberately does - // not report), but resolving the script reads its `#!` line and reports - // the interpreter path — a non-`/proc` path that is reported and is the - // observable proof the fexecve interposer ran and resolved the fd. - let interpreter = Path::new("/fspy-exec-fd-missing-interpreter"); - let script = dir.join("script"); - std::fs::write(&script, format!("#!{}\n", interpreter.display()))?; - - let accesses = track_fn!(script.to_str().unwrap().to_owned(), |script: String| { - let program_fd = nix::fcntl::open( - script.as_str(), - nix::fcntl::OFlag::O_RDONLY, - nix::sys::stat::Mode::empty(), - ) - .expect("failed to open the script fd"); - - // nix's fexecve goes through the libc symbol, so the interposer runs. - // It fails because the interpreter is missing; the child exits - // cleanly afterwards. - let err = nix::unistd::fexecve(program_fd, &[c"script"], &[c"X=1"]) - .expect_err("fexecve unexpectedly succeeded"); - assert_eq!(err, nix::Error::ENOENT, "unexpected fexecve error: {err}"); + let marker = std::fs::canonicalize(tmp_dir.path())?.join("marker"); + std::fs::write(&marker, "read by the exec'd shell\n")?; + + let accesses = track_fn!(marker.to_str().unwrap().to_owned(), |marker: String| { + use std::ffi::CString; + + let program_fd = + nix::fcntl::open("/bin/sh", nix::fcntl::OFlag::O_RDONLY, nix::sys::stat::Mode::empty()) + .expect("failed to open /bin/sh"); + + // `:` with an input redirect opens the marker and runs no subprocess. + let script = CString::new(format!(": < {marker}")).unwrap(); + // The environment passed here holds nothing of fspy's: no LD_PRELOAD, + // no payload. Re-injecting them is precisely the interposer's job, so + // if it does not run, the exec'd shell is untraced and the marker + // access below never arrives. nix's fexecve goes through the libc + // symbol, so the interposer does run. + let err = + nix::unistd::fexecve(program_fd, &[c"sh", c"-c", script.as_c_str()], &[c"UNRELATED=1"]) + .expect_err("fexecve returned instead of replacing the process"); + panic!("fexecve failed: {err}"); }) .await?; - assert_contains(&accesses, interpreter, AccessMode::READ); + // Reported by the shell that `fexecve` started — a different program + // image than the one the tracker was injected into. + assert_contains(&accesses, &marker, AccessMode::READ); Ok(()) }