Skip to main content

hardware_enclave/memory/
pool.rs

1// Copyright 2026 Jay Gowdy
2// SPDX-License-Identifier: MIT
3
4//! Tiered pool of mlock'd slab slots and standalone guard-paged buffers.
5//!
6//! # Initialization
7//! The global memory pool is lazily initialized on first use. For reliable startup-time
8//! error reporting, call [`init_pool()`] explicitly before using any [`MemoryEnclave`] or
9//! [`pool_acquire()`] operations.
10
11#![allow(unsafe_code)]
12
13use std::sync::{Condvar, Mutex, OnceLock};
14
15use super::memcall::page_size;
16use super::secure_buffer::SecureBuffer;
17use super::slab::{SecureSlab, DEFAULT_SLOT_SIZE, SLOT_WAIT_TIMEOUT};
18use crate::error::{Error, Result};
19
20// ── Pool slot origin ────────────────────────────────────────────────
21
22enum PoolSlotOrigin {
23    /// Slot lives in a tier's slab. `tier_index` identifies which `Tier` to return it to.
24    Slab {
25        tier_index: usize,
26        slot_index: usize,
27    },
28    /// Slot owns a standalone guard-paged buffer (no tier fits, or tier exhausted).
29    Standalone(SecureBuffer),
30}
31
32/// A handle to a locked memory region containing secret data.
33///
34/// The slot is either backed by a tier slab (mlock'd single-page pool) or a
35/// standalone `SecureBuffer` (guard pages + mlock, for larger allocations).
36///
37/// `PoolSlot` is `Send` but NOT `Sync`; exclusive reference semantics prevent
38/// concurrent mutation.
39///
40/// # Safety of the slab pointer
41///
42/// When origin is `Slab`, `ptr` points into a `TieredPool` tier's `SecureSlab`
43/// which lives in a `OnceLock<TieredPool>` and is never dropped for the
44/// process lifetime. The pointer therefore cannot dangle as long as the
45/// process is alive.
46///
47/// # Safety
48/// `PoolSlot` must not outlive the global pool. Only acquire via the module-level
49/// `pool_acquire()` and `coffer_view()` functions, not via a local `TieredPool` instance.
50/// The `TieredPool::acquire()` method is intentionally `pub(crate)` for this reason.
51pub struct PoolSlot {
52    ptr: *mut u8,
53    len: usize,
54    origin: PoolSlotOrigin,
55}
56
57// SAFETY: PoolSlot owns either:
58// (a) a slab slot index — the underlying memory is the global slab (OnceLock, process-lifetime).
59//     The raw pointer is valid for the process lifetime. No thread-local state.
60// (b) a standalone SecureBuffer — which is itself Send.
61// PoolSlot is NOT Sync because concurrent mutable access to the same slot is not protected.
62unsafe impl Send for PoolSlot {}
63
64impl std::fmt::Debug for PoolSlot {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        f.debug_struct("PoolSlot").field("len", &self.len).finish()
67    }
68}
69
70impl PoolSlot {
71    pub(crate) fn from_slab(
72        ptr: *mut u8,
73        len: usize,
74        tier_index: usize,
75        slot_index: usize,
76    ) -> Self {
77        Self {
78            ptr,
79            len,
80            origin: PoolSlotOrigin::Slab {
81                tier_index,
82                slot_index,
83            },
84        }
85    }
86
87    fn from_standalone(mut buf: SecureBuffer) -> Self {
88        // buf starts Mutable (freshly allocated); melt is a no-op if already mutable.
89        drop(buf.melt());
90        let ptr = buf.bytes().as_mut_ptr();
91        let len = buf.size();
92        Self {
93            ptr,
94            len,
95            origin: PoolSlotOrigin::Standalone(buf),
96        }
97    }
98
99    /// Mutable access to the slot's bytes.
100    pub fn bytes(&mut self) -> &mut [u8] {
101        // SAFETY: ptr is valid for len bytes (either in the global OnceLock slab, which is
102        // process-lifetime, or in the standalone SecureBuffer owned by this PoolSlot).
103        // &mut self guarantees exclusive access — PoolSlot is not Sync.
104        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
105    }
106
107    /// Read-only access to the slot's bytes.
108    pub fn as_slice(&self) -> &[u8] {
109        // SAFETY: ptr is valid for len bytes; no aliased mutable reference exists
110        // because PoolSlot is not Sync and we hold a shared reference.
111        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
112    }
113
114    /// Total capacity of this slot in bytes.
115    pub fn size(&self) -> usize {
116        self.len
117    }
118
119    /// Returns the slab slot index if this slot is backed by the global slab.
120    #[allow(dead_code)]
121    pub(crate) fn slab_index(&self) -> Option<usize> {
122        match &self.origin {
123            PoolSlotOrigin::Slab { slot_index, .. } => Some(*slot_index),
124            PoolSlotOrigin::Standalone(_) => None,
125        }
126    }
127
128    /// Returns the tier index if this slot is backed by the global slab.
129    #[allow(dead_code)]
130    pub(crate) fn tier_index(&self) -> Option<usize> {
131        match &self.origin {
132            PoolSlotOrigin::Slab { tier_index, .. } => Some(*tier_index),
133            PoolSlotOrigin::Standalone(_) => None,
134        }
135    }
136}
137
138impl Drop for PoolSlot {
139    fn drop(&mut self) {
140        match &mut self.origin {
141            PoolSlotOrigin::Slab {
142                tier_index,
143                slot_index,
144            } => {
145                // Zeroize before acquiring the lock so memory is clean even under contention.
146                // SAFETY: ptr points into the global TieredPool's slab page (OnceLock,
147                // process-lifetime). The PoolSlot doc requires slots to come from global_pool(),
148                // so this pointer is always valid for the process lifetime.
149                unsafe {
150                    use zeroize::Zeroize;
151                    std::slice::from_raw_parts_mut(self.ptr, self.len).zeroize();
152                }
153                let pool = global_pool();
154                if let Ok(mut slab) = pool.tiers[*tier_index].slab.lock() {
155                    slab.release(*slot_index);
156                }
157                // Wake any waiter blocked in `acquire`.
158                pool.tiers[*tier_index].cv.notify_one();
159            }
160            PoolSlotOrigin::Standalone(buf) => {
161                // Zeroize before buf's own drop, which also zeroizes.
162                drop(buf.melt());
163                // SAFETY: ptr points into buf's inner region (the SecureBuffer owned by this
164                // PoolSlot). buf is still alive at this point — we haven't dropped it yet.
165                unsafe {
166                    use zeroize::Zeroize;
167                    std::slice::from_raw_parts_mut(self.ptr, self.len).zeroize();
168                }
169                // buf drops here: zeroizes again + unmaps.
170            }
171        }
172    }
173}
174
175// ── Tiered pool ───────────────────────────────────────────────────
176
177/// One tier: a single mlock'd-page slab with a fixed slot size.
178struct Tier {
179    slot_size: usize,
180    slab: Mutex<SecureSlab>,
181    /// Notified on every slot release so waiting `acquire` calls can retry.
182    cv: Condvar,
183}
184
185impl std::fmt::Debug for Tier {
186    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
187        f.debug_struct("Tier")
188            .field("slot_size", &self.slot_size)
189            .finish()
190    }
191}
192
193/// Configuration for the tiered pool.
194///
195/// `tier_sizes` lists the slot sizes for each tier. Must be non-empty.
196/// Duplicate sizes are deduplicated; the list is sorted ascending internally.
197/// Each size must satisfy: `1 <= size <= page_size() / 3`.
198#[derive(Debug, Clone)]
199pub struct TieredPoolConfig {
200    /// Slot sizes for each tier, non-empty, each in `1..=page_size()/3`.
201    pub tier_sizes: Vec<usize>,
202}
203
204impl Default for TieredPoolConfig {
205    /// Single 32-byte tier (asherah-ffi compatible).
206    fn default() -> Self {
207        Self {
208            tier_sizes: vec![DEFAULT_SLOT_SIZE],
209        }
210    }
211}
212
213/// Statically-owned tiered pool.
214///
215/// Tier 0's slab is initialised with the Coffer (master key halves in slots 0+1).
216/// All other tiers have `init_coffer = false`.
217///
218/// `TieredPool` is `Send + Sync`: all mutable state is behind `Mutex`-guarded
219/// fields. The `Vec<Tier>` contains raw pointers inside `SecureSlab`, which
220/// requires explicit `Send`/`Sync` impls.
221#[derive(Debug)]
222pub struct TieredPool {
223    tiers: Vec<Tier>,
224}
225
226// SAFETY: TieredPool contains Vec<Tier> where Tier holds Mutex<SecureSlab> (the slab's
227// NonNull<u8> is Send but not Sync by default). All mutable access to each slab is
228// serialized by its per-tier Mutex. The Condvar is Sync. There is no unguarded shared
229// mutable state.
230unsafe impl Send for TieredPool {}
231unsafe impl Sync for TieredPool {}
232
233impl TieredPool {
234    /// Create a new `TieredPool` from `config`.
235    ///
236    /// Validates, deduplicates, and sorts `config.tier_sizes`, then allocates
237    /// one mlock'd page per tier. Tier 0's slab initialises the Coffer.
238    pub fn new(config: TieredPoolConfig) -> Result<Self> {
239        // Harden the process as early as possible (idempotent).
240        // Skipped in test builds: mitigations like ProcessStrictHandleCheckPolicy
241        // and ProcessExtensionPointDisablePolicy can interfere with the test runner
242        // and spawned threads, causing hangs on Windows CI.
243        #[cfg(not(test))]
244        crate::harden_process();
245
246        let ps = page_size();
247        let max_slot = ps / 3;
248
249        if config.tier_sizes.is_empty() {
250            return Err(Error::Memory(
251                "TieredPoolConfig: tier_sizes must be non-empty".into(),
252            ));
253        }
254
255        // Sort and dedup.
256        let mut sizes = config.tier_sizes;
257        sizes.sort_unstable();
258        sizes.dedup();
259
260        // Validate each size.
261        for &sz in &sizes {
262            if sz == 0 {
263                return Err(Error::Memory(
264                    "TieredPoolConfig: tier size 0 is invalid".into(),
265                ));
266            }
267            if sz > max_slot {
268                return Err(Error::Memory(format!(
269                    "TieredPoolConfig: tier size {sz} exceeds page_size/3 ({max_slot})"
270                )));
271            }
272        }
273
274        // Tier 0's slab initialises the Coffer (AES-256 key). Its slot_size must be >= 32.
275        if sizes[0] < 32 {
276            return Err(Error::Memory(format!(
277                "TieredPool: first tier slot_size must be >= 32 for coffer, got {}",
278                sizes[0]
279            )));
280        }
281
282        // Allocate one SecureSlab per tier.
283        // Tier 0 initialises the Coffer; all others do not.
284        let mut tiers = Vec::with_capacity(sizes.len());
285        for (i, sz) in sizes.into_iter().enumerate() {
286            let init_coffer = i == 0;
287            let slab = SecureSlab::new(sz, init_coffer)?;
288            tiers.push(Tier {
289                slot_size: sz,
290                slab: Mutex::new(slab),
291                cv: Condvar::new(),
292            });
293        }
294
295        Ok(Self { tiers })
296    }
297
298    /// Index of the smallest tier whose slot_size >= `size`, or `None` if all tiers are too small.
299    fn tier_for_size(&self, size: usize) -> Option<usize> {
300        self.tiers.iter().position(|t| t.slot_size >= size)
301    }
302
303    /// Acquire a slot from the appropriate tier.
304    ///
305    /// Routes to the smallest tier whose slot_size >= `size`. Waits up to
306    /// `SLOT_WAIT_TIMEOUT` for a free slot using a `Condvar` (no spin/sleep);
307    /// falls back to a standalone `SecureBuffer` if exhausted or no tier fits.
308    ///
309    /// # Safety note
310    /// The returned `PoolSlot` contains a raw pointer into this pool's slab. It is
311    /// only safe to use when `self` is the global pool (i.e. called via `pool_acquire()`).
312    /// Do not call this on a local `TieredPool` instance and let the `PoolSlot` outlive it.
313    pub(crate) fn acquire(&self, size: usize) -> Result<PoolSlot> {
314        if let Some(tier_idx) = self.tier_for_size(size) {
315            let deadline = std::time::Instant::now() + SLOT_WAIT_TIMEOUT;
316            let mut guard = self.tiers[tier_idx]
317                .slab
318                .lock()
319                .unwrap_or_else(|e| e.into_inner());
320            loop {
321                if let Some(slot_idx) = guard.acquire_transient() {
322                    let (ptr, len) = guard
323                        .slot_raw(slot_idx)
324                        .expect("slot_raw: index validated by acquire_transient");
325                    drop(guard);
326                    return Ok(PoolSlot::from_slab(ptr, len, tier_idx, slot_idx));
327                }
328                let timeout = deadline.saturating_duration_since(std::time::Instant::now());
329                if timeout.is_zero() {
330                    tracing::warn!(
331                        size,
332                        tier_idx,
333                        "pool acquire: all slab slots exhausted; using standalone SecureBuffer"
334                    );
335                    drop(guard);
336                    break;
337                }
338                let result = self.tiers[tier_idx]
339                    .cv
340                    .wait_timeout(guard, timeout)
341                    .unwrap_or_else(|e| e.into_inner());
342                guard = result.0;
343            }
344        }
345
346        // Standalone fallback (no tier fits, or tier exhausted).
347        Ok(PoolSlot::from_standalone(SecureBuffer::new(size)?))
348    }
349
350    /// Reconstruct the Coffer master key from tier 0's slab into a `PoolSlot`.
351    ///
352    /// # Safety note
353    /// Same lifetime constraint as `acquire()`: only safe when `self` is the global pool.
354    /// Use the module-level `coffer_view()` function instead of calling this directly.
355    pub(crate) fn coffer_view(&self) -> Result<PoolSlot> {
356        let mut guard = self.tiers[0].slab.lock().unwrap_or_else(|e| e.into_inner());
357        let slot_idx = guard
358            .coffer_view()
359            .ok_or_else(|| Error::Memory("coffer_view: no free slab slot".into()))?;
360        let (ptr, len) = guard
361            .slot_raw(slot_idx)
362            .expect("slot_raw: index validated by coffer_view");
363        drop(guard);
364        Ok(PoolSlot::from_slab(ptr, len, 0, slot_idx))
365    }
366
367    /// The largest slot size available in any tier.
368    pub fn max_slab_slot_size(&self) -> usize {
369        self.tiers.iter().map(|t| t.slot_size).max().unwrap_or(0)
370    }
371
372    /// Number of configured tiers.
373    pub fn tier_count(&self) -> usize {
374        self.tiers.len()
375    }
376
377    /// Slot size for tier `i`, or `None` if out of range.
378    pub fn tier_slot_size(&self, i: usize) -> Option<usize> {
379        self.tiers.get(i).map(|t| t.slot_size)
380    }
381}
382
383// ── Global pool ───────────────────────────────────────────────────
384
385static POOL: OnceLock<TieredPool> = OnceLock::new();
386
387/// Initialize the global pool with a custom config.
388///
389/// Must be called before any pool operation. If the pool is already initialized
390/// (via a prior call or via lazy default init), returns
391/// `Error::Memory("pool already initialized")`.
392pub fn init_pool(config: TieredPoolConfig) -> Result<()> {
393    let pool = TieredPool::new(config)?;
394    POOL.set(pool)
395        .map_err(|_| Error::Memory("pool already initialized".into()))
396}
397
398pub(crate) fn global_pool() -> &'static TieredPool {
399    POOL.get_or_init(|| {
400        TieredPool::new(TieredPoolConfig::default())
401            .expect("enclave: default tiered pool init failed — OsRng unavailable")
402    })
403}
404
405/// Acquire a pool slot for `size` bytes.
406///
407/// Routes to the smallest tier whose slot_size >= `size`. Waits up to 30 s
408/// for a free slot, then falls back to a standalone guard-paged `SecureBuffer`.
409pub fn pool_acquire(size: usize) -> Result<PoolSlot> {
410    global_pool().acquire(size)
411}
412
413/// Release a pool slot. The slot's contents are zeroized.
414/// Prefer dropping the `PoolSlot` directly; this is provided for explicit release.
415pub fn pool_release(slot: PoolSlot) {
416    drop(slot);
417}
418
419/// Get a `PoolSlot` containing the Coffer master key.
420/// Release promptly after use; the slot is from the pool and blocks that slot while held.
421pub fn coffer_view() -> Result<PoolSlot> {
422    global_pool().coffer_view()
423}
424
425// ── Slab-delegated hot cache helpers ─────────────────────────────
426
427/// Insert plaintext into the slab hot cache for `id`.
428/// Only caches if `data.len() == tier-0 slot_size` (exact fit).
429pub(super) fn hot_cache_insert(id: u64, data: &[u8]) {
430    let pool = global_pool();
431    let mut slab = pool.tiers[0].slab.lock().unwrap_or_else(|e| e.into_inner());
432    slab.cache_insert(id, data);
433}
434
435/// Look up plaintext from the slab hot cache.
436/// Returns a transient `PoolSlot` with a copy of the cached bytes, or `None` on miss.
437pub(super) fn hot_cache_get(id: u64) -> Option<PoolSlot> {
438    let pool = global_pool();
439    let mut guard = pool.tiers[0].slab.lock().unwrap_or_else(|e| e.into_inner());
440    let slot_idx = guard.cache_get(id)?;
441    // slot_idx was just returned by cache_get(), which only returns valid transient indices.
442    let (ptr, len) = guard.slot_raw(slot_idx)?;
443    drop(guard);
444    Some(PoolSlot::from_slab(ptr, len, 0, slot_idx))
445}
446
447/// Evict an entry from the slab hot cache and notify waiters.
448pub(super) fn hot_cache_evict(id: u64) {
449    let pool = global_pool();
450    {
451        let mut slab = pool.tiers[0].slab.lock().unwrap_or_else(|e| e.into_inner());
452        slab.cache_evict(id);
453    }
454    pool.tiers[0].cv.notify_one();
455}
456
457#[cfg(test)]
458#[allow(clippy::unwrap_used, clippy::panic)]
459mod tests {
460    use std::sync::Mutex;
461
462    use super::super::slab::FIRST_SHARED_SLOT;
463    use super::*;
464
465    /// Serializes tests that touch the global TieredPool to prevent interference.
466    static TEST_LOCK: Mutex<()> = Mutex::new(());
467
468    // ── Global pool tests ───────────────────────────────────────────
469
470    #[test]
471    fn pool_acquire_small_uses_slab() {
472        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
473        let slot = pool_acquire(16).unwrap();
474        assert!(slot.slab_index().is_some());
475        assert_eq!(slot.size(), DEFAULT_SLOT_SIZE);
476    }
477
478    #[test]
479    fn pool_acquire_large_uses_standalone() {
480        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
481        let slot = pool_acquire(8192).unwrap();
482        assert!(slot.slab_index().is_none());
483        assert_eq!(slot.size(), 8192);
484    }
485
486    #[test]
487    fn pool_acquire_zero_uses_slab() {
488        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
489        // size 0 <= slot_size, so it should use the slab.
490        let slot = pool_acquire(0).unwrap();
491        assert!(slot.slab_index().is_some());
492    }
493
494    #[test]
495    fn pool_slot_write_and_read() {
496        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
497        let mut slot = pool_acquire(16).unwrap();
498        let data = b"test data 12345!";
499        slot.bytes()[..data.len()].copy_from_slice(data);
500        assert_eq!(&slot.as_slice()[..data.len()], data);
501    }
502
503    #[test]
504    fn pool_slot_zeroized_on_drop() {
505        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
506        // Acquire a slab slot, write a pattern, drop it, re-acquire and verify zeroed.
507        let mut slot = pool_acquire(16).unwrap();
508        let sz = slot.size();
509        slot.bytes().iter_mut().for_each(|b| *b = 0xDE);
510        // The slab_index tells us which slot; after drop the same index should be free.
511        let slot_idx = slot.slab_index().unwrap();
512        drop(slot);
513        // Acquire again — we should get a zeroed slot.
514        let pool = global_pool();
515        let mut guard = pool.tiers[0].slab.lock().unwrap_or_else(|e| e.into_inner());
516        // Drain free list until we get the same index back.
517        let mut acquired = vec![];
518        while let Some(idx) = guard.acquire_transient() {
519            acquired.push(idx);
520            if idx == slot_idx {
521                break;
522            }
523        }
524        if acquired.last() == Some(&slot_idx) {
525            let (ptr, _) = guard
526                .slot_raw(slot_idx)
527                .expect("slot_raw: index just acquired from slab");
528            let s = unsafe { std::slice::from_raw_parts(ptr, sz) };
529            assert!(s.iter().all(|&b| b == 0), "slot must be zeroed after drop");
530        }
531        for idx in acquired {
532            guard.release(idx);
533        }
534    }
535
536    #[test]
537    fn coffer_view_returns_key_sized_slot() {
538        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
539        let slot = coffer_view().unwrap();
540        assert_eq!(slot.size(), DEFAULT_SLOT_SIZE);
541        // Slot must be from tier 0 (slab-backed).
542        assert_eq!(slot.tier_index(), Some(0));
543    }
544
545    #[test]
546    fn coffer_view_is_deterministic() {
547        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
548        let s1 = coffer_view().unwrap();
549        let key1 = s1.as_slice().to_vec();
550        drop(s1);
551        let s2 = coffer_view().unwrap();
552        let key2 = s2.as_slice().to_vec();
553        drop(s2);
554        assert_eq!(key1, key2, "coffer_view must return same key each call");
555        assert!(!key1.iter().all(|&b| b == 0));
556    }
557
558    #[test]
559    fn hot_cache_insert_get_evict() {
560        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
561        let data = [0xAB_u8; DEFAULT_SLOT_SIZE];
562        hot_cache_insert(1001, &data);
563        let slot = hot_cache_get(1001).unwrap();
564        assert_eq!(slot.as_slice(), &data);
565        drop(slot);
566        hot_cache_evict(1001);
567        assert!(hot_cache_get(1001).is_none());
568    }
569
570    #[test]
571    fn hot_cache_get_returns_pool_slot() {
572        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
573        let data = [0xCC_u8; DEFAULT_SLOT_SIZE];
574        hot_cache_insert(2002, &data);
575        let slot = hot_cache_get(2002).expect("should be a cache hit");
576        // Result is a slab-backed slot from tier 0.
577        assert_eq!(slot.tier_index(), Some(0));
578        assert!(slot
579            .slab_index()
580            .map(|i| i >= FIRST_SHARED_SLOT)
581            .unwrap_or(false));
582        drop(slot);
583        hot_cache_evict(2002);
584    }
585
586    // ── Local TieredPool tests ───────────────────────────────────────
587
588    #[test]
589    fn tiered_pool_routes_small_to_first_tier() {
590        // Use the global pool so PoolSlot::drop returns to the correct pool.
591        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
592        let slot = pool_acquire(16).unwrap();
593        assert_eq!(
594            slot.tier_index(),
595            Some(0),
596            "should route to tier 0 (32-byte)"
597        );
598        assert_eq!(slot.size(), DEFAULT_SLOT_SIZE);
599    }
600
601    #[test]
602    fn tiered_pool_routes_medium_to_second_tier() {
603        // A 48-byte request exceeds the default 32-byte tier — falls back to standalone.
604        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
605        let slot = pool_acquire(48).unwrap();
606        assert!(
607            slot.tier_index().is_none(),
608            "48-byte request exceeds default tier; should be standalone"
609        );
610        assert_eq!(slot.size(), 48);
611    }
612
613    #[test]
614    fn tiered_pool_routes_large_to_standalone() {
615        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
616        let slot = pool_acquire(8192).unwrap();
617        assert!(slot.tier_index().is_none(), "should be standalone");
618        assert_eq!(slot.size(), 8192);
619    }
620
621    #[test]
622    fn init_pool_default_config_has_one_tier() {
623        // Inspect the global pool (already initialised by other tests).
624        let pool = global_pool();
625        assert_eq!(pool.tier_count(), 1);
626        assert_eq!(pool.tier_slot_size(0), Some(DEFAULT_SLOT_SIZE));
627        assert_eq!(pool.max_slab_slot_size(), DEFAULT_SLOT_SIZE);
628    }
629
630    #[test]
631    fn tiered_pool_config_validates_ascending() {
632        // Dedup test: no PoolSlots acquired, so local pool is safe.
633        let pool = TieredPool::new(TieredPoolConfig {
634            tier_sizes: vec![32, 32],
635        })
636        .unwrap();
637        assert_eq!(pool.tier_count(), 1, "duplicates should be deduped");
638    }
639
640    #[test]
641    fn tiered_pool_config_validates_max_slot_size() {
642        let ps = page_size();
643        let too_large = ps / 3 + 1;
644        let err = TieredPool::new(TieredPoolConfig {
645            tier_sizes: vec![too_large],
646        });
647        assert!(err.is_err(), "slot size > page_size/3 must be rejected");
648    }
649
650    #[test]
651    fn local_pool_coffer_view_works() {
652        // Verify coffer_view on the global pool; all slab PoolSlots must come from global_pool.
653        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
654        let slot = coffer_view().unwrap();
655        assert_eq!(slot.size(), DEFAULT_SLOT_SIZE);
656        assert_eq!(slot.tier_index(), Some(0));
657    }
658
659    // ── New tests for review findings ────────────────────────────────
660
661    #[test]
662    fn tiered_pool_first_tier_must_be_32_bytes() {
663        // BLK-7: first tier slot_size < 32 must be rejected (coffer requires AES-256 key size).
664        let result = TieredPool::new(TieredPoolConfig {
665            tier_sizes: vec![16],
666        });
667        assert!(
668            result.is_err(),
669            "first tier < 32 should fail (coffer requires slot_size >= 32)"
670        );
671    }
672
673    #[test]
674    fn coffer_view_key_is_32_bytes_and_nonzero() {
675        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
676        let slot = coffer_view().unwrap();
677        assert_eq!(slot.size(), 32);
678        // Key must not be all zeros (OsRng ran during slab init).
679        assert!(
680            slot.as_slice().iter().any(|&b| b != 0),
681            "coffer key must not be all zeros"
682        );
683    }
684
685    #[test]
686    fn empty_tier_sizes_rejected() {
687        let result = TieredPool::new(TieredPoolConfig { tier_sizes: vec![] });
688        assert!(result.is_err(), "empty tier_sizes must be rejected");
689    }
690
691    #[test]
692    fn tier_sizes_sorted_ascending_internally() {
693        // Pass sizes in reverse order — should be sorted internally.
694        let pool = TieredPool::new(TieredPoolConfig {
695            tier_sizes: vec![64, 32],
696        })
697        .unwrap();
698        assert_eq!(pool.tier_count(), 2);
699        assert_eq!(pool.tier_slot_size(0), Some(32));
700        assert_eq!(pool.tier_slot_size(1), Some(64));
701    }
702
703    #[test]
704    fn multi_tier_routing_smallest_fit() {
705        // Two tiers: 32 and 64. Verify tier routing by slot_size inspection only
706        // (do not acquire PoolSlots from a local pool — drop would return to wrong pool).
707        let pool = TieredPool::new(TieredPoolConfig {
708            tier_sizes: vec![32, 64],
709        })
710        .unwrap();
711        assert_eq!(pool.tier_count(), 2);
712        // tier_for_size is private; verify via tier_slot_size that sizes are correct.
713        assert_eq!(pool.tier_slot_size(0), Some(32));
714        assert_eq!(pool.tier_slot_size(1), Some(64));
715        // Using the global pool: size 33 > 32 → standalone (single 32-byte tier in global pool).
716        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
717        let slot = pool_acquire(33).unwrap();
718        assert!(
719            slot.tier_index().is_none(),
720            "size 33 exceeds single 32-byte tier → standalone"
721        );
722        assert_eq!(slot.size(), 33);
723        // Size 32 uses tier 0.
724        let slot2 = pool_acquire(32).unwrap();
725        assert_eq!(
726            slot2.tier_index(),
727            Some(0),
728            "size 32 must use tier 0 (32-byte)"
729        );
730        drop(slot);
731        drop(slot2);
732    }
733
734    #[test]
735    fn pool_slot_tier_index_matches_acquisition_tier() {
736        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
737        // Default pool has one tier (32 bytes). A small acquisition must be tier 0.
738        let slot = pool_acquire(16).unwrap();
739        assert_eq!(slot.tier_index(), Some(0));
740        assert_eq!(
741            slot.slab_index().map(|i| i >= FIRST_SHARED_SLOT),
742            Some(true)
743        );
744    }
745
746    #[test]
747    fn standalone_slot_has_no_tier_or_slab_index() {
748        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
749        let slot = pool_acquire(9999).unwrap();
750        assert!(slot.tier_index().is_none());
751        assert!(slot.slab_index().is_none());
752        assert_eq!(slot.size(), 9999);
753    }
754
755    #[test]
756    fn pool_slot_zeroized_on_drop_standalone() {
757        // Standalone slot (large allocation) must also be zeroed on drop.
758        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
759        let mut slot = pool_acquire(512).unwrap();
760        slot.bytes().fill(0xBE);
761        // After drop, memory is in a guard-paged SecureBuffer that's zeroed in drop.
762        drop(slot);
763        // Can't inspect freed memory directly, but verify drop doesn't panic.
764    }
765
766    #[test]
767    fn hot_cache_not_populated_for_large_plaintext() {
768        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
769        // hot_cache_insert only caches data that fits in tier 0's slot size (32 bytes).
770        // A 64-byte payload should not be cached.
771        let big_data = [0x42_u8; 64];
772        hot_cache_insert(9876, &big_data);
773        // Should be a miss (not cached).
774        let result = hot_cache_get(9876);
775        assert!(result.is_none(), "oversized data must not be cached");
776    }
777
778    #[test]
779    fn hot_cache_multiple_ids_are_independent() {
780        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
781        let data_a = [0xAA_u8; DEFAULT_SLOT_SIZE];
782        let data_b = [0xBB_u8; DEFAULT_SLOT_SIZE];
783        hot_cache_insert(100, &data_a);
784        hot_cache_insert(101, &data_b);
785        let slot_a = hot_cache_get(100).unwrap();
786        let slot_b = hot_cache_get(101).unwrap();
787        assert_eq!(slot_a.as_slice(), &data_a, "id 100 must return data_a");
788        assert_eq!(slot_b.as_slice(), &data_b, "id 101 must return data_b");
789        drop(slot_a);
790        drop(slot_b);
791        hot_cache_evict(100);
792        hot_cache_evict(101);
793    }
794
795    #[test]
796    fn coffer_view_returns_same_key_every_time() {
797        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
798        let s1 = coffer_view().unwrap();
799        let k1 = s1.as_slice().to_vec();
800        drop(s1);
801        let s2 = coffer_view().unwrap();
802        let k2 = s2.as_slice().to_vec();
803        drop(s2);
804        let s3 = coffer_view().unwrap();
805        let k3 = s3.as_slice().to_vec();
806        drop(s3);
807        assert_eq!(k1, k2, "coffer key must be same on second call");
808        assert_eq!(k2, k3, "coffer key must be same on third call");
809        assert!(
810            k1.iter().any(|&b| b != 0),
811            "coffer key must not be all zeros"
812        );
813    }
814
815    #[test]
816    fn concurrent_pool_acquire_and_release() {
817        use std::sync::Arc;
818        use std::thread;
819        let _guard = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
820        // 8 threads each acquire a pool slot, write, verify, and release.
821        let barrier = Arc::new(std::sync::Barrier::new(8));
822        let handles: Vec<_> = (0..8_u8)
823            .map(|i| {
824                let b = Arc::clone(&barrier);
825                thread::spawn(move || {
826                    let mut slot = pool_acquire(16).unwrap();
827                    slot.bytes()[0] = i;
828                    b.wait(); // synchronize so all threads are in-flight simultaneously
829                    assert_eq!(slot.as_slice()[0], i, "thread {i}: slot content must match");
830                    drop(slot);
831                })
832            })
833            .collect();
834        for h in handles {
835            h.join().expect("thread panicked");
836        }
837    }
838}