chia_query/peer/pool.rs
1use std::net::SocketAddr;
2use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex as StdMutex};
4use std::time::{Duration, Instant};
5
6use chia_protocol::{CoinStateUpdate, Message, NewPeakWallet, ProtocolMessageTypes};
7use chia_traits::Streamable;
8use futures_util::stream::{FuturesUnordered, StreamExt};
9use tokio::sync::{mpsc, RwLock};
10
11use chia_wallet_sdk::client::Peer;
12use tokio_tungstenite::Connector;
13
14use crate::types::ChiaQueryError;
15use crate::NetworkType;
16
17use super::connect;
18use super::frames::{
19 FrameFanout, FrameSource, FrameSubscription, PoolFrame, SessionEndReason, SessionId,
20};
21use super::plurality::{CORROBORATION_FLOOR, PEER_LIFETIME, PRIORITY_SLOTS};
22
23/// How many dial rounds [`PeerPool::fill_toward_capacity`] may spend reaching capacity.
24///
25/// DERIVED, not chosen, for the same reason [`default_max_peers`](super::plurality::default_max_peers)
26/// is: the priority addresses are tried SEQUENTIALLY and a round admits at most one of them, so
27/// [`PRIORITY_SLOTS`] rounds can be consumed before a single dial reaches discovery at all. Two
28/// more are then owed — one that reaches discovery with every priority address excluded, and one
29/// for the ordinary attrition of a round whose dials did not all land.
30///
31/// A literal `3` was wrong on exactly the host this rule exists for: an operator with
32/// `TRUSTED_FULLNODE` who also runs a node spends round one on the trusted address and round two
33/// on the loopback, leaving ONE round for discovery and no slack. If that round admitted fewer
34/// than [`CORROBORATION_FLOOR`] independent peers the pool could never arm, and there was no
35/// fourth round in which to try. Deriving it means a third priority address widens the budget
36/// instead of silently consuming it.
37const FILL_ROUNDS: usize = PRIORITY_SLOTS + 2;
38
39// ---------------------------------------------------------------------------
40// Pool entry
41// ---------------------------------------------------------------------------
42
43struct PeerEntry {
44 peer: Peer,
45 address: SocketAddr,
46 /// The session this connection publishes its frames under.
47 ///
48 /// Held so an ejection can name a CONNECTION rather than an address: a peer whose session
49 /// ended is removed only if the entry at that address is still the one that ended, never
50 /// its freshly dialled replacement.
51 session: SessionId,
52 /// How this peer was reached. Held so a caller counting independent opinions can tell a
53 /// preferred local node from a discovered one — see [`connect::PeerOrigin`].
54 origin: connect::PeerOrigin,
55 /// When this connection entered the pool, so it can be rotated out on a TIMER.
56 ///
57 /// The pool's other eviction is failure-driven, and failure is not the risk this guards: a set
58 /// of peers that all keep answering is exactly the set an attacker only has to capture once
59 /// (NC-12). Age is the only signal that separates the two.
60 admitted_at: Instant,
61}
62
63// ---------------------------------------------------------------------------
64// PeerRequirement
65// ---------------------------------------------------------------------------
66
67/// Whether at least one peer must connect for the pool to be considered usable.
68///
69/// A client that can fall back to the coinset HTTP tier is still useful with zero
70/// peers, so failing construction on peer discovery would deny a keyless reader over a
71/// peer-tier problem it does not need (dig_ecosystem#2210).
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub enum PeerRequirement {
74 /// Peer discovery failing is fatal.
75 Required,
76 /// An empty pool is acceptable; it refills in the background.
77 Optional,
78}
79
80// ---------------------------------------------------------------------------
81// PeerPool
82// ---------------------------------------------------------------------------
83
84/// Whether the pool holds enough independent voices for a corroborated read.
85///
86/// A two-variant answer rather than a bare count, so a caller cannot accidentally proceed with
87/// "some" corroboration: the insufficient case names what it has AND what it needed, which is the
88/// information a log line or a user-facing message actually requires.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum CorroborationReadiness {
91 /// Enough independent peers, besides the one answering, to corroborate.
92 Armed { corroborators: usize },
93 /// Too few. The read must be REFUSED, never attempted with fewer voices.
94 Insufficient {
95 corroborators: usize,
96 required: usize,
97 },
98}
99
100pub struct PeerPool {
101 entries: RwLock<Vec<PeerEntry>>,
102 next_idx: AtomicUsize,
103 max_peers: usize,
104 tls: Connector,
105 network: NetworkType,
106 connect_timeout: Duration,
107 /// Latest peak height observed from any connected peer's NewPeakWallet
108 /// messages. Updated in the background by receiver handler tasks.
109 peak_height: Arc<AtomicU32>,
110 /// Fans every inbound frame out to the pool's subscribers.
111 ///
112 /// The atomic above answers "how high is the chain"; this carries the frames THEMSELVES, which
113 /// is what a consumer following coin states needs and what its absence forced into a second
114 /// dialled session (dig_ecosystem#2761).
115 fanout: Arc<FrameFanout>,
116 /// Sessions that have STOPPED, waiting to be ejected on the next maintenance pass.
117 ///
118 /// A receiver handler runs in its own task and cannot take the pool's write lock without
119 /// holding a reference to the pool, so it records the death here and lets
120 /// [`maintain`](Self::maintain) act on it. Without this the pool's only removals are
121 /// failure-driven — a dead connection stays in `entries`, counted as a held peer and
122 /// offered to `select_peer`, until a request happens to pick it or `PEER_LIFETIME`
123 /// elapses.
124 dead_sessions: Arc<StdMutex<Vec<FrameSource>>>,
125}
126
127impl PeerPool {
128 /// Spin up the pool by connecting to `max_peers` random full-node peers
129 /// concurrently. Under [`PeerRequirement::Required`] at least one peer must
130 /// succeed, otherwise we return [`ChiaQueryError::PeerDiscoveryFailed`]; under
131 /// [`PeerRequirement::Optional`] an empty pool is returned and refills later.
132 pub async fn new(
133 network: NetworkType,
134 tls: Connector,
135 max_peers: usize,
136 requirement: PeerRequirement,
137 connect_timeout: Duration,
138 ) -> Result<Self, ChiaQueryError> {
139 let pool = Self {
140 entries: RwLock::new(Vec::new()),
141 next_idx: AtomicUsize::new(0),
142 max_peers,
143 tls,
144 network,
145 connect_timeout,
146 peak_height: Arc::new(AtomicU32::new(0)),
147 fanout: Arc::new(FrameFanout::new()),
148 dead_sessions: Arc::new(StdMutex::new(Vec::new())),
149 };
150
151 pool.fill_toward_capacity().await;
152
153 if !pool.has_peers().await {
154 if requirement == PeerRequirement::Required {
155 return Err(ChiaQueryError::PeerDiscoveryFailed);
156 }
157 log::warn!("no peers connected; serving from the coinset fallback until one does");
158 }
159
160 Ok(pool)
161 }
162
163 /// Dial toward capacity in ROUNDS, each excluding what the earlier rounds admitted.
164 ///
165 /// One round of `max_peers` concurrent dials is not enough, and the reason is the priority
166 /// path: every dial offers `TRUSTED_FULLNODE` and the loopback ahead of discovery, and
167 /// concurrent dials know nothing of each other, so on a machine running a full node ALL of
168 /// them return the same local address and every one but the first is discarded as a duplicate.
169 /// A single round therefore leaves a pool of ONE peer on exactly the machines most likely to
170 /// have several available — and one peer is a pool that can never corroborate anything.
171 ///
172 /// A later round excludes what the earlier ones admitted, so the priority addresses are no
173 /// longer offered and the dial falls through to discovery. Rounds stop as soon as one admits
174 /// nothing: a round that admitted nothing is evidence that dialling again would not help
175 /// either, and that bound is what keeps a host with no reachable peers from spending
176 /// `FILL_ROUNDS` whole timeouts on the same answer.
177 async fn fill_toward_capacity(&self) {
178 for _ in 0..FILL_ROUNDS {
179 let held: Vec<SocketAddr> = {
180 let entries = self.entries.read().await;
181 entries.iter().map(|e| e.address).collect()
182 };
183 let wanted = self.max_peers.saturating_sub(held.len());
184 if wanted == 0 {
185 return;
186 }
187
188 let mut dials = FuturesUnordered::new();
189 for _ in 0..wanted {
190 let tls = self.tls.clone();
191 let held = held.clone();
192 let network = self.network;
193 let timeout = self.connect_timeout;
194 dials.push(async move {
195 connect::connect_random_peer_excluding(network, &tls, timeout, &held).await
196 });
197 }
198
199 let mut admitted = 0usize;
200 while let Some(result) = dials.next().await {
201 match result {
202 Ok((peer, addr, receiver, origin)) => {
203 if self.admit_and_follow(peer, addr, receiver, origin).await {
204 admitted += 1;
205 }
206 }
207 Err(e) => log::debug!("peer connect failed: {e}"),
208 }
209 }
210
211 if admitted == 0 {
212 return;
213 }
214 }
215 }
216
217 /// Admit a connection and, if it was admitted, start following its frames.
218 ///
219 /// The ONE path by which a session becomes visible to subscribers, so the ordering it holds
220 /// holds everywhere: the session is identified, then admitted, then ANNOUNCED, and only then
221 /// does its task begin publishing. Announcing before admission would name a session for a
222 /// duplicate that was discarded; announcing after the task started would race its first frame.
223 async fn admit_and_follow(
224 &self,
225 peer: Peer,
226 address: SocketAddr,
227 receiver: mpsc::Receiver<Message>,
228 origin: connect::PeerOrigin,
229 ) -> bool {
230 let source = self.fanout.allocate_session(address);
231 if !self.admit(peer, address, origin, source.session).await {
232 return false;
233 }
234 self.fanout.open_session(source).await;
235 self.spawn_receiver_handler(source, receiver);
236 true
237 }
238
239 /// Latest peak height observed across all connected peers.
240 /// Returns 0 if no peak has been received yet.
241 pub fn peak_height(&self) -> u32 {
242 self.peak_height.load(Ordering::Relaxed)
243 }
244
245 /// Round-robin select a peer from the pool.
246 /// Returns `None` when the pool is empty.
247 pub async fn select_peer(&self) -> Option<(Peer, SocketAddr)> {
248 let entries = self.entries.read().await;
249 if entries.is_empty() {
250 return None;
251 }
252 let idx = self.next_idx.fetch_add(1, Ordering::Relaxed) % entries.len();
253 let entry = &entries[idx];
254 Some((entry.peer.clone(), entry.address))
255 }
256
257 /// Every peer that could CORROBORATE an answer already given by the peer at `asked`.
258 ///
259 /// A corroborating peer must be two things at once, and neither alone is enough:
260 ///
261 /// - **A different address than `asked`.** Asking the same connection twice returns the same
262 /// opinion twice, which reads as agreement while being one voice.
263 /// - **[`PeerOrigin::Discovered`](connect::PeerOrigin).** A peer reached from a preferred
264 /// address — an operator's node, or one on this machine — is an excellent peer to READ from
265 /// and is not evidence about the chain independent of this host, exactly as
266 /// [`independent_peer_count`](Self::independent_peer_count) records.
267 ///
268 /// They are returned ALL AT ONCE, and there is deliberately no singular form of this. Asking
269 /// corroborators one at a time lets the first responder settle a claim about the chain, which
270 /// is exactly the power a hostile peer has (dig_ecosystem#2462) — and a single corroborator
271 /// cannot reach `CORROBORATION_FLOOR` at all, so a caller that took one would be building an
272 /// answer it is not allowed to report as corroborated.
273 ///
274 /// Returns an empty vector when the pool holds nobody who qualifies, which is the honest
275 /// answer that there is nobody to corroborate with.
276 pub async fn select_corroborating_peers(&self, asked: SocketAddr) -> Vec<(Peer, SocketAddr)> {
277 self.entries
278 .read()
279 .await
280 .iter()
281 .filter(|e| Self::is_corroborator(e, asked))
282 .map(|e| (e.peer.clone(), e.address))
283 .collect()
284 }
285
286 /// Remove a peer from the pool and asynchronously connect a replacement.
287 pub async fn eject_peer(&self, addr: SocketAddr) {
288 {
289 let mut entries = self.entries.write().await;
290 entries.retain(|e| e.address != addr);
291 }
292 log::debug!(
293 "peer ejected from pool; will refill on next request (network={:?})",
294 self.network,
295 );
296 }
297
298 /// Whether the pool has at least one usable peer.
299 pub async fn has_peers(&self) -> bool {
300 !self.entries.read().await.is_empty()
301 }
302
303 /// How many peers the pool HOLDS right now.
304 ///
305 /// This is a live count of the connections currently in the pool, not
306 /// [`max_peers`](Self::new)'s target: a pool that is still filling reports what it has, and
307 /// reports the target only once it has reached it. A caller showing this number to a user is
308 /// stating a fact about the machine, so a configured intention must never stand in for it.
309 ///
310 /// A peer is removed by [`eject_peer`](Self::eject_peer), which runs when a request to it
311 /// FAILS. So the count is of peers held and believed usable; a connection that has died
312 /// silently is still counted until something tries to use it. That is the same liveness
313 /// standard [`has_peers`](Self::has_peers) has always answered by, made countable.
314 pub async fn peer_count(&self) -> usize {
315 self.entries.read().await.len()
316 }
317
318 /// How many peers the pool holds that are INDEPENDENT opinions.
319 ///
320 /// [`peer_count`](Self::peer_count) answers "how many connections do I have"; this answers
321 /// "how many of them could corroborate each other". They differ by the peers reached from a
322 /// preferred address — an operator's trusted node or one on this machine — which are excellent
323 /// peers to READ from and are not evidence about the chain independent of this host. A caller
324 /// deciding whether enough separate sources agree MUST use this number, because counting a
325 /// co-resident node as an independent voice is the thing that made a single local process able
326 /// to look like a full peer set (dig_ecosystem#2648).
327 pub async fn independent_peer_count(&self) -> usize {
328 self.entries
329 .read()
330 .await
331 .iter()
332 .filter(|e| e.origin == connect::PeerOrigin::Discovered)
333 .count()
334 }
335
336 /// Whether a peer entry qualifies as a corroborator: not the answering peer, and discovered.
337 fn is_corroborator(entry: &PeerEntry, asked: SocketAddr) -> bool {
338 entry.address != asked && entry.origin == connect::PeerOrigin::Discovered
339 }
340
341 /// Whether the pool can honestly attempt a CORROBORATED read of an answer given by `asked`.
342 ///
343 /// The count is of the peers that WILL be asked to corroborate — precisely the set
344 /// [`select_corroborating_peers`](Self::select_corroborating_peers) returns for the same
345 /// address, because both use [`is_corroborator`](Self::is_corroborator) to decide the set.
346 ///
347 /// **It takes the answering address rather than subtracting one blindly.** An earlier version
348 /// charged the asker's slot against the independent set whatever the asker was, so a read from
349 /// the operator's own node — which is not in that set at all — silently spent an independent
350 /// voice it had never occupied. On the host this crate is sized for, two genuinely independent
351 /// peers agreeing with a co-resident node were downgraded to `Uncorroborated*` and pushed on
352 /// to the centralized coinset tier, which is the opposite of what NC-12 asks for.
353 ///
354 /// **This refuses; it never degrades.** Corroborating against however many peers happen to be
355 /// present turns a four-voice quorum into a two-voice one that still reports itself
356 /// corroborated, and no consumer downstream can tell those apart. A caller handed
357 /// [`CorroborationReadiness::Insufficient`] must decline the read, not proceed with fewer
358 /// voices.
359 pub async fn corroboration_readiness(&self, asked: SocketAddr) -> CorroborationReadiness {
360 let corroborators = self
361 .entries
362 .read()
363 .await
364 .iter()
365 .filter(|e| Self::is_corroborator(e, asked))
366 .count();
367 if corroborators >= CORROBORATION_FLOOR {
368 CorroborationReadiness::Armed { corroborators }
369 } else {
370 CorroborationReadiness::Insufficient {
371 corroborators,
372 required: CORROBORATION_FLOOR,
373 }
374 }
375 }
376
377 /// Rotate out the OLDEST discovered peer that has outlived [`PEER_LIFETIME`], if any.
378 ///
379 /// Returns the address ejected, so a caller can log or refill deliberately. This is NC-12's
380 /// cycling half and it is driven by AGE alone: a peer that has answered every request is
381 /// exactly the peer this removes, because a set that never fails is a set an attacker only has
382 /// to capture once. The pool's other eviction, [`eject_peer`](Self::eject_peer), fires on
383 /// request FAILURE and cannot substitute for this — a captured peer does not fail.
384 ///
385 /// Only [`PeerOrigin::Discovered`](connect::PeerOrigin) entries are rotated. A priority entry
386 /// is the operator's own node or one on this machine; cycling it would re-dial the same
387 /// address, spending a handshake to change nothing.
388 ///
389 /// One per call, so cycling can never empty the pool in a single sweep.
390 pub async fn cycle_expired_peers(&self) -> Option<SocketAddr> {
391 let mut entries = self.entries.write().await;
392 let now = Instant::now();
393
394 let oldest = entries
395 .iter()
396 .enumerate()
397 .filter(|(_, e)| {
398 e.origin == connect::PeerOrigin::Discovered
399 && now.duration_since(e.admitted_at) >= PEER_LIFETIME
400 })
401 .min_by_key(|(_, e)| e.admitted_at)
402 .map(|(idx, e)| (idx, e.address));
403
404 let (idx, address) = oldest?;
405 entries.remove(idx);
406 log::debug!("peer {address} rotated out after {PEER_LIFETIME:?} (NC-12 cycling)");
407 Some(address)
408 }
409
410 /// One maintenance pass: rotate out an over-age peer, then refill toward capacity.
411 ///
412 /// Cycling before refilling is deliberate. Refilling first would find the pool at capacity and
413 /// do nothing, so the rotation would leave a permanently smaller pool.
414 pub async fn maintain(&self) {
415 self.eject_dead_sessions().await;
416 self.cycle_expired_peers().await;
417 self.try_refill().await;
418 }
419
420 /// Remove the peers whose sessions have ENDED.
421 ///
422 /// Session death is the pool's third eviction reason and it is neither of the other two: a
423 /// disconnected or protocol-violating peer has not failed a request, so
424 /// [`eject_peer`](Self::eject_peer) never fires for it, and it need not be old, so cycling may
425 /// be minutes away. Until it is removed the pool counts it as held and `select_peer` keeps
426 /// offering it — a peer count that overstates what the pool can actually reach.
427 ///
428 /// Matched on address AND session, so a replacement already dialled to the same address is
429 /// never removed by its predecessor's death.
430 async fn eject_dead_sessions(&self) {
431 let dead = std::mem::take(&mut *self.dead_sessions_guard());
432 if dead.is_empty() {
433 return;
434 }
435
436 let mut entries = self.entries.write().await;
437 entries.retain(|entry| {
438 let died = dead
439 .iter()
440 .any(|d| d.address == entry.address && d.session == entry.session);
441 if died {
442 log::debug!("peer {} ejected: its session ended", entry.address);
443 }
444 !died
445 });
446 }
447
448 /// The dead-session list, recovering from a poisoned lock rather than panicking.
449 ///
450 /// A panic in one handler task must not take the pool's maintenance down with it, and the list
451 /// is a plain `Vec` with no invariant a poisoned lock could have left half-applied.
452 fn dead_sessions_guard(&self) -> std::sync::MutexGuard<'_, Vec<FrameSource>> {
453 self.dead_sessions
454 .lock()
455 .unwrap_or_else(|poisoned| poisoned.into_inner())
456 }
457
458 /// Admit a connection, or reject it, deciding under the WRITE lock.
459 ///
460 /// Returns whether it was admitted. Rejected because the pool is full, or because its address
461 /// is already held — a pool of N connections to one address reports itself healthy while being
462 /// a single point of both failure and deceit (dig_ecosystem#2648).
463 ///
464 /// **Both checks are made while HOLDING the write lock, and that placement is the whole
465 /// correctness of this.** Dials run concurrently, so any check made before acquiring the lock —
466 /// under the read lock, or by the caller — is a time-of-check/time-of-use gap: two fills of the
467 /// same address each observe it absent, then each pushes, and the duplicate is admitted by
468 /// exactly the code written to prevent it. The check and the push must be one critical section.
469 async fn admit(
470 &self,
471 peer: Peer,
472 address: SocketAddr,
473 origin: connect::PeerOrigin,
474 session: SessionId,
475 ) -> bool {
476 let mut entries = self.entries.write().await;
477
478 if entries.len() >= self.max_peers {
479 log::debug!("peer {address} not admitted: pool is at capacity");
480 return false;
481 }
482 if entries.iter().any(|e| e.address == address) {
483 log::debug!("peer {address} not admitted: already held");
484 return false;
485 }
486
487 entries.push(PeerEntry {
488 peer,
489 address,
490 origin,
491 session,
492 admitted_at: Instant::now(),
493 });
494 log::debug!("peer admitted: {address} ({origin:?})");
495 true
496 }
497
498 /// If the pool is under capacity, try to connect one new peer.
499 /// Also spawns a background task to handle its inbound `NewPeakWallet`
500 /// messages.
501 pub async fn try_refill(&self) {
502 let held: Vec<SocketAddr> = {
503 let entries = self.entries.read().await;
504 if entries.len() >= self.max_peers {
505 return;
506 }
507 entries.iter().map(|e| e.address).collect()
508 };
509
510 // `held` is a hint to the dial, not the guard: it saves dialling an address already in the
511 // pool (the local one is offered on every call), and it may be stale the moment it is read.
512 // `admit` re-decides under the write lock, which is where the invariant actually holds.
513 match connect::connect_random_peer_excluding(
514 self.network,
515 &self.tls,
516 self.connect_timeout,
517 &held,
518 )
519 .await
520 {
521 Ok((peer, addr, receiver, origin)) => {
522 if self.admit_and_follow(peer, addr, receiver, origin).await {
523 log::debug!("replacement peer connected: {addr}");
524 }
525 }
526 Err(e) => log::warn!("replacement peer connect failed: {e}"),
527 }
528 }
529
530 // -----------------------------------------------------------------------
531 // Receiver helpers (handle NewPeakWallet from peers)
532 // -----------------------------------------------------------------------
533
534 /// Spawn a background task that reads inbound messages from one peer session, updating the
535 /// shared peak height and fanning every recognised frame out to the pool's subscribers.
536 ///
537 /// `source` identifies the session. Every frame this task emits carries it, so a subscriber
538 /// can tell one held peer's frames from another's — which is what lets it follow the peer it
539 /// chose and eject one whose frames it rejected.
540 ///
541 /// **`pub(crate)`, so the one-path claim on [`admit_and_follow`](Self::admit_and_follow) holds
542 /// across the crate boundary too.** A caller outside the crate could otherwise supply its own
543 /// [`FrameSource`] and publish frames under a session the pool never allocated — attribution
544 /// that is unforgeable in-crate becomes forgeable the moment the constructor is exported.
545 ///
546 /// **The task never ends quietly.** Both ways a session can stop — the transport closing, and
547 /// a message this crate cannot decode — publish a [`PoolFrame::SessionEnded`] and record the
548 /// death for [`eject_dead_sessions`](Self::eject_dead_sessions). Returning silently would
549 /// leave a subscriber unable to distinguish a peer that stopped from a chain that is quiet,
550 /// and would leave the pool holding a connection nothing will ever read from.
551 pub(crate) fn spawn_receiver_handler(
552 &self,
553 source: FrameSource,
554 mut receiver: mpsc::Receiver<Message>,
555 ) {
556 let peak = Arc::clone(&self.peak_height);
557 let fanout = Arc::clone(&self.fanout);
558 let dead_sessions = Arc::clone(&self.dead_sessions);
559
560 tokio::spawn(async move {
561 let reason = loop {
562 let Some(msg) = receiver.recv().await else {
563 break SessionEndReason::Disconnected;
564 };
565
566 match msg.msg_type {
567 ProtocolMessageTypes::NewPeakWallet => {
568 let Ok(new_peak) = NewPeakWallet::from_bytes(&msg.data) else {
569 log::warn!(
570 "peer {} sent an undecodable NewPeakWallet; ending the session",
571 source.address
572 );
573 break SessionEndReason::UndecodableFrame;
574 };
575 let prev = peak.fetch_max(new_peak.height, Ordering::Relaxed);
576 if new_peak.height > prev {
577 log::debug!(
578 "new peak from peer {}: {}",
579 source.address,
580 new_peak.height
581 );
582 }
583 fanout
584 .publish(
585 source,
586 PoolFrame::Peak {
587 height: new_peak.height,
588 header_hash: new_peak.header_hash,
589 },
590 )
591 .await;
592 }
593 ProtocolMessageTypes::CoinStateUpdate => {
594 let Ok(update) = CoinStateUpdate::from_bytes(&msg.data) else {
595 log::warn!(
596 "peer {} sent an undecodable CoinStateUpdate; ending the session",
597 source.address
598 );
599 break SessionEndReason::UndecodableFrame;
600 };
601 fanout
602 .publish(
603 source,
604 PoolFrame::CoinStates {
605 height: update.height,
606 fork_height: update.fork_height,
607 items: update.items,
608 },
609 )
610 .await;
611 }
612 _ => {}
613 }
614 };
615
616 dead_sessions
617 .lock()
618 .unwrap_or_else(|poisoned| poisoned.into_inner())
619 .push(source);
620 fanout
621 .publish(source, PoolFrame::SessionEnded { reason })
622 .await;
623 });
624 }
625
626 /// Subscribe to this pool's frames, with room for `capacity` unread ones.
627 ///
628 /// Falling further behind than `capacity` ENDS the subscription — see
629 /// [`FrameSubscription`](super::frames::FrameSubscription) for why a gap is not an option.
630 pub async fn subscribe_frames(&self, capacity: usize) -> FrameSubscription {
631 self.fanout.subscribe(capacity).await
632 }
633}
634
635/// Construction and admission reachable from OTHER modules' tests.
636///
637/// [`PeerPool::new`] dials the network, so a test of anything built ON the pool — the backend's
638/// absence corroboration, for one — cannot use it. These wrap the private internals rather than
639/// widening them, so production code keeps exactly one admission path.
640#[cfg(test)]
641impl PeerPool {
642 pub(crate) fn for_tests(max_peers: usize) -> Self {
643 Self {
644 entries: RwLock::new(Vec::new()),
645 next_idx: AtomicUsize::new(0),
646 max_peers,
647 tls: connect::create_generated_tls().expect("generate a TLS identity"),
648 network: NetworkType::Mainnet,
649 connect_timeout: Duration::from_millis(1),
650 peak_height: Arc::new(AtomicU32::new(0)),
651 fanout: Arc::new(FrameFanout::new()),
652 dead_sessions: Arc::new(StdMutex::new(Vec::new())),
653 }
654 }
655
656 pub(crate) async fn admit_for_tests(
657 &self,
658 peer: Peer,
659 address: SocketAddr,
660 origin: connect::PeerOrigin,
661 ) -> bool {
662 self.admitted(peer, address, origin).await
663 }
664
665 /// Admit a connection under a freshly allocated session, reporting only whether it was taken.
666 ///
667 /// Production admits through [`admit_and_follow`](Self::admit_and_follow), which also starts
668 /// the session; a test that only cares about the pool's membership uses this so it does not
669 /// have to invent a receiver it will never feed.
670 pub(crate) async fn admitted(
671 &self,
672 peer: Peer,
673 address: SocketAddr,
674 origin: connect::PeerOrigin,
675 ) -> bool {
676 let session = self.fanout.allocate_session(address).session;
677 self.admit(peer, address, origin, session).await
678 }
679
680 /// Admit a connection AND follow `receiver`, exactly as a real dial would.
681 pub(crate) async fn admit_and_follow_for_tests(
682 &self,
683 peer: Peer,
684 address: SocketAddr,
685 receiver: mpsc::Receiver<Message>,
686 origin: connect::PeerOrigin,
687 ) -> bool {
688 self.admit_and_follow(peer, address, receiver, origin).await
689 }
690
691 /// Run only the dead-session half of [`maintain`](Self::maintain), which does not dial.
692 pub(crate) async fn eject_dead_sessions_for_tests(&self) {
693 self.eject_dead_sessions().await;
694 }
695
696 /// The addresses currently held, in pool order.
697 pub(crate) async fn held_addresses_for_tests(&self) -> Vec<SocketAddr> {
698 self.entries
699 .read()
700 .await
701 .iter()
702 .map(|e| e.address)
703 .collect()
704 }
705
706 /// Admit a connection as if it had entered the pool at `admitted_at`.
707 ///
708 /// Age is otherwise only reachable by waiting, and a test that waits five minutes is a test
709 /// nobody runs. Wrapping the private field rather than widening it keeps production code on
710 /// exactly one admission path.
711 pub(crate) async fn admit_at_for_tests(
712 &self,
713 peer: Peer,
714 address: SocketAddr,
715 origin: connect::PeerOrigin,
716 admitted_at: Instant,
717 ) -> bool {
718 if !self.admitted(peer, address, origin).await {
719 return false;
720 }
721 let mut entries = self.entries.write().await;
722 if let Some(entry) = entries.iter_mut().find(|e| e.address == address) {
723 entry.admitted_at = admitted_at;
724 }
725 true
726 }
727}
728
729#[cfg(test)]
730mod tests {
731 use super::*;
732 use crate::peer::connect::{create_generated_tls, PeerOrigin};
733 use crate::peer::plurality::{default_max_peers, QUORUM_SAMPLE};
734 use crate::peer::test_support::{address, loopback_peer};
735
736 use super::PeerPool as _Pool;
737 fn empty_pool(max_peers: usize) -> PeerPool {
738 _Pool::for_tests(max_peers)
739 }
740
741 /// **The defect, and the one shape that separates a locked re-check from a TOCTOU dedupe.**
742 ///
743 /// Eight fills of the SAME address are admitted CONCURRENTLY, which is how the pool fills in
744 /// production: `PeerPool::new` races `max_peers` dials with no knowledge of each other, and each
745 /// may return the same priority address. A dedupe that reads the entry list before taking the
746 /// write lock passes a sequential test and fails this one — every task observes the address
747 /// absent, then every task pushes.
748 ///
749 /// `max_peers` is 8, not 1, deliberately: a capacity of one would make the pool reject the
750 /// duplicates for being FULL rather than for being duplicates, and would stay green with the
751 /// distinctness check deleted entirely.
752 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
753 async fn one_address_cannot_fill_the_pool_however_many_fills_race() {
754 let pool = Arc::new(empty_pool(8));
755 let peer = loopback_peer().await;
756 let occupied = address(1);
757
758 let mut fills = Vec::new();
759 for _ in 0..8 {
760 let pool = Arc::clone(&pool);
761 let peer = peer.clone();
762 fills.push(tokio::spawn(async move {
763 pool.admitted(peer, occupied, PeerOrigin::Priority).await
764 }));
765 }
766
767 let admitted = futures_util::future::join_all(fills)
768 .await
769 .into_iter()
770 .filter(|r| *r.as_ref().expect("the admission task must not panic"))
771 .count();
772
773 assert_eq!(
774 admitted, 1,
775 "exactly one fill of an address may be admitted"
776 );
777 assert_eq!(
778 pool.peer_count().await,
779 1,
780 "eight concurrent fills of one address must leave one connection, not eight"
781 );
782 }
783
784 /// The control that keeps the test above honest: concurrency itself must not cost admissions.
785 ///
786 /// Without this, an `admit` that rejected everything after the first — or that lost racing
787 /// pushes — would satisfy the distinctness test while breaking the pool.
788 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
789 async fn distinct_addresses_all_fill_concurrently() {
790 let pool = Arc::new(empty_pool(8));
791 let peer = loopback_peer().await;
792
793 let mut fills = Vec::new();
794 for octet in 1..=8u8 {
795 let pool = Arc::clone(&pool);
796 let peer = peer.clone();
797 fills.push(tokio::spawn(async move {
798 pool.admitted(peer, address(octet), PeerOrigin::Discovered)
799 .await
800 }));
801 }
802 futures_util::future::join_all(fills).await;
803
804 assert_eq!(
805 pool.peer_count().await,
806 8,
807 "eight distinct addresses must all be admitted"
808 );
809 }
810
811 /// Capacity is enforced in the same critical section, so racing fills cannot overshoot it.
812 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
813 async fn concurrent_fills_never_exceed_max_peers() {
814 let pool = Arc::new(empty_pool(3));
815 let peer = loopback_peer().await;
816
817 let mut fills = Vec::new();
818 for octet in 1..=10u8 {
819 let pool = Arc::clone(&pool);
820 let peer = peer.clone();
821 fills.push(tokio::spawn(async move {
822 pool.admitted(peer, address(octet), PeerOrigin::Discovered)
823 .await
824 }));
825 }
826 futures_util::future::join_all(fills).await;
827
828 assert_eq!(pool.peer_count().await, 3, "max_peers is a hard ceiling");
829 }
830
831 /// **A preferred peer is not a corroborating one.**
832 ///
833 /// Two `Discovered` peers sit beside one `Priority` peer, so the two counts differ by exactly
834 /// the priority entry. A single-origin fixture cannot show that: all-priority or all-discovered
835 /// both make the two counts move together, which an implementation returning `peer_count` for
836 /// both would satisfy.
837 #[tokio::test]
838 async fn a_preferred_peer_is_held_but_not_counted_as_an_independent_opinion() {
839 let pool = empty_pool(5);
840 let peer = loopback_peer().await;
841
842 assert!(
843 pool.admitted(peer.clone(), address(1), PeerOrigin::Priority)
844 .await
845 );
846 assert!(
847 pool.admitted(peer.clone(), address(2), PeerOrigin::Discovered)
848 .await
849 );
850 assert!(
851 pool.admitted(peer, address(3), PeerOrigin::Discovered)
852 .await
853 );
854
855 assert_eq!(pool.peer_count().await, 3, "three connections are held");
856 assert_eq!(
857 pool.independent_peer_count().await,
858 2,
859 "the co-resident peer is held and read from, but is not an independent voice"
860 );
861 }
862
863 /// An ejected address is admissible again — distinctness must not become a permanent ban.
864 #[tokio::test]
865 async fn an_ejected_address_can_be_admitted_again() {
866 let pool = empty_pool(5);
867 let peer = loopback_peer().await;
868 let addr = address(1);
869
870 assert!(
871 pool.admitted(peer.clone(), addr, PeerOrigin::Discovered)
872 .await
873 );
874 assert!(
875 !pool
876 .admitted(peer.clone(), addr, PeerOrigin::Discovered)
877 .await,
878 "still held, so still a duplicate"
879 );
880
881 pool.eject_peer(addr).await;
882
883 assert!(
884 pool.admitted(peer, addr, PeerOrigin::Discovered).await,
885 "a re-dialled peer must be admissible after ejection"
886 );
887 assert_eq!(pool.peer_count().await, 1);
888 }
889
890 /// `max_peers: 0` attempts no connection at all, so the pool is deterministically
891 /// empty offline — an exact, network-free fixture for the empty-pool branch.
892 async fn pool_with_no_connection_attempts(
893 requirement: PeerRequirement,
894 ) -> Result<PeerPool, ChiaQueryError> {
895 PeerPool::new(
896 NetworkType::Mainnet,
897 create_generated_tls().expect("generate a TLS identity"),
898 0,
899 requirement,
900 Duration::from_millis(1),
901 )
902 .await
903 }
904
905 /// The control: an empty pool is still fatal when nothing can serve in its place.
906 #[tokio::test]
907 async fn empty_pool_is_fatal_when_peers_are_required() {
908 assert!(matches!(
909 pool_with_no_connection_attempts(PeerRequirement::Required).await,
910 Err(ChiaQueryError::PeerDiscoveryFailed)
911 ));
912 }
913
914 /// The fix: with a fallback able to serve, an empty pool must not deny the client.
915 #[tokio::test]
916 async fn empty_pool_is_tolerated_when_peers_are_optional() {
917 let pool = pool_with_no_connection_attempts(PeerRequirement::Optional)
918 .await
919 .expect("an optional peer pool must construct with zero peers");
920 assert!(!pool.has_peers().await);
921 }
922
923 /// **The count is what is HELD, never what was asked for.**
924 ///
925 /// Built by hand rather than through [`PeerPool::new`] so `max_peers` can be a realistic 5
926 /// while the pool provably holds nothing — the one shape that separates a measurement from a
927 /// configured intention. A `peer_count` that returned `max_peers` would satisfy every
928 /// assertion reachable through the offline constructor, whose `max_peers` is necessarily 0,
929 /// and would then report "5 peers" on a machine holding none.
930 #[tokio::test]
931 async fn an_unfilled_pool_counts_what_it_holds_not_the_target_it_was_given() {
932 let pool = empty_pool(5);
933
934 assert_eq!(
935 pool.peer_count().await,
936 0,
937 "held is 0 while the target is 5"
938 );
939 assert!(!pool.has_peers().await);
940 }
941
942 /// **NC-12 cycling: an over-age peer is rotated out by the TIMER, with nothing having failed.**
943 ///
944 /// The distinguishing fixture is the pair. One entry was admitted well past `PEER_LIFETIME`
945 /// ago, one moments ago, and BOTH are healthy — no request is made, so `eject_peer`, the
946 /// pool's only other eviction, is never reached. A cycler that fired on failure, or one that
947 /// swept indiscriminately, changes the outcome here: the first leaves both peers, the second
948 /// removes both.
949 #[tokio::test]
950 async fn an_over_age_peer_is_rotated_out_with_no_request_having_failed() {
951 let pool = empty_pool(7);
952 let peer = loopback_peer().await;
953
954 let stale = address(1);
955 let fresh = address(2);
956 let long_past = Instant::now() - (PEER_LIFETIME + Duration::from_secs(100));
957
958 assert!(
959 pool.admit_at_for_tests(peer.clone(), stale, PeerOrigin::Discovered, long_past)
960 .await
961 );
962 assert!(
963 pool.admit_at_for_tests(peer, fresh, PeerOrigin::Discovered, Instant::now())
964 .await
965 );
966
967 let rotated = pool.cycle_expired_peers().await;
968
969 assert_eq!(
970 rotated,
971 Some(stale),
972 "the peer past its lifetime must be rotated out on age alone"
973 );
974 assert_eq!(
975 pool.held_addresses_for_tests().await,
976 vec![fresh],
977 "cycling must remove the over-age peer and keep the fresh one"
978 );
979 }
980
981 /// The control: a pool of healthy, recent peers is left ALONE.
982 ///
983 /// Without it, a cycler that ejected the oldest entry unconditionally — no lifetime check at
984 /// all — would satisfy the test above while churning the pool on every pass.
985 #[tokio::test]
986 async fn peers_within_their_lifetime_are_not_rotated() {
987 let pool = empty_pool(7);
988 let peer = loopback_peer().await;
989
990 for octet in 1..=3u8 {
991 assert!(
992 pool.admit_at_for_tests(
993 peer.clone(),
994 address(octet),
995 PeerOrigin::Discovered,
996 Instant::now() - (PEER_LIFETIME - Duration::from_secs(30)),
997 )
998 .await
999 );
1000 }
1001
1002 assert_eq!(
1003 pool.cycle_expired_peers().await,
1004 None,
1005 "a peer inside its lifetime must not be rotated"
1006 );
1007 assert_eq!(pool.peer_count().await, 3);
1008 }
1009
1010 /// A priority entry is not rotated: cycling it would re-dial the same address.
1011 #[tokio::test]
1012 async fn an_over_age_priority_peer_is_not_rotated() {
1013 let pool = empty_pool(7);
1014 let peer = loopback_peer().await;
1015 let long_past = Instant::now() - (PEER_LIFETIME + Duration::from_secs(100));
1016
1017 assert!(
1018 pool.admit_at_for_tests(peer, address(1), PeerOrigin::Priority, long_past)
1019 .await
1020 );
1021
1022 assert_eq!(pool.cycle_expired_peers().await, None);
1023 assert_eq!(pool.peer_count().await, 1);
1024 }
1025
1026 /// **The sizing property: one priority entry must not cost the quorum.**
1027 ///
1028 /// A pool filled to the SHIPPED default — one priority entry, which is the ordinary case
1029 /// because the dialler tries the loopback first, and discovered peers for the rest — must
1030 /// still hold more independent voices than a sample needs. At the previous default of 5 this
1031 /// fixture yields four, and the assertion fails.
1032 #[tokio::test]
1033 async fn one_priority_entry_does_not_cost_the_quorum() {
1034 let pool = empty_pool(default_max_peers());
1035 let peer = loopback_peer().await;
1036
1037 assert!(
1038 pool.admitted(peer.clone(), address(1), PeerOrigin::Priority)
1039 .await
1040 );
1041 for octet in 2..=(default_max_peers() as u8) {
1042 assert!(
1043 pool.admitted(peer.clone(), address(octet), PeerOrigin::Discovered)
1044 .await
1045 );
1046 }
1047
1048 assert_eq!(pool.peer_count().await, default_max_peers());
1049 // A whole sample, plus the one session a subscriber is following and which therefore
1050 // cannot corroborate itself.
1051 let owed = QUORUM_SAMPLE + 1;
1052 let independent = pool.independent_peer_count().await;
1053 assert!(
1054 independent >= owed,
1055 "a full pool holding one priority entry still owes a whole sample plus the session being followed; it holds {independent}"
1056 );
1057 }
1058
1059 /// **Below the floor the pool REFUSES rather than corroborating with fewer voices.**
1060 ///
1061 /// Two discovered peers means exactly one corroborator once the answering peer is set aside —
1062 /// one short. The wrong implementation is not an error, it is a *degradation*: proceeding on
1063 /// that single second opinion and still calling the result corroborated. So the assertion is
1064 /// on the refusal AND on the count it reports, which a bare boolean could not distinguish from
1065 /// an empty pool.
1066 #[tokio::test]
1067 async fn a_pool_below_the_corroboration_floor_refuses_rather_than_degrading() {
1068 let pool = empty_pool(7);
1069 let peer = loopback_peer().await;
1070
1071 for octet in 1..=2u8 {
1072 assert!(
1073 pool.admitted(peer.clone(), address(octet), PeerOrigin::Discovered)
1074 .await
1075 );
1076 }
1077
1078 assert_eq!(
1079 pool.corroboration_readiness(address(1)).await,
1080 CorroborationReadiness::Insufficient {
1081 corroborators: 1,
1082 required: CORROBORATION_FLOOR,
1083 }
1084 );
1085 }
1086
1087 /// The control: at the floor exactly, corroboration ARMS.
1088 ///
1089 /// Pins the bound from the other side — a gate that refused everything would satisfy the test
1090 /// above on its own.
1091 #[tokio::test]
1092 async fn a_pool_at_the_corroboration_floor_arms() {
1093 let pool = empty_pool(7);
1094 let peer = loopback_peer().await;
1095
1096 for octet in 1..=(CORROBORATION_FLOOR as u8 + 1) {
1097 assert!(
1098 pool.admitted(peer.clone(), address(octet), PeerOrigin::Discovered)
1099 .await
1100 );
1101 }
1102
1103 assert_eq!(
1104 pool.corroboration_readiness(address(1)).await,
1105 CorroborationReadiness::Armed {
1106 corroborators: CORROBORATION_FLOOR
1107 }
1108 );
1109 }
1110
1111 /// A preferred peer is not a corroborator, so it cannot arm the gate.
1112 ///
1113 /// Same peer COUNT as the arming control above, different origins — the one fixture shape that
1114 /// separates "enough connections" from "enough independent voices" (dig_ecosystem#2648).
1115 #[tokio::test]
1116 async fn priority_peers_cannot_arm_the_corroboration_gate() {
1117 let pool = empty_pool(7);
1118 let peer = loopback_peer().await;
1119
1120 assert!(
1121 pool.admitted(peer.clone(), address(1), PeerOrigin::Discovered)
1122 .await
1123 );
1124 for octet in 2..=(CORROBORATION_FLOOR as u8 + 1) {
1125 assert!(
1126 pool.admitted(peer.clone(), address(octet), PeerOrigin::Priority)
1127 .await
1128 );
1129 }
1130
1131 assert!(matches!(
1132 pool.corroboration_readiness(address(1)).await,
1133 CorroborationReadiness::Insufficient { .. }
1134 ));
1135 }
1136
1137 /// **A PRIORITY peer answering does not spend an independent voice it never occupied.**
1138 ///
1139 /// The fixture varies exactly one thing against
1140 /// [`a_pool_at_the_corroboration_floor_arms`]: WHO was asked. The independent set is a floor's
1141 /// worth on its own, and the answer comes from a preferred peer that is not in that set — so
1142 /// charging the asker's slot against it, as a blind `- 1` does, reports a pool with two
1143 /// genuine corroborators as having one.
1144 ///
1145 /// That is not a missed opportunity, it is a downgrade with a destination: the answer becomes
1146 /// `Uncorroborated*` and the router settles it against the centralized coinset tier
1147 /// (`router.rs`), substituting one HTTPS source for the untrusted plurality NC-12 asks for. On
1148 /// a host with `TRUSTED_FULLNODE` or a co-resident node — the configuration this pool is sized
1149 /// for — that is the ordinary path, not an edge case.
1150 #[tokio::test]
1151 async fn a_priority_peer_answering_does_not_consume_an_independent_slot() {
1152 let pool = empty_pool(7);
1153 let peer = loopback_peer().await;
1154
1155 for octet in 1..=(CORROBORATION_FLOOR as u8) {
1156 assert!(
1157 pool.admitted(peer.clone(), address(octet), PeerOrigin::Discovered)
1158 .await
1159 );
1160 }
1161 let preferred = address(200);
1162 assert!(
1163 pool.admitted(peer.clone(), preferred, PeerOrigin::Priority)
1164 .await
1165 );
1166
1167 assert_eq!(
1168 pool.corroboration_readiness(preferred).await,
1169 CorroborationReadiness::Armed {
1170 corroborators: CORROBORATION_FLOOR
1171 },
1172 "a preferred peer is not an independent voice, so answering from one cannot cost the independent set a member"
1173 );
1174 }
1175
1176 #[tokio::test]
1177 async fn corroboration_readiness_and_select_use_the_same_predicate() {
1178 let pool = empty_pool(7);
1179 let peer = loopback_peer().await;
1180
1181 // Build a pool with mixed origins: Priority, Discovered, and the asker itself.
1182 let asked = address(100);
1183 let priority = address(200);
1184 let discovered_1 = address(1);
1185 let discovered_2 = address(2);
1186
1187 assert!(
1188 pool.admitted(peer.clone(), asked, PeerOrigin::Discovered)
1189 .await
1190 );
1191 assert!(
1192 pool.admitted(peer.clone(), priority, PeerOrigin::Priority)
1193 .await
1194 );
1195 assert!(
1196 pool.admitted(peer.clone(), discovered_1, PeerOrigin::Discovered)
1197 .await
1198 );
1199 assert!(
1200 pool.admitted(peer.clone(), discovered_2, PeerOrigin::Discovered)
1201 .await
1202 );
1203
1204 // Both should see exactly 2 corroborators: discovered_1 and discovered_2.
1205 // Not `asked` (excluded by address), not `priority` (excluded by origin).
1206 let readiness = pool.corroboration_readiness(asked).await;
1207 let selected = pool.select_corroborating_peers(asked).await;
1208
1209 let readiness_count = match readiness {
1210 CorroborationReadiness::Armed { corroborators } => corroborators,
1211 CorroborationReadiness::Insufficient { corroborators, .. } => corroborators,
1212 };
1213
1214 assert_eq!(
1215 readiness_count,
1216 selected.len(),
1217 "corroboration_readiness and select_corroborating_peers must use the same predicate"
1218 );
1219 assert_eq!(
1220 readiness_count, 2,
1221 "both should count exactly the two Discovered peers that are not the asker"
1222 );
1223 }
1224
1225 /// **`FILL_ROUNDS` is DERIVED from the dialler, so a new priority address cannot starve it.**
1226 ///
1227 /// Network-free arithmetic, in the shape of the pool-sizing derivations: the priority
1228 /// addresses are tried sequentially and a round admits at most one of them, so `PRIORITY_SLOTS`
1229 /// rounds can pass before any dial reaches discovery. What is left must still be enough for a
1230 /// discovery round AND one of attrition.
1231 ///
1232 /// The literal `3` this replaced satisfied that only while `PRIORITY_SLOTS` was 1. At 2 it
1233 /// left exactly one discovery round with no slack, on precisely the host the rounds exist for.
1234 #[test]
1235 fn fill_rounds_leaves_a_discovery_round_and_one_of_attrition() {
1236 let for_discovery = FILL_ROUNDS - PRIORITY_SLOTS;
1237
1238 assert_eq!(
1239 for_discovery, 2,
1240 "FILL_ROUNDS ({FILL_ROUNDS}) minus the {PRIORITY_SLOTS} rounds the priority addresses can consume must leave a discovery round and one of attrition"
1241 );
1242 assert_eq!(
1243 FILL_ROUNDS,
1244 PRIORITY_SLOTS + 2,
1245 "FILL_ROUNDS stays coupled to PRIORITY_SLOTS + 2"
1246 );
1247 // This assertion and the one above state the same mathematical fact: FILL_ROUNDS - PRIORITY_SLOTS == 2
1248 // and FILL_ROUNDS == PRIORITY_SLOTS + 2 are equivalent. Both are present because changing either
1249 // one invalidates FILL_ROUNDS' budget, but they do not add independent verification.
1250 }
1251 // -----------------------------------------------------------------------
1252 // Session lifecycle: attribution, loud endings, and the ejection they drive
1253 // -----------------------------------------------------------------------
1254
1255 use chia_protocol::Bytes32;
1256
1257 use super::super::frames::SourcedFrame;
1258
1259 /// A well-formed `NewPeakWallet` message at `height`.
1260 fn peak_message(height: u32) -> Message {
1261 let peak = NewPeakWallet::new(Bytes32::new([height as u8; 32]), height, 0, 0);
1262 Message {
1263 msg_type: ProtocolMessageTypes::NewPeakWallet,
1264 id: None,
1265 data: peak.to_bytes().expect("encode a peak").into(),
1266 }
1267 }
1268
1269 /// A `NewPeakWallet` message whose BODY cannot be decoded.
1270 ///
1271 /// The type byte is honest and the payload is one byte, far short of the
1272 /// `Bytes32 + u32 + u128 + u32` the body requires — which is what any peer can send at will,
1273 /// costing it nothing.
1274 fn undecodable_peak_message() -> Message {
1275 Message {
1276 msg_type: ProtocolMessageTypes::NewPeakWallet,
1277 id: None,
1278 data: vec![0x00].into(),
1279 }
1280 }
1281
1282 /// Drain the frames that have arrived, waiting briefly for the handler task to run.
1283 ///
1284 /// The handler is a separate task, so a bare `try_recv` races it. This yields until the
1285 /// expected number of frames has arrived or the budget runs out, and returns whatever it has —
1286 /// so a test asserting on the CONTENT fails on its own assertion rather than on a timeout.
1287 async fn drain_at_least(
1288 subscription: &mut FrameSubscription,
1289 wanted: usize,
1290 ) -> Vec<SourcedFrame> {
1291 let mut seen = Vec::new();
1292 for _ in 0..200 {
1293 while let Ok(frame) = subscription.try_recv() {
1294 seen.push(frame);
1295 }
1296 if seen.len() >= wanted {
1297 break;
1298 }
1299 tokio::time::sleep(Duration::from_millis(5)).await;
1300 }
1301 seen
1302 }
1303
1304 /// Admit a peer at `addr` and follow a channel the test itself feeds.
1305 async fn followed_session(
1306 pool: &PeerPool,
1307 addr: SocketAddr,
1308 ) -> (mpsc::Sender<Message>, FrameSource) {
1309 let (sender, receiver) = mpsc::channel(8);
1310 let peer = loopback_peer().await;
1311 let before = pool.held_addresses_for_tests().await.len();
1312 assert!(
1313 pool.admit_and_follow_for_tests(peer, addr, receiver, PeerOrigin::Discovered)
1314 .await,
1315 "the fixture peer must be admitted"
1316 );
1317 assert_eq!(pool.held_addresses_for_tests().await.len(), before + 1);
1318
1319 let source = pool
1320 .entries
1321 .read()
1322 .await
1323 .iter()
1324 .find(|e| e.address == addr)
1325 .map(|e| FrameSource {
1326 address: e.address,
1327 session: e.session,
1328 })
1329 .expect("the admitted entry");
1330 (sender, source)
1331 }
1332
1333 /// **A frame carries the address of the peer that sent it, all the way from the socket.**
1334 ///
1335 /// TWO sessions are followed and each is fed a peak of its own. A handler that published
1336 /// without attribution — or that attributed every frame to one session — gives both frames the
1337 /// same source and fails here; a one-session fixture cannot tell those apart from correct
1338 /// behaviour.
1339 ///
1340 /// This is the property whose absence let any held peer's `CoinStateUpdate` reach a subscriber
1341 /// as if it came from the peer that subscriber had chosen to follow.
1342 #[tokio::test]
1343 async fn a_frame_reaching_a_subscriber_names_the_session_it_came_from() {
1344 let pool = empty_pool(4);
1345 let mut subscription = pool.subscribe_frames(32).await;
1346
1347 let (first, first_source) = followed_session(&pool, address(1)).await;
1348 let (second, second_source) = followed_session(&pool, address(2)).await;
1349
1350 first.send(peak_message(100)).await.expect("send");
1351 second.send(peak_message(200)).await.expect("send");
1352
1353 let seen = drain_at_least(&mut subscription, 4).await;
1354
1355 let peaks: Vec<(SocketAddr, u32)> = seen
1356 .iter()
1357 .filter_map(|f| match f.frame {
1358 PoolFrame::Peak { height, .. } => Some((f.source.address, height)),
1359 _ => None,
1360 })
1361 .collect();
1362
1363 assert!(
1364 peaks.contains(&(address(1), 100)),
1365 "peer 1's peak must arrive under peer 1's address: {peaks:?}"
1366 );
1367 assert!(
1368 peaks.contains(&(address(2), 200)),
1369 "peer 2's peak must arrive under peer 2's address: {peaks:?}"
1370 );
1371 assert_ne!(
1372 first_source.session, second_source.session,
1373 "two sessions must not share an identity"
1374 );
1375 }
1376
1377 /// **An undecodable frame ENDS the session; it is never skipped.**
1378 ///
1379 /// The fixture is ordered so that skipping is distinguishable from ending: a valid peak, then
1380 /// an undecodable one, then a second valid peak. An implementation that ignores what it cannot
1381 /// decode — the `if let Ok(..)` this replaces — delivers BOTH peaks and no ending, which is a
1382 /// subscriber missing an update it will never learn it missed.
1383 #[tokio::test]
1384 async fn an_undecodable_frame_ends_the_session_rather_than_being_skipped() {
1385 let pool = empty_pool(4);
1386 let mut subscription = pool.subscribe_frames(32).await;
1387 let (sender, source) = followed_session(&pool, address(1)).await;
1388
1389 sender.send(peak_message(100)).await.expect("send");
1390 sender.send(undecodable_peak_message()).await.expect("send");
1391 sender.send(peak_message(101)).await.expect("send");
1392
1393 let seen = drain_at_least(&mut subscription, 3).await;
1394 let frames: Vec<&PoolFrame> = seen
1395 .iter()
1396 .filter(|f| f.source == source)
1397 .map(|f| &f.frame)
1398 .collect();
1399
1400 assert!(
1401 frames
1402 .iter()
1403 .any(|f| matches!(f, PoolFrame::Peak { height: 100, .. })),
1404 "the frames before the bad one are still delivered: {frames:?}"
1405 );
1406 assert!(
1407 frames.contains(&&PoolFrame::SessionEnded {
1408 reason: SessionEndReason::UndecodableFrame
1409 }),
1410 "an undecodable frame must END the session, loudly: {frames:?}"
1411 );
1412 assert!(
1413 !frames
1414 .iter()
1415 .any(|f| matches!(f, PoolFrame::Peak { height: 101, .. })),
1416 "nothing after the undecodable frame belongs to this session: {frames:?}"
1417 );
1418 }
1419
1420 /// The control: a session fed only VALID frames is not ended.
1421 ///
1422 /// Without it, a handler that ended every session on its first message would satisfy the test
1423 /// above.
1424 #[tokio::test]
1425 async fn a_session_fed_only_valid_frames_stays_open() {
1426 let pool = empty_pool(4);
1427 let mut subscription = pool.subscribe_frames(32).await;
1428 let (sender, source) = followed_session(&pool, address(1)).await;
1429
1430 sender.send(peak_message(100)).await.expect("send");
1431 sender.send(peak_message(101)).await.expect("send");
1432
1433 let seen = drain_at_least(&mut subscription, 3).await;
1434 let frames: Vec<&PoolFrame> = seen
1435 .iter()
1436 .filter(|f| f.source == source)
1437 .map(|f| &f.frame)
1438 .collect();
1439
1440 assert!(
1441 !frames
1442 .iter()
1443 .any(|f| matches!(f, PoolFrame::SessionEnded { .. })),
1444 "a well-behaved session must stay open: {frames:?}"
1445 );
1446 assert_eq!(
1447 frames
1448 .iter()
1449 .filter(|f| matches!(f, PoolFrame::Peak { .. }))
1450 .count(),
1451 2,
1452 "both valid peaks must be delivered: {frames:?}"
1453 );
1454 }
1455
1456 /// A transport that closes ends the session loudly rather than silently.
1457 #[tokio::test]
1458 async fn a_closed_transport_ends_the_session_loudly() {
1459 let pool = empty_pool(4);
1460 let mut subscription = pool.subscribe_frames(32).await;
1461 let (sender, source) = followed_session(&pool, address(1)).await;
1462
1463 drop(sender);
1464
1465 let seen = drain_at_least(&mut subscription, 2).await;
1466 assert!(
1467 seen.iter().any(|f| f.source == source
1468 && f.frame
1469 == PoolFrame::SessionEnded {
1470 reason: SessionEndReason::Disconnected
1471 }),
1472 "a dropped transport must be announced, not left as silence: {seen:?}"
1473 );
1474 }
1475
1476 /// **A peer whose session ended is EJECTED, without waiting for a failed request or the
1477 /// rotation timer.**
1478 ///
1479 /// Two peers are held and only ONE dies, which is the fixture shape that separates ejecting
1480 /// the right peer from ejecting on any death: a pass that removed both, or removed the wrong
1481 /// one, fails here. The surviving peer is the control and is never fed anything, so nothing
1482 /// about it changes except that its neighbour died.
1483 #[tokio::test]
1484 async fn a_peer_whose_session_ended_is_ejected_without_waiting_for_a_failure() {
1485 let pool = empty_pool(4);
1486 let mut subscription = pool.subscribe_frames(32).await;
1487
1488 let (dying, dying_source) = followed_session(&pool, address(1)).await;
1489 let (_surviving, _) = followed_session(&pool, address(2)).await;
1490
1491 drop(dying);
1492
1493 // Wait for the death to be announced, which is published after it is recorded.
1494 let seen = drain_at_least(&mut subscription, 3).await;
1495 assert!(
1496 seen.iter()
1497 .any(|f| f.source == dying_source
1498 && matches!(f.frame, PoolFrame::SessionEnded { .. })),
1499 "the fixture depends on the session actually ending: {seen:?}"
1500 );
1501
1502 assert_eq!(
1503 pool.held_addresses_for_tests().await.len(),
1504 2,
1505 "a dead session is still HELD until maintenance runs — which is the gap being closed"
1506 );
1507
1508 pool.eject_dead_sessions_for_tests().await;
1509
1510 assert_eq!(
1511 pool.held_addresses_for_tests().await,
1512 vec![address(2)],
1513 "exactly the peer whose session ended is removed"
1514 );
1515 }
1516
1517 /// A replacement dialled to the same address is not removed by its predecessor's death.
1518 ///
1519 /// The interleaving is a real one and it is the ONLY one that can exhibit this: a request to
1520 /// the dead connection fails, so `eject_peer` removes it by ADDRESS before maintenance runs;
1521 /// a refill then re-dials that same address; and only afterwards does maintenance drain the
1522 /// death that is still recorded against it. An ejection matching on address alone — the
1523 /// obvious implementation — removes the live replacement there and leaves the pool short, with
1524 /// nothing anywhere reporting a problem.
1525 ///
1526 /// Draining the dead list BEFORE re-admitting cannot show this, because the drain empties the
1527 /// list and the second pass then has nothing to match with. That ordering was this test's
1528 /// first shape and it passed against address-only matching, which is to say it proved nothing.
1529 #[tokio::test]
1530 async fn a_replacement_at_the_same_address_survives_its_predecessors_death() {
1531 let pool = empty_pool(4);
1532 let mut subscription = pool.subscribe_frames(32).await;
1533
1534 let (dying, dying_source) = followed_session(&pool, address(1)).await;
1535 drop(dying);
1536
1537 let seen = drain_at_least(&mut subscription, 2).await;
1538 assert!(
1539 seen.iter()
1540 .any(|f| f.source == dying_source
1541 && matches!(f.frame, PoolFrame::SessionEnded { .. })),
1542 "the fixture depends on the session actually ending: {seen:?}"
1543 );
1544
1545 // A request to the dead connection fails first, which is how it leaves `entries` while its
1546 // death is still recorded.
1547 pool.eject_peer(address(1)).await;
1548 assert!(pool.held_addresses_for_tests().await.is_empty());
1549
1550 let (_replacement, replacement_source) = followed_session(&pool, address(1)).await;
1551 assert_ne!(replacement_source.session, dying_source.session);
1552
1553 // Maintenance now drains a death recorded against an address the replacement holds.
1554 pool.eject_dead_sessions_for_tests().await;
1555
1556 assert_eq!(
1557 pool.held_addresses_for_tests().await,
1558 vec![address(1)],
1559 "the replacement session must survive its predecessor's death"
1560 );
1561 }
1562}