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}