Skip to main content

snarkos_node_network/
peering.rs

1// Copyright (c) 2019-2026 Provable Inc.
2// This file is part of the snarkOS library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use crate::{CandidatePeer, ConnectedPeer, ConnectionMode, NodeType, Peer, Resolver};
17
18use snarkos_node_tcp::{ConnectError, P2P, is_bogon_ip, is_unspecified_or_broadcast_ip};
19use snarkvm::prelude::{Address, Network};
20
21use anyhow::Result;
22#[cfg(feature = "locktick")]
23use locktick::parking_lot::RwLock;
24#[cfg(not(feature = "locktick"))]
25use parking_lot::RwLock;
26use std::{
27    cmp,
28    collections::{
29        HashSet,
30        hash_map::{Entry, HashMap},
31    },
32    fs,
33    io::{self, Write},
34    net::{IpAddr, SocketAddr},
35    path::Path,
36    str::FromStr,
37    time::Instant,
38};
39use tokio::task;
40use tracing::*;
41
42/// Application-level errors generated by the peering module.
43/// This is never returned directly, but only as the payload for a `ConnectError`.
44#[derive(Debug)]
45pub enum PeeringError {
46    NoExternalPeersAllowed,
47}
48
49impl snarkos_node_tcp::ApplicationError for PeeringError {}
50
51impl std::fmt::Display for PeeringError {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        match self {
54            Self::NoExternalPeersAllowed => write!(f, "no untrusted peers allowed"),
55        }
56    }
57}
58
59pub trait PeerPoolHandling<N: Network>: P2P {
60    const OWNER: &str;
61
62    /// The maximum number of peers permitted to be stored in the peer pool.
63    const MAXIMUM_POOL_SIZE: usize;
64
65    /// The number of candidate peers to be removed from the pool once `MAXIMUM_POOL_SIZE` is reached.
66    /// It must be lower than `MAXIMUM_POOL_SIZE`.
67    const PEER_SLASHING_COUNT: usize;
68
69    /// Returns the mapping of all known peers (connected or otherwise), keyed by their public listener address.
70    fn peer_pool(&self) -> &RwLock<HashMap<SocketAddr, Peer<N>>>;
71
72    /// Returns the resolver for translating between public listener addresses and connected addresses.
73    fn resolver(&self) -> &RwLock<Resolver<N>>;
74
75    /// Returns `true` if the owning node is in development mode.
76    fn is_dev(&self) -> bool;
77
78    /// Returns `true` if the node is in trusted peers only mode.
79    fn trusted_peers_only(&self) -> bool;
80
81    /// Returns the node type.
82    fn node_type(&self) -> NodeType;
83
84    /// Returns the listener address of this node.
85    fn local_ip(&self) -> SocketAddr {
86        self.tcp().listening_addr().expect("The TCP listener is not enabled")
87    }
88
89    /// Returns `true` if the given IP is this node.
90    fn is_local_ip(&self, addr: SocketAddr) -> bool {
91        addr == self.local_ip()
92            || (addr.ip().is_unspecified() || addr.ip().is_loopback()) && addr.port() == self.local_ip().port()
93    }
94
95    /// Returns `true` if the given IP is not this node, is not a bogon address, and is not unspecified.
96    fn is_valid_peer_ip(&self, ip: SocketAddr) -> bool {
97        !self.is_local_ip(ip) && !is_bogon_ip(ip.ip()) && !is_unspecified_or_broadcast_ip(ip.ip())
98    }
99
100    /// Returns the maximum number of connected peers.
101    fn max_connected_peers(&self) -> usize {
102        self.tcp().config().max_connections as usize
103    }
104
105    /// Ensure we can and are allowed to connect to the given listener address of a peer.
106    fn check_connection_attempt(&self, listener_addr: SocketAddr) -> Result<(), ConnectError> {
107        // Ensure the peer IP is not this node.
108        if self.is_local_ip(listener_addr) {
109            return Err(ConnectError::SelfConnect { address: listener_addr });
110        }
111        // Ensure the node does not surpass the maximum number of peer connections.
112        if self.number_of_connected_peers() >= self.max_connected_peers() {
113            return Err(ConnectError::MaximumConnectionsReached { limit: self.max_connected_peers() as u16 });
114        }
115        // Ensure the node is not already connected to this peer.
116        if self.is_connected(listener_addr) {
117            return Err(ConnectError::AlreadyConnected { address: listener_addr });
118        }
119        // Ensure the node is not already connecting to this peer.
120        if self.is_connecting(listener_addr) {
121            return Err(ConnectError::AlreadyConnecting { address: listener_addr });
122        }
123        // Ensure the peer IP is not banned.
124        if self.is_ip_banned(listener_addr.ip()) {
125            return Err(ConnectError::BannedIp { ip: listener_addr.ip() });
126        }
127        // If the node is in trusted peers only mode, ensure the peer is trusted.
128        if self.trusted_peers_only() && !self.is_trusted(listener_addr) {
129            return Err(ConnectError::application(PeeringError::NoExternalPeersAllowed));
130        }
131
132        Ok(())
133    }
134
135    /// Attempts to connect to the given peer's listener address.
136    ///
137    /// Returns an earlier error, if, for example, we are already connected to the peer.
138    /// Otherwise, it returns a handle to the tokio tasks that sets up the connection.
139    ///
140    /// # Concurrency
141    /// Only one task may call this function for a given listener address at a time.
142    fn connect(&self, listener_addr: SocketAddr) -> Result<task::JoinHandle<Result<(), ConnectError>>, ConnectError> {
143        // Return early if the attempt is against the protocol rules.
144        self.check_connection_attempt(listener_addr)?;
145
146        // Update the last connection attempt time for the peer.
147        if let Some(Peer::Candidate(peer)) = self.peer_pool().write().get_mut(&listener_addr) {
148            peer.last_connection_attempt = Some(Instant::now());
149            peer.total_connection_attempts += 1;
150        } else {
151            warn!("{} No candidate peer entry exists for '{listener_addr:?}' while connecting.", Self::OWNER);
152        }
153
154        let tcp = self.tcp().clone();
155        Ok(tokio::spawn(async move {
156            debug!("{} Connecting to {listener_addr}...", Self::OWNER);
157            tcp.connect(listener_addr).await
158        }))
159    }
160
161    /// Disconnects from the given peer IP, if the peer is connected. The returned boolean
162    /// indicates whether the peer was actually disconnected from, or if this was a noop.
163    fn disconnect(&self, listener_addr: SocketAddr) -> task::JoinHandle<bool> {
164        if let Some(connected_addr) = self.resolve_to_ambiguous(listener_addr) {
165            let tcp = self.tcp().clone();
166            tokio::spawn(async move { tcp.disconnect(connected_addr).await })
167        } else {
168            tokio::spawn(async { false })
169        }
170    }
171
172    /// Downgrades a connected peer to candidate status.
173    ///
174    /// Returns true if the peer was fully connected.
175    fn downgrade_peer_to_candidate(&self, listener_addr: SocketAddr) -> bool {
176        let mut peer_pool = self.peer_pool().write();
177        let Some(peer) = peer_pool.get_mut(&listener_addr) else {
178            trace!("{} Downgrade peer to candidate failed - peer not found", Self::OWNER);
179            return false;
180        };
181
182        if let Peer::Connected(conn_peer) = peer {
183            // Exception: the BootstrapClient only has a single Resolver,
184            // so it may only map a validator's Aleo address once, for its
185            // Gateway-mode connection. This also means that the Router-mode
186            // connection may not remove that mapping.
187            let aleo_addr = if self.node_type() == NodeType::BootstrapClient
188                && conn_peer.connection_mode == ConnectionMode::Router
189            {
190                None
191            } else {
192                Some(conn_peer.aleo_addr)
193            };
194            self.resolver().write().remove_peer(conn_peer.connected_addr, aleo_addr);
195            peer.downgrade_to_candidate(listener_addr);
196            true
197        } else {
198            peer.downgrade_to_candidate(listener_addr);
199            false
200        }
201    }
202
203    /// Adds new candidate peers to the peer pool, ensuring their validity and following the
204    /// limit on the number of peers in the pool. The listener addresses may be paired with
205    /// the last known block height of the associated peer.
206    fn insert_candidate_peers(&self, mut listener_addrs: Vec<(SocketAddr, Option<u32>)>) {
207        let trusted_peers = self.trusted_peers();
208
209        // Hold a write guard from now on, so as not to accidentally slash multiple times
210        // based on multiple batches of candidate peers, and to not overwrite any entries.
211        let mut peer_pool = self.peer_pool().write();
212
213        // Perform filtering to ensure candidate validity. Also count how many entries are updates.
214        let mut num_updates: usize = 0;
215        listener_addrs.retain(|&(addr, height)| {
216            !self.is_ip_banned(addr.ip())
217                && if self.is_dev() { !is_bogon_ip(addr.ip()) } else { self.is_valid_peer_ip(addr) }
218                && peer_pool
219                    .get(&addr)
220                    .map(|peer| peer.is_candidate() && height.is_some())
221                    .inspect(|is_valid_update| {
222                        if *is_valid_update {
223                            num_updates += 1
224                        }
225                    })
226                    .unwrap_or(true)
227        });
228
229        // If we've managed to filter out every entry, there's nothing to do.
230        if listener_addrs.is_empty() {
231            return;
232        }
233
234        // If we're about to exceed the peer pool size limit, apply candidate slashing.
235        if peer_pool.len() + listener_addrs.len() - num_updates >= Self::MAXIMUM_POOL_SIZE
236            && Self::PEER_SLASHING_COUNT != 0
237        {
238            // Collect the prospect peers, along with the key they are ranked by.
239            let now = Instant::now();
240            let mut peers_to_slash = peer_pool
241                .iter()
242                .filter_map(|(addr, peer)| match peer {
243                    Peer::Candidate(candidate) if !trusted_peers.contains(addr) => {
244                        Some((*addr, candidate.total_connection_attempts, candidate.last_connection_attempt))
245                    }
246                    _ => None,
247                })
248                .collect::<Vec<_>>();
249
250            // Sort by failed attempts (descending), then by the attempt's age (oldest first).
251            peers_to_slash.sort_unstable_by_key(|(_, attempts, last_attempt)| {
252                (cmp::Reverse(*attempts), last_attempt.unwrap_or(now))
253            });
254
255            // Retain the candidate peers with the most failed attempts and oldest timestamps.
256            peers_to_slash.truncate(Self::PEER_SLASHING_COUNT);
257
258            // Remove the peers to slash from the pool.
259            peer_pool.retain(|addr, _| !peers_to_slash.iter().any(|(slashed, _, _)| slashed == addr));
260        }
261
262        // Make sure that we won't breach the pool size limit in case the slashing didn't suffice.
263        listener_addrs.truncate(Self::MAXIMUM_POOL_SIZE.saturating_sub(peer_pool.len()));
264
265        // If we've managed to truncate to 0, exit.
266        if listener_addrs.is_empty() {
267            return;
268        }
269
270        // Insert or update the applicable candidate peers.
271        for (addr, height) in listener_addrs {
272            match peer_pool.entry(addr) {
273                Entry::Vacant(entry) => {
274                    entry.insert(Peer::new_candidate(addr, false));
275                }
276                Entry::Occupied(mut entry) => {
277                    if let Peer::Candidate(peer) = entry.get_mut() {
278                        peer.last_height_seen = height;
279                    }
280                }
281            }
282        }
283    }
284
285    /// Completely removes an entry from the peer pool.
286    fn remove_peer(&self, listener_addr: SocketAddr) {
287        self.peer_pool().write().remove(&listener_addr);
288    }
289
290    /// Returns the connected peer address from the listener IP address.
291    fn resolve_to_ambiguous(&self, listener_addr: SocketAddr) -> Option<SocketAddr> {
292        if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
293            Some(peer.connected_addr)
294        } else {
295            None
296        }
297    }
298
299    /// Returns the connected peer aleo address from the listener IP address.
300    fn resolve_to_aleo_addr(&self, listener_addr: SocketAddr) -> Option<Address<N>> {
301        if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
302            Some(peer.aleo_addr)
303        } else {
304            None
305        }
306    }
307
308    /// Returns `true` if the node is connecting to the given peer's listener address.
309    fn is_connecting(&self, listener_addr: SocketAddr) -> bool {
310        self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connecting())
311    }
312
313    /// Returns `true` if the node is connected to the given peer listener address.
314    fn is_connected(&self, listener_addr: SocketAddr) -> bool {
315        self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connected())
316    }
317
318    /// Returns `true` if the node is connected to the given Aleo address.
319    fn is_connected_address(&self, aleo_address: Address<N>) -> bool {
320        // The resolver only contains data on connected peers.
321        self.resolver().read().get_peer_ip_for_address(aleo_address).is_some()
322    }
323
324    /// Returns `true` if the node is connected or connecting to the given peer listener address.
325    fn is_connecting_or_connected(&self, listener_addr: SocketAddr) -> bool {
326        self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_connecting() || peer.is_connected())
327    }
328
329    /// Returns `true` if the given listener address is trusted.
330    fn is_trusted(&self, listener_addr: SocketAddr) -> bool {
331        self.peer_pool().read().get(&listener_addr).is_some_and(|peer| peer.is_trusted())
332    }
333
334    /// Returns the number of all peers.
335    fn number_of_peers(&self) -> usize {
336        self.peer_pool().read().len()
337    }
338
339    /// Returns the number of connected peers.
340    fn number_of_connected_peers(&self) -> usize {
341        self.peer_pool().read().values().filter(|peer| peer.is_connected()).count()
342    }
343
344    /// Returns the number of connected validators.
345    ///
346    /// Returns `None` if the peer pool is currently locked; only the metrics code cares about the
347    /// count, and it would rather skip a sample than contend for the lock.
348    fn number_of_connected_validators(&self) -> Option<usize> {
349        Some(
350            self.peer_pool()
351                .try_read()?
352                .values()
353                .filter(|peer| peer.as_connected().is_some_and(|peer| peer.is_validator()))
354                .count(),
355        )
356    }
357
358    /// Returns the number of connecting peers.
359    ///
360    /// Returns `None` if the peer pool is currently locked; see
361    /// [`PeerPoolHandling::number_of_connected_validators`].
362    fn number_of_connecting_peers(&self) -> Option<usize> {
363        Some(self.peer_pool().try_read()?.values().filter(|peer| peer.is_connecting()).count())
364    }
365
366    /// Returns the number of candidate peers.
367    fn number_of_candidate_peers(&self) -> usize {
368        self.peer_pool().read().values().filter(|peer| matches!(peer, Peer::Candidate(_))).count()
369    }
370
371    /// Returns the connected peer given the peer IP, if it exists.
372    fn get_connected_peer(&self, listener_addr: SocketAddr) -> Option<ConnectedPeer<N>> {
373        if let Some(Peer::Connected(peer)) = self.peer_pool().read().get(&listener_addr) {
374            Some(peer.clone())
375        } else {
376            None
377        }
378    }
379
380    /// Updates the connected peer - if it exists -  given the peer IP and a closure.
381    /// The returned status indicates whether the update was successful, i.e. the peer had existed.
382    fn update_connected_peer<F: FnMut(&mut ConnectedPeer<N>)>(
383        &self,
384        listener_addr: &SocketAddr,
385        mut update_fn: F,
386    ) -> bool {
387        if let Some(Peer::Connected(peer)) = self.peer_pool().write().get_mut(listener_addr) {
388            update_fn(peer);
389            true
390        } else {
391            false
392        }
393    }
394
395    /// Returns the list of all peers (connected, connecting, and candidate).
396    fn get_peers(&self) -> Vec<Peer<N>> {
397        self.peer_pool().read().values().cloned().collect()
398    }
399
400    /// Returns all connected peers.
401    fn get_connected_peers(&self) -> Vec<ConnectedPeer<N>> {
402        self.filter_connected_peers(|_| true)
403    }
404
405    /// Returns an optionally bounded list of all connected peers sorted by their
406    /// block height (highest first).
407    fn get_best_connected_peers(&self, max_entries: Option<usize>) -> Vec<ConnectedPeer<N>> {
408        // Get a snapshot of the currently connected peers.
409        let mut peers = self.get_connected_peers();
410
411        // Sort the prospect peers by greatest height.
412        peers.sort_unstable_by_key(|peer| cmp::Reverse(peer.last_height_seen));
413        if let Some(max) = max_entries {
414            peers.truncate(max);
415        }
416
417        peers
418    }
419
420    /// Returns all connected peers that satisify the given predicate.
421    fn filter_connected_peers<P: FnMut(&ConnectedPeer<N>) -> bool>(&self, mut predicate: P) -> Vec<ConnectedPeer<N>> {
422        self.peer_pool()
423            .read()
424            .values()
425            .filter_map(|p| {
426                if let Peer::Connected(peer) = p
427                    && predicate(peer)
428                {
429                    Some(peer)
430                } else {
431                    None
432                }
433            })
434            .cloned()
435            .collect()
436    }
437
438    /// Returns the list of connected peers.
439    fn connected_peers(&self) -> Vec<SocketAddr> {
440        self.peer_pool().read().iter().filter_map(|(addr, peer)| peer.is_connected().then_some(*addr)).collect()
441    }
442
443    /// Returns the list of trusted peers.
444    fn trusted_peers(&self) -> Vec<SocketAddr> {
445        self.peer_pool().read().iter().filter_map(|(addr, peer)| peer.is_trusted().then_some(*addr)).collect()
446    }
447
448    /// Returns the list of candidate peers.
449    fn get_candidate_peers(&self) -> Vec<CandidatePeer<N>> {
450        self.peer_pool()
451            .read()
452            .values()
453            .filter_map(|peer| if let Peer::Candidate(peer) = peer { Some(peer.clone()) } else { None })
454            .collect()
455    }
456
457    /// Returns the list of trusted candidate peers.
458    fn get_trusted_candidate_peers(&self) -> Vec<CandidatePeer<N>> {
459        self.peer_pool()
460            .read()
461            .values()
462            .filter_map(|peer| {
463                if let Peer::Candidate(peer) = peer
464                    && peer.trusted
465                {
466                    Some(peer.clone())
467                } else {
468                    None
469                }
470            })
471            .collect()
472    }
473
474    /// Loads any previously cached peer addresses so they can be introduced as initial
475    /// candidate peers to connect to.
476    fn load_cached_peers(path: &Path) -> Result<Vec<SocketAddr>> {
477        let peers = match fs::read_to_string(path) {
478            Ok(cached_peers_str) => {
479                let mut cached_peers = Vec::new();
480                for peer_addr_str in cached_peers_str.lines() {
481                    match SocketAddr::from_str(peer_addr_str) {
482                        Ok(addr) => cached_peers.push(addr),
483                        Err(error) => warn!("Couldn't parse the cached peer address '{peer_addr_str}': {error}"),
484                    }
485                }
486                cached_peers
487            }
488            Err(error) if error.kind() == io::ErrorKind::NotFound => {
489                // Not an issue - the cache may not exist yet.
490                Vec::new()
491            }
492            Err(error) => {
493                warn!("{} Couldn't load cached peers at {}: {error}", Self::OWNER, path.display());
494                Vec::new()
495            }
496        };
497
498        Ok(peers)
499    }
500
501    /// Preserve the peers who have the greatest known block heights, and the fewest
502    /// outstanding failed connection attempts.
503    ///
504    /// # Arguments
505    /// * `path` - The path to the file to save the peers to.
506    /// * `max_entries` - The maximum number of peers to save (if there are more, the extra ones are truncated).
507    /// * `store_ports` - Whether to store the ports of the peers, or just the IP addresses.
508    fn save_best_peers(&self, path: &Path, max_entries: Option<usize>, store_ports: bool) -> Result<()> {
509        // Collect all prospect peers.
510        let mut peers = self.get_peers();
511
512        // Sort the list of peers: greatest height first, then fewest failed connection attempts.
513        peers.sort_unstable_by_key(|peer| (cmp::Reverse(peer.last_height_seen()), peer.failed_connection_attempts()));
514        if let Some(max) = max_entries {
515            peers.truncate(max);
516        }
517
518        // Dump the connected and deduplicated peers to a file.
519        let addrs: HashSet<_> = peers
520            .iter()
521            .map(
522                |peer| {
523                    if store_ports { peer.listener_addr().to_string() } else { peer.listener_addr().ip().to_string() }
524                },
525            )
526            .collect();
527
528        let mut file = fs::File::create(path)?;
529        for addr in addrs {
530            writeln!(file, "{addr}")?;
531        }
532
533        Ok(())
534    }
535
536    // Introduces a new connecting peer into the peer pool if unknown, or promotes
537    // a known candidate peer to a connecting one. The returned boolean indicates
538    // whether the peer has been added/promoted, or rejected due to already being
539    // shaken hands with or connected.
540    fn add_connecting_peer(&self, listener_addr: SocketAddr) -> Result<(), ConnectError> {
541        match self.peer_pool().write().entry(listener_addr) {
542            Entry::Vacant(entry) => {
543                entry.insert(Peer::new_connecting(listener_addr, false));
544                Ok(())
545            }
546            Entry::Occupied(mut entry) => match entry.get() {
547                // Promote in place rather than replacing the entry, so the candidate's
548                // connection-attempt metadata survives the transition.
549                Peer::Candidate(_) => {
550                    entry.get_mut().promote_to_connecting();
551                    Ok(())
552                }
553                Peer::Connecting(_) => Err(ConnectError::AlreadyConnecting { address: listener_addr }),
554                Peer::Connected(_) => Err(ConnectError::AlreadyConnected { address: listener_addr }),
555            },
556        }
557    }
558
559    /// Temporarily IP-ban and disconnect from the peer with the given listener address and an
560    /// optional reason for the ban. This also removes the peer from the candidate pool.
561    fn ip_ban_peer(&self, listener_addr: SocketAddr, reason: Option<&str>) {
562        // Ignore IP-banning if we are in dev mode.
563        if self.is_dev() {
564            return;
565        }
566
567        let ip = listener_addr.ip();
568        debug!("IP-banning {ip}{}", reason.map(|r| format!(" reason: {r}")).unwrap_or_default());
569
570        // Insert/update the low-level IP ban list.
571        self.tcp().banned_peers().update_ip_ban(ip);
572
573        // Disconnect from the peer.
574        self.disconnect(listener_addr);
575        // Remove the peer from the pool.
576        self.remove_peer(listener_addr);
577    }
578
579    /// Check whether the given IP address is currently banned.
580    fn is_ip_banned(&self, ip: IpAddr) -> bool {
581        self.tcp().banned_peers().is_ip_banned(&ip)
582    }
583
584    /// Insert or update a banned IP.
585    fn update_ip_ban(&self, ip: IpAddr) {
586        self.tcp().banned_peers().update_ip_ban(ip);
587    }
588}
589
590#[cfg(test)]
591mod tests {
592    use super::*;
593    use crate::Peer;
594    use snarkos_node_tcp::{Config, P2P, Tcp};
595    use snarkvm::{prelude::Rng, utilities::TestRng};
596
597    use std::{collections::HashMap, net::SocketAddr, time::Instant};
598
599    type CurrentNetwork = snarkvm::prelude::MainnetV0;
600
601    struct MockPeerPool<N: Network> {
602        tcp: Tcp,
603        peer_pool: RwLock<HashMap<SocketAddr, Peer<N>>>,
604        resolver: RwLock<Resolver<N>>,
605    }
606
607    impl<N: Network> MockPeerPool<N> {
608        fn new() -> Self {
609            let config = Config { listener_ip: None, ..Default::default() };
610            Self { tcp: Tcp::new(config), peer_pool: Default::default(), resolver: Default::default() }
611        }
612    }
613
614    impl<N: Network> P2P for MockPeerPool<N> {
615        fn tcp(&self) -> &Tcp {
616            &self.tcp
617        }
618    }
619
620    impl<N: Network> PeerPoolHandling<N> for MockPeerPool<N> {
621        const MAXIMUM_POOL_SIZE: usize = 100;
622        const OWNER: &str = "MockPeerPool";
623        const PEER_SLASHING_COUNT: usize = 10;
624
625        fn peer_pool(&self) -> &RwLock<HashMap<SocketAddr, Peer<N>>> {
626            &self.peer_pool
627        }
628
629        fn resolver(&self) -> &RwLock<Resolver<N>> {
630            &self.resolver
631        }
632
633        fn is_dev(&self) -> bool {
634            false
635        }
636
637        fn trusted_peers_only(&self) -> bool {
638            false
639        }
640
641        fn node_type(&self) -> NodeType {
642            NodeType::Client
643        }
644    }
645
646    fn make_connected_peer(port: u16, node_type: NodeType, rng: &mut TestRng) -> (SocketAddr, Peer<CurrentNetwork>) {
647        use snarkvm::prelude::Address;
648        let listener_addr = SocketAddr::from(([127, 0, 0, 1], port));
649        let connected_addr = SocketAddr::from(([127, 0, 0, 1], port + 10000));
650        let now = Instant::now();
651        let peer = Peer::Connected(ConnectedPeer {
652            listener_addr,
653            connected_addr,
654            connection_mode: ConnectionMode::Router,
655            trusted: false,
656            aleo_addr: Address::<CurrentNetwork>::new(rng.random()),
657            node_type,
658            version: 1,
659            snarkos_sha: None,
660            last_height_seen: None,
661            first_seen: now,
662            last_seen: now,
663        });
664        (listener_addr, peer)
665    }
666
667    #[test]
668    fn test_peer_state_transitions() {
669        use snarkvm::prelude::Address;
670
671        let pool = MockPeerPool::<CurrentNetwork>::new();
672        let mut rng = TestRng::default();
673
674        let listener_addr = SocketAddr::from(([192, 0, 2, 1], 4000));
675        let connected_addr = SocketAddr::from(([192, 0, 2, 1], 14000));
676        let aleo_addr = Address::<CurrentNetwork>::new(rng.random());
677
678        // Step 1: insert as a candidate.
679        pool.peer_pool().write().insert(listener_addr, Peer::new_candidate(listener_addr, false));
680
681        assert_eq!(pool.number_of_candidate_peers(), 1);
682        assert_eq!(pool.number_of_connecting_peers(), Some(0));
683        assert_eq!(pool.number_of_connected_peers(), 0);
684        assert!(!pool.is_connecting(listener_addr));
685        assert!(!pool.is_connected(listener_addr));
686
687        // Step 2: promote to connecting.
688        assert!(pool.add_connecting_peer(listener_addr).is_ok());
689
690        assert_eq!(pool.number_of_candidate_peers(), 0);
691        assert_eq!(pool.number_of_connecting_peers(), Some(1));
692        assert_eq!(pool.number_of_connected_peers(), 0);
693        assert!(pool.is_connecting(listener_addr));
694        assert!(!pool.is_connected(listener_addr));
695
696        // Step 3: complete the handshake — upgrade to connected.
697        pool.peer_pool().write().get_mut(&listener_addr).unwrap().upgrade_to_connected(
698            connected_addr,
699            listener_addr.port(),
700            aleo_addr,
701            NodeType::Validator,
702            1,
703            None,
704            ConnectionMode::Router,
705        );
706
707        assert_eq!(pool.number_of_candidate_peers(), 0);
708        assert_eq!(pool.number_of_connecting_peers(), Some(0));
709        assert_eq!(pool.number_of_connected_peers(), 1);
710        assert!(!pool.is_connecting(listener_addr));
711        assert!(pool.is_connected(listener_addr));
712        assert_eq!(pool.number_of_connected_validators(), Some(1));
713
714        // Verify the connected peer's fields.
715        let connected = pool.get_connected_peer(listener_addr).expect("peer should be connected");
716        assert_eq!(connected.listener_addr, listener_addr);
717        assert_eq!(connected.connected_addr, connected_addr);
718        assert_eq!(connected.aleo_addr, aleo_addr);
719        assert_eq!(connected.node_type, NodeType::Validator);
720    }
721
722    #[test]
723    fn test_number_of_connected_validators() {
724        let pool = MockPeerPool::<CurrentNetwork>::new();
725        let mut rng = TestRng::default();
726
727        // Empty pool: no validators.
728        assert_eq!(pool.number_of_connected_validators(), Some(0));
729
730        // Insert 2 validators and 1 client.
731        let (addr1, peer1) = make_connected_peer(3000, NodeType::Validator, &mut rng);
732        let (addr2, peer2) = make_connected_peer(3001, NodeType::Validator, &mut rng);
733        let (addr3, peer3) = make_connected_peer(3002, NodeType::Client, &mut rng);
734        {
735            let mut pool_write = pool.peer_pool().write();
736            pool_write.insert(addr1, peer1);
737            pool_write.insert(addr2, peer2);
738            pool_write.insert(addr3, peer3);
739        }
740
741        assert_eq!(pool.number_of_connected_validators(), Some(2));
742        assert_eq!(pool.number_of_connected_peers(), 3);
743
744        // A candidate peer should not be counted as a validator.
745        let candidate_addr = SocketAddr::from(([127, 0, 0, 1], 3003));
746        pool.peer_pool().write().insert(candidate_addr, Peer::new_candidate(candidate_addr, false));
747
748        assert_eq!(pool.number_of_connected_validators(), Some(2));
749        assert_eq!(pool.number_of_connected_peers(), 3);
750    }
751}