Skip to main content

rust_sanitize/
store.rs

1//! Thread-safe, concurrent one-way replacement store.
2//!
3//! # Concurrency Model
4//!
5//! The store uses [`dashmap::DashMap`] — a concurrent hash map with shard-level
6//! locking (default 64 shards). This gives us:
7//!
8//! - **Lock-free reads** for lookups of already-mapped values.
9//! - **Shard-level write locks** that are held only while inserting a new entry.
10//!   With 64 shards and 8–16 threads, the probability of two threads contending
11//!   on the same shard is very low.
12//! - **Atomic get-or-insert** via the `entry()` API, which prevents TOCTOU races
13//!   and guarantees first-writer-wins semantics.
14//!
15//! # Structure
16//!
17//! The forward map is two-level: `Category → original → sanitized`.
18//!
19//! ```text
20//! DashMap<Category, Arc<DashMap<ZeroizingString, (CompactString, usize)>>>
21//!    outer (~20 entries, always hot in cache)
22//!               └── inner (one per category, holds the actual values)
23//! ```
24//!
25//! This lets the fast-path read call `inner.get(original: &str)` without
26//! constructing a temporary `String`, because `ZeroizingString: Borrow<str>`.
27//! For files where the same value appears thousands of times, this eliminates
28//! thousands of `malloc`/`free` cycles on the hot path.
29//!
30//! Replacements are **one-way only** — there is no reverse map, no mapping
31//! file, and no restore capability.
32//!
33//! # Memory Characteristics
34//!
35//! At 10M unique values with average key length 20 bytes and average value
36//! length 30 bytes:
37//! - Forward map: 10M × (20 + 30 + ~120 DashMap overhead) ≈ 1.7 GB
38//! - **Total: ~1.7 GB** — acceptable for server workloads.
39//!
40//! An optional `capacity_limit` can be set to prevent unbounded growth.
41
42use crate::allowlist::AllowlistMatcher;
43use crate::category::Category;
44use crate::error::{Result, SanitizeError};
45use crate::generator::ReplacementGenerator;
46use compact_str::CompactString;
47use dashmap::DashMap;
48use std::borrow::Borrow;
49use std::sync::atomic::{AtomicUsize, Ordering};
50use std::sync::Arc;
51use zeroize::Zeroize;
52
53/// An opaque cursor into the [`MappingStore`] insertion sequence.
54///
55/// Obtained from [`MappingStore::snapshot`] and passed to
56/// [`MappingStore::iter_since`]. Using a dedicated type prevents accidentally
57/// passing an unrelated `usize` (a count, an index, a capacity) to
58/// `iter_since`, which would silently yield the wrong subset of entries.
59///
60/// **Breaking change in 0.13.0:** `MappingStore::snapshot` previously returned
61/// `usize`; it now returns `StoreSnapshot`. Call sites should change
62/// `store.snapshot()` → the new type and `iter_since(n)` → `iter_since(snap)`.
63/// To iterate all entries use [`StoreSnapshot::start`] instead of the literal `0`.
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub struct StoreSnapshot(usize);
66
67impl StoreSnapshot {
68    /// A snapshot representing the beginning of the store (before any
69    /// insertions). Passing this to [`MappingStore::iter_since`] yields every
70    /// entry in the store — equivalent to the former `iter_since(0)`.
71    #[must_use]
72    pub fn start() -> Self {
73        Self(0)
74    }
75}
76
77impl Default for StoreSnapshot {
78    fn default() -> Self {
79        Self::start()
80    }
81}
82
83// ---------------------------------------------------------------------------
84// ZeroizingString — map key for the inner (per-category) DashMap
85// ---------------------------------------------------------------------------
86
87/// A `String` that zeroizes its heap buffer on drop.
88///
89/// `Zeroizing<String>` from the `zeroize` crate does not implement `Hash`,
90/// so it cannot be used as a `HashMap` key. This newtype adds `Hash` while
91/// keeping the zeroize-on-drop guarantee via an explicit `Drop` impl.
92///
93/// Implementing `Borrow<str>` allows `DashMap<ZeroizingString, _>::get(s: &str)`
94/// to work without constructing a temporary `ZeroizingString` — the key insight
95/// that makes the fast-path read allocation-free.
96#[derive(Debug, Clone, PartialEq, Eq)]
97struct ZeroizingString(String);
98
99impl std::hash::Hash for ZeroizingString {
100    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
101        self.0.hash(state);
102    }
103}
104
105impl Drop for ZeroizingString {
106    fn drop(&mut self) {
107        self.0.zeroize();
108    }
109}
110
111/// Enables `DashMap<ZeroizingString, _>::get(s: &str)` — zero allocation on
112/// cache hits. Correct because `ZeroizingString` delegates `Hash` and `Eq`
113/// to its inner `String`, which is consistent with `str`'s `Hash` and `Eq`.
114impl Borrow<str> for ZeroizingString {
115    fn borrow(&self) -> &str {
116        &self.0
117    }
118}
119
120// ---------------------------------------------------------------------------
121// Convenience type alias for the inner map
122// ---------------------------------------------------------------------------
123
124type InnerMap = DashMap<ZeroizingString, (CompactString, usize)>;
125
126// ---------------------------------------------------------------------------
127// MappingStore
128// ---------------------------------------------------------------------------
129
130/// Thread-safe concurrent one-way replacement store.
131///
132/// Caches forward mappings for per-run consistency (same input always
133/// produces the same output within a run). There is no reverse map,
134/// no journal, and no persistence — replacements are one-way only.
135///
136/// See the [module-level documentation](self) for concurrency and memory details.
137pub struct MappingStore {
138    /// `category → original → (sanitized, insertion_index)`
139    ///
140    /// Two-level map: outer is keyed by `Category` (tiny, always in cache),
141    /// inner is keyed by `ZeroizingString` (actual values). The inner map is
142    /// behind an `Arc` so it can be obtained without holding the outer shard
143    /// lock during inner map operations.
144    forward: DashMap<Category, Arc<InnerMap>>,
145    /// Replacement generator (HMAC deterministic or CSPRNG random).
146    generator: Arc<dyn ReplacementGenerator>,
147    /// Current number of mappings (atomic for lock-free reads).
148    len: AtomicUsize,
149    /// Optional upper bound on the number of mappings.
150    capacity_limit: Option<usize>,
151    /// Optional allowlist — matched values pass through unchanged and are
152    /// not recorded in the forward map.
153    allowlist: Option<Arc<AllowlistMatcher>>,
154}
155
156impl MappingStore {
157    // ---------------- Construction ----------------
158
159    /// Create a new, empty mapping store.
160    ///
161    /// # Arguments
162    ///
163    /// - `generator` — replacement strategy (HMAC or random).
164    /// - `capacity_limit` — optional max number of unique mappings.
165    #[must_use]
166    pub fn new(generator: Arc<dyn ReplacementGenerator>, capacity_limit: Option<usize>) -> Self {
167        Self {
168            forward: DashMap::with_capacity(32),
169            generator,
170            len: AtomicUsize::new(0),
171            capacity_limit,
172            allowlist: None,
173        }
174    }
175
176    /// Create a new store with an allowlist. Values matching the allowlist
177    /// are returned unchanged and never recorded in the forward map.
178    #[must_use]
179    pub fn new_with_allowlist(
180        generator: Arc<dyn ReplacementGenerator>,
181        capacity_limit: Option<usize>,
182        allowlist: Arc<AllowlistMatcher>,
183    ) -> Self {
184        Self {
185            forward: DashMap::with_capacity(32),
186            generator,
187            len: AtomicUsize::new(0),
188            capacity_limit,
189            allowlist: Some(allowlist),
190        }
191    }
192
193    /// Return the allowlist attached to this store, if any.
194    pub fn allowlist(&self) -> Option<&AllowlistMatcher> {
195        self.allowlist.as_deref()
196    }
197
198    // ---------------- Core API ----------------
199
200    /// Get or create the sanitized replacement for `(category, original)`.
201    ///
202    /// This is the primary API for one-way sanitization.
203    ///
204    /// **Hot-path allocation:** When the value is already cached, this method
205    /// is allocation-free. The inner `DashMap::get` accepts `&str` directly via
206    /// `ZeroizingString: Borrow<str>`, so no temporary `String` is constructed.
207    ///
208    /// **Thread-safety:** Uses `DashMap::entry()` which holds a shard-level
209    /// lock only for the duration of the insert closure. The generator is
210    /// called inside the lock, but generation is fast (one HMAC or one RNG
211    /// call). Capacity enforcement uses `compare_exchange` to prevent
212    /// TOCTOU over-insertion.
213    ///
214    /// **Per-run consistency:** Once a value is mapped, all subsequent
215    /// lookups return the same sanitized value (first-writer-wins).
216    ///
217    /// # Errors
218    ///
219    /// Returns [`SanitizeError::CapacityExceeded`] if the store has
220    /// reached its configured capacity limit.
221    pub fn get_or_insert(&self, category: &Category, original: &str) -> Result<CompactString> {
222        // Allowlist check: return the original value unchanged without recording it.
223        if let Some(al) = &self.allowlist {
224            if al.is_allowed(original) {
225                return Ok(CompactString::new(original));
226            }
227        }
228
229        // Fast path: already mapped — zero allocation.
230        // `inner.get(original)` accepts `&str` via `ZeroizingString: Borrow<str>`.
231        // Clone the Arc while we already hold the outer shard reference so the
232        // slow path below never needs to acquire the outer shard a second time.
233        let inner: Arc<InnerMap> = match self.forward.get(category) {
234            Some(outer) => {
235                if let Some(existing) = outer.value().get(original) {
236                    return Ok(existing.value().0.clone());
237                }
238                outer.value().clone()
239            }
240            None => self
241                .forward
242                .entry(category.clone())
243                .or_insert_with(|| Arc::new(DashMap::new()))
244                .value()
245                .clone(),
246        };
247
248        if let Some(limit) = self.capacity_limit {
249            // Atomically reserve a capacity slot *before* generating the value.
250            // This eliminates the TOCTOU race where multiple threads pass the
251            // capacity check and all insert.
252            //
253            // insertion_index is set to `current` (the pre-increment value) from
254            // the successful CAS — not from a separate load after the loop, which
255            // could observe a higher count from a concurrent inserter and assign
256            // the wrong monotonic position to this entry.
257            let insertion_index;
258            loop {
259                let current = self.len.load(Ordering::Acquire);
260                if current >= limit {
261                    // One more chance: key may have been inserted by another thread.
262                    if let Some(existing) = inner.get(original) {
263                        return Ok(existing.value().0.clone());
264                    }
265                    return Err(SanitizeError::CapacityExceeded { current, limit });
266                }
267                if self
268                    .len
269                    .compare_exchange_weak(
270                        current,
271                        current + 1,
272                        Ordering::AcqRel,
273                        Ordering::Acquire,
274                    )
275                    .is_ok()
276                {
277                    insertion_index = current;
278                    break;
279                }
280                // CAS failed → another thread incremented; retry.
281            }
282
283            // Slot reserved — generate and insert (first-writer-wins).
284            let mut was_inserted = false;
285            let result = inner
286                .entry(ZeroizingString(original.to_owned()))
287                .or_insert_with(|| {
288                    was_inserted = true;
289                    let val = self.generator.generate(category, original);
290                    (CompactString::new(val), insertion_index)
291                })
292                .value()
293                .0
294                .clone();
295
296            if !was_inserted {
297                // Another thread inserted first — release our reserved slot.
298                self.len.fetch_sub(1, Ordering::Release);
299            }
300
301            Ok(result)
302        } else {
303            // No capacity limit — generate inside the entry lock so only the
304            // first writer calls the generator (first-writer-wins semantics).
305            let result = inner
306                .entry(ZeroizingString(original.to_owned()))
307                .or_insert_with(|| {
308                    let insertion_index = self.len.fetch_add(1, Ordering::AcqRel);
309                    let val = self.generator.generate(category, original);
310                    (CompactString::new(val), insertion_index)
311                })
312                .value()
313                .0
314                .clone();
315
316            Ok(result)
317        }
318    }
319
320    /// Look up an existing forward mapping without creating one.
321    #[must_use]
322    pub fn forward_lookup(&self, category: &Category, original: &str) -> Option<CompactString> {
323        let inner = self.forward.get(category)?;
324        inner.value().get(original).map(|r| r.value().0.clone())
325    }
326
327    /// Register `alias` as an additional original that maps to the **same**
328    /// `sanitized` replacement under `category`.
329    ///
330    /// Structured processors use this to record the *source-escaped* form of a
331    /// discovered value — e.g. the JSON value `a"b` appears in the raw bytes as
332    /// `a\"b`. The format-preserving scanner matches against raw input bytes, so
333    /// without the alias the escaped occurrence would not be redacted. Aliasing
334    /// (rather than a fresh mapping) keeps the escaped occurrence consistent with
335    /// the parsed value's token.
336    ///
337    /// First-writer-wins: an existing mapping for `alias` is left unchanged.
338    /// Allowlisted or empty aliases are ignored. Like [`Self::get_or_insert`],
339    /// the new entry participates in [`Self::iter_since`].
340    pub fn register_alias(&self, category: &Category, alias: &str, sanitized: &str) {
341        if alias.is_empty() {
342            return;
343        }
344        if let Some(al) = &self.allowlist {
345            if al.is_allowed(alias) {
346                return;
347            }
348        }
349        let inner: Arc<InnerMap> = match self.forward.get(category) {
350            Some(outer) => outer.value().clone(),
351            None => self
352                .forward
353                .entry(category.clone())
354                .or_insert_with(|| Arc::new(DashMap::new()))
355                .value()
356                .clone(),
357        };
358        inner
359            .entry(ZeroizingString(alias.to_owned()))
360            .or_insert_with(|| {
361                let insertion_index = self.len.fetch_add(1, Ordering::AcqRel);
362                (CompactString::new(sanitized), insertion_index)
363            });
364    }
365
366    // ---------------- Metrics ----------------
367
368    /// Number of unique mappings in the store.
369    #[must_use]
370    pub fn len(&self) -> usize {
371        self.len.load(Ordering::Relaxed)
372    }
373
374    /// Whether the store is empty.
375    #[must_use]
376    pub fn is_empty(&self) -> bool {
377        self.len() == 0
378    }
379
380    /// Remove all mappings, zeroizing the original plaintexts.
381    ///
382    /// Takes `&self` so it is usable on a shared `Arc<MappingStore>`. Only
383    /// call this after all concurrent readers and writers have finished —
384    /// `DashMap::clear` acquires shard locks one at a time, so a concurrent
385    /// `get_or_insert` racing with `clear` will observe a partially-cleared
386    /// store.
387    ///
388    /// **Breaking change in 0.13.0:** the signature changed from
389    /// `clear(&mut self)` to `clear(&self)`.
390    pub fn clear(&self) {
391        // DashMap::clear() acquires each shard lock in turn and drops all
392        // entries, triggering ZeroizingString::drop for every key. Cloned
393        // Arc<InnerMap> refs held by concurrent threads survive until their
394        // last clone drops, but clear() is intended for post-run teardown
395        // only, so no concurrent access should be in flight.
396        self.forward.clear();
397        self.len.store(0, Ordering::Release);
398    }
399
400    // ---------------- Snapshot / diff (for format-preserving pass) ----------------
401
402    /// Snapshot the current insertion count.
403    ///
404    /// Returns a [`StoreSnapshot`] that can be passed to [`Self::iter_since`] to
405    /// iterate only the entries added *after* this point — useful for
406    /// finding which mappings a structured processor pass discovered without
407    /// building a full `HashSet` of all existing keys.
408    ///
409    /// O(1), no allocation.
410    ///
411    /// **Breaking change in 0.13.0:** previously returned `usize`;
412    /// now returns [`StoreSnapshot`].
413    #[must_use]
414    pub fn snapshot(&self) -> StoreSnapshot {
415        StoreSnapshot(self.len.load(Ordering::Acquire))
416    }
417
418    /// Iterate over entries added at or after the given snapshot.
419    ///
420    /// `since` is the value returned by a previous call to [`Self::snapshot`].
421    /// Entries whose insertion index is ≥ `since` are yielded; older entries
422    /// are skipped. Still O(n) in total store size, but avoids allocating a
423    /// `HashSet` of all prior keys. Use [`StoreSnapshot::start`] to iterate
424    /// all entries.
425    ///
426    /// **Breaking change in 0.13.0:** the parameter type changed from `usize`
427    /// to [`StoreSnapshot`].
428    ///
429    /// Implementation note: the inner `.collect::<Vec<_>>()` inside the
430    /// `flat_map` is required to release the DashMap shard lock before
431    /// yielding items — it allocates one `Vec` per category shard visited.
432    pub fn iter_since(
433        &self,
434        since: StoreSnapshot,
435    ) -> impl Iterator<Item = (Category, CompactString, CompactString)> + '_ {
436        self.forward.iter().flat_map(move |outer| {
437            let cat = outer.key().clone();
438            outer
439                .value()
440                .iter()
441                .filter_map(move |inner| {
442                    let (sanitized, idx) = inner.value();
443                    if *idx >= since.0 {
444                        Some((
445                            cat.clone(),
446                            CompactString::new(inner.key().0.as_str()),
447                            sanitized.clone(),
448                        ))
449                    } else {
450                        None
451                    }
452                })
453                .collect::<Vec<_>>()
454        })
455    }
456
457    // ---------------- Iteration (for external use) ----------------
458
459    /// Iterate over all mappings. Yields `(category, original, sanitized)`.
460    ///
461    /// Note: iteration over `DashMap` is not snapshot-consistent if concurrent
462    /// inserts are happening. Call this after all workers have finished.
463    ///
464    /// Implementation note: allocates one `Vec` per category shard to release
465    /// the DashMap shard lock between categories.
466    pub fn iter(&self) -> impl Iterator<Item = (Category, CompactString, CompactString)> + '_ {
467        self.forward.iter().flat_map(|outer| {
468            let cat = outer.key().clone();
469            outer
470                .value()
471                .iter()
472                .map(move |inner| {
473                    (
474                        cat.clone(),
475                        CompactString::new(inner.key().0.as_str()),
476                        inner.value().0.clone(),
477                    )
478                })
479                .collect::<Vec<_>>()
480        })
481    }
482}
483
484/// Zeroize original keys stored in the forward map on drop.
485impl Drop for MappingStore {
486    fn drop(&mut self) {
487        self.clear();
488    }
489}
490
491/// Compile-time assertion that a type is `Send + Sync`.
492macro_rules! static_assertions_send_sync {
493    ($t:ty) => {
494        const _: fn() = || {
495            fn assert_send<T: Send>() {}
496            fn assert_sync<T: Sync>() {}
497            assert_send::<$t>();
498            assert_sync::<$t>();
499        };
500    };
501}
502
503static_assertions_send_sync!(MappingStore);
504
505// ---------------------------------------------------------------------------
506// Unit tests
507// ---------------------------------------------------------------------------
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use crate::generator::{HmacGenerator, RandomGenerator};
513    use std::sync::Arc;
514
515    fn hmac_store(limit: Option<usize>) -> MappingStore {
516        let gen = Arc::new(HmacGenerator::new([42u8; 32]));
517        MappingStore::new(gen, limit)
518    }
519
520    fn random_store() -> MappingStore {
521        let gen = Arc::new(RandomGenerator::new());
522        MappingStore::new(gen, None)
523    }
524
525    // --- Basic operations ---
526
527    #[test]
528    fn insert_and_lookup() {
529        let store = hmac_store(None);
530        let s1 = store
531            .get_or_insert(&Category::Email, "alice@corp.com")
532            .unwrap();
533        assert!(!s1.is_empty());
534        assert!(s1.contains("@corp.com"), "domain must be preserved");
535        assert_eq!(s1.len(), "alice@corp.com".len(), "length must be preserved");
536        assert_eq!(store.len(), 1);
537    }
538
539    #[test]
540    fn same_input_same_output() {
541        let store = hmac_store(None);
542        let s1 = store
543            .get_or_insert(&Category::Email, "alice@corp.com")
544            .unwrap();
545        let s2 = store
546            .get_or_insert(&Category::Email, "alice@corp.com")
547            .unwrap();
548        assert_eq!(s1, s2, "repeated insert must return cached value");
549        assert_eq!(store.len(), 1, "no duplicate entry");
550    }
551
552    #[test]
553    fn different_inputs_different_outputs() {
554        let store = hmac_store(None);
555        let s1 = store
556            .get_or_insert(&Category::Email, "alice@corp.com")
557            .unwrap();
558        let s2 = store
559            .get_or_insert(&Category::Email, "bob@corp.com")
560            .unwrap();
561        assert_ne!(s1, s2);
562        assert_eq!(store.len(), 2);
563    }
564
565    #[test]
566    fn different_categories_different_outputs() {
567        let store = hmac_store(None);
568        let s1 = store.get_or_insert(&Category::Email, "test").unwrap();
569        let s2 = store.get_or_insert(&Category::Name, "test").unwrap();
570        assert_ne!(s1, s2);
571    }
572
573    #[test]
574    fn forward_lookup_works() {
575        let store = hmac_store(None);
576        let sanitized = store.get_or_insert(&Category::IpV4, "192.168.1.1").unwrap();
577        let found = store.forward_lookup(&Category::IpV4, "192.168.1.1");
578        assert_eq!(found, Some(sanitized));
579    }
580
581    #[test]
582    fn forward_lookup_missing() {
583        let store = hmac_store(None);
584        assert!(store.forward_lookup(&Category::Email, "nope").is_none());
585    }
586
587    // --- Capacity limit ---
588
589    #[test]
590    fn capacity_limit_enforced() {
591        let store = hmac_store(Some(2));
592        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
593        store.get_or_insert(&Category::Email, "b@b.com").unwrap();
594        let result = store.get_or_insert(&Category::Email, "c@c.com");
595        assert!(result.is_err());
596        match result.unwrap_err() {
597            SanitizeError::CapacityExceeded {
598                current: 2,
599                limit: 2,
600            } => {}
601            other => panic!("unexpected error: {:?}", other),
602        }
603    }
604
605    #[test]
606    fn capacity_limit_allows_duplicate() {
607        let store = hmac_store(Some(1));
608        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
609        // Re-inserting same value should succeed (fast path).
610        let s2 = store.get_or_insert(&Category::Email, "a@a.com").unwrap();
611        assert!(!s2.is_empty());
612    }
613
614    // --- Random generator within store ---
615
616    #[test]
617    fn random_store_caches() {
618        let store = random_store();
619        let s1 = store
620            .get_or_insert(&Category::Email, "alice@corp.com")
621            .unwrap();
622        let s2 = store
623            .get_or_insert(&Category::Email, "alice@corp.com")
624            .unwrap();
625        assert_eq!(s1, s2, "random store must still cache the first result");
626    }
627
628    // --- Iteration ---
629
630    #[test]
631    fn iter_yields_all_mappings() {
632        let store = hmac_store(None);
633        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
634        store.get_or_insert(&Category::IpV4, "1.2.3.4").unwrap();
635        let collected: Vec<_> = store.iter().collect();
636        assert_eq!(collected.len(), 2);
637    }
638
639    // --- Concurrent inserts (basic smoke test) ---
640
641    #[test]
642    fn concurrent_inserts_no_panic() {
643        use std::sync::Arc;
644        use std::thread;
645
646        let gen = Arc::new(HmacGenerator::new([99u8; 32]));
647        let store = Arc::new(MappingStore::new(gen, None));
648
649        let mut handles = vec![];
650        for t in 0..8 {
651            let store = Arc::clone(&store);
652            handles.push(thread::spawn(move || {
653                for i in 0..1000 {
654                    let val = format!("thread{}-val{}", t, i);
655                    store.get_or_insert(&Category::Email, &val).unwrap();
656                }
657            }));
658        }
659
660        for h in handles {
661            h.join().unwrap();
662        }
663
664        assert_eq!(store.len(), 8000);
665    }
666
667    #[test]
668    fn concurrent_inserts_same_key_idempotent() {
669        use std::sync::Arc;
670        use std::thread;
671
672        let gen = Arc::new(HmacGenerator::new([7u8; 32]));
673        let store = Arc::new(MappingStore::new(gen, None));
674
675        let mut handles = vec![];
676        for _ in 0..8 {
677            let store = Arc::clone(&store);
678            handles.push(thread::spawn(move || {
679                let mut results = Vec::new();
680                for i in 0..100 {
681                    let val = format!("shared-{}", i);
682                    let r = store.get_or_insert(&Category::Email, &val).unwrap();
683                    results.push((val, r));
684                }
685                results
686            }));
687        }
688
689        let mut all_results: Vec<Vec<(String, CompactString)>> = vec![];
690        for h in handles {
691            all_results.push(h.join().unwrap());
692        }
693
694        // All threads must agree on every mapping.
695        assert_eq!(store.len(), 100);
696        for i in 0..100 {
697            let val = format!("shared-{}", i);
698            let expected = store.forward_lookup(&Category::Email, &val).unwrap();
699            for thread_results in &all_results {
700                let (_, got) = &thread_results[i];
701                assert_eq!(
702                    got, &expected,
703                    "all threads must see the same mapping for {}",
704                    val
705                );
706            }
707        }
708    }
709
710    // --- is_empty / clear ---
711
712    #[test]
713    fn is_empty_on_new_store() {
714        let store = hmac_store(None);
715        assert!(store.is_empty());
716    }
717
718    #[test]
719    fn is_empty_false_after_insert() {
720        let store = hmac_store(None);
721        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
722        assert!(!store.is_empty());
723    }
724
725    #[test]
726    fn clear_via_arc_shares_state() {
727        // The primary motivation for clear(&self) over clear(&mut self) is
728        // usability on Arc<MappingStore>. Verify that calling clear through a
729        // second Arc handle empties the store seen by the first handle.
730        let store = Arc::new(hmac_store(None));
731        let clone = Arc::clone(&store);
732        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
733        assert_eq!(store.len(), 1);
734        clone.clear();
735        assert_eq!(store.len(), 0, "clear via Arc must empty the shared store");
736        assert!(store.is_empty());
737    }
738
739    #[test]
740    fn clear_resets_store() {
741        let store = hmac_store(None);
742        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
743        store.get_or_insert(&Category::IpV4, "1.2.3.4").unwrap();
744        assert_eq!(store.len(), 2);
745        store.clear();
746        assert_eq!(store.len(), 0);
747        assert!(store.is_empty());
748    }
749
750    #[test]
751    fn clear_then_reinsert_works() {
752        let store = hmac_store(None);
753        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
754        store.clear();
755        let result = store.get_or_insert(&Category::Email, "a@a.com");
756        assert!(result.is_ok());
757        assert_eq!(store.len(), 1);
758    }
759
760    // --- snapshot / iter_since ---
761
762    #[test]
763    fn snapshot_and_iter_since_yields_only_new() {
764        let store = hmac_store(None);
765        store.get_or_insert(&Category::Email, "old@a.com").unwrap();
766        let snap = store.snapshot();
767        store.get_or_insert(&Category::IpV4, "1.2.3.4").unwrap();
768        store.get_or_insert(&Category::Name, "Alice").unwrap();
769
770        let new_entries: Vec<_> = store.iter_since(snap).collect();
771        assert_eq!(new_entries.len(), 2);
772        // None of the new entries should be the pre-snapshot email.
773        assert!(!new_entries
774            .iter()
775            .any(|(cat, orig, _)| { *cat == Category::Email && orig.as_str() == "old@a.com" }));
776    }
777
778    #[test]
779    fn snapshot_default_and_start_are_equivalent() {
780        let store = hmac_store(None);
781        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
782        store.get_or_insert(&Category::IpV4, "1.2.3.4").unwrap();
783        let via_start: Vec<_> = store.iter_since(StoreSnapshot::start()).collect();
784        let via_default: Vec<_> = store.iter_since(StoreSnapshot::default()).collect();
785        assert_eq!(
786            via_start.len(),
787            via_default.len(),
788            "default() must yield identical results to start()"
789        );
790    }
791
792    #[test]
793    fn iter_since_zero_yields_all() {
794        let store = hmac_store(None);
795        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
796        store.get_or_insert(&Category::IpV4, "1.2.3.4").unwrap();
797        let all: Vec<_> = store.iter_since(StoreSnapshot::start()).collect();
798        assert_eq!(all.len(), 2);
799    }
800
801    #[test]
802    fn iter_since_at_end_yields_nothing() {
803        let store = hmac_store(None);
804        store.get_or_insert(&Category::Email, "a@a.com").unwrap();
805        let snap = store.snapshot();
806        let new: Vec<_> = store.iter_since(snap).collect();
807        assert!(new.is_empty());
808    }
809
810    // --- new_with_allowlist ---
811
812    #[test]
813    fn allowlist_passes_value_through_unchanged() {
814        use crate::allowlist::AllowlistMatcher;
815        let matcher =
816            AllowlistMatcher::new(vec!["localhost".to_string(), "127.0.0.1".to_string()]).matcher;
817        let gen = Arc::new(HmacGenerator::new([42u8; 32]));
818        let store = MappingStore::new_with_allowlist(gen, None, Arc::new(matcher));
819
820        assert!(store.allowlist().is_some());
821
822        // Allowlisted value must be returned verbatim.
823        let result = store
824            .get_or_insert(&Category::Hostname, "localhost")
825            .unwrap();
826        assert_eq!(result.as_str(), "localhost");
827    }
828
829    #[test]
830    fn allowlist_still_replaces_non_listed() {
831        use crate::allowlist::AllowlistMatcher;
832        let matcher = AllowlistMatcher::new(vec!["localhost".to_string()]).matcher;
833        let gen = Arc::new(HmacGenerator::new([42u8; 32]));
834        let store = MappingStore::new_with_allowlist(gen, None, Arc::new(matcher));
835
836        let result = store
837            .get_or_insert(&Category::Hostname, "prod.corp.com")
838            .unwrap();
839        assert_ne!(result.as_str(), "prod.corp.com");
840    }
841}