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