Skip to main content

dig_dht/
provider_store.rs

1//! [`ProviderStore`] — the local key→providers map a node serves on `find_providers` / `add_provider`.
2//!
3//! Every DHT node keeps a small store of provider records it has been told about (via
4//! `add_provider`, because it is one of the `k` closest to those content keys) plus the records for
5//! content **it itself holds and announces**. The store is:
6//!
7//! - **keyed by content key** (the 64-hex [`Key`](crate::Key)) → a set of [`ProviderRecord`]s (one
8//!   per distinct provider `peer_id`);
9//! - **TTL'd** — [`get`](ProviderStore::get) never returns expired records, and
10//!   [`gc`](ProviderStore::gc) drops them so the store does not grow without bound;
11//! - **dedup-on-provider** — re-announcing from the same provider replaces that provider's record
12//!   (refreshing its `expires_at` + addresses), it does not accumulate duplicates;
13//! - **bounded** — [`put`](ProviderStore::put) enforces a per-content-key cap
14//!   ([`ProviderStoreLimits::max_providers_per_key`]) and a global record ceiling
15//!   ([`ProviderStoreLimits::max_total_records`]); an inbound record from an untrusted peer can
16//!   never grow the store without bound (SPEC §6.3, §14).
17//!
18//! It also tracks the set of content keys **this node announces** (content it holds) so the
19//! maintenance loop can republish them before their TTL elapses ([`local_announcements`]).
20//!
21//! [`local_announcements`]: ProviderStore::local_announcements
22
23use std::collections::{HashMap, HashSet};
24
25use crate::record::ProviderRecord;
26
27/// Bounds enforced by [`ProviderStore::put`] — the admission control that keeps the store from
28/// growing without bound under inbound `add_provider` traffic from untrusted peers.
29///
30/// Both caps are enforced **on every `put`**, not just at GC time: a single misbehaving peer that
31/// floods `add_provider` for many distinct content keys (or many distinct providers per key) is
32/// rejected once a cap is hit, rather than accepted and relying on TTL expiry to eventually free
33/// memory (SPEC §6.3, §14 "Unbounded provider store").
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct ProviderStoreLimits {
36    /// Maximum distinct provider records kept **per content key**. When a `put` for a new provider
37    /// would exceed this, an existing record is evicted to make room: an EXPIRED record if the key
38    /// holds one, otherwise the soonest-to-expire among the key's NEWEST slots, leaving its
39    /// longest-established LIVE providers reserved (see [`ProviderStore::eviction_victim`]).
40    pub max_providers_per_key: usize,
41    /// Maximum total records across **all** content keys. When a `put` for a genuinely new
42    /// (content_key, provider) pair would exceed this, the request is rejected outright (no
43    /// eviction across keys — that would let one attacker evict another key's legitimate holders).
44    pub max_total_records: usize,
45}
46
47impl Default for ProviderStoreLimits {
48    /// Conservative defaults: `k` (20, the Kademlia replication parameter) providers per key is
49    /// already generous replication, and a global ceiling that comfortably covers a node
50    /// participating in many lookups while still bounding worst-case memory from a single
51    /// misbehaving peer.
52    fn default() -> Self {
53        ProviderStoreLimits {
54            max_providers_per_key: 20,
55            max_total_records: 100_000,
56        }
57    }
58}
59
60/// The outcome of a [`ProviderStore::put`] — whether the record was admitted.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub enum PutOutcome {
63    /// The record was stored (fresh insert or refresh of an existing provider's record).
64    Accepted,
65    /// The record was rejected: the store is at capacity and the record did not qualify for
66    /// eviction-based admission (a new provider would exceed
67    /// [`ProviderStoreLimits::max_total_records`], or the per-key cap is full of records that all
68    /// expire no sooner than the incoming one).
69    RejectedOverCapacity,
70}
71
72/// Share of a content key's slots reserved for its longest-established providers — the divisor is
73/// applied to [`ProviderStoreLimits::max_providers_per_key`], so half the slots are protected from
74/// eviction and the newest half form the "churn zone" where eviction happens (#1434).
75///
76/// Half is chosen so the floor is always strictly smaller than the cap: a newcomer can therefore
77/// ALWAYS be admitted by evicting inside the churn zone, and the protection never turns into a
78/// refusal to learn about new honest holders.
79const ESTABLISHED_FLOOR_DIVISOR: usize = 2;
80
81/// One stored provider record plus **when this node first admitted it** — its establishment.
82///
83/// Establishment is an admission SEQUENCE number, not a timestamp: the store needs only the relative
84/// order in which providers were first learned, and an ordinal cannot be manipulated by an attacker
85/// choosing when to announce, nor does it need a clock threaded through [`ProviderStore::put`].
86#[derive(Debug)]
87struct ProviderEntry {
88    record: ProviderRecord,
89    /// Admission order — assigned once, on first admission, and PRESERVED across refreshes so
90    /// republishing (how an honest holder stays findable) never costs a holder its establishment.
91    admitted_seq: u64,
92}
93
94/// One content key in a [`ProviderSnapshot`]: the key, and how many live providers this node knows
95/// for it. Deliberately carries NO provider identity — see [`ProviderStore::snapshot`].
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct ProviderSnapshotEntry {
98    /// The 64-hex content key.
99    pub content_key: String,
100    /// How many non-expired providers this node holds a record for.
101    pub providers: usize,
102}
103
104/// A bounded, aggregated view of a node's provider store — see [`ProviderStore::snapshot`].
105#[derive(Debug, Clone, PartialEq, Eq)]
106pub struct ProviderSnapshot {
107    /// Content keys with at least one live provider, sorted by key, capped at the requested maximum.
108    pub entries: Vec<ProviderSnapshotEntry>,
109    /// How many keys had a live provider BEFORE the cap was applied, so a consumer can report
110    /// "showing N of M" rather than presenting a truncated view as complete.
111    pub total_keys: usize,
112    /// Whether the cap dropped entries.
113    pub truncated: bool,
114}
115
116/// A node's local provider records + the set of content keys it announces itself.
117#[derive(Debug)]
118pub struct ProviderStore {
119    /// content_key (64-hex) → provider_peer_id (64-hex) → entry.
120    by_key: HashMap<String, HashMap<String, ProviderEntry>>,
121    /// content keys (64-hex) this node holds + announces (for republish).
122    announced: HashSet<String>,
123    /// Admission-control bounds enforced by [`put`](Self::put).
124    limits: ProviderStoreLimits,
125    /// Monotonic source of [`ProviderEntry::admitted_seq`] — the next admission's ordinal.
126    next_admitted_seq: u64,
127}
128
129impl Default for ProviderStore {
130    fn default() -> Self {
131        ProviderStore::new()
132    }
133}
134
135impl ProviderStore {
136    /// A new empty store with the default [`ProviderStoreLimits`].
137    pub fn new() -> Self {
138        ProviderStore::with_limits(ProviderStoreLimits::default())
139    }
140
141    /// A new empty store enforcing `limits` on every [`put`](Self::put).
142    pub fn with_limits(limits: ProviderStoreLimits) -> Self {
143        ProviderStore {
144            by_key: HashMap::new(),
145            announced: HashSet::new(),
146            limits,
147            next_admitted_seq: 0,
148        }
149    }
150
151    /// Store (or refresh) a provider record, subject to [`ProviderStoreLimits`].
152    ///
153    /// Keyed by (content_key, provider_peer_id): a second record from the same provider for the
154    /// same key REPLACES the first (refreshes expiry + addresses) rather than duplicating — this
155    /// always succeeds regardless of capacity, since it does not grow the store.
156    ///
157    /// A genuinely new (content_key, provider) pair is admission-controlled:
158    /// - if the key already holds [`ProviderStoreLimits::max_providers_per_key`] *other* providers,
159    ///   one is evicted to make room — chosen by [`eviction_victim`], which reserves the key's
160    ///   longest-established slots so a Sybil flood cannot displace an incumbent holder (#1434);
161    /// - if the store is at [`ProviderStoreLimits::max_total_records`] globally, the new record is
162    ///   rejected — [`PutOutcome::RejectedOverCapacity`] — rather than evicting another key's
163    ///   records (which would let one attacker's flood evict another key's legitimate holders).
164    ///
165    /// [`eviction_victim`]: Self::eviction_victim
166    pub fn put(&mut self, record: ProviderRecord) -> PutOutcome {
167        self.put_at(record, crate::clock::now_secs())
168    }
169
170    /// [`put`](Self::put) with an explicit `now` (absolute Unix seconds) — the same admission
171    /// decision, taking the caller's clock instead of reading the system one.
172    ///
173    /// `now` is what lets eviction tell a LIVE provider from an expired one, which is the difference
174    /// between reclaiming a dead slot and evicting a real holder (see [`eviction_victim`]). A caller
175    /// that already has a timestamp — the serving side computes one for the TTL clamp — SHOULD pass
176    /// it, so the clamp and the admission decision are made against a single instant.
177    ///
178    /// [`eviction_victim`]: Self::eviction_victim
179    pub fn put_at(&mut self, record: ProviderRecord, now: u64) -> PutOutcome {
180        if let Some(existing) = self
181            .by_key
182            .get_mut(&record.content_key)
183            .and_then(|providers| providers.get_mut(&record.provider_peer_id))
184        {
185            // Refresh: same provider, same key. It does not grow the store, so no capacity check —
186            // and `admitted_seq` is deliberately left untouched (see [`ProviderEntry`]).
187            existing.record = record;
188            return PutOutcome::Accepted;
189        }
190
191        // Global ceiling check FIRST, before touching this key's entry, so a rejected record never
192        // leaves a stray empty entry behind and so the check reads the true pre-insert total (not
193        // skewed by an entry we are about to create).
194        if self.len() >= self.limits.max_total_records {
195            return PutOutcome::RejectedOverCapacity;
196        }
197        if let Some(providers) = self.by_key.get_mut(&record.content_key) {
198            if providers.len() >= self.limits.max_providers_per_key {
199                let Some(evict_id) =
200                    Self::eviction_victim(providers, self.limits.max_providers_per_key, now)
201                else {
202                    // Every slot is established — admitting would breach the per-key cap, so the
203                    // cap wins. Unreachable while the floor stays a strict fraction of the cap; kept
204                    // as the explicit guard that the per-key invariant is never violated.
205                    return PutOutcome::RejectedOverCapacity;
206                };
207                providers.remove(&evict_id);
208            }
209        }
210
211        let admitted_seq = self.next_admitted_seq;
212        self.next_admitted_seq += 1;
213        self.by_key
214            .entry(record.content_key.clone())
215            .or_default()
216            .insert(
217                record.provider_peer_id.clone(),
218                ProviderEntry {
219                    record,
220                    admitted_seq,
221                },
222            );
223        PutOutcome::Accepted
224    }
225
226    /// Pick which of a full key's providers to evict, or `None` if none may be.
227    ///
228    /// **Why not simply soonest-to-expire (#1434).** Every inbound record has its `expires_at`
229    /// clamped to `now + provider_ttl` at admission, so a provider that announces LATER necessarily
230    /// carries a strictly LATER expiry. Pure soonest-to-expire eviction therefore made the honest
231    /// incumbent the deterministic victim of anyone announcing after it: `max_providers_per_key`
232    /// Sybil identities — free, since a `ProviderRecord` is unsigned self-assertion — could evict
233    /// the ONLY real holder of a capsule and replace it with peers that fail the fetch, making that
234    /// content undiscoverable through this node. Repeated across the k-closest nodes that is
235    /// network-wide censorship of a key.
236    ///
237    /// **The policy, in two steps.**
238    ///
239    /// 1. **An EXPIRED record is the victim, wherever it sits — the floor included.** A record past
240    ///    its `expires_at` is already invisible to [`get`](Self::get) and merely awaits the next
241    ///    [`gc`](Self::gc), so reclaiming its slot costs nothing. Liveness therefore OUTRANKS
242    ///    establishment. Were the floor allowed to protect a dead record, a live holder in the churn
243    ///    zone would be evicted to keep a corpse — and that needs no attacker, because a node's GC
244    ///    tick is coarser than the provider TTL: a key whose earliest providers have gone offline
245    ///    (ordinary churn — shutdown, cache eviction) carries expired records inside its floor for a
246    ///    whole GC period, and during that window every new announcement would evict a LIVE
247    ///    provider, making a capsule LESS discoverable the more holders announce it. That is the
248    ///    replication flywheel running backwards.
249    /// 2. **Otherwise every record is live, and the establishment floor governs.** The
250    ///    `max_providers_per_key / ESTABLISHED_FLOOR_DIVISOR` longest-established providers are
251    ///    RESERVED; the victim is the soonest-to-expire among the newest slots (the churn zone),
252    ///    that being the least valuable LIVE record to keep. This mirrors the k-bucket policy this
253    ///    crate already applies to contacts — long-lived entries resist eviction attacks — and
254    ///    bounds what a flood can achieve: an attacker may churn the unreserved slots at will but
255    ///    cannot displace an ALREADY-ESTABLISHED holder, however many identities it spends or
256    ///    however it times its expiries.
257    ///
258    /// Ties break on `admitted_seq` in both steps, so the choice is deterministic rather than
259    /// hash-order dependent.
260    ///
261    /// **Residual, NOT closed here.** The floor protects an incumbent, not a latecomer: an attacker
262    /// that establishes BEFORE the honest holder retains the full pre-#1434 eviction primitive, and
263    /// because this store is in-memory only, every restart resets the floor to first-come. See the
264    /// caveat in `SPEC.md` §6.3/§14 — closing it needs signed provider records (#1573).
265    fn eviction_victim(
266        providers: &HashMap<String, ProviderEntry>,
267        max_providers_per_key: usize,
268        now: u64,
269    ) -> Option<String> {
270        let mut by_establishment: Vec<&ProviderEntry> = providers.values().collect();
271        by_establishment.sort_by_key(|e| e.admitted_seq);
272
273        // Step 1 — reclaim a dead slot in preference to ANY live record, the floor included.
274        let expired = by_establishment
275            .iter()
276            .filter(|e| e.record.is_expired(now))
277            .min_by_key(|e| (e.record.expires_at, e.admitted_seq));
278        if let Some(dead) = expired {
279            return Some(dead.record.provider_peer_id.clone());
280        }
281
282        // Step 2 — every record is live: reserve the established floor, evict inside the churn zone.
283        let established_floor = max_providers_per_key / ESTABLISHED_FLOOR_DIVISOR;
284        by_establishment
285            .into_iter()
286            .skip(established_floor)
287            .min_by_key(|e| (e.record.expires_at, e.admitted_seq))
288            .map(|e| e.record.provider_peer_id.clone())
289    }
290
291    /// Remove exactly the record for `(content_key, provider_peer_id)`, if present. Returns whether
292    /// a record was removed.
293    ///
294    /// This is the store half of an **authenticated retract** (SPEC §6.6): a caller that has
295    /// verified a signed retract from `provider_peer_id` removes only that provider's record for
296    /// that key. It MUST NOT touch any OTHER provider of the same key — a retract signed by one
297    /// holder can never evict another holder's record (censorship-resistance). A content key left
298    /// with no remaining providers is dropped so the store does not accumulate empty entries.
299    pub fn remove(&mut self, content_key: &str, provider_peer_id: &str) -> bool {
300        let Some(providers) = self.by_key.get_mut(content_key) else {
301            return false;
302        };
303        let removed = providers.remove(provider_peer_id).is_some();
304        if providers.is_empty() {
305            self.by_key.remove(content_key);
306        }
307        removed
308    }
309
310    /// Drop EVERY record for `content_key`, returning how many were removed.
311    ///
312    /// Unlike [`remove`](Self::remove) — the authenticated per-holder retract — this is a
313    /// whole-key wipe, so it MUST NOT be reachable from any wire path: a peer able to drive it
314    /// against the authoritative store would hold a censorship primitive over any key it names.
315    /// Its one caller is the node's own decision to forget a DISCOVERY-CACHE entry whose holders
316    /// all turned out to be undialable (`DhtService::forget_discovered`), where the records being
317    /// dropped are this node's own unverified hearsay and nobody else can see them.
318    pub fn remove_key(&mut self, content_key: &str) -> usize {
319        self.by_key
320            .remove(content_key)
321            .map(|providers| providers.len())
322            .unwrap_or(0)
323    }
324
325    /// The live (non-expired at `now`) provider records for `content_key`. Expired records are
326    /// skipped (and cleaned up by [`gc`](Self::gc)); returns an empty vec if none are known/live.
327    pub fn get(&self, content_key: &str, now: u64) -> Vec<ProviderRecord> {
328        self.by_key
329            .get(content_key)
330            .map(|providers| {
331                providers
332                    .values()
333                    .map(|e| &e.record)
334                    .filter(|r| !r.is_expired(now))
335                    .cloned()
336                    .collect()
337            })
338            .unwrap_or_default()
339    }
340
341    /// Drop every expired record (and any content key left with no live providers) as of `now`.
342    /// Returns the number of records removed. Call periodically from the maintenance loop.
343    pub fn gc(&mut self, now: u64) -> usize {
344        let mut removed = 0;
345        self.by_key.retain(|_key, providers| {
346            let before = providers.len();
347            providers.retain(|_pid, e| !e.record.is_expired(now));
348            removed += before - providers.len();
349            !providers.is_empty()
350        });
351        removed
352    }
353
354    /// Record that this node holds + announces `content_key` (so the maintenance loop republishes
355    /// it). Idempotent.
356    pub fn mark_announced(&mut self, content_key: String) {
357        self.announced.insert(content_key);
358    }
359
360    /// Stop announcing `content_key` (this node no longer holds the content). Returns whether it was
361    /// being announced.
362    pub fn unmark_announced(&mut self, content_key: &str) -> bool {
363        self.announced.remove(content_key)
364    }
365
366    /// The content keys this node announces (holds) — the republish work list.
367    pub fn local_announcements(&self) -> Vec<String> {
368        self.announced.iter().cloned().collect()
369    }
370
371    /// A bounded, AGGREGATED view of what this node holds in its DHT provider store — content keys
372    /// and how many live providers each has, with no provider identities (dig_ecosystem #1935).
373    ///
374    /// This is what lets the relay show the network's content layer without joining the DHT: a
375    /// Kademlia node stores records for keys near its OWN `peer_id`, so these are records about
376    /// MANY OTHER peers' content, not a self-report of what this node caches. The union across
377    /// several nodes is a broad slice of the real DHT.
378    ///
379    /// # Why counts and not identities
380    ///
381    /// A provider record IS a `(peer_id, content_key)` pair — exactly the linkage the relay's `/map`
382    /// refuses to publish (its tests assert no `peer_id` and no raw IP ever appear). Returning
383    /// counts keeps that contract intact rather than carving an exception into it. A caller that
384    /// genuinely needs identities can still use [`get`](Self::get) per key.
385    ///
386    /// Expired records are excluded as of `now`, so the counts match what [`get`](Self::get) would
387    /// return rather than including records the store has not GC'd yet.
388    ///
389    /// `max_keys` bounds the result: the store is attacker-influenced (any peer can announce), so an
390    /// unbounded snapshot would let a Sybil dictate the response size. When the cap truncates,
391    /// [`ProviderSnapshot::truncated`] is set and `total_keys` still reports the true total, so a
392    /// consumer can say "showing N of M" instead of silently presenting a partial view as complete.
393    /// `max_keys == 0` yields no entries but still reports `total_keys`.
394    pub fn snapshot(&self, now: u64, max_keys: usize) -> ProviderSnapshot {
395        let mut entries: Vec<ProviderSnapshotEntry> = self
396            .by_key
397            .iter()
398            .filter_map(|(content_key, providers)| {
399                let live = providers
400                    .values()
401                    .filter(|e| !e.record.is_expired(now))
402                    .count();
403                // A key whose every record has expired is not part of the view.
404                (live > 0).then(|| ProviderSnapshotEntry {
405                    content_key: content_key.clone(),
406                    providers: live,
407                })
408            })
409            .collect();
410
411        // Deterministic order so the same store yields the same snapshot, and so truncation takes a
412        // stable subset rather than an arbitrary one from HashMap iteration order.
413        entries.sort_by(|a, b| a.content_key.cmp(&b.content_key));
414
415        let total_keys = entries.len();
416        let truncated = total_keys > max_keys;
417        entries.truncate(max_keys);
418
419        ProviderSnapshot {
420            entries,
421            total_keys,
422            truncated,
423        }
424    }
425
426    /// Total live+stale records across all keys (diagnostics / tests).
427    pub fn len(&self) -> usize {
428        self.by_key.values().map(|p| p.len()).sum()
429    }
430
431    /// Whether the store holds no records.
432    pub fn is_empty(&self) -> bool {
433        self.len() == 0
434    }
435}
436
437#[cfg(test)]
438mod tests {
439    use super::*;
440    use crate::key::Key;
441    use crate::record::CandidateAddr;
442    use dig_nat::PeerId;
443
444    /// The instant the eviction tests reason at. Every `expires_at` they use is in the FUTURE
445    /// relative to this, so their records are LIVE and the assertions are about establishment —
446    /// not about a record that had silently already expired.
447    const NOW: u64 = 0;
448
449    fn rec(content: &Key, provider: u8, expires_at: u64) -> ProviderRecord {
450        ProviderRecord::new(
451            content,
452            &PeerId::from_bytes([provider; 32]),
453            vec![CandidateAddr::direct("h", 9444)],
454            expires_at,
455        )
456    }
457
458    // -- #1935: the aggregated snapshot the relay's /dht endpoint is built on -----------------
459
460    #[test]
461    fn snapshot_counts_live_providers_per_key_and_never_leaks_an_identity() {
462        // The privacy property is the point: a provider record IS (peer_id, content_key), which is
463        // exactly the linkage the relay's /map refuses to publish. The snapshot must carry counts.
464        let mut s = ProviderStore::new();
465        let k1 = Key::from_bytes([1u8; 32]);
466        let k2 = Key::from_bytes([2u8; 32]);
467        s.put(rec(&k1, 10, NOW + 100));
468        s.put(rec(&k1, 11, NOW + 100));
469        s.put(rec(&k2, 12, NOW + 100));
470
471        let snap = s.snapshot(NOW, 100);
472
473        assert_eq!(snap.total_keys, 2);
474        assert!(!snap.truncated);
475        let counts: Vec<usize> = snap.entries.iter().map(|e| e.providers).collect();
476        assert_eq!(counts, vec![2, 1], "two providers for k1, one for k2");
477
478        // Nothing in the snapshot may be a provider peer_id. Assert structurally rather than by
479        // string-matching, so the property cannot rot when a field is added.
480        let rendered = format!("{snap:?}");
481        for provider in [10u8, 11, 12] {
482            let pid = PeerId::from_bytes([provider; 32]).to_hex();
483            assert!(
484                !rendered.contains(&pid),
485                "provider identity {pid} must never appear in a snapshot"
486            );
487        }
488    }
489
490    #[test]
491    fn snapshot_excludes_expired_records_and_keys_left_with_none() {
492        // Must agree with `get`, which also filters on expiry — otherwise the relay would advertise
493        // providers the node would not actually return.
494        let mut s = ProviderStore::new();
495        let live = Key::from_bytes([1u8; 32]);
496        let dead = Key::from_bytes([2u8; 32]);
497        s.put(rec(&live, 10, NOW + 100));
498        s.put(rec(&dead, 11, NOW + 1));
499
500        let snap = s.snapshot(NOW + 50, 100);
501
502        assert_eq!(
503            snap.total_keys, 1,
504            "the fully-expired key drops out entirely"
505        );
506        assert_eq!(snap.entries[0].providers, 1);
507        assert_eq!(
508            snap.entries[0].content_key,
509            live.to_hex(),
510            "the surviving key is the live one"
511        );
512    }
513
514    #[test]
515    fn snapshot_is_bounded_and_reports_the_true_total_when_truncated() {
516        // The store is attacker-influenced — any peer can announce — so an unbounded snapshot would
517        // let a Sybil dictate the response size. Truncation must be VISIBLE, not silent.
518        let mut s = ProviderStore::new();
519        for i in 0..10u8 {
520            s.put(rec(&Key::from_bytes([i; 32]), 100 + i, NOW + 100));
521        }
522
523        let snap = s.snapshot(NOW, 3);
524
525        assert_eq!(snap.entries.len(), 3);
526        assert!(snap.truncated);
527        assert_eq!(snap.total_keys, 10, "the true total survives truncation");
528    }
529
530    #[test]
531    fn snapshot_is_deterministic_so_truncation_takes_a_stable_subset() {
532        // HashMap iteration order is arbitrary; without sorting, two calls could return different
533        // subsets and a consumer polling the relay would see content flicker in and out.
534        let mut s = ProviderStore::new();
535        for i in 0..8u8 {
536            s.put(rec(&Key::from_bytes([i; 32]), 100 + i, NOW + 100));
537        }
538        assert_eq!(s.snapshot(NOW, 4), s.snapshot(NOW, 4));
539    }
540
541    #[test]
542    fn a_zero_cap_yields_no_entries_but_still_reports_the_total() {
543        let mut s = ProviderStore::new();
544        s.put(rec(&Key::from_bytes([1u8; 32]), 10, NOW + 100));
545        let snap = s.snapshot(NOW, 0);
546        assert!(snap.entries.is_empty());
547        assert!(snap.truncated);
548        assert_eq!(snap.total_keys, 1);
549    }
550
551    #[test]
552    fn put_then_get_returns_live_record() {
553        let mut s = ProviderStore::new();
554        let key = Key::from_bytes([0xAA; 32]);
555        s.put(rec(&key, 1, 100));
556        let got = s.get(&key.to_hex(), 50);
557        assert_eq!(got.len(), 1);
558        assert_eq!(
559            got[0].provider_peer_id,
560            PeerId::from_bytes([1u8; 32]).to_hex()
561        );
562    }
563
564    #[test]
565    fn get_hides_expired_records() {
566        let mut s = ProviderStore::new();
567        let key = Key::from_bytes([0xAA; 32]);
568        s.put(rec(&key, 1, 100));
569        assert!(
570            s.get(&key.to_hex(), 100).is_empty(),
571            "expired at exactly TTL"
572        );
573        assert!(s.get(&key.to_hex(), 200).is_empty());
574    }
575
576    #[test]
577    fn same_provider_dedups_and_refreshes() {
578        let mut s = ProviderStore::new();
579        let key = Key::from_bytes([0xAA; 32]);
580        s.put(rec(&key, 1, 100));
581        s.put(rec(&key, 1, 500)); // same provider, later expiry
582        assert_eq!(s.len(), 1, "same provider must not duplicate");
583        // The refreshed expiry wins.
584        assert_eq!(s.get(&key.to_hex(), 300).len(), 1);
585    }
586
587    #[test]
588    fn distinct_providers_for_same_key_coexist() {
589        let mut s = ProviderStore::new();
590        let key = Key::from_bytes([0xAA; 32]);
591        s.put(rec(&key, 1, 100));
592        s.put(rec(&key, 2, 100));
593        assert_eq!(s.get(&key.to_hex(), 50).len(), 2);
594    }
595
596    // ---- Admission control (HIGH #1: unbounded provider store, SECURITY_AUDIT_P2P.md #179) ----
597
598    #[test]
599    fn put_returns_accepted_under_capacity() {
600        let mut s = ProviderStore::new();
601        let key = Key::from_bytes([0xAA; 32]);
602        assert_eq!(s.put(rec(&key, 1, 100)), PutOutcome::Accepted);
603    }
604
605    #[test]
606    fn refreshing_same_provider_always_succeeds_even_at_per_key_cap() {
607        // A refresh (same provider, same key) never counts as "new" so it must never be blocked by
608        // the per-key cap even when the key is already full.
609        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
610            max_providers_per_key: 1,
611            max_total_records: 1000,
612        });
613        let key = Key::from_bytes([0xAA; 32]);
614        assert_eq!(s.put(rec(&key, 1, 100)), PutOutcome::Accepted);
615        assert_eq!(s.put(rec(&key, 1, 999)), PutOutcome::Accepted, "refresh");
616        assert_eq!(s.len(), 1);
617    }
618
619    #[test]
620    fn per_key_cap_evicts_soonest_to_expire_within_the_churn_zone() {
621        // One malicious/heavy peer announcing many DISTINCT providers for the SAME content key must
622        // not grow that key's provider set past `max_providers_per_key` — the audit's "no cap on
623        // providers-per-key" finding.
624        // Cap 4 → the two longest-established slots are reserved (#1434), so the eviction choice
625        // is made among the two newest — the churn zone. Within that zone the soonest-to-expire
626        // record is still the least valuable one to keep.
627        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
628            max_providers_per_key: 4,
629            max_total_records: 1000,
630        });
631        let key = Key::from_bytes([0xAA; 32]);
632        assert_eq!(s.put_at(rec(&key, 1, 100), NOW), PutOutcome::Accepted); // established
633        assert_eq!(s.put_at(rec(&key, 2, 200), NOW), PutOutcome::Accepted); // established
634        assert_eq!(s.put_at(rec(&key, 3, 900), NOW), PutOutcome::Accepted); // churn zone
635        assert_eq!(s.put_at(rec(&key, 4, 800), NOW), PutOutcome::Accepted); // churn zone, expires sooner
636        assert_eq!(s.put_at(rec(&key, 5, 999), NOW), PutOutcome::Accepted);
637        assert_eq!(
638            s.get(&key.to_hex(), 0).len(),
639            4,
640            "per-key cap must not be exceeded"
641        );
642        assert!(
643            !live_provider_ids(&s, &key).contains(&PeerId::from_bytes([4u8; 32]).to_hex()),
644            "the soonest-to-expire record in the churn zone must be the one evicted"
645        );
646    }
647
648    /// The live provider peer_ids for `key` (order-independent membership assertions).
649    fn live_provider_ids(s: &ProviderStore, key: &Key) -> std::collections::HashSet<String> {
650        s.get(&key.to_hex(), 0)
651            .into_iter()
652            .map(|r| r.provider_peer_id)
653            .collect()
654    }
655
656    // ---- Sybil-resistant eviction (#1434) ----
657
658    #[test]
659    fn sustained_sybil_flood_cannot_evict_the_lone_established_holder() {
660        // #1434: every record clamps its expiry to `now + provider_ttl` at put time, so an attacker
661        // who announces LATER always holds a strictly-later `expires_at` than an honest incumbent.
662        // Under pure soonest-to-expire eviction that made the honest holder the deterministic
663        // victim, and 20 Sybil identities could make the only real holder of a capsule
664        // undiscoverable at this node — content-discovery censorship. Stated over the CLASS: no
665        // volume of later-expiring newcomers may evict a provider inside the established floor.
666        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
667            max_providers_per_key: 20,
668            max_total_records: 100_000,
669        });
670        let key = Key::from_bytes([0xAA; 32]);
671        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
672        assert_eq!(s.put_at(rec(&key, 1, 100), NOW), PutOutcome::Accepted);
673
674        // A sustained flood of distinct Sybil providers, each expiring strictly later than the last
675        // — the worst case for expiry-ordered eviction.
676        for i in 0..500u64 {
677            let sybil = ProviderRecord::new(
678                &key,
679                &PeerId::from_bytes(sybil_id(i)),
680                vec![CandidateAddr::direct("h", 9444)],
681                1_000 + i,
682            );
683            s.put_at(sybil, NOW);
684        }
685
686        assert!(
687            live_provider_ids(&s, &key).contains(&honest),
688            "the lone honest holder must survive a sustained Sybil flood"
689        );
690        assert_eq!(
691            s.get(&key.to_hex(), 0).len(),
692            20,
693            "the per-key cap still bounds the set"
694        );
695    }
696
697    #[test]
698    fn established_floor_protects_the_earliest_admitted_providers() {
699        // The one-off variant: exactly one provider beyond the cap. Eviction must fall inside the
700        // churn zone and never touch the reserved, longest-established slots.
701        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
702            max_providers_per_key: 4,
703            max_total_records: 1000,
704        });
705        let key = Key::from_bytes([0xAA; 32]);
706        // Established slots deliberately hold the SOONEST expiries — under the old policy they
707        // would have been evicted first.
708        s.put_at(rec(&key, 1, 10), NOW);
709        s.put_at(rec(&key, 2, 20), NOW);
710        s.put_at(rec(&key, 3, 900), NOW);
711        s.put_at(rec(&key, 4, 800), NOW);
712        s.put_at(rec(&key, 5, 999), NOW);
713
714        let live = live_provider_ids(&s, &key);
715        assert!(
716            live.contains(&PeerId::from_bytes([1u8; 32]).to_hex()),
717            "the first-admitted provider is inside the established floor"
718        );
719        assert!(
720            live.contains(&PeerId::from_bytes([2u8; 32]).to_hex()),
721            "the second-admitted provider is inside the established floor"
722        );
723    }
724
725    #[test]
726    fn republish_does_not_reset_a_holders_establishment() {
727        // A holder stays findable by republishing before its TTL elapses. If a refresh reset the
728        // record's establishment, republishing — the very act that keeps an honest holder alive —
729        // would drop it into the churn zone and hand the attacker the eviction it wanted.
730        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
731            max_providers_per_key: 4,
732            max_total_records: 1000,
733        });
734        let key = Key::from_bytes([0xAA; 32]);
735        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
736        s.put_at(rec(&key, 1, 100), NOW);
737        for i in 0..3u64 {
738            s.put_at(rec(&key, 10 + i as u8, 500 + i), NOW);
739        }
740        s.put_at(rec(&key, 1, 5_000), NOW); // the honest holder republishes
741        for i in 0..50u64 {
742            s.put_at(
743                ProviderRecord::new(
744                    &key,
745                    &PeerId::from_bytes(sybil_id(i)),
746                    vec![CandidateAddr::direct("h", 9444)],
747                    9_000 + i,
748                ),
749                NOW,
750            );
751        }
752        assert!(
753            live_provider_ids(&s, &key).contains(&honest),
754            "a republished record keeps its establishment"
755        );
756    }
757
758    // ---- Liveness outranks establishment (#1434 follow-up) ----
759
760    #[test]
761    fn an_expired_record_in_the_floor_is_evicted_before_a_live_one() {
762        // The pre-#1434 policy evicted the soonest-to-expire record, so an EXPIRED record was always
763        // the first victim. The establishment floor must not invert that: a dead record inside the
764        // reserved floor cannot outrank a live provider in the churn zone. Without a liveness check
765        // this needs NO attacker — a node's GC tick is coarser than the provider TTL, so whenever the
766        // earliest-admitted half of a key goes offline, every new announcement for that key evicts a
767        // LIVE holder and announcing more holders makes the capsule LESS discoverable.
768        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
769            max_providers_per_key: 4,
770            max_total_records: 1000,
771        });
772        let key = Key::from_bytes([0xAA; 32]);
773        let now = 10_000;
774        // The reserved floor (seq 0, 1) is long expired...
775        s.put_at(rec(&key, 1, 100), now);
776        s.put_at(rec(&key, 2, 200), now);
777        // ...while the churn zone (seq 2, 3) holds two LIVE honest providers.
778        s.put_at(rec(&key, 3, now + 5_000), now);
779        s.put_at(rec(&key, 4, now + 6_000), now);
780
781        s.put_at(rec(&key, 5, now + 7_000), now);
782
783        let live = live_provider_ids_at(&s, &key, now);
784        assert!(
785            live.contains(&PeerId::from_bytes([3u8; 32]).to_hex())
786                && live.contains(&PeerId::from_bytes([4u8; 32]).to_hex()),
787            "both LIVE providers must survive; an expired record in the floor is the victim"
788        );
789    }
790
791    #[test]
792    fn one_expired_record_anywhere_is_the_victim_before_any_live_record() {
793        // The one-off variant: exactly ONE expired record, sitting inside the reserved floor, with
794        // every other slot live. It must still be the one evicted.
795        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
796            max_providers_per_key: 4,
797            max_total_records: 1000,
798        });
799        let key = Key::from_bytes([0xAA; 32]);
800        let now = 10_000;
801        s.put_at(rec(&key, 1, 100), now); // expired, seq 0 → inside the floor
802        s.put_at(rec(&key, 2, now + 1_000), now);
803        s.put_at(rec(&key, 3, now + 2_000), now);
804        s.put_at(rec(&key, 4, now + 3_000), now);
805
806        s.put_at(rec(&key, 5, now + 4_000), now);
807
808        assert_eq!(
809            live_provider_ids_at(&s, &key, now).len(),
810            4,
811            "reclaiming the dead slot leaves every live provider intact"
812        );
813    }
814
815    #[test]
816    fn the_floor_still_protects_an_established_holder_when_every_record_is_live() {
817        // Liveness must take precedence WITHOUT weakening #1434: with no dead slot to reclaim, the
818        // establishment floor governs again and a sustained flood cannot displace the incumbent.
819        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
820            max_providers_per_key: 20,
821            max_total_records: 100_000,
822        });
823        let key = Key::from_bytes([0xAA; 32]);
824        let now = 10_000;
825        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
826        s.put_at(rec(&key, 1, now + 1_000), now);
827        for i in 0..500u64 {
828            s.put_at(
829                ProviderRecord::new(
830                    &key,
831                    &PeerId::from_bytes(sybil_id(i)),
832                    vec![CandidateAddr::direct("h", 9444)],
833                    now + 2_000 + i,
834                ),
835                now,
836            );
837        }
838        assert!(
839            live_provider_ids_at(&s, &key, now).contains(&honest),
840            "an all-live key keeps the #1434 protection"
841        );
842    }
843
844    #[test]
845    fn put_delegates_to_put_at_with_the_wall_clock() {
846        // `put` is the compatibility wrapper (its signature is public API): same admission decision,
847        // with `now` read from the system clock.
848        let mut wall = ProviderStore::new();
849        let key = Key::from_bytes([0xAA; 32]);
850        assert_eq!(wall.put(rec(&key, 1, u64::MAX)), PutOutcome::Accepted);
851        assert_eq!(wall.len(), 1);
852    }
853
854    /// The live provider peer_ids for `key` as of `now`.
855    fn live_provider_ids_at(
856        s: &ProviderStore,
857        key: &Key,
858        now: u64,
859    ) -> std::collections::HashSet<String> {
860        s.get(&key.to_hex(), now)
861            .into_iter()
862            .map(|r| r.provider_peer_id)
863            .collect()
864    }
865
866    /// A distinct Sybil peer_id per index (varying the high bytes so ids stay distinct past 255).
867    fn sybil_id(i: u64) -> [u8; 32] {
868        let mut b = [0xEE; 32];
869        b[0..8].copy_from_slice(&i.to_be_bytes());
870        b
871    }
872
873    #[test]
874    fn global_cap_rejects_new_content_keys_over_ceiling() {
875        // Many DISTINCT content keys (not just many providers per key) must also be bounded — the
876        // audit's "no cap on distinct content keys ... no global record ceiling" finding.
877        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
878            max_providers_per_key: 20,
879            max_total_records: 2,
880        });
881        let k1 = Key::from_bytes([0x01; 32]);
882        let k2 = Key::from_bytes([0x02; 32]);
883        let k3 = Key::from_bytes([0x03; 32]);
884        assert_eq!(s.put(rec(&k1, 1, 100)), PutOutcome::Accepted);
885        assert_eq!(s.put(rec(&k2, 1, 100)), PutOutcome::Accepted);
886        assert_eq!(
887            s.put(rec(&k3, 1, 100)),
888            PutOutcome::RejectedOverCapacity,
889            "third distinct record must be rejected once the global ceiling is hit"
890        );
891        assert_eq!(s.len(), 2, "rejected record must not be stored");
892        assert!(
893            s.get(&k3.to_hex(), 0).is_empty(),
894            "rejected key must not appear in the store at all"
895        );
896    }
897
898    #[test]
899    fn global_cap_does_not_evict_a_different_key_to_make_room() {
900        // A single attacker flooding new keys must not be able to evict a DIFFERENT (legitimate)
901        // key's providers just by hitting the global ceiling.
902        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
903            max_providers_per_key: 20,
904            max_total_records: 1,
905        });
906        let legit = Key::from_bytes([0xAA; 32]);
907        s.put(rec(&legit, 1, 100));
908        let attacker_key = Key::from_bytes([0xBB; 32]);
909        assert_eq!(
910            s.put(rec(&attacker_key, 2, 100)),
911            PutOutcome::RejectedOverCapacity
912        );
913        assert_eq!(
914            s.get(&legit.to_hex(), 0).len(),
915            1,
916            "the legitimate key's record must survive"
917        );
918    }
919
920    #[test]
921    fn remove_deletes_only_the_named_provider_record() {
922        // Authenticated retract (SPEC §6.6): removing (key, provider-1) must leave provider-2 of the
923        // SAME key untouched — a retract signed by one holder cannot censor another holder.
924        let mut s = ProviderStore::new();
925        let key = Key::from_bytes([0xAA; 32]);
926        s.put(rec(&key, 1, 100));
927        s.put(rec(&key, 2, 100));
928        let pid1 = PeerId::from_bytes([1u8; 32]).to_hex();
929        let pid2 = PeerId::from_bytes([2u8; 32]).to_hex();
930        assert!(
931            s.remove(&key.to_hex(), &pid1),
932            "the named record was removed"
933        );
934        let survivors: std::collections::HashSet<String> = s
935            .get(&key.to_hex(), 0)
936            .into_iter()
937            .map(|r| r.provider_peer_id)
938            .collect();
939        assert_eq!(survivors.len(), 1, "the other provider must survive");
940        assert!(survivors.contains(&pid2));
941        assert!(!survivors.contains(&pid1));
942    }
943
944    #[test]
945    fn remove_of_absent_record_returns_false() {
946        let mut s = ProviderStore::new();
947        let key = Key::from_bytes([0xAA; 32]);
948        s.put(rec(&key, 1, 100));
949        let absent = PeerId::from_bytes([9u8; 32]).to_hex();
950        assert!(!s.remove(&key.to_hex(), &absent), "no such provider");
951        assert!(!s.remove(&"00".repeat(32), &absent), "no such content key");
952        assert_eq!(s.len(), 1, "nothing removed");
953    }
954
955    #[test]
956    fn remove_drops_content_key_when_last_provider_leaves() {
957        let mut s = ProviderStore::new();
958        let key = Key::from_bytes([0xAA; 32]);
959        s.put(rec(&key, 1, 100));
960        let pid1 = PeerId::from_bytes([1u8; 32]).to_hex();
961        assert!(s.remove(&key.to_hex(), &pid1));
962        assert!(
963            s.is_empty(),
964            "the now-empty content key must be dropped entirely"
965        );
966    }
967
968    #[test]
969    fn gc_removes_expired_and_empty_keys() {
970        let mut s = ProviderStore::new();
971        let k1 = Key::from_bytes([0x01; 32]);
972        let k2 = Key::from_bytes([0x02; 32]);
973        s.put(rec(&k1, 1, 100)); // expires at 100
974        s.put(rec(&k2, 1, 500)); // expires at 500
975        let removed = s.gc(200);
976        assert_eq!(removed, 1);
977        assert!(s.get(&k1.to_hex(), 200).is_empty());
978        assert_eq!(s.get(&k2.to_hex(), 200).len(), 1);
979    }
980
981    #[test]
982    fn announcements_track_and_untrack() {
983        let mut s = ProviderStore::new();
984        let key = Key::from_bytes([0x07; 32]).to_hex();
985        s.mark_announced(key.clone());
986        s.mark_announced(key.clone()); // idempotent
987        assert_eq!(s.local_announcements(), vec![key.clone()]);
988        assert!(s.unmark_announced(&key));
989        assert!(!s.unmark_announced(&key));
990        assert!(s.local_announcements().is_empty());
991    }
992}