Skip to main content

dig_pex/
engine.rs

1//! The [`PexEngine`] — the transport-agnostic, sans-IO core both a DIG Node and the relay embed
2//! (SPEC Appendix A).
3//!
4//! You feed the engine four kinds of input and it returns the messages to send + the events to act
5//! on; it does no I/O itself (the node/relay do the actual dig-nat mux / WebSocket reads and writes):
6//!
7//! - **link events** — [`link_up`](PexEngine::link_up) (produces our outgoing handshake + snapshot)
8//!   and [`link_down`](PexEngine::link_down) (discards all per-link state, SPEC §5.5);
9//! - **inbound messages** — [`on_message`](PexEngine::on_message) validates + advances the receiver
10//!   state machine, returning verified-candidate / dropped events and any `pex_error` replies, and
11//!   penalizing misbehavior (SPEC §5, §6.4, §7, §11);
12//! - **local peer-set changes** — [`upsert_known`](PexEngine::upsert_known) /
13//!   [`remove_known`](PexEngine::remove_known) maintain the first-hand set PEX advertises (SPEC §9.3);
14//! - **clock ticks** — [`tick`](PexEngine::tick) (~1/s) emits per-link `pex_delta`s for pending
15//!   changes, spaced by the effective interval (SPEC §6).
16//!
17//! Timestamps are **Unix epoch milliseconds**. See [`crate`] docs for the node vs relay embedding.
18
19use std::collections::HashMap;
20
21use crate::caps::{
22    PEX_MAX_ADDED, PEX_MAX_DROPPED, PEX_MAX_HINTS, PEX_MAX_INTERVAL, PEX_MAX_RECEIVED_PER_LINK,
23    PEX_MAX_SNAPSHOT, PEX_VERSION, PEX_VIOLATION_LIMIT,
24};
25use crate::entry::{PeerEntry, ValidateCtx};
26use crate::error::PexErrorCode;
27use crate::state::{LinkState, RecvPhase};
28use crate::timer::{arrival_floor_ms, clamp_interval, effective_interval_secs, jitter_ms};
29use crate::wire::PexMessage;
30
31/// Configuration for a [`PexEngine`] (SPEC Appendix A).
32#[derive(Debug, Clone)]
33pub struct PexConfig {
34    /// This participant's own transport identity (`peer_id`, `<64hex>`) — excluded from every
35    /// advertisement and used to skip self-entries on receive (SPEC §5.4).
36    pub local_peer_id: String,
37    /// The network this participant serves — every handshake declares it and every entry MUST match
38    /// it (SPEC §5.2, §7.3).
39    pub network_id: String,
40    /// This participant's own capability flags, sent in its handshake (SPEC §4.2). For the relay
41    /// introducer this is `["introducer"]`.
42    pub flags: Vec<String>,
43    /// The declared send interval (seconds) — clamped into `[30, 3600]` (SPEC §6.2). Default `60`.
44    pub interval: u32,
45    /// Whether to add SPEC §6.3 send jitter. Default `true`; tests may disable it for deterministic
46    /// scheduling (0% jitter is within the allowed `0..+10%`).
47    pub jitter: bool,
48}
49
50impl PexConfig {
51    /// A new config for `local_peer_id` on `network_id`, with the default 60 s interval, no flags,
52    /// and jitter enabled.
53    #[must_use]
54    pub fn new(local_peer_id: impl Into<String>, network_id: impl Into<String>) -> Self {
55        PexConfig {
56            local_peer_id: local_peer_id.into(),
57            network_id: network_id.into(),
58            flags: Vec::new(),
59            interval: crate::caps::PEX_DEFAULT_INTERVAL,
60            jitter: true,
61        }
62    }
63
64    /// Builder: set this participant's own capability flags.
65    #[must_use]
66    pub fn with_flags(mut self, flags: Vec<String>) -> Self {
67        self.flags = flags;
68        self
69    }
70
71    /// Builder: set the declared send interval (seconds), clamped into `[30, 3600]`.
72    #[must_use]
73    pub fn with_interval(mut self, secs: u32) -> Self {
74        self.interval = clamp_interval(secs);
75        self
76    }
77
78    /// Builder: enable/disable send jitter (SPEC §6.3).
79    #[must_use]
80    pub fn with_jitter(mut self, jitter: bool) -> Self {
81        self.jitter = jitter;
82        self
83    }
84}
85
86/// An event the engine surfaces from an inbound message for the host to act on (SPEC Appendix A).
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub enum PexEvent {
89    /// Validated, verified-candidate peer hints — feed to the address manager as new-table
90    /// candidates to dial + verify (SPEC §9.3). These are hints, never authenticated facts (§11.1).
91    Candidates(Vec<PeerEntry>),
92    /// The link's sender dropped these `peer_id`s (SPEC §8.3) — advisory. Unlist the sender as a
93    /// source for them; never delete a first-hand-verified peer on this alone.
94    Dropped {
95        /// The dropped ids this link had previously told us (ids it never told us are ignored).
96        peer_ids: Vec<String>,
97    },
98    /// The link's sender committed a violation (SPEC §11.2). `mute` is `true` once the direction is
99    /// muted — either at the strike limit (misbehavior: `code` 1/3/4/6 → the host MAY penalize /
100    /// disconnect) or immediately for a version/network mismatch (`code` 2/5 → benign; the host MUST
101    /// NOT tear down the underlying connection for that alone, SPEC §5.2).
102    Violation {
103        /// The SPEC §4.5 error code.
104        code: u16,
105        /// Whether the incoming direction is now muted.
106        mute: bool,
107    },
108}
109
110/// The result of feeding the engine an inbound message or transport error: the messages to send back
111/// on the link, plus the events for the host to act on.
112#[derive(Debug, Clone, Default, PartialEq, Eq)]
113pub struct PexOutcome {
114    /// Messages to write back on the link (e.g. a `pex_error`). Best-effort / advisory (SPEC §4.5).
115    pub replies: Vec<PexMessage>,
116    /// Events for the host (candidates / dropped / violation).
117    pub events: Vec<PexEvent>,
118}
119
120/// A deduplicated inbound hint (SPEC §9.2) — the current best entry for a `peer_id` across all links.
121#[derive(Debug, Clone)]
122struct ReceivedHint {
123    /// The link (`peer_id`) that is currently the source for this hint.
124    source: String,
125    /// The `last_seen` of the current hint — newer wins across senders.
126    last_seen: u64,
127}
128
129/// The transport-agnostic PEX engine (SPEC Appendix A). One instance per participant; it multiplexes
130/// all of that participant's links.
131#[derive(Debug)]
132pub struct PexEngine {
133    cfg: PexConfig,
134    /// The first-hand known-peer set PEX advertises, keyed by `peer_id` (SPEC §9.3 outbound).
135    known: HashMap<String, PeerEntry>,
136    /// Per-link state, keyed by the transport `peer_id`.
137    links: HashMap<String, LinkState>,
138    /// Global inbound dedup (SPEC §9.2): `peer_id → current best hint`.
139    hints: HashMap<String, ReceivedHint>,
140    /// Bumped on every `known` mutation ([`upsert_known`](Self::upsert_known) /
141    /// [`remove_known`](Self::remove_known)) — invalidates `advertisable_cache` (#179 MED
142    /// optimization: the freshest-first, self-excluded base list is identical across every link
143    /// within one tick, so it is computed once and reused rather than per-link).
144    known_epoch: u64,
145    /// The cached partner-independent advertisable base list (self excluded, stale dropped,
146    /// freshest-first) — `(epoch it was built at, the `now_secs` it was built for, the list)`.
147    /// Recomputed only when `known_epoch` or `now_secs` has moved since the cached build.
148    advertisable_cache: std::cell::RefCell<Option<(u64, u64, Vec<PeerEntry>)>>,
149    /// Test-only instrumentation (#179 MED): counts actual `advertisable_base` rebuilds, so a test
150    /// can assert the cache is hit (O(1) rebuild per tick) rather than only checking behavioral
151    /// equivalence, which a naive per-link recompute would also satisfy.
152    #[cfg(test)]
153    advertisable_rebuilds: std::cell::Cell<u64>,
154}
155
156impl PexEngine {
157    /// Create an engine from `cfg`.
158    #[must_use]
159    pub fn new(cfg: PexConfig) -> Self {
160        PexEngine {
161            cfg,
162            known: HashMap::new(),
163            links: HashMap::new(),
164            hints: HashMap::new(),
165            known_epoch: 0,
166            advertisable_cache: std::cell::RefCell::new(None),
167            #[cfg(test)]
168            advertisable_rebuilds: std::cell::Cell::new(0),
169        }
170    }
171
172    // ----- local first-hand set (SPEC §9.3 outbound) -----
173
174    /// Add or update a **first-hand-known** peer in the advertise set (SPEC §8.1, §9.3). The caller
175    /// supplies the honest `via` + a fresh `last_seen`; the [`Provenance`](crate::Provenance) type
176    /// structurally forbids a `"pex"` provenance, so a PEX-learned entry can never be re-advertised
177    /// unverified. The change surfaces as `added` in the next [`tick`](Self::tick) delta on each link.
178    pub fn upsert_known(&mut self, entry: PeerEntry) {
179        // Never advertise ourselves (the link is our own advertisement — SPEC §5.4).
180        if entry.peer_id == self.cfg.local_peer_id {
181            return;
182        }
183        self.known.insert(entry.peer_id.clone(), entry);
184        self.known_epoch = self.known_epoch.wrapping_add(1);
185    }
186
187    /// Remove a peer from the advertise set — it disconnected or went stale (SPEC §9.3). It surfaces
188    /// as `dropped` in the next delta on links that were told it.
189    pub fn remove_known(&mut self, peer_id: &str) {
190        self.known.remove(peer_id);
191        self.known_epoch = self.known_epoch.wrapping_add(1);
192    }
193
194    // ----- link lifecycle (SPEC §5) -----
195
196    /// A link came up: register it and produce our outgoing direction — the `pex_handshake` followed
197    /// (back-to-back) by the `pex_snapshot` of our current first-hand set (SPEC §5.1, §6.1). Write
198    /// the returned messages on our sending stream. Preserves any receiver-side state if the link
199    /// already exists (e.g. an inbound message arrived first).
200    pub fn link_up(&mut self, peer_id: &str, now_ms: u64) -> Vec<PexMessage> {
201        let interval = self.cfg.interval;
202        let handshake = PexMessage::PexHandshake {
203            version: PEX_VERSION,
204            network_id: self.cfg.network_id.clone(),
205            interval,
206            flags: self.cfg.flags.clone(),
207        };
208
209        // Build the snapshot from the current advertisable set (freshest-first, capped, self+partner
210        // excluded) before mutating the link, so `known` isn't borrowed across the link mutation. Uses
211        // the same cached partner-independent base as `tick`'s deltas (#179 MED optimization).
212        let now_secs = now_ms / 1000;
213        let mut peers = self.advertisable_for(peer_id, now_secs);
214        peers.truncate(PEX_MAX_SNAPSHOT);
215
216        let remote_declared = self.links.get(peer_id).and_then(|l| l.remote_declared_secs);
217        let jitter = self.draw_jitter(effective_interval_secs(interval, remote_declared));
218
219        let link = self
220            .links
221            .entry(peer_id.to_string())
222            .or_insert_with(|| LinkState::new(interval));
223        link.self_interval_secs = interval;
224        link.handshake_sent = true;
225        for e in &peers {
226            link.told.insert(e.peer_id.clone(), e.fingerprint_hash());
227        }
228        link.snapshot_sent = true;
229        link.last_data_send_ms = Some(now_ms);
230        link.send_jitter_ms = jitter;
231
232        vec![handshake, PexMessage::PexSnapshot { peers }]
233    }
234
235    /// A link went down: discard all per-link state, and unlist it as the source of any current hints
236    /// (SPEC §5.5, §9.2). A new connection starts fresh.
237    pub fn link_down(&mut self, peer_id: &str) {
238        self.links.remove(peer_id);
239        self.hints.retain(|_, h| h.source != peer_id);
240    }
241
242    // ----- inbound (SPEC §5.3, §6.4, §7, §11) -----
243
244    /// Feed one decoded inbound message from `peer_id`. Returns replies to send + events to act on.
245    /// A malformed *entry* inside a valid message is skipped silently; a malformed *message* /
246    /// rate / oversize / state violation is discarded with a strike (SPEC §7.3, §11.2).
247    pub fn on_message(&mut self, peer_id: &str, msg: PexMessage, now_ms: u64) -> PexOutcome {
248        // Ensure a link exists (an inbound message may precede our own `link_up`).
249        let interval = self.cfg.interval;
250        let muted = {
251            let link = self
252                .links
253                .entry(peer_id.to_string())
254                .or_insert_with(|| LinkState::new(interval));
255            link.muted
256        };
257        if muted {
258            // Direction muted — ignore all further inbound PEX (SPEC §5.2, §11.2).
259            return PexOutcome::default();
260        }
261
262        match msg {
263            PexMessage::PexError { code, .. } => self.on_pex_error(peer_id, code, now_ms),
264            PexMessage::PexHandshake {
265                version,
266                network_id,
267                interval: declared,
268                ..
269            } => self.on_handshake(peer_id, version, &network_id, declared),
270            PexMessage::PexSnapshot { peers } => self.on_snapshot(peer_id, peers, now_ms),
271            PexMessage::PexDelta { added, dropped } => {
272                self.on_delta(peer_id, added, dropped, now_ms)
273            }
274        }
275    }
276
277    /// Record a transport-detected violation the engine could not see itself: a frame-size overrun
278    /// (`Oversized`) or an undecodable/malformed frame (`BadMessage`) — SPEC §7.2, §7.3. Counts a
279    /// strike and mutes at the limit, exactly like an engine-detected violation.
280    pub fn record_violation(
281        &mut self,
282        peer_id: &str,
283        code: PexErrorCode,
284        _now_ms: u64,
285    ) -> PexOutcome {
286        let interval = self.cfg.interval;
287        self.links
288            .entry(peer_id.to_string())
289            .or_insert_with(|| LinkState::new(interval));
290        self.strike(peer_id, code)
291    }
292
293    // ----- clock (SPEC §6.1) -----
294
295    /// Drive the send cadence (call ~1/s). For each link whose effective interval has elapsed since
296    /// its last data message and that has pending changes, emits a `pex_delta` (SPEC §4.4, §6). A
297    /// delta with no changes is suppressed (empty deltas are never sent). Returns `(peer_id,
298    /// message)` pairs to write to the matching links.
299    pub fn tick(&mut self, now_ms: u64) -> Vec<(String, PexMessage)> {
300        let now_secs = now_ms / 1000;
301        let mut out = Vec::new();
302        // Snapshot the link keys to avoid borrowing `self.links` while mutating per-link below.
303        let peer_ids: Vec<String> = self.links.keys().cloned().collect();
304        for peer_id in peer_ids {
305            let (eligible, effective) = {
306                let link = &self.links[&peer_id];
307                if !link.snapshot_sent {
308                    continue; // we are receive-only on this link
309                }
310                let effective =
311                    effective_interval_secs(link.self_interval_secs, link.remote_declared_secs);
312                let base = link.last_data_send_ms.unwrap_or(0);
313                let eligible = now_ms >= base + u64::from(effective) * 1000 + link.send_jitter_ms;
314                (eligible, effective)
315            };
316            if !eligible {
317                continue;
318            }
319
320            let (added, dropped) = self.build_delta(&peer_id, now_secs);
321            if added.is_empty() && dropped.is_empty() {
322                continue; // suppress empty deltas (SPEC §4.4)
323            }
324
325            // Commit the told-state for exactly what we send (SPEC §9.1); the capped remainder recurs.
326            let link = self.links.get_mut(&peer_id).expect("link exists");
327            for e in &added {
328                link.told.insert(e.peer_id.clone(), e.fingerprint_hash());
329            }
330            for id in &dropped {
331                link.told.remove(id);
332            }
333            link.last_data_send_ms = Some(now_ms);
334            let jitter = self.draw_jitter(effective);
335            self.links
336                .get_mut(&peer_id)
337                .expect("link exists")
338                .send_jitter_ms = jitter;
339
340            out.push((peer_id, PexMessage::PexDelta { added, dropped }));
341        }
342        out
343    }
344
345    // ----- read-only accessors (observability / tests) -----
346
347    /// Number of peers in the first-hand advertise set.
348    #[must_use]
349    pub fn known_count(&self) -> usize {
350        self.known.len()
351    }
352
353    /// Number of live links.
354    #[must_use]
355    pub fn link_count(&self) -> usize {
356        self.links.len()
357    }
358
359    /// Whether the incoming direction of `peer_id`'s link is muted (SPEC §5.2, §11.2).
360    #[must_use]
361    pub fn is_muted(&self, peer_id: &str) -> bool {
362        self.links.get(peer_id).is_some_and(|l| l.muted)
363    }
364
365    /// The violation strike count on `peer_id`'s incoming direction (SPEC §11.2).
366    #[must_use]
367    pub fn strikes(&self, peer_id: &str) -> u32 {
368        self.links.get(peer_id).map_or(0, |l| l.strikes)
369    }
370
371    /// How many peer ids we have currently told `peer_id`'s link (told-state size, SPEC §9.1).
372    #[must_use]
373    pub fn told_count(&self, peer_id: &str) -> usize {
374        self.links.get(peer_id).map_or(0, |l| l.told.len())
375    }
376
377    /// The current deduplicated hint for `peer_id` (SPEC §9.2): `(source link, last_seen)`, if any.
378    #[must_use]
379    pub fn current_hint(&self, peer_id: &str) -> Option<(&str, u64)> {
380        self.hints
381            .get(peer_id)
382            .map(|h| (h.source.as_str(), h.last_seen))
383    }
384
385    /// How many `peer_id`s `peer_id`'s link has told us are currently tracked in its `received`
386    /// accumulator (SPEC §9.2, §11.3) — bounded by [`crate::caps::PEX_MAX_RECEIVED_PER_LINK`].
387    #[must_use]
388    pub fn received_count(&self, peer_id: &str) -> usize {
389        self.links.get(peer_id).map_or(0, |l| l.received.len())
390    }
391
392    /// The total number of deduplicated hints currently held across all links (SPEC §9.2, §11.3) —
393    /// bounded by [`crate::caps::PEX_MAX_HINTS`].
394    #[must_use]
395    pub fn hints_count(&self) -> usize {
396        self.hints.len()
397    }
398
399    // ----- internals -----
400
401    fn draw_jitter(&self, effective_secs: u32) -> u64 {
402        if self.cfg.jitter {
403            jitter_ms(effective_secs)
404        } else {
405            0
406        }
407    }
408
409    /// `pex_error` is acceptable in any state and never changes the receiver state (SPEC §5.3). A
410    /// sender receiving code `3` SHOULD back off — double its effective interval, capped (SPEC §6.4).
411    ///
412    /// `pex_error` is advisory and **unauthenticated** (SPEC §4.5): any non-muted peer can send it at
413    /// will. Two gates bound how far/fast a spoofed code-3 flood can push us (LOW #179 fix):
414    ///
415    /// 1. **Plausibility** — only honored if we actually sent a data message to this peer recently
416    ///    enough that a rate violation is plausible: `now_ms` must fall within
417    ///    `arrival_floor_ms(self_interval_secs)` of `last_data_send_ms`. A code-3 arriving long after
418    ///    our last send (or before we have ever sent anything) cannot correspond to a real violation
419    ///    of *our* sends, so it is ignored.
420    /// 2. **Rate limit** — even a plausible code-3 is honored at most once per (pre-doubling)
421    ///    effective interval: a flood of code-3 frames right after a legitimate one cannot keep
422    ///    doubling toward `PEX_MAX_INTERVAL` faster than one genuine violation could.
423    fn on_pex_error(&mut self, peer_id: &str, code: u16, now_ms: u64) -> PexOutcome {
424        if code == PexErrorCode::RateViolation.as_u16() {
425            if let Some(link) = self.links.get_mut(peer_id) {
426                let plausible = link.last_data_send_ms.is_some_and(|sent| {
427                    now_ms.saturating_sub(sent) < arrival_floor_ms(link.self_interval_secs)
428                });
429                let effective_ms = u64::from(link.self_interval_secs) * 1000;
430                let rate_limited = link
431                    .last_backoff_applied_ms
432                    .is_some_and(|applied| now_ms.saturating_sub(applied) < effective_ms);
433                if plausible && !rate_limited {
434                    link.self_interval_secs = clamp_interval(
435                        (link.self_interval_secs.saturating_mul(2)).min(PEX_MAX_INTERVAL),
436                    );
437                    link.last_backoff_applied_ms = Some(now_ms);
438                }
439            }
440        }
441        PexOutcome::default()
442    }
443
444    fn on_handshake(
445        &mut self,
446        peer_id: &str,
447        version: u32,
448        network_id: &str,
449        declared: u32,
450    ) -> PexOutcome {
451        let phase = self.links[peer_id].phase;
452        if phase != RecvPhase::AwaitingHandshake {
453            // A repeat handshake once past the handshake state is a protocol violation (SPEC §5.3).
454            return self.strike(peer_id, PexErrorCode::ProtocolViolation);
455        }
456        if version != PEX_VERSION {
457            return self.mute_mismatch(peer_id, PexErrorCode::UnsupportedVersion);
458        }
459        if network_id != self.cfg.network_id {
460            return self.mute_mismatch(peer_id, PexErrorCode::NetworkMismatch);
461        }
462        let link = self.links.get_mut(peer_id).expect("link exists");
463        link.remote_declared_secs = Some(clamp_interval(declared));
464        link.phase = RecvPhase::AwaitingSnapshot;
465        PexOutcome::default()
466    }
467
468    fn on_snapshot(&mut self, peer_id: &str, peers: Vec<PeerEntry>, now_ms: u64) -> PexOutcome {
469        match self.links[peer_id].phase {
470            RecvPhase::AwaitingHandshake => self.strike(peer_id, PexErrorCode::ProtocolViolation),
471            RecvPhase::Streaming => self.strike(peer_id, PexErrorCode::ProtocolViolation),
472            RecvPhase::AwaitingSnapshot => {
473                if peers.len() > PEX_MAX_SNAPSHOT {
474                    return self.strike(peer_id, PexErrorCode::Oversized);
475                }
476                let link = self.links.get_mut(peer_id).expect("link exists");
477                link.phase = RecvPhase::Streaming;
478                link.last_arrival_ms = Some(now_ms); // the snapshot starts the arrival clock
479                self.ingest_added(peer_id, peers, now_ms)
480            }
481        }
482    }
483
484    fn on_delta(
485        &mut self,
486        peer_id: &str,
487        added: Vec<PeerEntry>,
488        dropped: Vec<String>,
489        now_ms: u64,
490    ) -> PexOutcome {
491        match self.links[peer_id].phase {
492            RecvPhase::AwaitingHandshake | RecvPhase::AwaitingSnapshot => {
493                // Data before handshake, or a delta before the snapshot (SPEC §5.3).
494                return self.strike(peer_id, PexErrorCode::ProtocolViolation);
495            }
496            RecvPhase::Streaming => {}
497        }
498
499        // Rate enforcement (SPEC §6.4): a delta arriving under the floor is discarded + struck.
500        let (floor, last) = {
501            let link = &self.links[peer_id];
502            (
503                arrival_floor_ms(link.remote_declared_secs.unwrap_or(0)),
504                link.last_arrival_ms,
505            )
506        };
507        if let Some(last) = last {
508            if now_ms.saturating_sub(last) < floor {
509                return self.strike(peer_id, PexErrorCode::RateViolation);
510            }
511        }
512
513        // List caps: reject the whole message, never truncate (SPEC §7.2).
514        if added.len() > PEX_MAX_ADDED || dropped.len() > PEX_MAX_DROPPED {
515            return self.strike(peer_id, PexErrorCode::Oversized);
516        }
517        // Structural MUST: a peer_id may not appear in both `added` and `dropped` (SPEC §4.4).
518        let added_ids: std::collections::HashSet<&str> =
519            added.iter().map(|e| e.peer_id.as_str()).collect();
520        if dropped.iter().any(|d| added_ids.contains(d.as_str())) {
521            return self.strike(peer_id, PexErrorCode::BadMessage);
522        }
523
524        self.links
525            .get_mut(peer_id)
526            .expect("link exists")
527            .last_arrival_ms = Some(now_ms);
528
529        let mut outcome = self.ingest_added(peer_id, added, now_ms);
530        outcome
531            .events
532            .extend(self.ingest_dropped(peer_id, dropped).events);
533        outcome
534    }
535
536    /// Validate + dedup a batch of inbound entries into `Candidates` (SPEC §3.3, §9.2). Malformed
537    /// entries are skipped silently.
538    fn ingest_added(&mut self, peer_id: &str, entries: Vec<PeerEntry>, now_ms: u64) -> PexOutcome {
539        let now_secs = now_ms / 1000;
540        let mut candidates = Vec::new();
541        for e in entries {
542            let ctx = ValidateCtx {
543                receiver_peer_id: &self.cfg.local_peer_id,
544                sender_peer_id: peer_id,
545                network_id: &self.cfg.network_id,
546                now_secs,
547            };
548            if e.validate(&ctx).is_err() {
549                continue; // malformed entry — skip silently (SPEC §3.3, §7.3)
550            }
551            let ce = e.clamped(now_secs);
552            // Attribute the hint to this link so a later `dropped` can be matched (SPEC §8.3). Bound
553            // the accumulator first (HIGH #179): a single authenticated peer must not be able to grow
554            // this link's `received` map without limit by streaming many distinct fresh peer_ids.
555            let link = self.links.get_mut(peer_id).expect("link exists");
556            if !link.received.contains_key(&ce.peer_id)
557                && link.received.len() >= PEX_MAX_RECEIVED_PER_LINK
558            {
559                evict_oldest(&mut link.received, |last_seen| *last_seen);
560            }
561            link.received.insert(ce.peer_id.clone(), ce.last_seen);
562            // Dedup: newest `last_seen` wins as the current hint (SPEC §9.2); only surface an entry
563            // that is new or fresher than what we already hold, to avoid re-dialing stale duplicates.
564            let fresher = match self.hints.get(&ce.peer_id) {
565                Some(h) => ce.last_seen > h.last_seen,
566                None => true,
567            };
568            if fresher {
569                // Bound the global hints map the same way (HIGH #179): many links each contributing
570                // distinct peer_ids must not grow this map without limit.
571                if !self.hints.contains_key(&ce.peer_id) && self.hints.len() >= PEX_MAX_HINTS {
572                    evict_oldest(&mut self.hints, |h| h.last_seen);
573                }
574                self.hints.insert(
575                    ce.peer_id.clone(),
576                    ReceivedHint {
577                        source: peer_id.to_string(),
578                        last_seen: ce.last_seen,
579                    },
580                );
581                candidates.push(ce);
582            }
583        }
584        let mut outcome = PexOutcome::default();
585        if !candidates.is_empty() {
586            outcome.events.push(PexEvent::Candidates(candidates));
587        }
588        outcome
589    }
590
591    /// Attribute `dropped` ids: only those this link previously told us are acted on (SPEC §4.4,
592    /// §8.3). If a dropped id's current hint was sourced from this link, clear it (unlist the source).
593    fn ingest_dropped(&mut self, peer_id: &str, dropped: Vec<String>) -> PexOutcome {
594        let mut attributed = Vec::new();
595        for id in dropped {
596            let told_us = self
597                .links
598                .get_mut(peer_id)
599                .expect("link exists")
600                .received
601                .remove(&id)
602                .is_some();
603            if told_us {
604                if let Some(h) = self.hints.get(&id) {
605                    if h.source == peer_id {
606                        self.hints.remove(&id);
607                    }
608                }
609                attributed.push(id);
610            }
611        }
612        let mut outcome = PexOutcome::default();
613        if !attributed.is_empty() {
614            outcome.events.push(PexEvent::Dropped {
615                peer_ids: attributed,
616            });
617        }
618        outcome
619    }
620
621    /// Count a misbehavior strike (SPEC §11.2): discard the message, reply `pex_error` (advisory),
622    /// mute at the limit, and surface a `Violation` event. Version/network mismatch use
623    /// [`mute_mismatch`](Self::mute_mismatch) instead (immediate, non-strike mute).
624    fn strike(&mut self, peer_id: &str, code: PexErrorCode) -> PexOutcome {
625        let link = self.links.get_mut(peer_id).expect("link exists");
626        link.strikes += 1;
627        let mute = link.strikes >= PEX_VIOLATION_LIMIT;
628        if mute {
629            link.muted = true;
630            self.free_muted_link_state(peer_id);
631        }
632        PexOutcome {
633            replies: vec![PexMessage::PexError {
634                code: code.as_u16(),
635                message: code.message().to_string(),
636            }],
637            events: vec![PexEvent::Violation {
638                code: code.as_u16(),
639                mute,
640            }],
641        }
642    }
643
644    /// Immediately mute the direction for a version/network mismatch (SPEC §5.2). This is NOT a
645    /// strike (the peer is simply on a different version/network) and MUST NOT tear down the
646    /// underlying connection — PEX is an optional overlay.
647    fn mute_mismatch(&mut self, peer_id: &str, code: PexErrorCode) -> PexOutcome {
648        self.links.get_mut(peer_id).expect("link exists").muted = true;
649        self.free_muted_link_state(peer_id);
650        PexOutcome {
651            replies: vec![PexMessage::PexError {
652                code: code.as_u16(),
653                message: code.message().to_string(),
654            }],
655            events: vec![PexEvent::Violation {
656                code: code.as_u16(),
657                mute: true,
658            }],
659        }
660    }
661
662    /// Free accumulated state for a link whose incoming direction was just muted (SPEC §9.2, §11.3 —
663    /// HIGH #179 fix): treat mute like a soft `link_down` for the `received`/`hints` accumulators,
664    /// since a muted direction accepts no further inbound PEX (`on_message` early-returns) and so can
665    /// never again reference or grow that state. This bounds memory promptly rather than waiting for
666    /// the real `link_down` (which may be much later, or never, if the transport itself stays open).
667    fn free_muted_link_state(&mut self, peer_id: &str) {
668        if let Some(link) = self.links.get_mut(peer_id) {
669            link.received.clear();
670        }
671        self.hints.retain(|_, h| h.source != peer_id);
672    }
673}
674
675/// Evict the single oldest entry (by `last_seen`, ascending) from a bounded accumulator (SPEC §9.2,
676/// §11.3 — HIGH #179). Called once, immediately before an insert that would otherwise exceed the
677/// map's cardinality bound, so the map never grows past its cap. Ties break on key order for
678/// determinism. A no-op on an empty map (the caller only reaches the cap check when non-empty).
679fn evict_oldest<V>(map: &mut HashMap<String, V>, last_seen: impl Fn(&V) -> u64) {
680    if let Some(oldest_key) = map
681        .iter()
682        .min_by(|(ka, va), (kb, vb)| last_seen(va).cmp(&last_seen(vb)).then_with(|| ka.cmp(kb)))
683        .map(|(k, _)| k.clone())
684    {
685        map.remove(&oldest_key);
686    }
687}
688
689/// The **partner-independent** advertisable base for `known` at `now_secs`: self excluded (SPEC
690/// §5.4), stale entries dropped (SPEC §8.2), sorted **freshest-first** then by `peer_id` for a
691/// deterministic order (SPEC §4.3, §9.1). This ordering is identical for every link in a given tick
692/// (only the partner exclusion differs per link) — see [`PexEngine::advertisable_for`], which caches
693/// this and applies the cheap per-link partner exclusion (#179 MED optimization).
694fn advertisable_base(
695    known: &HashMap<String, PeerEntry>,
696    local_peer_id: &str,
697    now_secs: u64,
698) -> Vec<PeerEntry> {
699    let mut out: Vec<PeerEntry> = known
700        .values()
701        .filter(|e| e.peer_id != local_peer_id)
702        .filter(|e| {
703            // Not stale: within PEX_MAX_ENTRY_AGE (a future last_seen is treated as fresh).
704            e.last_seen >= now_secs || now_secs - e.last_seen <= crate::caps::PEX_MAX_ENTRY_AGE
705        })
706        .cloned()
707        .collect();
708    out.sort_by(|a, b| {
709        b.last_seen
710            .cmp(&a.last_seen)
711            .then_with(|| a.peer_id.cmp(&b.peer_id))
712    });
713    out
714}
715
716impl PexEngine {
717    /// The partner-independent advertisable base list for `now_secs`, computed once and reused for
718    /// every link (#179 MED optimization): freshest-first, self excluded, stale dropped. Cached in
719    /// `advertisable_cache` and only recomputed when `known_epoch` (bumped by
720    /// [`upsert_known`](Self::upsert_known)/[`remove_known`](Self::remove_known)) or `now_secs` has
721    /// moved since the cached build — so `L` links in one `tick` share a single O(K log K) build
722    /// instead of each paying it, where `K = known.len()`.
723    fn advertisable_cached(&self, now_secs: u64) -> std::cell::Ref<'_, Vec<PeerEntry>> {
724        {
725            let cache = self.advertisable_cache.borrow();
726            if let Some((epoch, cached_secs, _)) = cache.as_ref() {
727                if *epoch == self.known_epoch && *cached_secs == now_secs {
728                    drop(cache);
729                    return std::cell::Ref::map(self.advertisable_cache.borrow(), |c| {
730                        &c.as_ref().unwrap().2
731                    });
732                }
733            }
734        }
735        let fresh = advertisable_base(&self.known, &self.cfg.local_peer_id, now_secs);
736        #[cfg(test)]
737        self.advertisable_rebuilds
738            .set(self.advertisable_rebuilds.get() + 1);
739        *self.advertisable_cache.borrow_mut() = Some((self.known_epoch, now_secs, fresh));
740        std::cell::Ref::map(self.advertisable_cache.borrow(), |c| &c.as_ref().unwrap().2)
741    }
742
743    /// Test-only: how many times the advertisable base list has actually been rebuilt (#179 MED) —
744    /// used to assert the per-tick cache is hit rather than recomputed per link.
745    #[cfg(test)]
746    fn advertisable_rebuild_count(&self) -> u64 {
747        self.advertisable_rebuilds.get()
748    }
749
750    /// The advertisable subset for a link to `partner` at `now_secs`: the cached partner-independent
751    /// base (see [`advertisable_cached`](Self::advertisable_cached)) with `partner` excluded (SPEC
752    /// §5.4) cheaply during iteration — no additional clone or sort of the shared list.
753    fn advertisable_for(&self, partner: &str, now_secs: u64) -> Vec<PeerEntry> {
754        self.advertisable_cached(now_secs)
755            .iter()
756            .filter(|e| e.peer_id != partner)
757            .cloned()
758            .collect()
759    }
760
761    /// Compute the delta for a link relative to its told-state (SPEC §9.1): `added` = advertisable
762    /// entries not yet told (or told with a changed fingerprint), freshest-first, capped at
763    /// [`PEX_MAX_ADDED`]; `dropped` = told ids no longer advertisable, capped at [`PEX_MAX_DROPPED`].
764    fn build_delta(&self, peer_id: &str, now_secs: u64) -> (Vec<PeerEntry>, Vec<String>) {
765        let link = &self.links[peer_id];
766        let base = self.advertisable_cached(now_secs);
767
768        let mut added = Vec::new();
769        let mut advert_ids: std::collections::HashSet<&str> = std::collections::HashSet::new();
770        for e in base.iter().filter(|e| e.peer_id != peer_id) {
771            advert_ids.insert(e.peer_id.as_str());
772            if added.len() >= PEX_MAX_ADDED {
773                continue; // keep collecting ids for the dropped-set below; added is already capped
774            }
775            match link.told.get(&e.peer_id) {
776                // Cheap, allocation-free u64 equality — the hot-path check (#179 MED optimization).
777                Some(fp) if *fp == e.fingerprint_hash() => {} // unchanged — never re-advertise (SPEC §9.1)
778                _ => added.push(e.clone()),
779            }
780        }
781
782        let mut dropped = Vec::new();
783        for id in link.told.keys() {
784            if dropped.len() >= PEX_MAX_DROPPED {
785                break;
786            }
787            if !advert_ids.contains(id.as_str()) {
788                dropped.push(id.clone());
789            }
790        }
791        dropped.sort(); // deterministic order
792
793        (added, dropped)
794    }
795}
796
797#[cfg(test)]
798mod cap_tests {
799    use super::*;
800    use crate::caps::{PEX_MAX_HINTS, PEX_MAX_RECEIVED_PER_LINK};
801    use crate::entry::{Address, Provenance};
802
803    fn hex_id(n: u32) -> String {
804        // A deterministic, distinct 64-hex peer_id for index `n`.
805        format!("{n:064x}")
806    }
807
808    fn eng(local: &str) -> PexEngine {
809        PexEngine::new(PexConfig::new(local.to_string(), "mainnet".to_string()).with_jitter(false))
810    }
811
812    /// Ensure a link entry exists for `peer_id` (mirrors what `on_message` does before dispatch) so
813    /// the private `ingest_added`/`strike` helpers can be exercised directly in these tests.
814    fn ensure_link(e: &mut PexEngine, peer_id: &str) {
815        let interval = e.cfg.interval;
816        e.links
817            .entry(peer_id.to_string())
818            .or_insert_with(|| LinkState::new(interval));
819    }
820
821    fn distinct_entry(n: u32, now_secs: u64) -> PeerEntry {
822        PeerEntry::new(hex_id(n), "mainnet", now_secs, Provenance::Direct)
823            .with_address(Address::direct("203.0.113.7", 9444))
824    }
825
826    /// HIGH finding (#179): a single authenticated peer streaming far more than
827    /// `PEX_MAX_RECEIVED_PER_LINK` distinct fresh `peer_id`s over the life of a link must not grow
828    /// that link's `received` accumulator without bound — it must stay capped, with the oldest
829    /// entries evicted to make room for newer ones.
830    #[test]
831    fn received_map_is_capped_per_link_with_eviction() {
832        let local = hex_id(0);
833        let sender = hex_id(1);
834        let mut e = eng(&local);
835        let now_ms = 1_000_000_000_u64;
836
837        ensure_link(&mut e, &sender);
838        // Stream well past the cap in batches (ingest_added has no per-call size limit of its own —
839        // the message-level cap is enforced by on_delta/on_snapshot before this point).
840        let total = PEX_MAX_RECEIVED_PER_LINK + 500;
841        for n in 2..2 + total as u32 {
842            let entry = distinct_entry(n, now_ms / 1000);
843            e.ingest_added(&sender, vec![entry], now_ms);
844        }
845
846        assert!(
847            e.received_count(&sender) <= PEX_MAX_RECEIVED_PER_LINK,
848            "received map must stay bounded at PEX_MAX_RECEIVED_PER_LINK, got {}",
849            e.received_count(&sender)
850        );
851        // The oldest ids (evicted first) must be gone; the newest must remain.
852        assert!(
853            !e.links[&sender].received.contains_key(&hex_id(2)),
854            "the oldest entry should have been evicted"
855        );
856        let newest = hex_id(1 + total as u32);
857        assert!(
858            e.links[&sender].received.contains_key(&newest),
859            "the newest entry must survive eviction"
860        );
861    }
862
863    /// HIGH finding (#179): the engine-global `hints` map must stay bounded even when many distinct
864    /// links each contribute distinct fresh `peer_id`s, with oldest-`last_seen` eviction.
865    #[test]
866    fn hints_map_is_capped_globally_with_eviction() {
867        let local = hex_id(0);
868        let mut e = eng(&local);
869        let now_ms = 1_000_000_000_u64;
870
871        let total = PEX_MAX_HINTS + 500;
872        for n in 0..total as u32 {
873            // A distinct sender per entry so every hint is a genuinely new peer_id from a live link.
874            let sender = hex_id(1_000_000 + n);
875            ensure_link(&mut e, &sender);
876            let entry = distinct_entry(2_000_000 + n, now_ms / 1000 + u64::from(n));
877            e.ingest_added(&sender, vec![entry], now_ms);
878        }
879
880        assert!(
881            e.hints_count() <= PEX_MAX_HINTS,
882            "hints map must stay bounded at PEX_MAX_HINTS, got {}",
883            e.hints_count()
884        );
885        // The oldest (lowest last_seen) hint must have been evicted; the newest must remain.
886        assert!(
887            e.current_hint(&hex_id(2_000_000)).is_none(),
888            "the oldest hint should have been evicted"
889        );
890        let newest_peer = hex_id(2_000_000 + total as u32 - 1);
891        assert!(
892            e.current_hint(&newest_peer).is_some(),
893            "the newest hint must survive eviction"
894        );
895    }
896
897    /// Muting a direction is treated like a soft `link_down` for accumulated state (#179 fix note):
898    /// the link's `received` entries and any global `hints` sourced from it are freed immediately,
899    /// not left to accumulate until the real `link_down`.
900    #[test]
901    fn muting_a_direction_frees_its_received_and_sourced_hints() {
902        let local = hex_id(0);
903        let sender = hex_id(1);
904        let mut e = eng(&local);
905        let now_ms = 1_000_000_000_u64;
906
907        ensure_link(&mut e, &sender);
908        e.ingest_added(&sender, vec![distinct_entry(2, now_ms / 1000)], now_ms);
909        assert_eq!(e.received_count(&sender), 1);
910        assert!(e.current_hint(&hex_id(2)).is_some());
911
912        // Force three strikes to mute the incoming direction.
913        for _ in 0..3 {
914            e.strike(&sender, PexErrorCode::ProtocolViolation);
915        }
916        assert!(e.is_muted(&sender));
917
918        assert_eq!(
919            e.received_count(&sender),
920            0,
921            "received state must be freed when the direction is muted"
922        );
923        assert!(
924            e.current_hint(&hex_id(2)).is_none(),
925            "hints sourced from a now-muted link must be cleared"
926        );
927    }
928}
929
930#[cfg(test)]
931mod advertisable_cache_tests {
932    use super::*;
933    use crate::entry::{Address, Provenance};
934
935    fn hex_id(n: u32) -> String {
936        format!("{n:064x}")
937    }
938
939    fn eng(local: &str) -> PexEngine {
940        PexEngine::new(PexConfig::new(local.to_string(), "mainnet".to_string()).with_jitter(false))
941    }
942
943    fn known_entry(n: u32, last_seen: u64) -> PeerEntry {
944        PeerEntry::new(hex_id(n), "mainnet", last_seen, Provenance::Direct)
945            .with_address(Address::direct("203.0.113.7", 9444))
946    }
947
948    /// MEDIUM finding (#179): a single `tick` covering many links must build the partner-independent
949    /// advertisable base list ONCE and reuse it across every link, not clone+re-sort per link.
950    #[test]
951    fn tick_rebuilds_advertisable_base_once_for_many_links() {
952        let mut e = eng(&hex_id(0));
953        for n in 100..110 {
954            e.upsert_known(known_entry(n, 1_000));
955        }
956        for n in 0..20u32 {
957            e.link_up(&hex_id(n), 1_000_000);
958        }
959        // link_up itself uses the cache; each link_up call is at the SAME now_secs, so all 20 share
960        // one rebuild.
961        assert_eq!(
962            e.advertisable_rebuild_count(),
963            1,
964            "20 link_ups at the same now_secs must share a single advertisable rebuild, got {}",
965            e.advertisable_rebuild_count()
966        );
967
968        // Advance past every link's interval and tick: the per-tick delta computation for all 20
969        // links must again share a single rebuild (a fresh one, since now_secs moved). Deltas may be
970        // empty (nothing changed since link_up already told everything) — the rebuild count is what
971        // this test asserts, not the delta contents.
972        let _out = e.tick(1_000_000 + 61_000);
973        assert_eq!(
974            e.advertisable_rebuild_count(),
975            2,
976            "one tick covering 20 links must add exactly one more rebuild (now_secs changed once), got {}",
977            e.advertisable_rebuild_count()
978        );
979    }
980
981    /// The cache must not go stale: a `known` mutation between two calls at the same `now_secs` MUST
982    /// force a rebuild so a newly upserted (or removed) peer is reflected immediately.
983    #[test]
984    fn cache_invalidates_on_known_mutation_even_at_same_now_secs() {
985        let mut e = eng(&hex_id(0));
986        e.upsert_known(known_entry(1, 1_000));
987        let now_ms = 1_000_000;
988
989        let first = e.advertisable_for(&hex_id(99), now_ms / 1000);
990        assert_eq!(first.len(), 1);
991        assert_eq!(e.advertisable_rebuild_count(), 1);
992
993        // Same now_secs, but known changed — must rebuild, not serve the stale cached list.
994        e.upsert_known(known_entry(2, 1_000));
995        let second = e.advertisable_for(&hex_id(99), now_ms / 1000);
996        assert_eq!(second.len(), 2, "the newly upserted peer must appear");
997        assert_eq!(
998            e.advertisable_rebuild_count(),
999            2,
1000            "a known mutation must invalidate the cache even at the same now_secs"
1001        );
1002
1003        // Same now_secs, no mutation — must reuse the cache (no third rebuild).
1004        let third = e.advertisable_for(&hex_id(98), now_ms / 1000);
1005        assert_eq!(third.len(), 2);
1006        assert_eq!(
1007            e.advertisable_rebuild_count(),
1008            2,
1009            "an unchanged known set at the same now_secs must reuse the cached build"
1010        );
1011    }
1012
1013    /// The cached base is still correctly filtered per-partner: excluding one link's partner must
1014    /// never leak into another link's view, even though they share the same cached base list.
1015    #[test]
1016    fn cached_base_still_excludes_each_links_own_partner() {
1017        let mut e = eng(&hex_id(0));
1018        e.upsert_known(known_entry(1, 1_000));
1019        e.upsert_known(known_entry(2, 1_000));
1020        let now_secs = 1_000;
1021
1022        let for_1 = e.advertisable_for(&hex_id(1), now_secs);
1023        assert!(
1024            for_1.iter().all(|p| p.peer_id != hex_id(1)),
1025            "peer 1's own link must never be advertised back to it"
1026        );
1027        assert!(for_1.iter().any(|p| p.peer_id == hex_id(2)));
1028
1029        let for_2 = e.advertisable_for(&hex_id(2), now_secs);
1030        assert!(for_2.iter().all(|p| p.peer_id != hex_id(2)));
1031        assert!(for_2.iter().any(|p| p.peer_id == hex_id(1)));
1032    }
1033}