vdso_rng/
pool.rs

1use core::{ffi::c_void, ptr::NonNull};
2
3use alloc::vec::Vec;
4use crossbeam_queue::SegQueue;
5use lamlock::Lock;
6
7use crate::{config::Config, utils};
8
9#[repr(transparent)]
10#[derive(Debug, Clone, Copy)]
11pub struct Ptr(pub(crate) NonNull<c_void>);
12
13unsafe impl Send for Ptr {}
14
15/// A shared state block pool for `vDSO`-based `getrandom` operations.
16///
17/// This pool maintains a free list of opaque state blocks. Threads can rent a block
18/// from the pool and use it to generate random bytes. The pool is [`Sync`] and supports
19/// concurrent access. Under debug assertions, [`Pool::drop`] detects potential misuse,
20/// such as double drops.
21///
22/// The pool itself offers only [`Pool::new`]. To interact with the pool, see [`crate::LocalState`],
23/// which wraps individual state blocks for random generation.
24///
25/// ```rust
26/// use vdso_rng::Pool;
27/// let _pool = Pool::new().unwrap();
28/// ```
29///
30/// ### Memory Behavior
31/// The pool grows monotonically with system parallelism. Opaque state blocks are stored
32/// in memory-mapped pages that are not backed by swap. When the system is under memory pressure,
33/// the OS may reclaim these pages, which is generally safe.
34///
35/// ## Safety
36/// - **Not async-signal-safe**: Using the pool in signal handlers may cause deadlocks.
37/// - **Fork safety**: After `fork`, the kernel wipes the random states to avoid leaks.
38///   However, we do **not** guarantee correctness of pool usage across forks.
39pub struct Pool {
40    pub(crate) config: Config,
41    mmaps: Lock<Vec<Ptr>>,
42    freelist: SegQueue<Ptr>,
43}
44
45impl Pool {
46    pub fn new() -> Result<Self, crate::Error> {
47        let (function, page_size) =
48            crate::vdso::get_function_and_page_size().ok_or(crate::Error::NotSupported)?;
49        let config = unsafe { Config::new(function, page_size) };
50        let mmaps = Lock::new(Vec::new());
51        let freelist = SegQueue::new();
52        Ok(Self {
53            config,
54            mmaps,
55            freelist,
56        })
57    }
58    fn grow(
59        mmaps: &mut Vec<Ptr>,
60        config: &Config,
61        freelist: &SegQueue<Ptr>,
62    ) -> Result<(), crate::Error> {
63        let page = utils::mmap(
64            config.page_size * config.pages_per_block,
65            config.params.mmap_prot,
66            config.params.mmap_flags,
67        )
68        .ok_or(crate::Error::AllocationFailure)?;
69        mmaps.push(Ptr(page));
70        unsafe {
71            for p in 0..config.pages_per_block {
72                let page_ptr = page.byte_add(p * config.page_size);
73                for s in 0..config.states_per_page {
74                    let state_ptr =
75                        page_ptr.byte_add(s * config.params.size_of_opaque_states as usize);
76                    freelist.push(Ptr(state_ptr));
77                }
78            }
79        }
80        Ok(())
81    }
82    pub(crate) fn get(&self) -> Result<Ptr, crate::Error> {
83        if let Some(ptr) = self.freelist.pop() {
84            return Ok(ptr);
85        }
86        self.mmaps
87            .run(|mmaps| {
88                // Since the mmaps is locked, this loop should terminates in finite amount of time.
89                loop {
90                    match self.freelist.pop() {
91                        Some(ptr) => return Ok(ptr),
92                        None => {
93                            Self::grow(mmaps, &self.config, &self.freelist)?;
94                            continue;
95                        }
96                    }
97                }
98            })
99            .unwrap_or(Err(crate::Error::PoolPoisoned))
100    }
101    pub(crate) fn recycle(&self, ptr: Ptr) {
102        self.freelist.push(ptr);
103    }
104}
105
106impl Drop for Pool {
107    fn drop(&mut self) {
108        _ = self.mmaps.poison();
109        #[cfg(debug_assertions)]
110        let mut counter = 0;
111        while self.freelist.pop().is_some() {
112            #[cfg(debug_assertions)]
113            {
114                counter += 1;
115            }
116        }
117        _ = self.mmaps.inspect_poison(|mmaps| {
118            #[cfg(debug_assertions)]
119            debug_assert_eq!(
120                counter,
121                self.config.pages_per_block * self.config.states_per_page * mmaps.len(),
122                "Freelist should contain all states from all mmaps"
123            );
124            for ptr in mmaps.drain(..) {
125                unsafe {
126                    utils::munmap(ptr.0, self.config.page_size * self.config.pages_per_block)
127                };
128            }
129            core::ops::ControlFlow::Continue(())
130        });
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    extern crate std;
137    use super::*;
138
139    #[test]
140    fn pool_smoke_test() {
141        let pool = Pool::new().expect("Failed to create pool");
142        let ptr = pool.get().expect("Failed to get pointer from pool");
143        pool.recycle(ptr);
144    }
145
146    #[test]
147    fn pool_multi_thread_test() {
148        let parallelism = std::thread::available_parallelism().unwrap();
149        let pool = Pool::new().expect("Failed to create pool with VDSO function and page size");
150        std::thread::scope(|scope| {
151            for _ in 0..parallelism.get() {
152                scope.spawn(|| {
153                    let ptrs = (0..16)
154                        .map(|_| pool.get().expect("Failed to get pointer from pool"))
155                        .collect::<Vec<_>>();
156                    for ptr in ptrs {
157                        pool.recycle(ptr);
158                    }
159                });
160            }
161        });
162    }
163}
vdso_rng/pool.rs

vdso_rng/
pool.rs
