Skip to main content

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}