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    /// The live (non-expired at `now`) provider records for `content_key`. Expired records are
311    /// skipped (and cleaned up by [`gc`](Self::gc)); returns an empty vec if none are known/live.
312    pub fn get(&self, content_key: &str, now: u64) -> Vec<ProviderRecord> {
313        self.by_key
314            .get(content_key)
315            .map(|providers| {
316                providers
317                    .values()
318                    .map(|e| &e.record)
319                    .filter(|r| !r.is_expired(now))
320                    .cloned()
321                    .collect()
322            })
323            .unwrap_or_default()
324    }
325
326    /// Drop every expired record (and any content key left with no live providers) as of `now`.
327    /// Returns the number of records removed. Call periodically from the maintenance loop.
328    pub fn gc(&mut self, now: u64) -> usize {
329        let mut removed = 0;
330        self.by_key.retain(|_key, providers| {
331            let before = providers.len();
332            providers.retain(|_pid, e| !e.record.is_expired(now));
333            removed += before - providers.len();
334            !providers.is_empty()
335        });
336        removed
337    }
338
339    /// Record that this node holds + announces `content_key` (so the maintenance loop republishes
340    /// it). Idempotent.
341    pub fn mark_announced(&mut self, content_key: String) {
342        self.announced.insert(content_key);
343    }
344
345    /// Stop announcing `content_key` (this node no longer holds the content). Returns whether it was
346    /// being announced.
347    pub fn unmark_announced(&mut self, content_key: &str) -> bool {
348        self.announced.remove(content_key)
349    }
350
351    /// The content keys this node announces (holds) — the republish work list.
352    pub fn local_announcements(&self) -> Vec<String> {
353        self.announced.iter().cloned().collect()
354    }
355
356    /// A bounded, AGGREGATED view of what this node holds in its DHT provider store — content keys
357    /// and how many live providers each has, with no provider identities (dig_ecosystem #1935).
358    ///
359    /// This is what lets the relay show the network's content layer without joining the DHT: a
360    /// Kademlia node stores records for keys near its OWN `peer_id`, so these are records about
361    /// MANY OTHER peers' content, not a self-report of what this node caches. The union across
362    /// several nodes is a broad slice of the real DHT.
363    ///
364    /// # Why counts and not identities
365    ///
366    /// A provider record IS a `(peer_id, content_key)` pair — exactly the linkage the relay's `/map`
367    /// refuses to publish (its tests assert no `peer_id` and no raw IP ever appear). Returning
368    /// counts keeps that contract intact rather than carving an exception into it. A caller that
369    /// genuinely needs identities can still use [`get`](Self::get) per key.
370    ///
371    /// Expired records are excluded as of `now`, so the counts match what [`get`](Self::get) would
372    /// return rather than including records the store has not GC'd yet.
373    ///
374    /// `max_keys` bounds the result: the store is attacker-influenced (any peer can announce), so an
375    /// unbounded snapshot would let a Sybil dictate the response size. When the cap truncates,
376    /// [`ProviderSnapshot::truncated`] is set and `total_keys` still reports the true total, so a
377    /// consumer can say "showing N of M" instead of silently presenting a partial view as complete.
378    /// `max_keys == 0` yields no entries but still reports `total_keys`.
379    pub fn snapshot(&self, now: u64, max_keys: usize) -> ProviderSnapshot {
380        let mut entries: Vec<ProviderSnapshotEntry> = self
381            .by_key
382            .iter()
383            .filter_map(|(content_key, providers)| {
384                let live = providers
385                    .values()
386                    .filter(|e| !e.record.is_expired(now))
387                    .count();
388                // A key whose every record has expired is not part of the view.
389                (live > 0).then(|| ProviderSnapshotEntry {
390                    content_key: content_key.clone(),
391                    providers: live,
392                })
393            })
394            .collect();
395
396        // Deterministic order so the same store yields the same snapshot, and so truncation takes a
397        // stable subset rather than an arbitrary one from HashMap iteration order.
398        entries.sort_by(|a, b| a.content_key.cmp(&b.content_key));
399
400        let total_keys = entries.len();
401        let truncated = total_keys > max_keys;
402        entries.truncate(max_keys);
403
404        ProviderSnapshot {
405            entries,
406            total_keys,
407            truncated,
408        }
409    }
410
411    /// Total live+stale records across all keys (diagnostics / tests).
412    pub fn len(&self) -> usize {
413        self.by_key.values().map(|p| p.len()).sum()
414    }
415
416    /// Whether the store holds no records.
417    pub fn is_empty(&self) -> bool {
418        self.len() == 0
419    }
420}
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425    use crate::key::Key;
426    use crate::record::CandidateAddr;
427    use dig_nat::PeerId;
428
429    /// The instant the eviction tests reason at. Every `expires_at` they use is in the FUTURE
430    /// relative to this, so their records are LIVE and the assertions are about establishment —
431    /// not about a record that had silently already expired.
432    const NOW: u64 = 0;
433
434    fn rec(content: &Key, provider: u8, expires_at: u64) -> ProviderRecord {
435        ProviderRecord::new(
436            content,
437            &PeerId::from_bytes([provider; 32]),
438            vec![CandidateAddr::direct("h", 9444)],
439            expires_at,
440        )
441    }
442
443    // -- #1935: the aggregated snapshot the relay's /dht endpoint is built on -----------------
444
445    #[test]
446    fn snapshot_counts_live_providers_per_key_and_never_leaks_an_identity() {
447        // The privacy property is the point: a provider record IS (peer_id, content_key), which is
448        // exactly the linkage the relay's /map refuses to publish. The snapshot must carry counts.
449        let mut s = ProviderStore::new();
450        let k1 = Key::from_bytes([1u8; 32]);
451        let k2 = Key::from_bytes([2u8; 32]);
452        s.put(rec(&k1, 10, NOW + 100));
453        s.put(rec(&k1, 11, NOW + 100));
454        s.put(rec(&k2, 12, NOW + 100));
455
456        let snap = s.snapshot(NOW, 100);
457
458        assert_eq!(snap.total_keys, 2);
459        assert!(!snap.truncated);
460        let counts: Vec<usize> = snap.entries.iter().map(|e| e.providers).collect();
461        assert_eq!(counts, vec![2, 1], "two providers for k1, one for k2");
462
463        // Nothing in the snapshot may be a provider peer_id. Assert structurally rather than by
464        // string-matching, so the property cannot rot when a field is added.
465        let rendered = format!("{snap:?}");
466        for provider in [10u8, 11, 12] {
467            let pid = PeerId::from_bytes([provider; 32]).to_hex();
468            assert!(
469                !rendered.contains(&pid),
470                "provider identity {pid} must never appear in a snapshot"
471            );
472        }
473    }
474
475    #[test]
476    fn snapshot_excludes_expired_records_and_keys_left_with_none() {
477        // Must agree with `get`, which also filters on expiry — otherwise the relay would advertise
478        // providers the node would not actually return.
479        let mut s = ProviderStore::new();
480        let live = Key::from_bytes([1u8; 32]);
481        let dead = Key::from_bytes([2u8; 32]);
482        s.put(rec(&live, 10, NOW + 100));
483        s.put(rec(&dead, 11, NOW + 1));
484
485        let snap = s.snapshot(NOW + 50, 100);
486
487        assert_eq!(
488            snap.total_keys, 1,
489            "the fully-expired key drops out entirely"
490        );
491        assert_eq!(snap.entries[0].providers, 1);
492        assert_eq!(
493            snap.entries[0].content_key,
494            live.to_hex(),
495            "the surviving key is the live one"
496        );
497    }
498
499    #[test]
500    fn snapshot_is_bounded_and_reports_the_true_total_when_truncated() {
501        // The store is attacker-influenced — any peer can announce — so an unbounded snapshot would
502        // let a Sybil dictate the response size. Truncation must be VISIBLE, not silent.
503        let mut s = ProviderStore::new();
504        for i in 0..10u8 {
505            s.put(rec(&Key::from_bytes([i; 32]), 100 + i, NOW + 100));
506        }
507
508        let snap = s.snapshot(NOW, 3);
509
510        assert_eq!(snap.entries.len(), 3);
511        assert!(snap.truncated);
512        assert_eq!(snap.total_keys, 10, "the true total survives truncation");
513    }
514
515    #[test]
516    fn snapshot_is_deterministic_so_truncation_takes_a_stable_subset() {
517        // HashMap iteration order is arbitrary; without sorting, two calls could return different
518        // subsets and a consumer polling the relay would see content flicker in and out.
519        let mut s = ProviderStore::new();
520        for i in 0..8u8 {
521            s.put(rec(&Key::from_bytes([i; 32]), 100 + i, NOW + 100));
522        }
523        assert_eq!(s.snapshot(NOW, 4), s.snapshot(NOW, 4));
524    }
525
526    #[test]
527    fn a_zero_cap_yields_no_entries_but_still_reports_the_total() {
528        let mut s = ProviderStore::new();
529        s.put(rec(&Key::from_bytes([1u8; 32]), 10, NOW + 100));
530        let snap = s.snapshot(NOW, 0);
531        assert!(snap.entries.is_empty());
532        assert!(snap.truncated);
533        assert_eq!(snap.total_keys, 1);
534    }
535
536    #[test]
537    fn put_then_get_returns_live_record() {
538        let mut s = ProviderStore::new();
539        let key = Key::from_bytes([0xAA; 32]);
540        s.put(rec(&key, 1, 100));
541        let got = s.get(&key.to_hex(), 50);
542        assert_eq!(got.len(), 1);
543        assert_eq!(
544            got[0].provider_peer_id,
545            PeerId::from_bytes([1u8; 32]).to_hex()
546        );
547    }
548
549    #[test]
550    fn get_hides_expired_records() {
551        let mut s = ProviderStore::new();
552        let key = Key::from_bytes([0xAA; 32]);
553        s.put(rec(&key, 1, 100));
554        assert!(
555            s.get(&key.to_hex(), 100).is_empty(),
556            "expired at exactly TTL"
557        );
558        assert!(s.get(&key.to_hex(), 200).is_empty());
559    }
560
561    #[test]
562    fn same_provider_dedups_and_refreshes() {
563        let mut s = ProviderStore::new();
564        let key = Key::from_bytes([0xAA; 32]);
565        s.put(rec(&key, 1, 100));
566        s.put(rec(&key, 1, 500)); // same provider, later expiry
567        assert_eq!(s.len(), 1, "same provider must not duplicate");
568        // The refreshed expiry wins.
569        assert_eq!(s.get(&key.to_hex(), 300).len(), 1);
570    }
571
572    #[test]
573    fn distinct_providers_for_same_key_coexist() {
574        let mut s = ProviderStore::new();
575        let key = Key::from_bytes([0xAA; 32]);
576        s.put(rec(&key, 1, 100));
577        s.put(rec(&key, 2, 100));
578        assert_eq!(s.get(&key.to_hex(), 50).len(), 2);
579    }
580
581    // ---- Admission control (HIGH #1: unbounded provider store, SECURITY_AUDIT_P2P.md #179) ----
582
583    #[test]
584    fn put_returns_accepted_under_capacity() {
585        let mut s = ProviderStore::new();
586        let key = Key::from_bytes([0xAA; 32]);
587        assert_eq!(s.put(rec(&key, 1, 100)), PutOutcome::Accepted);
588    }
589
590    #[test]
591    fn refreshing_same_provider_always_succeeds_even_at_per_key_cap() {
592        // A refresh (same provider, same key) never counts as "new" so it must never be blocked by
593        // the per-key cap even when the key is already full.
594        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
595            max_providers_per_key: 1,
596            max_total_records: 1000,
597        });
598        let key = Key::from_bytes([0xAA; 32]);
599        assert_eq!(s.put(rec(&key, 1, 100)), PutOutcome::Accepted);
600        assert_eq!(s.put(rec(&key, 1, 999)), PutOutcome::Accepted, "refresh");
601        assert_eq!(s.len(), 1);
602    }
603
604    #[test]
605    fn per_key_cap_evicts_soonest_to_expire_within_the_churn_zone() {
606        // One malicious/heavy peer announcing many DISTINCT providers for the SAME content key must
607        // not grow that key's provider set past `max_providers_per_key` — the audit's "no cap on
608        // providers-per-key" finding.
609        // Cap 4 → the two longest-established slots are reserved (#1434), so the eviction choice
610        // is made among the two newest — the churn zone. Within that zone the soonest-to-expire
611        // record is still the least valuable one to keep.
612        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
613            max_providers_per_key: 4,
614            max_total_records: 1000,
615        });
616        let key = Key::from_bytes([0xAA; 32]);
617        assert_eq!(s.put_at(rec(&key, 1, 100), NOW), PutOutcome::Accepted); // established
618        assert_eq!(s.put_at(rec(&key, 2, 200), NOW), PutOutcome::Accepted); // established
619        assert_eq!(s.put_at(rec(&key, 3, 900), NOW), PutOutcome::Accepted); // churn zone
620        assert_eq!(s.put_at(rec(&key, 4, 800), NOW), PutOutcome::Accepted); // churn zone, expires sooner
621        assert_eq!(s.put_at(rec(&key, 5, 999), NOW), PutOutcome::Accepted);
622        assert_eq!(
623            s.get(&key.to_hex(), 0).len(),
624            4,
625            "per-key cap must not be exceeded"
626        );
627        assert!(
628            !live_provider_ids(&s, &key).contains(&PeerId::from_bytes([4u8; 32]).to_hex()),
629            "the soonest-to-expire record in the churn zone must be the one evicted"
630        );
631    }
632
633    /// The live provider peer_ids for `key` (order-independent membership assertions).
634    fn live_provider_ids(s: &ProviderStore, key: &Key) -> std::collections::HashSet<String> {
635        s.get(&key.to_hex(), 0)
636            .into_iter()
637            .map(|r| r.provider_peer_id)
638            .collect()
639    }
640
641    // ---- Sybil-resistant eviction (#1434) ----
642
643    #[test]
644    fn sustained_sybil_flood_cannot_evict_the_lone_established_holder() {
645        // #1434: every record clamps its expiry to `now + provider_ttl` at put time, so an attacker
646        // who announces LATER always holds a strictly-later `expires_at` than an honest incumbent.
647        // Under pure soonest-to-expire eviction that made the honest holder the deterministic
648        // victim, and 20 Sybil identities could make the only real holder of a capsule
649        // undiscoverable at this node — content-discovery censorship. Stated over the CLASS: no
650        // volume of later-expiring newcomers may evict a provider inside the established floor.
651        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
652            max_providers_per_key: 20,
653            max_total_records: 100_000,
654        });
655        let key = Key::from_bytes([0xAA; 32]);
656        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
657        assert_eq!(s.put_at(rec(&key, 1, 100), NOW), PutOutcome::Accepted);
658
659        // A sustained flood of distinct Sybil providers, each expiring strictly later than the last
660        // — the worst case for expiry-ordered eviction.
661        for i in 0..500u64 {
662            let sybil = ProviderRecord::new(
663                &key,
664                &PeerId::from_bytes(sybil_id(i)),
665                vec![CandidateAddr::direct("h", 9444)],
666                1_000 + i,
667            );
668            s.put_at(sybil, NOW);
669        }
670
671        assert!(
672            live_provider_ids(&s, &key).contains(&honest),
673            "the lone honest holder must survive a sustained Sybil flood"
674        );
675        assert_eq!(
676            s.get(&key.to_hex(), 0).len(),
677            20,
678            "the per-key cap still bounds the set"
679        );
680    }
681
682    #[test]
683    fn established_floor_protects_the_earliest_admitted_providers() {
684        // The one-off variant: exactly one provider beyond the cap. Eviction must fall inside the
685        // churn zone and never touch the reserved, longest-established slots.
686        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
687            max_providers_per_key: 4,
688            max_total_records: 1000,
689        });
690        let key = Key::from_bytes([0xAA; 32]);
691        // Established slots deliberately hold the SOONEST expiries — under the old policy they
692        // would have been evicted first.
693        s.put_at(rec(&key, 1, 10), NOW);
694        s.put_at(rec(&key, 2, 20), NOW);
695        s.put_at(rec(&key, 3, 900), NOW);
696        s.put_at(rec(&key, 4, 800), NOW);
697        s.put_at(rec(&key, 5, 999), NOW);
698
699        let live = live_provider_ids(&s, &key);
700        assert!(
701            live.contains(&PeerId::from_bytes([1u8; 32]).to_hex()),
702            "the first-admitted provider is inside the established floor"
703        );
704        assert!(
705            live.contains(&PeerId::from_bytes([2u8; 32]).to_hex()),
706            "the second-admitted provider is inside the established floor"
707        );
708    }
709
710    #[test]
711    fn republish_does_not_reset_a_holders_establishment() {
712        // A holder stays findable by republishing before its TTL elapses. If a refresh reset the
713        // record's establishment, republishing — the very act that keeps an honest holder alive —
714        // would drop it into the churn zone and hand the attacker the eviction it wanted.
715        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
716            max_providers_per_key: 4,
717            max_total_records: 1000,
718        });
719        let key = Key::from_bytes([0xAA; 32]);
720        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
721        s.put_at(rec(&key, 1, 100), NOW);
722        for i in 0..3u64 {
723            s.put_at(rec(&key, 10 + i as u8, 500 + i), NOW);
724        }
725        s.put_at(rec(&key, 1, 5_000), NOW); // the honest holder republishes
726        for i in 0..50u64 {
727            s.put_at(
728                ProviderRecord::new(
729                    &key,
730                    &PeerId::from_bytes(sybil_id(i)),
731                    vec![CandidateAddr::direct("h", 9444)],
732                    9_000 + i,
733                ),
734                NOW,
735            );
736        }
737        assert!(
738            live_provider_ids(&s, &key).contains(&honest),
739            "a republished record keeps its establishment"
740        );
741    }
742
743    // ---- Liveness outranks establishment (#1434 follow-up) ----
744
745    #[test]
746    fn an_expired_record_in_the_floor_is_evicted_before_a_live_one() {
747        // The pre-#1434 policy evicted the soonest-to-expire record, so an EXPIRED record was always
748        // the first victim. The establishment floor must not invert that: a dead record inside the
749        // reserved floor cannot outrank a live provider in the churn zone. Without a liveness check
750        // this needs NO attacker — a node's GC tick is coarser than the provider TTL, so whenever the
751        // earliest-admitted half of a key goes offline, every new announcement for that key evicts a
752        // LIVE holder and announcing more holders makes the capsule LESS discoverable.
753        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
754            max_providers_per_key: 4,
755            max_total_records: 1000,
756        });
757        let key = Key::from_bytes([0xAA; 32]);
758        let now = 10_000;
759        // The reserved floor (seq 0, 1) is long expired...
760        s.put_at(rec(&key, 1, 100), now);
761        s.put_at(rec(&key, 2, 200), now);
762        // ...while the churn zone (seq 2, 3) holds two LIVE honest providers.
763        s.put_at(rec(&key, 3, now + 5_000), now);
764        s.put_at(rec(&key, 4, now + 6_000), now);
765
766        s.put_at(rec(&key, 5, now + 7_000), now);
767
768        let live = live_provider_ids_at(&s, &key, now);
769        assert!(
770            live.contains(&PeerId::from_bytes([3u8; 32]).to_hex())
771                && live.contains(&PeerId::from_bytes([4u8; 32]).to_hex()),
772            "both LIVE providers must survive; an expired record in the floor is the victim"
773        );
774    }
775
776    #[test]
777    fn one_expired_record_anywhere_is_the_victim_before_any_live_record() {
778        // The one-off variant: exactly ONE expired record, sitting inside the reserved floor, with
779        // every other slot live. It must still be the one evicted.
780        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
781            max_providers_per_key: 4,
782            max_total_records: 1000,
783        });
784        let key = Key::from_bytes([0xAA; 32]);
785        let now = 10_000;
786        s.put_at(rec(&key, 1, 100), now); // expired, seq 0 → inside the floor
787        s.put_at(rec(&key, 2, now + 1_000), now);
788        s.put_at(rec(&key, 3, now + 2_000), now);
789        s.put_at(rec(&key, 4, now + 3_000), now);
790
791        s.put_at(rec(&key, 5, now + 4_000), now);
792
793        assert_eq!(
794            live_provider_ids_at(&s, &key, now).len(),
795            4,
796            "reclaiming the dead slot leaves every live provider intact"
797        );
798    }
799
800    #[test]
801    fn the_floor_still_protects_an_established_holder_when_every_record_is_live() {
802        // Liveness must take precedence WITHOUT weakening #1434: with no dead slot to reclaim, the
803        // establishment floor governs again and a sustained flood cannot displace the incumbent.
804        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
805            max_providers_per_key: 20,
806            max_total_records: 100_000,
807        });
808        let key = Key::from_bytes([0xAA; 32]);
809        let now = 10_000;
810        let honest = PeerId::from_bytes([1u8; 32]).to_hex();
811        s.put_at(rec(&key, 1, now + 1_000), now);
812        for i in 0..500u64 {
813            s.put_at(
814                ProviderRecord::new(
815                    &key,
816                    &PeerId::from_bytes(sybil_id(i)),
817                    vec![CandidateAddr::direct("h", 9444)],
818                    now + 2_000 + i,
819                ),
820                now,
821            );
822        }
823        assert!(
824            live_provider_ids_at(&s, &key, now).contains(&honest),
825            "an all-live key keeps the #1434 protection"
826        );
827    }
828
829    #[test]
830    fn put_delegates_to_put_at_with_the_wall_clock() {
831        // `put` is the compatibility wrapper (its signature is public API): same admission decision,
832        // with `now` read from the system clock.
833        let mut wall = ProviderStore::new();
834        let key = Key::from_bytes([0xAA; 32]);
835        assert_eq!(wall.put(rec(&key, 1, u64::MAX)), PutOutcome::Accepted);
836        assert_eq!(wall.len(), 1);
837    }
838
839    /// The live provider peer_ids for `key` as of `now`.
840    fn live_provider_ids_at(
841        s: &ProviderStore,
842        key: &Key,
843        now: u64,
844    ) -> std::collections::HashSet<String> {
845        s.get(&key.to_hex(), now)
846            .into_iter()
847            .map(|r| r.provider_peer_id)
848            .collect()
849    }
850
851    /// A distinct Sybil peer_id per index (varying the high bytes so ids stay distinct past 255).
852    fn sybil_id(i: u64) -> [u8; 32] {
853        let mut b = [0xEE; 32];
854        b[0..8].copy_from_slice(&i.to_be_bytes());
855        b
856    }
857
858    #[test]
859    fn global_cap_rejects_new_content_keys_over_ceiling() {
860        // Many DISTINCT content keys (not just many providers per key) must also be bounded — the
861        // audit's "no cap on distinct content keys ... no global record ceiling" finding.
862        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
863            max_providers_per_key: 20,
864            max_total_records: 2,
865        });
866        let k1 = Key::from_bytes([0x01; 32]);
867        let k2 = Key::from_bytes([0x02; 32]);
868        let k3 = Key::from_bytes([0x03; 32]);
869        assert_eq!(s.put(rec(&k1, 1, 100)), PutOutcome::Accepted);
870        assert_eq!(s.put(rec(&k2, 1, 100)), PutOutcome::Accepted);
871        assert_eq!(
872            s.put(rec(&k3, 1, 100)),
873            PutOutcome::RejectedOverCapacity,
874            "third distinct record must be rejected once the global ceiling is hit"
875        );
876        assert_eq!(s.len(), 2, "rejected record must not be stored");
877        assert!(
878            s.get(&k3.to_hex(), 0).is_empty(),
879            "rejected key must not appear in the store at all"
880        );
881    }
882
883    #[test]
884    fn global_cap_does_not_evict_a_different_key_to_make_room() {
885        // A single attacker flooding new keys must not be able to evict a DIFFERENT (legitimate)
886        // key's providers just by hitting the global ceiling.
887        let mut s = ProviderStore::with_limits(ProviderStoreLimits {
888            max_providers_per_key: 20,
889            max_total_records: 1,
890        });
891        let legit = Key::from_bytes([0xAA; 32]);
892        s.put(rec(&legit, 1, 100));
893        let attacker_key = Key::from_bytes([0xBB; 32]);
894        assert_eq!(
895            s.put(rec(&attacker_key, 2, 100)),
896            PutOutcome::RejectedOverCapacity
897        );
898        assert_eq!(
899            s.get(&legit.to_hex(), 0).len(),
900            1,
901            "the legitimate key's record must survive"
902        );
903    }
904
905    #[test]
906    fn remove_deletes_only_the_named_provider_record() {
907        // Authenticated retract (SPEC §6.6): removing (key, provider-1) must leave provider-2 of the
908        // SAME key untouched — a retract signed by one holder cannot censor another holder.
909        let mut s = ProviderStore::new();
910        let key = Key::from_bytes([0xAA; 32]);
911        s.put(rec(&key, 1, 100));
912        s.put(rec(&key, 2, 100));
913        let pid1 = PeerId::from_bytes([1u8; 32]).to_hex();
914        let pid2 = PeerId::from_bytes([2u8; 32]).to_hex();
915        assert!(
916            s.remove(&key.to_hex(), &pid1),
917            "the named record was removed"
918        );
919        let survivors: std::collections::HashSet<String> = s
920            .get(&key.to_hex(), 0)
921            .into_iter()
922            .map(|r| r.provider_peer_id)
923            .collect();
924        assert_eq!(survivors.len(), 1, "the other provider must survive");
925        assert!(survivors.contains(&pid2));
926        assert!(!survivors.contains(&pid1));
927    }
928
929    #[test]
930    fn remove_of_absent_record_returns_false() {
931        let mut s = ProviderStore::new();
932        let key = Key::from_bytes([0xAA; 32]);
933        s.put(rec(&key, 1, 100));
934        let absent = PeerId::from_bytes([9u8; 32]).to_hex();
935        assert!(!s.remove(&key.to_hex(), &absent), "no such provider");
936        assert!(!s.remove(&"00".repeat(32), &absent), "no such content key");
937        assert_eq!(s.len(), 1, "nothing removed");
938    }
939
940    #[test]
941    fn remove_drops_content_key_when_last_provider_leaves() {
942        let mut s = ProviderStore::new();
943        let key = Key::from_bytes([0xAA; 32]);
944        s.put(rec(&key, 1, 100));
945        let pid1 = PeerId::from_bytes([1u8; 32]).to_hex();
946        assert!(s.remove(&key.to_hex(), &pid1));
947        assert!(
948            s.is_empty(),
949            "the now-empty content key must be dropped entirely"
950        );
951    }
952
953    #[test]
954    fn gc_removes_expired_and_empty_keys() {
955        let mut s = ProviderStore::new();
956        let k1 = Key::from_bytes([0x01; 32]);
957        let k2 = Key::from_bytes([0x02; 32]);
958        s.put(rec(&k1, 1, 100)); // expires at 100
959        s.put(rec(&k2, 1, 500)); // expires at 500
960        let removed = s.gc(200);
961        assert_eq!(removed, 1);
962        assert!(s.get(&k1.to_hex(), 200).is_empty());
963        assert_eq!(s.get(&k2.to_hex(), 200).len(), 1);
964    }
965
966    #[test]
967    fn announcements_track_and_untrack() {
968        let mut s = ProviderStore::new();
969        let key = Key::from_bytes([0x07; 32]).to_hex();
970        s.mark_announced(key.clone());
971        s.mark_announced(key.clone()); // idempotent
972        assert_eq!(s.local_announcements(), vec![key.clone()]);
973        assert!(s.unmark_announced(&key));
974        assert!(!s.unmark_announced(&key));
975        assert!(s.local_announcements().is_empty());
976    }
977}