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