Skip to main content

reth_network/
peers.rs

1//! Peer related implementations
2
3use crate::{
4    error::SessionError,
5    session::{Direction, PendingSessionHandshakeError},
6    swarm::NetworkConnectionState,
7    trusted_peers_resolver::TrustedPeersResolver,
8};
9use alloy_primitives::map::{hash_map::Entry, FbBuildHasher, HashMap, HashSet};
10use futures::StreamExt;
11
12use rand::Rng;
13use reth_eth_wire::{errors::EthStreamError, DisconnectReason};
14use reth_ethereum_forks::ForkId;
15use reth_net_banlist::BanList;
16use reth_network_api::test_utils::{PeerCommand, PeersHandle};
17use reth_network_peers::{NodeRecord, PeerId, TrustedPeer};
18use reth_network_types::{
19    is_connection_failed_reputation,
20    peers::{
21        config::{PeerBackoffDurations, PEER_ROTATION_MIN_UPTIME},
22        reputation::{DEFAULT_REPUTATION, MAX_TRUSTED_PEER_REPUTATION_CHANGE},
23    },
24    ConnectionsConfig, Peer, PeerAddr, PeerConnectionState, PeerKind, PeersConfig,
25    PersistedPeerInfo, ReputationChangeKind, ReputationChangeOutcome, ReputationChangeWeights,
26};
27use std::{
28    collections::VecDeque,
29    fmt::Display,
30    io::{self},
31    net::{IpAddr, SocketAddr},
32    pin::Pin,
33    task::{Context, Poll},
34    time::Duration,
35};
36use thiserror::Error;
37use tokio::{
38    sync::mpsc,
39    time::{Instant, Interval, Sleep},
40};
41use tokio_stream::wrappers::UnboundedReceiverStream;
42use tracing::{trace, warn};
43
44/// Maintains the state of _all_ the peers known to the network.
45///
46/// This is supposed to be owned by the network itself, but can be reached via the [`PeersHandle`].
47/// From this type, connections to peers are established or disconnected, see [`PeerAction`].
48///
49/// The [`PeersManager`] will be notified on peer related changes
50#[derive(Debug)]
51pub struct PeersManager {
52    /// All peers known to the network
53    peers: HashMap<PeerId, Peer, FbBuildHasher<64>>,
54    /// The set of trusted peer ids.
55    ///
56    /// This tracks peer ids that are considered trusted, but for which we don't necessarily have
57    /// an address: [`Self::add_trusted_peer_id`]
58    trusted_peer_ids: HashSet<PeerId, FbBuildHasher<64>>,
59    /// A resolver used to periodically resolve DNS names for trusted peers. This updates the
60    /// peer's address when the DNS records change.
61    trusted_peers_resolver: TrustedPeersResolver,
62    /// Copy of the sender half, so new [`PeersHandle`] can be created on demand.
63    manager_tx: mpsc::UnboundedSender<PeerCommand>,
64    /// Receiver half of the command channel.
65    handle_rx: UnboundedReceiverStream<PeerCommand>,
66    /// Buffered actions until the manager is polled.
67    queued_actions: VecDeque<PeerAction>,
68    /// Interval for triggering connections if there are free slots.
69    refill_slots_interval: Interval,
70    /// How to weigh reputation changes
71    reputation_weights: ReputationChangeWeights,
72    /// Tracks current slot stats.
73    connection_info: ConnectionInfo,
74    /// Tracks unwanted ips/peer ids.
75    ban_list: BanList,
76    /// Tracks currently backed off peers.
77    backed_off_peers: HashMap<PeerId, std::time::Instant, FbBuildHasher<64>>,
78    /// Interval at which to check for peers to unban and release from the backoff map.
79    release_interval: Interval,
80    /// How long to ban bad peers.
81    ban_duration: Duration,
82    /// How long peers to which we could not connect for non-fatal reasons, e.g.
83    /// [`DisconnectReason::TooManyPeers`], are put in time out.
84    backoff_durations: PeerBackoffDurations,
85    /// If non-trusted peers should be connected to, or the connection from non-trusted
86    /// incoming peers should be accepted.
87    trusted_nodes_only: bool,
88    /// Timestamp of the last time [`Self::tick`] was called.
89    last_tick: Instant,
90    /// Maximum number of backoff attempts before we give up on a peer and dropping.
91    max_backoff_count: u8,
92    /// Tracks the connection state of the node
93    net_connection_state: NetworkConnectionState,
94    /// How long to temporarily ban ip on an incoming connection attempt.
95    incoming_ip_throttle_duration: Duration,
96    /// IP address filter for restricting network connections to specific IP ranges.
97    ip_filter: reth_net_banlist::IpFilter,
98    /// If true, discovered peers without a confirmed ENR fork ID will not be added until their
99    /// fork ID is verified via EIP-868.
100    enforce_enr_fork_id: bool,
101    /// One-shot sleep that fires when it's time to rotate a peer; reset with jitter after each
102    /// fire. `None` when rotation is disabled.
103    peer_rotation_sleep: Option<Pin<Box<Sleep>>>,
104    /// Mean duration for computing jittered rotation intervals. `None` when rotation is disabled.
105    peer_rotation_mean: Option<Duration>,
106}
107
108impl PeersManager {
109    /// Create a new instance with the given config
110    pub fn new(config: PeersConfig) -> Self {
111        let PeersConfig {
112            refill_slots_interval,
113            connection_info,
114            reputation_weights,
115            ban_list,
116            ban_duration,
117            backoff_durations,
118            trusted_nodes,
119            trusted_nodes_only,
120            trusted_nodes_resolution_interval,
121            basic_nodes,
122            persisted_peers,
123            max_backoff_count,
124            incoming_ip_throttle_duration,
125            ip_filter,
126            enforce_enr_fork_id,
127            peer_rotation_interval,
128        } = config;
129        let (manager_tx, handle_rx) = mpsc::unbounded_channel();
130        let now = Instant::now();
131
132        // We use half of the interval to decrease the max duration to `150%` in worst case
133        let unban_interval = ban_duration.min(backoff_durations.low) / 2;
134
135        let mut peers: HashMap<PeerId, Peer, FbBuildHasher<64>> = HashMap::with_capacity_and_hasher(
136            trusted_nodes.len() + basic_nodes.len() + persisted_peers.len(),
137            Default::default(),
138        );
139        let mut trusted_peer_ids: HashSet<PeerId, FbBuildHasher<64>> =
140            HashSet::with_capacity_and_hasher(trusted_nodes.len(), Default::default());
141
142        for trusted_peer in &trusted_nodes {
143            match trusted_peer.resolve_blocking() {
144                Ok(NodeRecord { address, tcp_port, udp_port, id }) => {
145                    trusted_peer_ids.insert(id);
146                    peers.entry(id).or_insert_with(|| {
147                        Peer::trusted(PeerAddr::new_with_ports(address, tcp_port, Some(udp_port)))
148                    });
149                }
150                Err(err) => {
151                    warn!(target: "net::peers", ?err, "Failed to resolve trusted peer");
152                }
153            }
154        }
155
156        for PersistedPeerInfo { record, kind, fork_id, reputation } in persisted_peers {
157            // When enforce_enr_fork_id is enabled, skip persisted peers that don't have a
158            // confirmed fork ID. These were likely accumulated from a different network during
159            // a prior run without the flag.
160            if enforce_enr_fork_id && fork_id.is_none() {
161                continue
162            }
163            let NodeRecord { address, tcp_port, udp_port, id } = record;
164            peers.entry(id).or_insert_with(|| {
165                let mut peer = Peer::with_kind(
166                    PeerAddr::new_with_ports(address, tcp_port, Some(udp_port)),
167                    kind,
168                );
169                peer.fork_id = fork_id.map(Box::new);
170                peer.reputation = reputation;
171                peer
172            });
173        }
174
175        for NodeRecord { address, tcp_port, udp_port, id } in basic_nodes {
176            peers.entry(id).or_insert_with(|| {
177                Peer::new(PeerAddr::new_with_ports(address, tcp_port, Some(udp_port)))
178            });
179        }
180
181        trace!(target: "net::peers", trusted_peers=?trusted_peer_ids, "Initialized peers manager");
182
183        Self {
184            peers,
185            trusted_peer_ids,
186            trusted_peers_resolver: TrustedPeersResolver::new(
187                trusted_nodes,
188                tokio::time::interval(trusted_nodes_resolution_interval), // 1 hour
189            ),
190            manager_tx,
191            handle_rx: UnboundedReceiverStream::new(handle_rx),
192            queued_actions: Default::default(),
193            reputation_weights,
194            refill_slots_interval: tokio::time::interval(refill_slots_interval),
195            release_interval: tokio::time::interval_at(now + unban_interval, unban_interval),
196            connection_info: ConnectionInfo::new(connection_info),
197            ban_list,
198            backed_off_peers: Default::default(),
199            ban_duration,
200            backoff_durations,
201            trusted_nodes_only,
202            last_tick: Instant::now(),
203            max_backoff_count,
204            net_connection_state: NetworkConnectionState::default(),
205            incoming_ip_throttle_duration,
206            ip_filter,
207            enforce_enr_fork_id,
208            peer_rotation_sleep: peer_rotation_interval
209                .map(|mean| Box::pin(tokio::time::sleep(jitter_rotation_interval(mean)))),
210            peer_rotation_mean: peer_rotation_interval,
211        }
212    }
213
214    /// Returns a new [`PeersHandle`] that can send commands to this type.
215    pub(crate) fn handle(&self) -> PeersHandle {
216        PeersHandle::new(self.manager_tx.clone())
217    }
218
219    /// Returns `true` if discovered peers must have a confirmed ENR fork ID before being added.
220    pub(crate) const fn enforce_enr_fork_id(&self) -> bool {
221        self.enforce_enr_fork_id
222    }
223
224    /// Returns the number of peers in the peer set
225    #[inline]
226    pub(crate) fn num_known_peers(&self) -> usize {
227        self.peers.len()
228    }
229
230    /// Returns an iterator over all peers as [`NodeRecord`]s.
231    pub(crate) fn iter_peers(&self) -> impl Iterator<Item = NodeRecord> + '_ {
232        self.peers.iter().map(|(peer_id, v)| {
233            NodeRecord::new_with_ports(
234                v.addr.tcp().ip(),
235                v.addr.tcp().port(),
236                v.addr.udp().map(|addr| addr.port()),
237                *peer_id,
238            )
239        })
240    }
241
242    /// Returns an iterator over peers suitable for persisting to disk.
243    ///
244    /// Filters out backed-off and banned peers, and includes metadata like kind, fork ID, and
245    /// reputation.
246    pub(crate) fn persistable_peers(&self) -> impl Iterator<Item = PersistedPeerInfo> + '_ {
247        self.peers.iter().filter(|(_, peer)| !peer.is_backed_off() && !peer.is_banned()).map(
248            |(peer_id, peer)| PersistedPeerInfo {
249                record: NodeRecord::new_with_ports(
250                    peer.addr.tcp().ip(),
251                    peer.addr.tcp().port(),
252                    peer.addr.udp().map(|addr| addr.port()),
253                    *peer_id,
254                ),
255                kind: peer.kind,
256                fork_id: peer.fork_id.as_deref().copied(),
257                reputation: peer.reputation,
258            },
259        )
260    }
261
262    /// Returns the `NodeRecord` and `PeerKind` for the given peer id
263    pub(crate) fn peer_by_id(&self, peer_id: PeerId) -> Option<(NodeRecord, PeerKind)> {
264        self.peers.get(&peer_id).map(|v| {
265            (
266                NodeRecord::new_with_ports(
267                    v.addr.tcp().ip(),
268                    v.addr.tcp().port(),
269                    v.addr.udp().map(|addr| addr.port()),
270                    peer_id,
271                ),
272                v.kind,
273            )
274        })
275    }
276
277    /// Returns `true` if the given peer is connected via an inbound session.
278    pub(crate) fn is_inbound_peer(&self, peer_id: &PeerId) -> bool {
279        self.peers.get(peer_id).is_some_and(|p| {
280            matches!(p.state, PeerConnectionState::In | PeerConnectionState::DisconnectingIn)
281        })
282    }
283
284    /// Returns an iterator over all peer ids for peers with the given kind
285    pub(crate) fn peers_by_kind(&self, kind: PeerKind) -> impl Iterator<Item = PeerId> + '_ {
286        self.peers.iter().filter_map(move |(peer_id, peer)| (peer.kind == kind).then_some(*peer_id))
287    }
288
289    /// Returns the number of currently active inbound connections.
290    #[inline]
291    pub(crate) const fn num_inbound_connections(&self) -> usize {
292        self.connection_info.num_inbound
293    }
294
295    /// Returns the number of currently __active__ outbound connections.
296    #[inline]
297    pub(crate) const fn num_outbound_connections(&self) -> usize {
298        self.connection_info.num_outbound
299    }
300
301    /// Returns the number of currently pending outbound connections.
302    #[inline]
303    pub(crate) const fn num_pending_outbound_connections(&self) -> usize {
304        self.connection_info.num_pending_out
305    }
306
307    /// Returns the number of currently backed off peers.
308    #[inline]
309    pub(crate) fn num_backed_off_peers(&self) -> usize {
310        self.backed_off_peers.len()
311    }
312
313    /// Returns the number of idle trusted peers.
314    fn num_idle_trusted_peers(&self) -> usize {
315        self.peers.iter().filter(|(_, peer)| peer.kind.is_trusted() && peer.state.is_idle()).count()
316    }
317
318    /// Invoked when a new _incoming_ tcp connection is accepted.
319    ///
320    /// returns an error if the inbound ip address is on the ban list
321    pub(crate) fn on_incoming_pending_session(
322        &mut self,
323        addr: IpAddr,
324    ) -> Result<(), InboundConnectionError> {
325        // Check if the IP is in the allowed ranges (netrestrict)
326        if !self.ip_filter.is_allowed(&addr) {
327            trace!(target: "net", ?addr, "Rejecting connection from IP not in allowed ranges");
328            return Err(InboundConnectionError::IpBanned)
329        }
330
331        if self.ban_list.is_banned_ip(&addr) {
332            return Err(InboundConnectionError::IpBanned)
333        }
334
335        // check if we even have slots for a new incoming connection
336        if !self.connection_info.has_in_capacity() {
337            if self.trusted_peer_ids.is_empty() {
338                // if we don't have any incoming slots and no trusted peers, we don't accept any new
339                // connections
340                return Err(InboundConnectionError::ExceedsCapacity)
341            }
342
343            // there's an edge case here where no incoming connections besides from trusted peers
344            // are allowed (max_inbound == 0), in which case we still need to allow new pending
345            // incoming connections until all trusted peers are connected.
346            let num_idle_trusted_peers = self.num_idle_trusted_peers();
347            if num_idle_trusted_peers <= self.trusted_peer_ids.len() {
348                // we still want to limit concurrent pending connections
349                let max_inbound =
350                    self.trusted_peer_ids.len().max(self.connection_info.config.max_inbound);
351                if self.connection_info.num_pending_in < max_inbound {
352                    self.connection_info.inc_pending_in();
353                    return Ok(())
354                }
355            }
356
357            // all trusted peers are either connected or connecting
358            return Err(InboundConnectionError::ExceedsCapacity)
359        }
360
361        // also cap the incoming connections we can process at once
362        if !self.connection_info.has_in_pending_capacity() {
363            return Err(InboundConnectionError::ExceedsCapacity)
364        }
365
366        // apply the rate limit
367        self.throttle_incoming_ip(addr);
368
369        self.connection_info.inc_pending_in();
370        Ok(())
371    }
372
373    /// Invoked when a previous call to [`Self::on_incoming_pending_session`] succeeded but it was
374    /// rejected.
375    pub(crate) const fn on_incoming_pending_session_rejected_internally(&mut self) {
376        self.connection_info.decr_pending_in();
377    }
378
379    /// Invoked when a pending session was closed.
380    pub(crate) const fn on_incoming_pending_session_gracefully_closed(&mut self) {
381        self.connection_info.decr_pending_in()
382    }
383
384    /// Invoked when a pending session was closed.
385    pub(crate) fn on_incoming_pending_session_dropped(
386        &mut self,
387        remote_addr: SocketAddr,
388        err: &PendingSessionHandshakeError,
389    ) {
390        if err.is_fatal_protocol_error() {
391            self.ban_ip(remote_addr.ip());
392
393            if err.merits_discovery_ban() {
394                self.queued_actions
395                    .push_back(PeerAction::DiscoveryBanIp { ip_addr: remote_addr.ip() })
396            }
397        }
398
399        self.connection_info.decr_pending_in();
400    }
401
402    /// Called when a new _incoming_ active session was established to the given peer.
403    ///
404    /// This will update the state of the peer if not yet tracked.
405    ///
406    /// If the reputation of the peer is below the `BANNED_REPUTATION` threshold, a disconnect will
407    /// be scheduled.
408    pub(crate) fn on_incoming_session_established(&mut self, peer_id: PeerId, addr: SocketAddr) {
409        self.connection_info.decr_pending_in();
410
411        // we only need to check the peer id here as the ip address will have been checked at
412        // on_incoming_pending_session. We also check if the peer is in the backoff list here.
413        if self.ban_list.is_banned_peer(&peer_id) {
414            self.queued_actions.push_back(PeerAction::DisconnectBannedIncoming { peer_id });
415            return
416        }
417
418        // check if the peer is trustable or not
419        let mut is_trusted = self.trusted_peer_ids.contains(&peer_id);
420        if self.trusted_nodes_only && !is_trusted {
421            self.queued_actions.push_back(PeerAction::DisconnectUntrustedIncoming { peer_id });
422            return
423        }
424
425        // start a new tick, so the peer is not immediately rewarded for the time since last tick
426        self.tick();
427
428        match self.peers.entry(peer_id) {
429            Entry::Occupied(mut entry) => {
430                let peer = entry.get_mut();
431                if peer.is_banned() {
432                    self.queued_actions.push_back(PeerAction::DisconnectBannedIncoming { peer_id });
433                    return
434                }
435                // it might be the case that we're also trying to connect to this peer at the same
436                // time, so we need to adjust the state here
437                if peer.state.is_pending_out() {
438                    self.connection_info.decr_state(peer.state);
439                }
440
441                peer.state = PeerConnectionState::In;
442                peer.mark_connected();
443
444                is_trusted = is_trusted || peer.is_trusted();
445            }
446            Entry::Vacant(entry) => {
447                // peer is missing in the table, we add it but mark it as to be removed after
448                // disconnect, because we only know the outgoing port
449                let mut peer = Peer::with_state(PeerAddr::from_tcp(addr), PeerConnectionState::In);
450                peer.mark_connected();
451                peer.remove_after_disconnect = true;
452                entry.insert(peer);
453                self.queued_actions.push_back(PeerAction::PeerAdded(peer_id));
454            }
455        }
456
457        let has_in_capacity = self.connection_info.has_in_capacity();
458        // increment new incoming connection
459        self.connection_info.inc_in();
460
461        // disconnect the peer if we don't have capacity for more inbound connections
462        if !is_trusted && !has_in_capacity {
463            self.queued_actions.push_back(PeerAction::Disconnect {
464                peer_id,
465                reason: Some(DisconnectReason::TooManyPeers),
466            });
467        }
468    }
469
470    /// Bans the peer temporarily with the configured ban timeout
471    fn ban_peer(&mut self, peer_id: PeerId) {
472        let ban_duration = if let Some(peer) = self.peers.get(&peer_id) &&
473            (peer.is_trusted() || peer.is_static())
474        {
475            // For misbehaving trusted or static peers, we provide a bit more leeway when
476            // penalizing them.
477            self.backoff_durations.low / 2
478        } else {
479            self.ban_duration
480        };
481
482        self.ban_list.ban_peer_until(peer_id, std::time::Instant::now() + ban_duration);
483        self.queued_actions.push_back(PeerAction::BanPeer { peer_id });
484    }
485
486    /// Bans the IP temporarily with the configured ban timeout
487    fn ban_ip(&mut self, ip: IpAddr) {
488        self.ban_list.ban_ip_until(ip, std::time::Instant::now() + self.ban_duration);
489    }
490
491    /// Bans the IP temporarily to rate limit inbound connection attempts per IP.
492    fn throttle_incoming_ip(&mut self, ip: IpAddr) {
493        self.ban_list
494            .ban_ip_until(ip, std::time::Instant::now() + self.incoming_ip_throttle_duration);
495    }
496
497    /// Temporarily puts the peer in timeout by inserting it into the backedoff peers set
498    fn backoff_peer_until(&mut self, peer_id: PeerId, until: std::time::Instant) {
499        trace!(target: "net::peers", ?peer_id, "backing off");
500
501        if let Some(peer) = self.peers.get_mut(&peer_id) {
502            peer.backed_off = true;
503            self.backed_off_peers.insert(peer_id, until);
504        }
505    }
506
507    /// Unbans the peer
508    fn unban_peer(&mut self, peer_id: PeerId) {
509        self.ban_list.unban_peer(&peer_id);
510        self.queued_actions.push_back(PeerAction::UnBanPeer { peer_id });
511    }
512
513    /// Tick function to update reputation of all connected peers.
514    /// Peers are rewarded with reputation increases for the time they are connected since the last
515    /// tick. This is to prevent peers from being disconnected eventually due to slashed
516    /// reputation because of some bad messages (most likely transaction related)
517    fn tick(&mut self) {
518        let now = Instant::now();
519        // Determine the number of seconds since the last tick.
520        // Ensuring that now is always greater than last_tick to account for issues with system
521        // time.
522        let secs_since_last_tick =
523            if self.last_tick > now { 0 } else { (now - self.last_tick).as_secs() as i32 };
524        self.last_tick = now;
525
526        // update reputation via seconds connected
527        for peer in self.peers.iter_mut().filter(|(_, peer)| peer.state.is_connected()) {
528            // update reputation via seconds connected, but keep the target _around_ the default
529            // reputation.
530            if peer.1.reputation < DEFAULT_REPUTATION {
531                peer.1.reputation += secs_since_last_tick;
532            }
533        }
534    }
535
536    /// Returns the tracked reputation for a peer.
537    pub(crate) fn get_reputation(&self, peer_id: &PeerId) -> Option<i32> {
538        self.peers.get(peer_id).map(|peer| peer.reputation)
539    }
540
541    /// Apply the corresponding reputation change to the given peer.
542    ///
543    /// If the peer is a trusted peer, it will be exempt from reputation slashing for certain
544    /// reputation changes that can be attributed to network conditions. If the peer is a
545    /// trusted peer, it will also be less strict with the reputation slashing.
546    pub(crate) fn apply_reputation_change(&mut self, peer_id: &PeerId, rep: ReputationChangeKind) {
547        trace!(target: "net::peers", ?peer_id, reputation=?rep, "applying reputation change");
548
549        let reputation_change = if rep.is_reset() {
550            None
551        } else {
552            let reputation_change = self.reputation_weights.change(rep).as_i32();
553            if reputation_change == 0 {
554                return
555            }
556            Some(reputation_change)
557        };
558
559        let outcome = if let Some(peer) = self.peers.get_mut(peer_id) {
560            if let Some(mut reputation_change) = reputation_change {
561                if peer.is_trusted() || peer.is_static() {
562                    // exempt trusted and static peers from reputation slashing for
563                    if matches!(
564                        rep,
565                        ReputationChangeKind::Dropped |
566                            ReputationChangeKind::BadAnnouncement |
567                            ReputationChangeKind::Timeout |
568                            ReputationChangeKind::AlreadySeenTransaction
569                    ) {
570                        return
571                    }
572
573                    // also be less strict with the reputation slashing for trusted peers
574                    if reputation_change < MAX_TRUSTED_PEER_REPUTATION_CHANGE {
575                        // this caps the reputation change to the maximum allowed for trusted peers
576                        reputation_change = MAX_TRUSTED_PEER_REPUTATION_CHANGE;
577                    }
578                }
579                peer.apply_reputation(reputation_change, rep)
580            } else {
581                peer.reset_reputation()
582            }
583        } else {
584            return
585        };
586
587        match outcome {
588            ReputationChangeOutcome::None => {}
589            ReputationChangeOutcome::Ban => {
590                self.ban_peer(*peer_id);
591            }
592            ReputationChangeOutcome::Unban => self.unban_peer(*peer_id),
593            ReputationChangeOutcome::DisconnectAndBan => {
594                self.queued_actions.push_back(PeerAction::Disconnect {
595                    peer_id: *peer_id,
596                    reason: Some(DisconnectReason::DisconnectRequested),
597                });
598                self.ban_peer(*peer_id);
599            }
600        }
601    }
602
603    /// Gracefully disconnected a pending _outgoing_ session
604    pub(crate) fn on_outgoing_pending_session_gracefully_closed(&mut self, peer_id: &PeerId) {
605        if let Some(peer) = self.peers.get_mut(peer_id) {
606            self.connection_info.decr_state(peer.state);
607            peer.state = PeerConnectionState::Idle;
608            peer.mark_disconnected();
609        }
610    }
611
612    /// Invoked when an _outgoing_ pending session was closed during authentication or the
613    /// handshake.
614    pub(crate) fn on_outgoing_pending_session_dropped(
615        &mut self,
616        remote_addr: &SocketAddr,
617        peer_id: &PeerId,
618        err: &PendingSessionHandshakeError,
619    ) {
620        self.on_connection_failure(remote_addr, peer_id, err, ReputationChangeKind::FailedToConnect)
621    }
622
623    /// Gracefully disconnected an active session
624    pub(crate) fn on_active_session_gracefully_closed(&mut self, peer_id: PeerId) {
625        match self.peers.entry(peer_id) {
626            Entry::Occupied(mut entry) => {
627                trace!(target: "net::peers", ?peer_id, direction=?entry.get().state, "active session gracefully closed");
628                self.connection_info.decr_state(entry.get().state);
629
630                if entry.get().remove_after_disconnect && !entry.get().is_trusted() {
631                    // this peer should be removed from the set
632                    entry.remove();
633                    self.queued_actions.push_back(PeerAction::PeerRemoved(peer_id));
634                } else {
635                    let peer = entry.get_mut();
636                    // reset the peer's state
637                    // we reset the backoff counter since we're able to establish a successful
638                    // session to that peer
639                    peer.severe_backoff_counter = 0;
640                    peer.state = PeerConnectionState::Idle;
641                    peer.mark_disconnected();
642
643                    // but we're backing off slightly to avoid dialing the peer again right away, to
644                    // give the remote time to also properly register the closed session and clean
645                    // up and to avoid any issues with ip throttling on the remote in case this
646                    // session was terminated right away.
647                    peer.backed_off = true;
648                    self.backed_off_peers.insert(
649                        peer_id,
650                        std::time::Instant::now() + self.incoming_ip_throttle_duration,
651                    );
652                    trace!(target: "net::peers", ?peer_id, kind=?peer.kind, duration=?self.incoming_ip_throttle_duration, "backing off on gracefully closed session");
653                }
654            }
655            Entry::Vacant(_) => return,
656        }
657
658        self.fill_outbound_slots();
659    }
660
661    /// Called when a _pending_ outbound connection is successful.
662    pub(crate) fn on_active_outgoing_established(&mut self, peer_id: PeerId) {
663        if let Some(peer) = self.peers.get_mut(&peer_id) {
664            trace!(target: "net::peers", ?peer_id, "established active outgoing connection");
665            self.connection_info.decr_state(peer.state);
666            self.connection_info.inc_out();
667            peer.state = PeerConnectionState::Out;
668            peer.mark_connected();
669        }
670    }
671
672    /// Called when an _active_ session to a peer was forcefully dropped due to an error.
673    ///
674    /// Depending on whether the error is fatal, the peer will be removed from the peer set
675    /// otherwise its reputation is slashed.
676    pub(crate) fn on_active_session_dropped(
677        &mut self,
678        remote_addr: &SocketAddr,
679        peer_id: &PeerId,
680        err: &EthStreamError,
681    ) {
682        self.on_connection_failure(remote_addr, peer_id, err, ReputationChangeKind::Dropped)
683    }
684
685    /// Called when an attempt to create an _outgoing_ pending session failed while setting up a tcp
686    /// connection.
687    pub(crate) fn on_outgoing_connection_failure(
688        &mut self,
689        remote_addr: &SocketAddr,
690        peer_id: &PeerId,
691        err: &io::Error,
692    ) {
693        // there's a race condition where we accepted an incoming connection while we were trying to
694        // connect to the same peer at the same time. if the outgoing connection failed
695        // after the incoming connection was accepted, we can ignore this error
696        if let Some(peer) = self.peers.get(peer_id) {
697            if peer.state.is_incoming() {
698                // we already have an active connection to the peer, so we can ignore this error
699                return
700            }
701
702            if peer.is_trusted() && is_connection_failed_reputation(peer.reputation) {
703                // trigger resolution task for trusted peer since multiple connection failures
704                // occurred
705                self.trusted_peers_resolver.interval.reset_immediately();
706            }
707        }
708
709        self.on_connection_failure(remote_addr, peer_id, err, ReputationChangeKind::FailedToConnect)
710    }
711
712    fn on_connection_failure(
713        &mut self,
714        remote_addr: &SocketAddr,
715        peer_id: &PeerId,
716        err: impl SessionError,
717        reputation_change: ReputationChangeKind,
718    ) {
719        trace!(target: "net::peers", ?remote_addr, ?peer_id, %err, "handling failed connection");
720
721        if err.is_fatal_protocol_error() {
722            trace!(target: "net::peers", ?remote_addr, ?peer_id, %err, "fatal connection error");
723            // remove the peer to which we can't establish a connection due to protocol related
724            // issues.
725            if let Entry::Occupied(mut entry) = self.peers.entry(*peer_id) {
726                self.connection_info.decr_state(entry.get().state);
727                // only remove if the peer is not trusted
728                if entry.get().is_trusted() {
729                    let peer = entry.get_mut();
730                    peer.state = PeerConnectionState::Idle;
731                    peer.mark_disconnected();
732                } else {
733                    entry.remove();
734                    self.queued_actions.push_back(PeerAction::PeerRemoved(*peer_id));
735                    // If the error is caused by a peer that should be banned from discovery
736                    if err.merits_discovery_ban() {
737                        self.queued_actions.push_back(PeerAction::DiscoveryBanPeerId {
738                            peer_id: *peer_id,
739                            ip_addr: remote_addr.ip(),
740                        })
741                    }
742                }
743            }
744
745            // ban the peer
746            self.ban_peer(*peer_id);
747        } else {
748            let mut backoff_until = None;
749            let mut remove_peer = false;
750
751            if let Some(peer) = self.peers.get_mut(peer_id) {
752                if let Some(kind) = err.should_backoff() {
753                    if peer.is_trusted() || peer.is_static() {
754                        // provide a bit more leeway for trusted peers and use a lower backoff so
755                        // that we keep re-trying them after backing off shortly, but we should at
756                        // least backoff for the low duration to not violate the ip based inbound
757                        // connection throttle that peer has in place, because this peer might not
758                        // have us registered as a trusted peer.
759                        let backoff = self.backoff_durations.low;
760                        backoff_until = Some(std::time::Instant::now() + backoff);
761                        trace!(target: "net::peers", ?peer_id, ?backoff, "backing off trusted peer");
762                    } else {
763                        // Increment peer.backoff_counter
764                        if kind.is_severe() {
765                            peer.severe_backoff_counter =
766                                peer.severe_backoff_counter.saturating_add(1);
767                        }
768                        trace!(target: "net::peers", ?peer_id, ?kind, severe_backoff_counter=peer.severe_backoff_counter, "backing off basic peer");
769
770                        let backoff_time =
771                            self.backoff_durations.backoff_until(kind, peer.severe_backoff_counter);
772
773                        // The peer has signaled that it is currently unable to process any more
774                        // connections, so we will hold off on attempting any new connections for a
775                        // while
776                        backoff_until = Some(backoff_time);
777                    }
778                } else {
779                    // If the error was not a backoff error, we reduce the peer's reputation
780                    let reputation_change = self.reputation_weights.change(reputation_change);
781                    peer.reputation = peer.reputation.saturating_add(reputation_change.as_i32());
782                };
783
784                self.connection_info.decr_state(peer.state);
785                peer.state = PeerConnectionState::Idle;
786                peer.mark_disconnected();
787
788                if peer.severe_backoff_counter > self.max_backoff_count &&
789                    !peer.is_trusted() &&
790                    !peer.is_static()
791                {
792                    // mark peer for removal if it has been backoff too many times and is _not_
793                    // trusted or static
794                    remove_peer = true;
795                }
796            }
797
798            // remove peer if it has been marked for removal
799            if remove_peer {
800                trace!(target: "net", ?peer_id, "removed peer after exceeding backoff counter");
801                let (peer_id, _) = self.peers.remove_entry(peer_id).expect("peer must exist");
802                self.queued_actions.push_back(PeerAction::PeerRemoved(peer_id));
803            } else if let Some(backoff_until) = backoff_until {
804                // otherwise, backoff the peer if marked as such
805                self.backoff_peer_until(*peer_id, backoff_until);
806            }
807        }
808
809        self.fill_outbound_slots();
810    }
811
812    /// Invoked if a pending session was disconnected because there's already a connection to the
813    /// peer.
814    ///
815    /// If the session was an outgoing connection, this means that the peer initiated a connection
816    /// to us at the same time and this connection is already established.
817    pub(crate) const fn on_already_connected(&mut self, direction: Direction) {
818        match direction {
819            Direction::Incoming => {
820                // need to decrement the ingoing counter
821                self.connection_info.decr_pending_in();
822            }
823            Direction::Outgoing(_) => {
824                // cleanup is handled when the incoming active session is activated in
825                // `on_incoming_session_established`
826            }
827        }
828    }
829
830    /// Called for a newly discovered peer.
831    ///
832    /// If the peer already exists, then the address, kind and `fork_id` will be updated.
833    pub(crate) fn add_peer(&mut self, peer_id: PeerId, addr: PeerAddr, fork_id: Option<ForkId>) {
834        self.add_peer_kind(peer_id, None, addr, fork_id)
835    }
836
837    /// Marks the given peer as trusted.
838    pub(crate) fn add_trusted_peer_id(&mut self, peer_id: PeerId) {
839        self.trusted_peer_ids.insert(peer_id);
840        if let Some(peer) = self.peers.get_mut(&peer_id) {
841            peer.kind = PeerKind::Trusted;
842        }
843    }
844
845    /// Called for a newly discovered trusted peer.
846    ///
847    /// If the peer already exists, then the address and kind will be updated.
848    #[cfg_attr(not(test), expect(dead_code))]
849    pub(crate) fn add_trusted_peer(&mut self, peer_id: PeerId, addr: PeerAddr) {
850        self.add_peer_kind(peer_id, Some(PeerKind::Trusted), addr, None)
851    }
852
853    /// Adds a trusted peer that may use a hostname instead of an IP address.
854    pub(crate) fn add_trusted_peer_node(&mut self, trusted: TrustedPeer) {
855        let peer_id = trusted.id;
856        self.trusted_peer_ids.insert(peer_id);
857        self.trusted_peers_resolver.remove(peer_id);
858        self.trusted_peers_resolver.trusted_peers.push(trusted);
859        self.trusted_peers_resolver.interval.reset_immediately();
860    }
861
862    /// Called for a newly discovered peer.
863    ///
864    /// If the peer already exists, then the address, kind and `fork_id` will be updated.
865    /// If the peer exists and a [`PeerKind`] is provided then the peer's kind is updated
866    pub(crate) fn add_peer_kind(
867        &mut self,
868        peer_id: PeerId,
869        kind: Option<PeerKind>,
870        addr: PeerAddr,
871        fork_id: Option<ForkId>,
872    ) {
873        let ip_addr = addr.tcp().ip();
874
875        // Check if the IP is in the allowed ranges (netrestrict)
876        if !self.ip_filter.is_allowed(&ip_addr) {
877            trace!(target: "net", ?peer_id, ?ip_addr, "Skipping peer from IP not in allowed ranges");
878            return
879        }
880
881        if self.ban_list.is_banned(&peer_id, &ip_addr) {
882            return
883        }
884
885        match self.peers.entry(peer_id) {
886            Entry::Occupied(mut entry) => {
887                let peer = entry.get_mut();
888                peer.fork_id = fork_id.map(Box::new);
889                peer.addr = addr;
890
891                if let Some(kind) = kind {
892                    peer.kind = kind;
893                }
894
895                if peer.state.is_incoming() {
896                    // now that we have an actual discovered address, for that peer and not just the
897                    // ip of the incoming connection, we don't need to remove the peer after
898                    // disconnecting, See `on_incoming_session_established`
899                    peer.remove_after_disconnect = false;
900                }
901            }
902            Entry::Vacant(entry) => {
903                trace!(target: "net::peers", ?peer_id, addr=?addr.tcp(), "discovered new node");
904                let mut peer = Peer::with_kind(addr, kind.unwrap_or(PeerKind::Basic));
905                peer.fork_id = fork_id.map(Box::new);
906                entry.insert(peer);
907                self.queued_actions.push_back(PeerAction::PeerAdded(peer_id));
908            }
909        }
910
911        if kind.is_some_and(|kind| kind.is_trusted()) {
912            // also track the peer in the peer id set
913            self.trusted_peer_ids.insert(peer_id);
914        }
915    }
916
917    /// Removes the tracked node from the set.
918    pub(crate) fn remove_peer(&mut self, peer_id: PeerId) {
919        let Entry::Occupied(entry) = self.peers.entry(peer_id) else { return };
920        if entry.get().is_trusted() {
921            return
922        }
923        let mut peer = entry.remove();
924
925        trace!(target: "net::peers", ?peer_id, "remove discovered node");
926        self.queued_actions.push_back(PeerAction::PeerRemoved(peer_id));
927
928        if peer.state.is_connected() {
929            trace!(target: "net::peers", ?peer_id, "disconnecting on remove from discovery");
930            // we terminate the active session here, but only remove the peer after the session
931            // was disconnected, this prevents the case where the session is scheduled for
932            // disconnect but the node is immediately rediscovered, See also
933            // [`Self::on_disconnected()`]
934            peer.remove_after_disconnect = true;
935            peer.state.disconnect();
936            self.peers.insert(peer_id, peer);
937            self.queued_actions.push_back(PeerAction::Disconnect {
938                peer_id,
939                reason: Some(DisconnectReason::DisconnectRequested),
940            })
941        }
942    }
943
944    /// Bans the peer indefinitely and removes it from the peer set.
945    ///
946    /// This follows [`Self::remove_peer`] and does not override trusted status. Remove trusted
947    /// peers from the trusted set before banning them.
948    pub(crate) fn ban_peer_by_admin(&mut self, peer_id: PeerId) {
949        if self.trusted_peer_ids.contains(&peer_id) ||
950            self.peers.get(&peer_id).is_some_and(Peer::is_trusted)
951        {
952            return
953        }
954
955        self.remove_peer(peer_id);
956        self.ban_list.ban_peer(peer_id);
957    }
958
959    /// Removes the peer from the ban list and resets its reputation.
960    pub(crate) fn unban_peer_by_admin(&mut self, peer_id: PeerId) {
961        self.ban_list.unban_peer(&peer_id);
962        if let Some(peer) = self.peers.get_mut(&peer_id) &&
963            peer.is_banned()
964        {
965            peer.unban();
966            self.queued_actions.push_back(PeerAction::UnBanPeer { peer_id });
967        }
968    }
969
970    /// Connect to the given peer. NOTE: if the maximum number of outbound sessions is reached,
971    /// this won't do anything. See `reth_network::SessionManager::dial_outbound`.
972    #[cfg_attr(not(test), expect(dead_code))]
973    pub(crate) fn add_and_connect(
974        &mut self,
975        peer_id: PeerId,
976        addr: PeerAddr,
977        fork_id: Option<ForkId>,
978    ) {
979        self.add_and_connect_kind(peer_id, PeerKind::Basic, addr, fork_id)
980    }
981
982    /// Connects a peer and its address with the given kind.
983    ///
984    /// Note: This is invoked on demand via an external command received by the manager
985    pub(crate) fn add_and_connect_kind(
986        &mut self,
987        peer_id: PeerId,
988        kind: PeerKind,
989        addr: PeerAddr,
990        fork_id: Option<ForkId>,
991    ) {
992        let ip_addr = addr.tcp().ip();
993
994        // Check if the IP is in the allowed ranges (netrestrict)
995        if !self.ip_filter.is_allowed(&ip_addr) {
996            trace!(target: "net", ?peer_id, ?ip_addr, "Skipping outbound connection to IP not in allowed ranges");
997            return
998        }
999
1000        if self.ban_list.is_banned(&peer_id, &ip_addr) {
1001            return
1002        }
1003
1004        match self.peers.entry(peer_id) {
1005            Entry::Occupied(mut entry) => {
1006                let peer = entry.get_mut();
1007                peer.kind = kind;
1008                peer.fork_id = fork_id.map(Box::new);
1009                peer.addr = addr;
1010
1011                if peer.state == PeerConnectionState::Idle {
1012                    // Try connecting again.
1013                    peer.state = PeerConnectionState::PendingOut;
1014                    self.connection_info.inc_pending_out();
1015                    self.queued_actions
1016                        .push_back(PeerAction::Connect { peer_id, remote_addr: addr.tcp() });
1017                }
1018            }
1019            Entry::Vacant(entry) => {
1020                trace!(target: "net::peers", ?peer_id, addr=?addr.tcp(), "connects new node");
1021                let mut peer = Peer::with_kind(addr, kind);
1022                peer.state = PeerConnectionState::PendingOut;
1023                peer.fork_id = fork_id.map(Box::new);
1024                entry.insert(peer);
1025                self.connection_info.inc_pending_out();
1026                self.queued_actions
1027                    .push_back(PeerAction::Connect { peer_id, remote_addr: addr.tcp() });
1028            }
1029        }
1030
1031        if kind.is_trusted() {
1032            self.trusted_peer_ids.insert(peer_id);
1033        }
1034    }
1035
1036    /// Removes the tracked node from the trusted set.
1037    pub(crate) fn remove_peer_from_trusted_set(&mut self, peer_id: PeerId) {
1038        self.trusted_peers_resolver.remove(peer_id);
1039
1040        let Entry::Occupied(mut entry) = self.peers.entry(peer_id) else {
1041            self.trusted_peer_ids.remove(&peer_id);
1042            return
1043        };
1044        if !entry.get().is_trusted() {
1045            return
1046        }
1047
1048        let peer = entry.get_mut();
1049        peer.kind = PeerKind::Basic;
1050
1051        self.trusted_peer_ids.remove(&peer_id);
1052    }
1053
1054    /// Returns the best idle peer to connect to.
1055    ///
1056    /// Peers that are `trusted` or `static`, see [`PeerKind`], are prioritized as long as they're
1057    /// not currently marked as banned or backed off.
1058    ///
1059    /// Among remaining peers, the one with the highest reputation is selected. When reputation is
1060    /// equal, a peer with a discovered `fork_id` is preferred since it indicates a compatible fork.
1061    ///
1062    /// If `trusted_nodes_only` is enabled, see [`PeersConfig`], then this will only consider
1063    /// `trusted` peers.
1064    ///
1065    /// Returns `None` if no peer is available.
1066    fn best_unconnected(&mut self) -> Option<(PeerId, &mut Peer)> {
1067        let mut unconnected = self.peers.iter_mut().filter(|(_, peer)| {
1068            !peer.is_backed_off() &&
1069                !peer.is_banned() &&
1070                peer.state.is_unconnected() &&
1071                (!self.trusted_nodes_only || peer.is_trusted())
1072        });
1073
1074        // keep track of the best peer, if there's one
1075        let mut best_peer = unconnected.next()?;
1076
1077        if best_peer.1.is_trusted() || best_peer.1.is_static() {
1078            return Some((*best_peer.0, best_peer.1))
1079        }
1080
1081        for maybe_better in unconnected {
1082            // if the peer is trusted or static, return it immediately
1083            if maybe_better.1.is_trusted() || maybe_better.1.is_static() {
1084                return Some((*maybe_better.0, maybe_better.1))
1085            }
1086
1087            // prefer higher reputation, break ties by fork_id presence
1088            match maybe_better.1.reputation.cmp(&best_peer.1.reputation) {
1089                std::cmp::Ordering::Greater => best_peer = maybe_better,
1090                std::cmp::Ordering::Equal
1091                    if maybe_better.1.fork_id.is_some() && best_peer.1.fork_id.is_none() =>
1092                {
1093                    best_peer = maybe_better
1094                }
1095                _ => {}
1096            }
1097        }
1098        Some((*best_peer.0, best_peer.1))
1099    }
1100
1101    /// Disconnects one eligible peer (inbound or outbound) to open a slot for new nodes,
1102    /// mirroring Geth's peer dropper logic:
1103    /// <https://github.com/ethereum/go-ethereum/blob/c5c75977ab55e4d7ea6147cc0e221b588e5e3754/eth/dropper.go#L106-L138>
1104    ///
1105    /// A peer is eligible when its pool (inbound or outbound) is at capacity, the peer is not
1106    /// trusted or static, and it has been connected for at least [`PEER_ROTATION_MIN_UPTIME`].
1107    fn try_rotate_peer(&mut self) {
1108        let outbound_at_capacity = self.connection_info.is_outbound_at_capacity();
1109        let inbound_at_capacity = self.connection_info.is_inbound_at_capacity();
1110
1111        if !outbound_at_capacity && !inbound_at_capacity {
1112            return
1113        }
1114
1115        let now = std::time::Instant::now();
1116
1117        let candidates = self
1118            .peers
1119            .iter()
1120            .filter_map(|(peer_id, peer)| {
1121                let eligible = match peer.state {
1122                    PeerConnectionState::Out => outbound_at_capacity,
1123                    PeerConnectionState::In => inbound_at_capacity,
1124                    _ => false,
1125                };
1126                (eligible &&
1127                    !peer.is_trusted() &&
1128                    !peer.is_static() &&
1129                    !self.trusted_peer_ids.contains(peer_id) &&
1130                    peer.connected_for_at_least(now, PEER_ROTATION_MIN_UPTIME))
1131                .then_some(*peer_id)
1132            })
1133            .collect::<Vec<_>>();
1134
1135        if candidates.is_empty() {
1136            return
1137        }
1138        let peer_id = candidates[rand::rng().random_range(0..candidates.len())];
1139
1140        trace!(target: "net::peers", ?peer_id, "rotating peer to open slot for new nodes");
1141
1142        self.queued_actions.push_back(PeerAction::Disconnect {
1143            peer_id,
1144            reason: Some(DisconnectReason::UselessPeer),
1145        });
1146    }
1147
1148    /// If there's capacity for new outbound connections, this will queue new
1149    /// [`PeerAction::Connect`] actions.
1150    ///
1151    /// New connections are only initiated, if slots are available and appropriate peers are
1152    /// available.
1153    fn fill_outbound_slots(&mut self) {
1154        self.tick();
1155
1156        if !self.net_connection_state.is_active() {
1157            // nothing to fill
1158            return
1159        }
1160
1161        // as long as there are slots available fill them with the best peers
1162        while self.connection_info.has_out_capacity() {
1163            let action = {
1164                let (peer_id, peer) = match self.best_unconnected() {
1165                    Some(peer) => peer,
1166                    _ => break,
1167                };
1168
1169                trace!(target: "net::peers", ?peer_id, addr=?peer.addr, "schedule outbound connection");
1170
1171                peer.state = PeerConnectionState::PendingOut;
1172                PeerAction::Connect { peer_id, remote_addr: peer.addr.tcp() }
1173            };
1174
1175            self.connection_info.inc_pending_out();
1176
1177            self.queued_actions.push_back(action);
1178        }
1179    }
1180
1181    fn on_resolved_peer(&mut self, peer_id: PeerId, new_record: NodeRecord) {
1182        if !self.trusted_peer_ids.contains(&peer_id) {
1183            trace!(target: "net::peers", ?peer_id, "Ignoring resolved trusted peer after removal");
1184            return
1185        }
1186
1187        let new_addr = PeerAddr::new_with_ports(
1188            new_record.address,
1189            new_record.tcp_port,
1190            Some(new_record.udp_port),
1191        );
1192
1193        if let Some(peer) = self.peers.get_mut(&peer_id) {
1194            if peer.addr != new_addr {
1195                peer.addr = new_addr;
1196                trace!(target: "net::peers", ?peer_id, addr=?peer.addr, "Updated resolved trusted peer address");
1197            }
1198        } else {
1199            trace!(target: "net::peers", ?peer_id, ?new_addr, "Adding trusted peer after first successful resolution");
1200            self.add_peer_kind(peer_id, Some(PeerKind::Trusted), new_addr, None);
1201        }
1202    }
1203
1204    /// Keeps track of network state changes.
1205    pub const fn on_network_state_change(&mut self, state: NetworkConnectionState) {
1206        self.net_connection_state = state;
1207    }
1208
1209    /// Returns the current network connection state.
1210    pub const fn connection_state(&self) -> &NetworkConnectionState {
1211        &self.net_connection_state
1212    }
1213
1214    /// Sets `net_connection_state` to `ShuttingDown`.
1215    pub const fn on_shutdown(&mut self) {
1216        self.net_connection_state = NetworkConnectionState::ShuttingDown;
1217    }
1218
1219    /// Advances the state.
1220    ///
1221    /// Event hooks invoked externally may trigger a new [`PeerAction`] that are buffered until
1222    /// [`PeersManager`] is polled.
1223    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<PeerAction> {
1224        loop {
1225            // drain buffered actions
1226            if let Some(action) = self.queued_actions.pop_front() {
1227                return Poll::Ready(action)
1228            }
1229
1230            while let Poll::Ready(Some(cmd)) = self.handle_rx.poll_next_unpin(cx) {
1231                match cmd {
1232                    PeerCommand::Add(peer_id, addr) => {
1233                        self.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
1234                    }
1235                    PeerCommand::Remove(peer) => self.remove_peer(peer),
1236                    PeerCommand::ReputationChange(peer_id, rep) => {
1237                        self.apply_reputation_change(&peer_id, rep)
1238                    }
1239                    PeerCommand::GetPeer(peer, tx) => {
1240                        let _ = tx.send(self.peers.get(&peer).cloned());
1241                    }
1242                    PeerCommand::GetPeers(tx) => {
1243                        let _ = tx.send(self.iter_peers().collect());
1244                    }
1245                }
1246            }
1247
1248            if self.release_interval.poll_tick(cx).is_ready() {
1249                let now = std::time::Instant::now();
1250                let (_, unbanned_peers) = self.ban_list.evict(now);
1251
1252                for peer_id in unbanned_peers {
1253                    if let Some(peer) = self.peers.get_mut(&peer_id) {
1254                        peer.unban();
1255                        self.queued_actions.push_back(PeerAction::UnBanPeer { peer_id });
1256                    }
1257                }
1258
1259                // clear the backoff list of expired backoffs, and mark the relevant peers as
1260                // ready to be dialed
1261                self.backed_off_peers.retain(|peer_id, until| {
1262                    if now > *until {
1263                        if let Some(peer) = self.peers.get_mut(peer_id) {
1264                            peer.backed_off = false;
1265                        }
1266                        return false
1267                    }
1268                    true
1269                })
1270            }
1271
1272            while self.refill_slots_interval.poll_tick(cx).is_ready() {
1273                self.fill_outbound_slots();
1274            }
1275
1276            if let Poll::Ready((peer_id, new_record)) = self.trusted_peers_resolver.poll(cx) {
1277                self.on_resolved_peer(peer_id, new_record);
1278            }
1279
1280            // Poll the jittered rotation timer and rotate one eligible peer when ready.
1281            let rotation_ready = self
1282                .peer_rotation_sleep
1283                .as_mut()
1284                .is_some_and(|sleep| sleep.as_mut().poll(cx).is_ready());
1285            if rotation_ready {
1286                self.try_rotate_peer();
1287                if let Some((mean, sleep)) =
1288                    self.peer_rotation_mean.zip(self.peer_rotation_sleep.as_mut())
1289                {
1290                    sleep
1291                        .as_mut()
1292                        .reset(tokio::time::Instant::now() + jitter_rotation_interval(mean));
1293                }
1294            }
1295
1296            if self.queued_actions.is_empty() {
1297                return Poll::Pending
1298            }
1299        }
1300    }
1301}
1302
1303impl Default for PeersManager {
1304    fn default() -> Self {
1305        Self::new(Default::default())
1306    }
1307}
1308
1309/// Tracks stats about connected nodes
1310#[derive(Debug, Clone, PartialEq, Eq, Default)]
1311pub struct ConnectionInfo {
1312    /// Counter for currently occupied slots for active outbound connections.
1313    num_outbound: usize,
1314    /// Counter for pending outbound connections.
1315    num_pending_out: usize,
1316    /// Counter for currently occupied slots for active inbound connections.
1317    num_inbound: usize,
1318    /// Counter for pending inbound connections.
1319    num_pending_in: usize,
1320    /// Restrictions on number of connections.
1321    config: ConnectionsConfig,
1322}
1323
1324// === impl ConnectionInfo ===
1325
1326impl ConnectionInfo {
1327    /// Returns a new [`ConnectionInfo`] with the given config.
1328    const fn new(config: ConnectionsConfig) -> Self {
1329        Self { config, num_outbound: 0, num_pending_out: 0, num_inbound: 0, num_pending_in: 0 }
1330    }
1331
1332    ///  Returns `true` if there's still capacity to perform an outgoing connection.
1333    const fn has_out_capacity(&self) -> bool {
1334        self.num_pending_out < self.config.max_concurrent_outbound_dials &&
1335            self.num_outbound < self.config.max_outbound
1336    }
1337
1338    /// Returns `true` if all active outbound slots are occupied (ignoring pending dials).
1339    const fn is_outbound_at_capacity(&self) -> bool {
1340        self.num_outbound >= self.config.max_outbound
1341    }
1342
1343    /// Returns `true` if all active inbound slots are occupied.
1344    const fn is_inbound_at_capacity(&self) -> bool {
1345        self.num_inbound >= self.config.max_inbound
1346    }
1347
1348    ///  Returns `true` if there's still capacity to accept a new incoming connection.
1349    const fn has_in_capacity(&self) -> bool {
1350        self.num_inbound < self.config.max_inbound
1351    }
1352
1353    /// Returns `true` if we can handle an additional incoming pending connection.
1354    const fn has_in_pending_capacity(&self) -> bool {
1355        self.num_pending_in < self.config.max_inbound
1356    }
1357
1358    const fn decr_state(&mut self, state: PeerConnectionState) {
1359        match state {
1360            PeerConnectionState::Idle => {}
1361            PeerConnectionState::DisconnectingIn | PeerConnectionState::In => self.decr_in(),
1362            PeerConnectionState::DisconnectingOut | PeerConnectionState::Out => self.decr_out(),
1363            PeerConnectionState::PendingOut => self.decr_pending_out(),
1364        }
1365    }
1366
1367    const fn decr_out(&mut self) {
1368        self.num_outbound -= 1;
1369    }
1370
1371    const fn inc_out(&mut self) {
1372        self.num_outbound += 1;
1373    }
1374
1375    const fn inc_pending_out(&mut self) {
1376        self.num_pending_out += 1;
1377    }
1378
1379    const fn inc_in(&mut self) {
1380        self.num_inbound += 1;
1381    }
1382
1383    const fn inc_pending_in(&mut self) {
1384        self.num_pending_in += 1;
1385    }
1386
1387    const fn decr_in(&mut self) {
1388        self.num_inbound -= 1;
1389    }
1390
1391    const fn decr_pending_out(&mut self) {
1392        self.num_pending_out -= 1;
1393    }
1394
1395    const fn decr_pending_in(&mut self) {
1396        self.num_pending_in -= 1;
1397    }
1398}
1399
1400/// Actions the peer manager can trigger.
1401#[derive(Debug)]
1402pub enum PeerAction {
1403    /// Start a new connection to a peer.
1404    Connect {
1405        /// The peer to connect to.
1406        peer_id: PeerId,
1407        /// Where to reach the node
1408        remote_addr: SocketAddr,
1409    },
1410    /// Disconnect an existing connection.
1411    Disconnect {
1412        /// The peer ID of the established connection.
1413        peer_id: PeerId,
1414        /// An optional reason for the disconnect.
1415        reason: Option<DisconnectReason>,
1416    },
1417    /// Disconnect an existing incoming connection, because the peers reputation is below the
1418    /// banned threshold or is on the [`BanList`]
1419    DisconnectBannedIncoming {
1420        /// The peer ID of the established connection.
1421        peer_id: PeerId,
1422    },
1423    /// Disconnect an untrusted incoming connection when trust-node-only is enabled.
1424    DisconnectUntrustedIncoming {
1425        /// The peer ID.
1426        peer_id: PeerId,
1427    },
1428    /// Ban the peer in discovery.
1429    DiscoveryBanPeerId {
1430        /// The peer ID.
1431        peer_id: PeerId,
1432        /// The IP address.
1433        ip_addr: IpAddr,
1434    },
1435    /// Ban the IP in discovery.
1436    DiscoveryBanIp {
1437        /// The IP address.
1438        ip_addr: IpAddr,
1439    },
1440    /// Ban the peer temporarily
1441    BanPeer {
1442        /// The peer ID.
1443        peer_id: PeerId,
1444    },
1445    /// Unban the peer temporarily
1446    UnBanPeer {
1447        /// The peer ID.
1448        peer_id: PeerId,
1449    },
1450    /// Emit peerAdded event
1451    PeerAdded(PeerId),
1452    /// Emit peerRemoved event
1453    PeerRemoved(PeerId),
1454}
1455
1456/// Error thrown when an incoming connection is rejected right away
1457#[derive(Debug, Error, PartialEq, Eq)]
1458pub enum InboundConnectionError {
1459    /// The remote's ip address is banned
1460    IpBanned,
1461    /// No capacity for new inbound connections
1462    ExceedsCapacity,
1463}
1464
1465impl Display for InboundConnectionError {
1466    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1467        write!(f, "{self:?}")
1468    }
1469}
1470
1471/// The reason a peer was backed off.
1472#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1473pub enum BackoffReason {
1474    /// The remote peer responded with `TooManyPeers` (0x04).
1475    TooManyPeers,
1476    /// The session was gracefully closed and we're backing off briefly.
1477    GracefulClose,
1478    /// A connection or protocol-level error occurred.
1479    ConnectionError,
1480}
1481
1482impl BackoffReason {
1483    /// Derives the backoff reason from an optional [`DisconnectReason`].
1484    pub const fn from_disconnect(reason: Option<DisconnectReason>) -> Self {
1485        match reason {
1486            Some(DisconnectReason::TooManyPeers) => Self::TooManyPeers,
1487            _ => Self::ConnectionError,
1488        }
1489    }
1490}
1491
1492/// Returns a random duration uniformly distributed in `[mean * 3/5, mean * 7/5]`.
1493///
1494/// With the default 5-minute mean this gives `[3 min, 7 min]`, matching Geth's peer dropper:
1495/// <https://github.com/ethereum/go-ethereum/blob/c5c75977ab55e4d7ea6147cc0e221b588e5e3754/eth/dropper.go#L32-L36>
1496fn jitter_rotation_interval(mean: Duration) -> Duration {
1497    let min_nanos = (mean * 3 / 5).as_nanos() as u64;
1498    let max_nanos = (mean * 7 / 5).as_nanos() as u64;
1499    Duration::from_nanos(rand::rng().random_range(min_nanos..=max_nanos))
1500}
1501
1502#[cfg(test)]
1503mod tests {
1504    use alloy_primitives::B512;
1505    use reth_eth_wire::{
1506        errors::{EthHandshakeError, EthStreamError, P2PHandshakeError, P2PStreamError},
1507        DisconnectReason,
1508    };
1509    use reth_ethereum_forks::{ForkHash, ForkId};
1510    use reth_net_banlist::BanList;
1511    use reth_network_api::Direction;
1512    use reth_network_peers::{NodeRecord, PeerId, TrustedPeer};
1513    use reth_network_types::{
1514        peers::reputation::DEFAULT_REPUTATION, BackoffKind, Peer, ReputationChangeKind,
1515    };
1516    use std::{
1517        future::{poll_fn, Future},
1518        io,
1519        net::{IpAddr, Ipv4Addr, SocketAddr},
1520        pin::Pin,
1521        task::{Context, Poll},
1522        time::Duration,
1523    };
1524    use url::Host;
1525
1526    use super::PeersManager;
1527    use crate::{
1528        error::SessionError,
1529        peers::{
1530            ConnectionInfo, InboundConnectionError, PeerAction, PeerAddr, PeerBackoffDurations,
1531            PeerConnectionState,
1532        },
1533        session::PendingSessionHandshakeError,
1534        PeersConfig,
1535    };
1536
1537    struct PeerActionFuture<'a> {
1538        peers: &'a mut PeersManager,
1539    }
1540
1541    impl Future for PeerActionFuture<'_> {
1542        type Output = PeerAction;
1543
1544        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1545            self.get_mut().peers.poll(cx)
1546        }
1547    }
1548
1549    macro_rules! event {
1550        ($peers:expr) => {
1551            PeerActionFuture { peers: &mut $peers }.await
1552        };
1553    }
1554
1555    fn set_connected_at(
1556        peers: &mut PeersManager,
1557        peer_id: PeerId,
1558        connected_at: std::time::Instant,
1559    ) {
1560        peers.peers.get_mut(&peer_id).expect("peer exists").connected_at = Some(connected_at);
1561    }
1562
1563    #[tokio::test]
1564    async fn test_insert() {
1565        let peer = PeerId::random();
1566        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1567        let mut peers = PeersManager::default();
1568        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1569
1570        match event!(peers) {
1571            PeerAction::PeerAdded(peer_id) => {
1572                assert_eq!(peer_id, peer);
1573            }
1574            _ => unreachable!(),
1575        }
1576        match event!(peers) {
1577            PeerAction::Connect { peer_id, remote_addr } => {
1578                assert_eq!(peer_id, peer);
1579                assert_eq!(remote_addr, socket_addr);
1580            }
1581            _ => unreachable!(),
1582        }
1583
1584        let (record, _) = peers.peer_by_id(peer).unwrap();
1585        assert_eq!(record.tcp_addr(), socket_addr);
1586        assert_eq!(record.udp_addr(), socket_addr);
1587    }
1588
1589    #[tokio::test]
1590    async fn test_insert_udp() {
1591        let peer = PeerId::random();
1592        let tcp_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1593        let udp_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
1594        let mut peers = PeersManager::default();
1595        peers.add_peer(peer, PeerAddr::new(tcp_addr, Some(udp_addr)), None);
1596
1597        match event!(peers) {
1598            PeerAction::PeerAdded(peer_id) => {
1599                assert_eq!(peer_id, peer);
1600            }
1601            _ => unreachable!(),
1602        }
1603        match event!(peers) {
1604            PeerAction::Connect { peer_id, remote_addr } => {
1605                assert_eq!(peer_id, peer);
1606                assert_eq!(remote_addr, tcp_addr);
1607            }
1608            _ => unreachable!(),
1609        }
1610
1611        let (record, _) = peers.peer_by_id(peer).unwrap();
1612        assert_eq!(record.tcp_addr(), tcp_addr);
1613        assert_eq!(record.udp_addr(), udp_addr);
1614    }
1615
1616    #[tokio::test]
1617    async fn test_ban() {
1618        let peer = PeerId::random();
1619        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1620        let mut peers = PeersManager::default();
1621        peers.ban_peer(peer);
1622        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1623
1624        match event!(peers) {
1625            PeerAction::BanPeer { peer_id } => {
1626                assert_eq!(peer_id, peer);
1627            }
1628            _ => unreachable!(),
1629        }
1630
1631        poll_fn(|cx| {
1632            assert!(peers.poll(cx).is_pending());
1633            Poll::Ready(())
1634        })
1635        .await;
1636    }
1637
1638    #[tokio::test]
1639    async fn test_unban() {
1640        let peer = PeerId::random();
1641        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1642        let mut peers = PeersManager::default();
1643        peers.ban_peer(peer);
1644        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1645
1646        match event!(peers) {
1647            PeerAction::BanPeer { peer_id } => {
1648                assert_eq!(peer_id, peer);
1649            }
1650            _ => unreachable!(),
1651        }
1652
1653        poll_fn(|cx| {
1654            assert!(peers.poll(cx).is_pending());
1655            Poll::Ready(())
1656        })
1657        .await;
1658
1659        peers.unban_peer(peer);
1660
1661        match event!(peers) {
1662            PeerAction::UnBanPeer { peer_id } => {
1663                assert_eq!(peer_id, peer);
1664            }
1665            _ => unreachable!(),
1666        }
1667
1668        poll_fn(|cx| {
1669            assert!(peers.poll(cx).is_pending());
1670            Poll::Ready(())
1671        })
1672        .await;
1673    }
1674
1675    #[tokio::test]
1676    async fn test_admin_ban_removes_peer_and_bans_indefinitely() {
1677        let peer = PeerId::random();
1678        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1679        let mut peers = PeersManager::default();
1680        peers.peers.insert(peer, Peer::new(PeerAddr::from_tcp(socket_addr)));
1681
1682        peers.ban_peer_by_admin(peer);
1683
1684        assert!(peers.ban_list.is_banned_peer(&peer));
1685        assert!(peers.peer_by_id(peer).is_none());
1686
1687        match peers.queued_actions.pop_front() {
1688            Some(PeerAction::PeerRemoved(peer_id)) => assert_eq!(peer_id, peer),
1689            other => panic!("unexpected action: {other:?}"),
1690        }
1691
1692        let (_, unbanned_peers) =
1693            peers.ban_list.evict(std::time::Instant::now() + Duration::from_secs(1));
1694        assert!(unbanned_peers.is_empty());
1695    }
1696
1697    #[tokio::test]
1698    async fn test_admin_ban_does_not_override_trusted_peer() {
1699        let peer = PeerId::random();
1700        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1701        let mut peers = PeersManager::default();
1702        peers.peers.insert(peer, Peer::trusted(PeerAddr::from_tcp(socket_addr)));
1703
1704        peers.ban_peer_by_admin(peer);
1705
1706        assert!(!peers.ban_list.is_banned_peer(&peer));
1707        assert!(peers.peer_by_id(peer).is_some());
1708        assert!(peers.queued_actions.is_empty());
1709    }
1710
1711    #[tokio::test]
1712    async fn test_admin_unban_resets_reputation() {
1713        let peer = PeerId::random();
1714        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1715        let mut peers = PeersManager::default();
1716        let mut peer_info = Peer::new(PeerAddr::from_tcp(socket_addr));
1717        peer_info.reputation = i32::MIN;
1718        peers.peers.insert(peer, peer_info);
1719        peers.ban_list.ban_peer(peer);
1720        assert!(peers.peers.get(&peer).is_some_and(Peer::is_banned));
1721
1722        peers.unban_peer_by_admin(peer);
1723
1724        assert!(!peers.ban_list.is_banned_peer(&peer));
1725        assert!(!peers.peers.get(&peer).is_some_and(Peer::is_banned));
1726        assert!(matches!(
1727            peers.queued_actions.pop_front(),
1728            Some(PeerAction::UnBanPeer { peer_id }) if peer_id == peer
1729        ));
1730    }
1731
1732    #[tokio::test]
1733    async fn test_backoff_on_busy() {
1734        let peer = PeerId::random();
1735        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1736
1737        let mut peers = PeersManager::new(PeersConfig::test());
1738        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1739
1740        match event!(peers) {
1741            PeerAction::PeerAdded(peer_id) => {
1742                assert_eq!(peer_id, peer);
1743            }
1744            _ => unreachable!(),
1745        }
1746        match event!(peers) {
1747            PeerAction::Connect { peer_id, .. } => {
1748                assert_eq!(peer_id, peer);
1749            }
1750            _ => unreachable!(),
1751        }
1752
1753        poll_fn(|cx| {
1754            assert!(peers.poll(cx).is_pending());
1755            Poll::Ready(())
1756        })
1757        .await;
1758
1759        peers.on_active_session_dropped(
1760            &socket_addr,
1761            &peer,
1762            &EthStreamError::P2PStreamError(P2PStreamError::Disconnected(
1763                DisconnectReason::TooManyPeers,
1764            )),
1765        );
1766
1767        poll_fn(|cx| {
1768            assert!(peers.poll(cx).is_pending());
1769            Poll::Ready(())
1770        })
1771        .await;
1772
1773        assert!(peers.backed_off_peers.contains_key(&peer));
1774        assert!(peers.peers.get(&peer).unwrap().is_backed_off());
1775
1776        tokio::time::sleep(peers.backoff_durations.low).await;
1777
1778        match event!(peers) {
1779            PeerAction::Connect { peer_id, .. } => {
1780                assert_eq!(peer_id, peer);
1781            }
1782            _ => unreachable!(),
1783        }
1784
1785        assert!(!peers.backed_off_peers.contains_key(&peer));
1786        assert!(!peers.peers.get(&peer).unwrap().is_backed_off());
1787    }
1788
1789    #[tokio::test]
1790    async fn test_backoff_on_no_response() {
1791        let peer = PeerId::random();
1792        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1793
1794        let backoff_durations = PeerBackoffDurations::test();
1795        let config = PeersConfig { backoff_durations, ..PeersConfig::test() };
1796        let mut peers = PeersManager::new(config);
1797        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1798
1799        match event!(peers) {
1800            PeerAction::PeerAdded(peer_id) => {
1801                assert_eq!(peer_id, peer);
1802            }
1803            _ => unreachable!(),
1804        }
1805        match event!(peers) {
1806            PeerAction::Connect { peer_id, .. } => {
1807                assert_eq!(peer_id, peer);
1808            }
1809            _ => unreachable!(),
1810        }
1811
1812        poll_fn(|cx| {
1813            assert!(peers.poll(cx).is_pending());
1814            Poll::Ready(())
1815        })
1816        .await;
1817
1818        peers.on_outgoing_pending_session_dropped(
1819            &socket_addr,
1820            &peer,
1821            &PendingSessionHandshakeError::Eth(EthStreamError::EthHandshakeError(
1822                EthHandshakeError::NoResponse,
1823            )),
1824        );
1825
1826        poll_fn(|cx| {
1827            assert!(peers.poll(cx).is_pending());
1828            Poll::Ready(())
1829        })
1830        .await;
1831
1832        assert!(peers.backed_off_peers.contains_key(&peer));
1833        assert!(peers.peers.get(&peer).unwrap().is_backed_off());
1834
1835        tokio::time::sleep(backoff_durations.high).await;
1836
1837        match event!(peers) {
1838            PeerAction::Connect { peer_id, .. } => {
1839                assert_eq!(peer_id, peer);
1840            }
1841            _ => unreachable!(),
1842        }
1843
1844        assert!(!peers.backed_off_peers.contains_key(&peer));
1845        assert!(!peers.peers.get(&peer).unwrap().is_backed_off());
1846    }
1847
1848    #[tokio::test]
1849    async fn test_low_backoff() {
1850        let peer = PeerId::random();
1851        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1852        let config = PeersConfig::test();
1853        let mut peers = PeersManager::new(config);
1854        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1855        let peer_struct = peers.peers.get_mut(&peer).unwrap();
1856
1857        let backoff_timestamp = peers
1858            .backoff_durations
1859            .backoff_until(BackoffKind::Low, peer_struct.severe_backoff_counter);
1860
1861        let expected = std::time::Instant::now() + peers.backoff_durations.low;
1862        assert!(backoff_timestamp <= expected);
1863    }
1864
1865    #[tokio::test]
1866    async fn test_multiple_backoff_calculations() {
1867        let peer = PeerId::random();
1868        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1869        let config = PeersConfig::default();
1870        let mut peers = PeersManager::new(config);
1871        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1872        let peer_struct = peers.peers.get_mut(&peer).unwrap();
1873
1874        // Simulate a peer that was already backed off once
1875        peer_struct.severe_backoff_counter = 1;
1876
1877        let now = std::time::Instant::now();
1878
1879        // Simulate the increment that happens in on_connection_failure
1880        peer_struct.severe_backoff_counter += 1;
1881        // Get official backoff time
1882        let backoff_time = peers
1883            .backoff_durations
1884            .backoff_until(BackoffKind::High, peer_struct.severe_backoff_counter);
1885
1886        // Duration of the backoff should be 2 * 15 minutes = 30 minutes
1887        let backoff_duration = std::time::Duration::new(30 * 60, 0);
1888
1889        // We can't use assert_eq! since there is a very small diff in the nano secs
1890        // Usually it is 1800s != 1799.9999996s
1891        assert!(backoff_time.duration_since(now) > backoff_duration);
1892    }
1893
1894    #[tokio::test]
1895    async fn test_ban_on_active_drop() {
1896        let peer = PeerId::random();
1897        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1898        let mut peers = PeersManager::default();
1899        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1900
1901        match event!(peers) {
1902            PeerAction::PeerAdded(peer_id) => {
1903                assert_eq!(peer_id, peer);
1904            }
1905            _ => unreachable!(),
1906        }
1907        match event!(peers) {
1908            PeerAction::Connect { peer_id, .. } => {
1909                assert_eq!(peer_id, peer);
1910            }
1911            _ => unreachable!(),
1912        }
1913
1914        poll_fn(|cx| {
1915            assert!(peers.poll(cx).is_pending());
1916            Poll::Ready(())
1917        })
1918        .await;
1919
1920        peers.on_active_session_dropped(
1921            &socket_addr,
1922            &peer,
1923            &EthStreamError::P2PStreamError(P2PStreamError::Disconnected(
1924                DisconnectReason::UselessPeer,
1925            )),
1926        );
1927
1928        match event!(peers) {
1929            PeerAction::PeerRemoved(peer_id) => {
1930                assert_eq!(peer_id, peer);
1931            }
1932            _ => unreachable!(),
1933        }
1934        match event!(peers) {
1935            PeerAction::BanPeer { peer_id } => {
1936                assert_eq!(peer_id, peer);
1937            }
1938            _ => unreachable!(),
1939        }
1940
1941        poll_fn(|cx| {
1942            assert!(peers.poll(cx).is_pending());
1943            Poll::Ready(())
1944        })
1945        .await;
1946
1947        assert!(!peers.peers.contains_key(&peer));
1948    }
1949
1950    #[tokio::test]
1951    async fn test_remove_on_max_backoff_count() {
1952        let peer = PeerId::random();
1953        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
1954        let config = PeersConfig::test();
1955        let mut peers = PeersManager::new(config.clone());
1956        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
1957        let peer_struct = peers.peers.get_mut(&peer).unwrap();
1958
1959        // Simulate a peer that was already backed off once
1960        peer_struct.severe_backoff_counter = config.max_backoff_count;
1961
1962        match event!(peers) {
1963            PeerAction::PeerAdded(peer_id) => {
1964                assert_eq!(peer_id, peer);
1965            }
1966            _ => unreachable!(),
1967        }
1968        match event!(peers) {
1969            PeerAction::Connect { peer_id, .. } => {
1970                assert_eq!(peer_id, peer);
1971            }
1972            _ => unreachable!(),
1973        }
1974
1975        poll_fn(|cx| {
1976            assert!(peers.poll(cx).is_pending());
1977            Poll::Ready(())
1978        })
1979        .await;
1980
1981        peers.on_outgoing_pending_session_dropped(
1982            &socket_addr,
1983            &peer,
1984            &PendingSessionHandshakeError::Eth(
1985                io::Error::new(io::ErrorKind::ConnectionRefused, "peer unreachable").into(),
1986            ),
1987        );
1988
1989        match event!(peers) {
1990            PeerAction::PeerRemoved(peer_id) => {
1991                assert_eq!(peer_id, peer);
1992            }
1993            _ => unreachable!(),
1994        }
1995
1996        poll_fn(|cx| {
1997            assert!(peers.poll(cx).is_pending());
1998            Poll::Ready(())
1999        })
2000        .await;
2001
2002        assert!(!peers.peers.contains_key(&peer));
2003    }
2004
2005    #[tokio::test]
2006    async fn test_ban_on_pending_drop() {
2007        let peer = PeerId::random();
2008        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2009        let mut peers = PeersManager::default();
2010        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2011
2012        match event!(peers) {
2013            PeerAction::PeerAdded(peer_id) => {
2014                assert_eq!(peer_id, peer);
2015            }
2016            _ => unreachable!(),
2017        }
2018        match event!(peers) {
2019            PeerAction::Connect { peer_id, .. } => {
2020                assert_eq!(peer_id, peer);
2021            }
2022            _ => unreachable!(),
2023        }
2024
2025        poll_fn(|cx| {
2026            assert!(peers.poll(cx).is_pending());
2027            Poll::Ready(())
2028        })
2029        .await;
2030
2031        peers.on_outgoing_pending_session_dropped(
2032            &socket_addr,
2033            &peer,
2034            &PendingSessionHandshakeError::Eth(EthStreamError::P2PStreamError(
2035                P2PStreamError::Disconnected(DisconnectReason::UselessPeer),
2036            )),
2037        );
2038
2039        match event!(peers) {
2040            PeerAction::PeerRemoved(peer_id) => {
2041                assert_eq!(peer_id, peer);
2042            }
2043            _ => unreachable!(),
2044        }
2045        match event!(peers) {
2046            PeerAction::BanPeer { peer_id } => {
2047                assert_eq!(peer_id, peer);
2048            }
2049            _ => unreachable!(),
2050        }
2051
2052        poll_fn(|cx| {
2053            assert!(peers.poll(cx).is_pending());
2054            Poll::Ready(())
2055        })
2056        .await;
2057
2058        assert!(!peers.peers.contains_key(&peer));
2059    }
2060
2061    #[tokio::test]
2062    async fn test_internally_closed_incoming() {
2063        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2064        let mut peers = PeersManager::default();
2065
2066        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2067        assert_eq!(peers.connection_info.num_pending_in, 1);
2068        peers.on_incoming_pending_session_rejected_internally();
2069        assert_eq!(peers.connection_info.num_pending_in, 0);
2070    }
2071
2072    #[tokio::test]
2073    async fn test_reject_incoming_at_pending_capacity() {
2074        let mut peers = PeersManager::default();
2075
2076        for count in 1..=peers.connection_info.config.max_inbound {
2077            let socket_addr =
2078                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, count as u8)), 8008);
2079            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2080            assert_eq!(peers.connection_info.num_pending_in, count);
2081        }
2082        assert!(peers.connection_info.has_in_capacity());
2083        assert!(!peers.connection_info.has_in_pending_capacity());
2084
2085        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 100)), 8008);
2086        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_err());
2087    }
2088
2089    #[tokio::test]
2090    async fn test_reject_incoming_at_pending_capacity_trusted_peers() {
2091        let mut peers = PeersManager::new(PeersConfig::test().with_max_inbound(2));
2092        let trusted = PeerId::random();
2093        peers.add_trusted_peer_id(trusted);
2094
2095        // connect the trusted peer
2096        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 0)), 8008);
2097        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
2098        peers.on_incoming_session_established(trusted, addr);
2099
2100        match event!(peers) {
2101            PeerAction::PeerAdded(id) => {
2102                assert_eq!(id, trusted);
2103            }
2104            _ => unreachable!(),
2105        }
2106
2107        // saturate the remaining inbound slots with untrusted peers
2108        let mut connected_untrusted_peer_ids = Vec::new();
2109        for i in 0..(peers.connection_info.config.max_inbound - 1) {
2110            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, (i + 1) as u8)), 8008);
2111            assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
2112            let peer_id = PeerId::random();
2113            peers.on_incoming_session_established(peer_id, addr);
2114            connected_untrusted_peer_ids.push(peer_id);
2115
2116            match event!(peers) {
2117                PeerAction::PeerAdded(id) => {
2118                    assert_eq!(id, peer_id);
2119                }
2120                _ => unreachable!(),
2121            }
2122        }
2123
2124        let mut pending_addrs = Vec::new();
2125
2126        // saturate available slots
2127        for i in 0..2 {
2128            let socket_addr =
2129                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, (i + 10) as u8)), 8008);
2130            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2131
2132            pending_addrs.push(socket_addr);
2133        }
2134
2135        assert_eq!(peers.connection_info.num_pending_in, 2);
2136
2137        // try to handle additional incoming connections at capacity
2138        for i in 0..2 {
2139            let socket_addr =
2140                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, (i + 20) as u8)), 8008);
2141            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_err());
2142        }
2143
2144        let err = PendingSessionHandshakeError::Eth(EthStreamError::P2PStreamError(
2145            P2PStreamError::HandshakeError(P2PHandshakeError::Disconnected(
2146                DisconnectReason::UselessPeer,
2147            )),
2148        ));
2149
2150        // Remove all pending peers
2151        for pending_addr in pending_addrs {
2152            peers.on_incoming_pending_session_dropped(pending_addr, &err);
2153        }
2154
2155        println!("num_pending_in: {}", peers.connection_info.num_pending_in);
2156
2157        println!(
2158            "num_inbound: {}, has_in_capacity: {}",
2159            peers.connection_info.num_inbound,
2160            peers.connection_info.has_in_capacity()
2161        );
2162
2163        // disconnect a connected peer
2164        peers.on_active_session_gracefully_closed(connected_untrusted_peer_ids[0]);
2165
2166        println!(
2167            "num_inbound: {}, has_in_capacity: {}",
2168            peers.connection_info.num_inbound,
2169            peers.connection_info.has_in_capacity()
2170        );
2171
2172        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 99)), 8008);
2173        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2174    }
2175
2176    #[tokio::test]
2177    async fn test_closed_incoming() {
2178        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2179        let mut peers = PeersManager::default();
2180
2181        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2182        assert_eq!(peers.connection_info.num_pending_in, 1);
2183        peers.on_incoming_pending_session_gracefully_closed();
2184        assert_eq!(peers.connection_info.num_pending_in, 0);
2185    }
2186
2187    #[tokio::test]
2188    async fn test_dropped_incoming() {
2189        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 0, 1, 2)), 8008);
2190        let ban_duration = Duration::from_millis(500);
2191        let config = PeersConfig { ban_duration, ..PeersConfig::test() };
2192        let mut peers = PeersManager::new(config);
2193
2194        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2195        assert_eq!(peers.connection_info.num_pending_in, 1);
2196        let err = PendingSessionHandshakeError::Eth(EthStreamError::P2PStreamError(
2197            P2PStreamError::HandshakeError(P2PHandshakeError::Disconnected(
2198                DisconnectReason::UselessPeer,
2199            )),
2200        ));
2201
2202        peers.on_incoming_pending_session_dropped(socket_addr, &err);
2203        assert_eq!(peers.connection_info.num_pending_in, 0);
2204        assert!(peers.ban_list.is_banned_ip(&socket_addr.ip()));
2205
2206        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_err());
2207
2208        // unbanned after timeout
2209        tokio::time::sleep(ban_duration).await;
2210
2211        poll_fn(|cx| {
2212            let _ = peers.poll(cx);
2213            Poll::Ready(())
2214        })
2215        .await;
2216
2217        assert!(!peers.ban_list.is_banned_ip(&socket_addr.ip()));
2218        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2219    }
2220
2221    #[tokio::test]
2222    async fn test_reputation_change_connected() {
2223        let peer = PeerId::random();
2224        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2225        let mut peers = PeersManager::default();
2226        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2227
2228        match event!(peers) {
2229            PeerAction::PeerAdded(peer_id) => {
2230                assert_eq!(peer_id, peer);
2231            }
2232            _ => unreachable!(),
2233        }
2234        match event!(peers) {
2235            PeerAction::Connect { peer_id, remote_addr } => {
2236                assert_eq!(peer_id, peer);
2237                assert_eq!(remote_addr, socket_addr);
2238            }
2239            _ => unreachable!(),
2240        }
2241
2242        let p = peers.peers.get_mut(&peer).unwrap();
2243        assert_eq!(p.state, PeerConnectionState::PendingOut);
2244
2245        peers.apply_reputation_change(&peer, ReputationChangeKind::BadProtocol);
2246
2247        let p = peers.peers.get(&peer).unwrap();
2248        assert_eq!(p.state, PeerConnectionState::PendingOut);
2249        assert!(p.is_banned());
2250
2251        peers.on_active_session_gracefully_closed(peer);
2252
2253        let p = peers.peers.get(&peer).unwrap();
2254        assert_eq!(p.state, PeerConnectionState::Idle);
2255        assert!(p.is_banned());
2256
2257        match event!(peers) {
2258            PeerAction::Disconnect { peer_id, .. } => {
2259                assert_eq!(peer_id, peer);
2260            }
2261            _ => unreachable!(),
2262        }
2263    }
2264
2265    #[tokio::test]
2266    async fn retain_trusted_status() {
2267        let _socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 99)), 8008);
2268        let trusted = PeerId::random();
2269        let mut peers =
2270            PeersManager::new(PeersConfig::test().with_trusted_nodes(vec![TrustedPeer {
2271                host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2272                tcp_port: 8008,
2273                udp_port: 8008,
2274                id: trusted,
2275            }]));
2276        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2277        peers.add_peer(trusted, PeerAddr::from_tcp(socket_addr), None);
2278        assert!(peers.peers.get(&trusted).unwrap().is_trusted());
2279    }
2280
2281    #[tokio::test]
2282    async fn accept_incoming_trusted_unknown_peer_address() {
2283        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 99)), 8008);
2284        let mut peers = PeersManager::new(PeersConfig::test().with_max_inbound(2));
2285        // try to connect trusted peer
2286        let trusted = PeerId::random();
2287        peers.add_trusted_peer_id(trusted);
2288
2289        // saturate the inbound slots
2290        for i in 0..peers.connection_info.config.max_inbound {
2291            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, i as u8)), 8008);
2292            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2293            let peer_id = PeerId::random();
2294            peers.on_incoming_session_established(peer_id, addr);
2295
2296            match event!(peers) {
2297                PeerAction::PeerAdded(id) => {
2298                    assert_eq!(id, peer_id);
2299                }
2300                _ => unreachable!(),
2301            }
2302        }
2303
2304        // try to connect untrusted peer
2305        let untrusted = PeerId::random();
2306        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 99)), 8008);
2307        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2308        peers.on_incoming_session_established(untrusted, socket_addr);
2309
2310        match event!(peers) {
2311            PeerAction::PeerAdded(id) => {
2312                assert_eq!(id, untrusted);
2313            }
2314            _ => unreachable!(),
2315        }
2316
2317        match event!(peers) {
2318            PeerAction::Disconnect { peer_id, reason } => {
2319                assert_eq!(peer_id, untrusted);
2320                assert_eq!(reason, Some(DisconnectReason::TooManyPeers));
2321            }
2322            _ => unreachable!(),
2323        }
2324
2325        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 100)), 8008);
2326        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2327        peers.on_incoming_session_established(trusted, socket_addr);
2328
2329        match event!(peers) {
2330            PeerAction::PeerAdded(id) => {
2331                assert_eq!(id, trusted);
2332            }
2333            _ => unreachable!(),
2334        }
2335
2336        poll_fn(|cx| {
2337            assert!(peers.poll(cx).is_pending());
2338            Poll::Ready(())
2339        })
2340        .await;
2341
2342        let peer = peers.peers.get(&trusted).unwrap();
2343        assert_eq!(peer.state, PeerConnectionState::In);
2344    }
2345
2346    #[tokio::test]
2347    async fn test_already_connected() {
2348        let peer = PeerId::random();
2349        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2350        let mut peers = PeersManager::default();
2351
2352        // Attempt to establish an incoming session, expecting `num_pending_in` to increase by 1
2353        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2354        assert_eq!(peers.connection_info.num_pending_in, 1);
2355
2356        // Establish a session with the peer, expecting the peer to be added and the `num_inbound`
2357        // to increase by 1
2358        peers.on_incoming_session_established(peer, socket_addr);
2359        let p = peers.peers.get_mut(&peer).expect("peer not found");
2360        assert_eq!(p.addr.tcp(), socket_addr);
2361        assert_eq!(peers.connection_info.num_pending_in, 0);
2362        assert_eq!(peers.connection_info.num_inbound, 1);
2363
2364        // Attempt to establish another incoming session, expecting the `num_pending_in` to increase
2365        // by 1
2366        assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
2367        assert_eq!(peers.connection_info.num_pending_in, 1);
2368
2369        // Simulate a rejection due to an already established connection, expecting the
2370        // `num_pending_in` to decrease by 1. The peer should remain connected and the `num_inbound`
2371        // should not be changed.
2372        peers.on_already_connected(Direction::Incoming);
2373
2374        let p = peers.peers.get_mut(&peer).expect("peer not found");
2375        assert_eq!(p.addr.tcp(), socket_addr);
2376        assert_eq!(peers.connection_info.num_pending_in, 0);
2377        assert_eq!(peers.connection_info.num_inbound, 1);
2378    }
2379
2380    #[tokio::test]
2381    async fn test_reputation_change_trusted_peer() {
2382        let peer = PeerId::random();
2383        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2384        let mut peers = PeersManager::default();
2385        peers.add_trusted_peer(peer, PeerAddr::from_tcp(socket_addr));
2386
2387        match event!(peers) {
2388            PeerAction::PeerAdded(peer_id) => {
2389                assert_eq!(peer_id, peer);
2390            }
2391            _ => unreachable!(),
2392        }
2393        match event!(peers) {
2394            PeerAction::Connect { peer_id, remote_addr } => {
2395                assert_eq!(peer_id, peer);
2396                assert_eq!(remote_addr, socket_addr);
2397            }
2398            _ => unreachable!(),
2399        }
2400
2401        assert_eq!(peers.peers.get_mut(&peer).unwrap().state, PeerConnectionState::PendingOut);
2402        peers.on_active_outgoing_established(peer);
2403        assert_eq!(peers.peers.get_mut(&peer).unwrap().state, PeerConnectionState::Out);
2404
2405        peers.apply_reputation_change(&peer, ReputationChangeKind::BadMessage);
2406
2407        {
2408            let p = peers.peers.get(&peer).unwrap();
2409            assert_eq!(p.state, PeerConnectionState::Out);
2410            // not banned yet
2411            assert!(!p.is_banned());
2412        }
2413
2414        // ensure peer is banned eventually
2415        loop {
2416            peers.apply_reputation_change(&peer, ReputationChangeKind::BadMessage);
2417
2418            let p = peers.peers.get(&peer).unwrap();
2419            if p.is_banned() {
2420                break
2421            }
2422        }
2423
2424        match event!(peers) {
2425            PeerAction::Disconnect { peer_id, .. } => {
2426                assert_eq!(peer_id, peer);
2427            }
2428            _ => unreachable!(),
2429        }
2430    }
2431
2432    #[tokio::test]
2433    async fn test_reputation_management() {
2434        let peer = PeerId::random();
2435        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2436        let mut peers = PeersManager::default();
2437        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2438        assert_eq!(peers.get_reputation(&peer), Some(0));
2439
2440        peers.apply_reputation_change(&peer, ReputationChangeKind::Other(1024));
2441        assert_eq!(peers.get_reputation(&peer), Some(1024));
2442
2443        peers.apply_reputation_change(&peer, ReputationChangeKind::Reset);
2444        assert_eq!(peers.get_reputation(&peer), Some(0));
2445    }
2446
2447    #[tokio::test]
2448    async fn test_remove_discovered_active() {
2449        let peer = PeerId::random();
2450        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2451        let mut peers = PeersManager::default();
2452        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2453
2454        match event!(peers) {
2455            PeerAction::PeerAdded(peer_id) => {
2456                assert_eq!(peer_id, peer);
2457            }
2458            _ => unreachable!(),
2459        }
2460        match event!(peers) {
2461            PeerAction::Connect { peer_id, remote_addr } => {
2462                assert_eq!(peer_id, peer);
2463                assert_eq!(remote_addr, socket_addr);
2464            }
2465            _ => unreachable!(),
2466        }
2467
2468        let p = peers.peers.get(&peer).unwrap();
2469        assert_eq!(p.state, PeerConnectionState::PendingOut);
2470
2471        peers.remove_peer(peer);
2472
2473        match event!(peers) {
2474            PeerAction::PeerRemoved(peer_id) => {
2475                assert_eq!(peer_id, peer);
2476            }
2477            _ => unreachable!(),
2478        }
2479        match event!(peers) {
2480            PeerAction::Disconnect { peer_id, .. } => {
2481                assert_eq!(peer_id, peer);
2482            }
2483            _ => unreachable!(),
2484        }
2485
2486        let p = peers.peers.get(&peer).unwrap();
2487        assert_eq!(p.state, PeerConnectionState::PendingOut);
2488
2489        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2490        let p = peers.peers.get(&peer).unwrap();
2491        assert_eq!(p.state, PeerConnectionState::PendingOut);
2492
2493        peers.on_active_session_gracefully_closed(peer);
2494        assert!(!peers.peers.contains_key(&peer));
2495    }
2496
2497    #[tokio::test]
2498    async fn test_fatal_outgoing_connection_error_trusted() {
2499        let peer = PeerId::random();
2500        let config = PeersConfig::test()
2501            .with_trusted_nodes(vec![TrustedPeer {
2502                host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2503                tcp_port: 8008,
2504                udp_port: 8008,
2505                id: peer,
2506            }])
2507            .with_trusted_nodes_only(true);
2508        let mut peers = PeersManager::new(config);
2509        let socket_addr = peers.peers.get(&peer).unwrap().addr.tcp();
2510
2511        match event!(peers) {
2512            PeerAction::Connect { peer_id, remote_addr } => {
2513                assert_eq!(peer_id, peer);
2514                assert_eq!(remote_addr, socket_addr);
2515            }
2516            _ => unreachable!(),
2517        }
2518
2519        let p = peers.peers.get(&peer).unwrap();
2520        assert_eq!(p.state, PeerConnectionState::PendingOut);
2521
2522        assert_eq!(peers.num_outbound_connections(), 0);
2523
2524        let err = PendingSessionHandshakeError::Eth(EthStreamError::EthHandshakeError(
2525            EthHandshakeError::NonStatusMessageInHandshake,
2526        ));
2527        assert!(err.is_fatal_protocol_error());
2528
2529        peers.on_outgoing_pending_session_dropped(&socket_addr, &peer, &err);
2530        assert_eq!(peers.num_outbound_connections(), 0);
2531
2532        // try tmp ban peer
2533        match event!(peers) {
2534            PeerAction::BanPeer { peer_id } => {
2535                assert_eq!(peer_id, peer);
2536            }
2537            err => unreachable!("{err:?}"),
2538        }
2539
2540        // ensure we still have trusted peer
2541        assert!(peers.peers.contains_key(&peer));
2542
2543        // await for the ban to expire
2544        tokio::time::sleep(peers.backoff_durations.medium).await;
2545
2546        match event!(peers) {
2547            PeerAction::Connect { peer_id, remote_addr } => {
2548                assert_eq!(peer_id, peer);
2549                assert_eq!(remote_addr, socket_addr);
2550            }
2551            err => unreachable!("{err:?}"),
2552        }
2553    }
2554
2555    #[tokio::test]
2556    async fn test_outgoing_connection_error() {
2557        let peer = PeerId::random();
2558        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2559        let mut peers = PeersManager::default();
2560        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2561
2562        match event!(peers) {
2563            PeerAction::PeerAdded(peer_id) => {
2564                assert_eq!(peer_id, peer);
2565            }
2566            _ => unreachable!(),
2567        }
2568        match event!(peers) {
2569            PeerAction::Connect { peer_id, remote_addr } => {
2570                assert_eq!(peer_id, peer);
2571                assert_eq!(remote_addr, socket_addr);
2572            }
2573            _ => unreachable!(),
2574        }
2575
2576        let p = peers.peers.get(&peer).unwrap();
2577        assert_eq!(p.state, PeerConnectionState::PendingOut);
2578
2579        assert_eq!(peers.num_outbound_connections(), 0);
2580
2581        peers.on_outgoing_connection_failure(
2582            &socket_addr,
2583            &peer,
2584            &io::Error::new(io::ErrorKind::ConnectionRefused, ""),
2585        );
2586
2587        assert_eq!(peers.num_outbound_connections(), 0);
2588    }
2589
2590    #[tokio::test]
2591    async fn test_outgoing_connection_gracefully_closed() {
2592        let peer = PeerId::random();
2593        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2594        let mut peers = PeersManager::default();
2595        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
2596
2597        match event!(peers) {
2598            PeerAction::PeerAdded(peer_id) => {
2599                assert_eq!(peer_id, peer);
2600            }
2601            _ => unreachable!(),
2602        }
2603        match event!(peers) {
2604            PeerAction::Connect { peer_id, remote_addr } => {
2605                assert_eq!(peer_id, peer);
2606                assert_eq!(remote_addr, socket_addr);
2607            }
2608            _ => unreachable!(),
2609        }
2610
2611        let p = peers.peers.get(&peer).unwrap();
2612        assert_eq!(p.state, PeerConnectionState::PendingOut);
2613
2614        assert_eq!(peers.num_outbound_connections(), 0);
2615
2616        peers.on_outgoing_pending_session_gracefully_closed(&peer);
2617
2618        assert_eq!(peers.num_outbound_connections(), 0);
2619        assert_eq!(peers.connection_info.num_pending_out, 0);
2620    }
2621
2622    #[tokio::test]
2623    async fn test_discovery_ban_list() {
2624        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
2625        let socket_addr = SocketAddr::new(ip, 8008);
2626        let ban_list = BanList::new(vec![], vec![ip]);
2627        let config = PeersConfig::default().with_ban_list(ban_list);
2628        let mut peer_manager = PeersManager::new(config);
2629        peer_manager.add_peer(B512::default(), PeerAddr::from_tcp(socket_addr), None);
2630
2631        assert!(peer_manager.peers.is_empty());
2632    }
2633
2634    #[tokio::test]
2635    async fn test_on_pending_ban_list() {
2636        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
2637        let socket_addr = SocketAddr::new(ip, 8008);
2638        let ban_list = BanList::new(vec![], vec![ip]);
2639        let config = PeersConfig::test().with_ban_list(ban_list);
2640        let mut peer_manager = PeersManager::new(config);
2641        let a = peer_manager.on_incoming_pending_session(socket_addr.ip());
2642        // because we have no active peers this should be fine for testings
2643        match a {
2644            Ok(_) => panic!(),
2645            Err(err) => match err {
2646                InboundConnectionError::IpBanned => {
2647                    assert_eq!(peer_manager.connection_info.num_pending_in, 0)
2648                }
2649                _ => unreachable!(),
2650            },
2651        }
2652    }
2653
2654    #[tokio::test]
2655    async fn test_on_active_inbound_ban_list() {
2656        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
2657        let socket_addr = SocketAddr::new(ip, 8008);
2658        let given_peer_id = PeerId::random();
2659        let ban_list = BanList::new(vec![given_peer_id], vec![]);
2660        let config = PeersConfig::test().with_ban_list(ban_list);
2661        let mut peer_manager = PeersManager::new(config);
2662        assert!(peer_manager.on_incoming_pending_session(socket_addr.ip()).is_ok());
2663        // non-trusted nodes should also increase pending_in
2664        assert_eq!(peer_manager.connection_info.num_pending_in, 1);
2665        peer_manager.on_incoming_session_established(given_peer_id, socket_addr);
2666        // after the connection is established, the peer should be removed, the num_pending_in
2667        // should be decreased, and the num_inbound should not be increased
2668        assert_eq!(peer_manager.connection_info.num_pending_in, 0);
2669        assert_eq!(peer_manager.connection_info.num_inbound, 0);
2670
2671        let Some(PeerAction::DisconnectBannedIncoming { peer_id }) =
2672            peer_manager.queued_actions.pop_front()
2673        else {
2674            panic!()
2675        };
2676
2677        assert_eq!(peer_id, given_peer_id)
2678    }
2679
2680    #[test]
2681    fn test_connection_limits() {
2682        let mut info = ConnectionInfo::default();
2683        info.inc_in();
2684        assert_eq!(info.num_inbound, 1);
2685        assert_eq!(info.num_outbound, 0);
2686        assert!(info.has_in_capacity());
2687
2688        info.decr_in();
2689        assert_eq!(info.num_inbound, 0);
2690        assert_eq!(info.num_outbound, 0);
2691
2692        info.inc_out();
2693        assert_eq!(info.num_inbound, 0);
2694        assert_eq!(info.num_outbound, 1);
2695        assert!(info.has_out_capacity());
2696
2697        info.decr_out();
2698        assert_eq!(info.num_inbound, 0);
2699        assert_eq!(info.num_outbound, 0);
2700    }
2701
2702    #[test]
2703    fn test_connection_peer_state() {
2704        let mut info = ConnectionInfo::default();
2705        info.inc_in();
2706
2707        info.decr_state(PeerConnectionState::In);
2708        assert_eq!(info.num_inbound, 0);
2709        assert_eq!(info.num_outbound, 0);
2710
2711        info.inc_out();
2712
2713        info.decr_state(PeerConnectionState::Out);
2714        assert_eq!(info.num_inbound, 0);
2715        assert_eq!(info.num_outbound, 0);
2716    }
2717
2718    #[tokio::test]
2719    async fn test_trusted_peers_are_prioritized() {
2720        let trusted_peer = PeerId::random();
2721        let trusted_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2722        let config = PeersConfig::test().with_trusted_nodes(vec![TrustedPeer {
2723            host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2724            tcp_port: 8008,
2725            udp_port: 8008,
2726            id: trusted_peer,
2727        }]);
2728        let mut peers = PeersManager::new(config);
2729
2730        let basic_peer = PeerId::random();
2731        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2732        peers.add_peer(basic_peer, PeerAddr::from_tcp(basic_sock), None);
2733
2734        match event!(peers) {
2735            PeerAction::PeerAdded(peer_id) => {
2736                assert_eq!(peer_id, basic_peer);
2737            }
2738            _ => unreachable!(),
2739        }
2740        match event!(peers) {
2741            PeerAction::Connect { peer_id, remote_addr } => {
2742                assert_eq!(peer_id, trusted_peer);
2743                assert_eq!(remote_addr, trusted_sock);
2744            }
2745            _ => unreachable!(),
2746        }
2747        match event!(peers) {
2748            PeerAction::Connect { peer_id, remote_addr } => {
2749                assert_eq!(peer_id, basic_peer);
2750                assert_eq!(remote_addr, basic_sock);
2751            }
2752            _ => unreachable!(),
2753        }
2754    }
2755
2756    #[tokio::test]
2757    async fn test_connect_trusted_nodes_only() {
2758        let trusted_peer = PeerId::random();
2759        let trusted_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
2760        let config = PeersConfig::test()
2761            .with_trusted_nodes(vec![TrustedPeer {
2762                host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2763                tcp_port: 8008,
2764                udp_port: 8008,
2765                id: trusted_peer,
2766            }])
2767            .with_trusted_nodes_only(true);
2768        let mut peers = PeersManager::new(config);
2769
2770        let basic_peer = PeerId::random();
2771        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2772        peers.add_peer(basic_peer, PeerAddr::from_tcp(basic_sock), None);
2773
2774        match event!(peers) {
2775            PeerAction::PeerAdded(peer_id) => {
2776                assert_eq!(peer_id, basic_peer);
2777            }
2778            _ => unreachable!(),
2779        }
2780        match event!(peers) {
2781            PeerAction::Connect { peer_id, remote_addr } => {
2782                assert_eq!(peer_id, trusted_peer);
2783                assert_eq!(remote_addr, trusted_sock);
2784            }
2785            _ => unreachable!(),
2786        }
2787        poll_fn(|cx| {
2788            assert!(peers.poll(cx).is_pending());
2789            Poll::Ready(())
2790        })
2791        .await;
2792    }
2793
2794    #[tokio::test]
2795    async fn test_incoming_with_trusted_nodes_only() {
2796        let trusted_peer = PeerId::random();
2797        let config = PeersConfig::test()
2798            .with_trusted_nodes(vec![TrustedPeer {
2799                host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2800                tcp_port: 8008,
2801                udp_port: 8008,
2802                id: trusted_peer,
2803            }])
2804            .with_trusted_nodes_only(true);
2805        let mut peers = PeersManager::new(config);
2806
2807        let basic_peer = PeerId::random();
2808        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2809        assert!(peers.on_incoming_pending_session(basic_sock.ip()).is_ok());
2810        // non-trusted nodes should also increase pending_in
2811        assert_eq!(peers.connection_info.num_pending_in, 1);
2812        peers.on_incoming_session_established(basic_peer, basic_sock);
2813        // after the connection is established, the peer should be removed, the num_pending_in
2814        // should be decreased, and the num_inbound mut not be increased
2815        assert_eq!(peers.connection_info.num_pending_in, 0);
2816        assert_eq!(peers.connection_info.num_inbound, 0);
2817
2818        let Some(PeerAction::DisconnectUntrustedIncoming { peer_id }) =
2819            peers.queued_actions.pop_front()
2820        else {
2821            panic!()
2822        };
2823        assert_eq!(basic_peer, peer_id);
2824        assert!(!peers.peers.contains_key(&basic_peer));
2825    }
2826
2827    #[tokio::test]
2828    async fn test_incoming_without_trusted_nodes_only() {
2829        let trusted_peer = PeerId::random();
2830        let config = PeersConfig::test()
2831            .with_trusted_nodes(vec![TrustedPeer {
2832                host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 2)),
2833                tcp_port: 8008,
2834                udp_port: 8008,
2835                id: trusted_peer,
2836            }])
2837            .with_trusted_nodes_only(false);
2838        let mut peers = PeersManager::new(config);
2839
2840        let basic_peer = PeerId::random();
2841        let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2842        assert!(peers.on_incoming_pending_session(basic_sock.ip()).is_ok());
2843
2844        // non-trusted nodes should also increase pending_in
2845        assert_eq!(peers.connection_info.num_pending_in, 1);
2846        peers.on_incoming_session_established(basic_peer, basic_sock);
2847        // after the connection is established, the peer should be removed, the num_pending_in
2848        // should be decreased, and the num_inbound must be increased
2849        assert_eq!(peers.connection_info.num_pending_in, 0);
2850        assert_eq!(peers.connection_info.num_inbound, 1);
2851        assert!(peers.peers.contains_key(&basic_peer));
2852    }
2853
2854    #[tokio::test]
2855    async fn test_incoming_at_capacity() {
2856        let mut config = PeersConfig::test();
2857        config.connection_info.max_inbound = 1;
2858        let mut peers = PeersManager::new(config);
2859
2860        let peer = PeerId::random();
2861        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2862        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
2863
2864        peers.on_incoming_session_established(peer, addr);
2865
2866        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2867        assert_eq!(
2868            peers.on_incoming_pending_session(addr.ip()).unwrap_err(),
2869            InboundConnectionError::ExceedsCapacity
2870        );
2871    }
2872
2873    #[tokio::test]
2874    async fn test_incoming_rate_limit() {
2875        let config = PeersConfig {
2876            incoming_ip_throttle_duration: Duration::from_millis(100),
2877            ..PeersConfig::test()
2878        };
2879        let mut peers = PeersManager::new(config);
2880
2881        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(168, 0, 1, 2)), 8009);
2882        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
2883        assert_eq!(
2884            peers.on_incoming_pending_session(addr.ip()).unwrap_err(),
2885            InboundConnectionError::IpBanned
2886        );
2887
2888        peers.release_interval.reset_immediately();
2889        tokio::time::sleep(peers.incoming_ip_throttle_duration).await;
2890
2891        // await unban
2892        poll_fn(|cx| loop {
2893            if peers.poll(cx).is_pending() {
2894                return Poll::Ready(());
2895            }
2896        })
2897        .await;
2898
2899        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
2900        assert_eq!(
2901            peers.on_incoming_pending_session(addr.ip()).unwrap_err(),
2902            InboundConnectionError::IpBanned
2903        );
2904    }
2905
2906    #[tokio::test]
2907    async fn test_tick() {
2908        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
2909        let socket_addr = SocketAddr::new(ip, 8008);
2910        let config = PeersConfig::test();
2911        let mut peer_manager = PeersManager::new(config);
2912        let peer_id = PeerId::random();
2913        peer_manager.add_peer(peer_id, PeerAddr::from_tcp(socket_addr), None);
2914
2915        tokio::time::sleep(Duration::from_secs(1)).await;
2916        peer_manager.tick();
2917
2918        // still unconnected
2919        assert_eq!(peer_manager.peers.get_mut(&peer_id).unwrap().reputation, DEFAULT_REPUTATION);
2920
2921        // mark as connected
2922        peer_manager.peers.get_mut(&peer_id).unwrap().state = PeerConnectionState::Out;
2923
2924        tokio::time::sleep(Duration::from_secs(1)).await;
2925        peer_manager.tick();
2926
2927        // still at default reputation
2928        assert_eq!(peer_manager.peers.get_mut(&peer_id).unwrap().reputation, DEFAULT_REPUTATION);
2929
2930        peer_manager.peers.get_mut(&peer_id).unwrap().reputation -= 1;
2931
2932        tokio::time::sleep(Duration::from_secs(1)).await;
2933        peer_manager.tick();
2934
2935        // tick applied
2936        assert!(peer_manager.peers.get_mut(&peer_id).unwrap().reputation >= DEFAULT_REPUTATION);
2937    }
2938
2939    #[tokio::test]
2940    async fn test_remove_incoming_after_disconnect() {
2941        let peer_id = PeerId::random();
2942        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2943        let mut peers = PeersManager::default();
2944
2945        peers.on_incoming_pending_session(addr.ip()).unwrap();
2946        peers.on_incoming_session_established(peer_id, addr);
2947        let peer = peers.peers.get(&peer_id).unwrap();
2948        assert_eq!(peer.state, PeerConnectionState::In);
2949        assert!(peer.remove_after_disconnect);
2950
2951        peers.on_active_session_gracefully_closed(peer_id);
2952        assert!(!peers.peers.contains_key(&peer_id))
2953    }
2954
2955    #[tokio::test]
2956    async fn test_keep_incoming_after_disconnect_if_discovered() {
2957        let peer_id = PeerId::random();
2958        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2959        let mut peers = PeersManager::default();
2960
2961        peers.on_incoming_pending_session(addr.ip()).unwrap();
2962        peers.on_incoming_session_established(peer_id, addr);
2963        let peer = peers.peers.get(&peer_id).unwrap();
2964        assert_eq!(peer.state, PeerConnectionState::In);
2965        assert!(peer.remove_after_disconnect);
2966
2967        // trigger discovery manually while the peer is still connected
2968        peers.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
2969
2970        peers.on_active_session_gracefully_closed(peer_id);
2971
2972        let peer = peers.peers.get(&peer_id).unwrap();
2973        assert_eq!(peer.state, PeerConnectionState::Idle);
2974        assert!(!peer.remove_after_disconnect);
2975    }
2976
2977    #[tokio::test]
2978    async fn test_peer_reconnect_after_graceful_close_respects_throttle() {
2979        let throttle_duration = Duration::from_millis(100);
2980        let config =
2981            PeersConfig { incoming_ip_throttle_duration: throttle_duration, ..PeersConfig::test() };
2982        let mut peers = PeersManager::new(config);
2983
2984        let peer_id = PeerId::random();
2985        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
2986
2987        // Add as regular peer
2988        peers.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
2989
2990        match event!(peers) {
2991            PeerAction::PeerAdded(id) => assert_eq!(id, peer_id),
2992            _ => unreachable!(),
2993        }
2994
2995        match event!(peers) {
2996            PeerAction::Connect { .. } => {}
2997            _ => unreachable!(),
2998        }
2999
3000        // Simulate outbound connection established
3001        peers.on_active_outgoing_established(peer_id);
3002        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::Out);
3003
3004        // Gracefully close the session
3005        peers.on_active_session_gracefully_closed(peer_id);
3006
3007        let peer = peers.peers.get(&peer_id).unwrap();
3008        assert_eq!(peer.state, PeerConnectionState::Idle);
3009        assert!(peer.backed_off);
3010
3011        // Verify the peer is in the backed_off_peers set
3012        assert!(peers.backed_off_peers.contains_key(&peer_id));
3013
3014        // Immediately try to poll - should not trigger any actions yet
3015        poll_fn(|cx| {
3016            assert!(peers.poll(cx).is_pending());
3017            Poll::Ready(())
3018        })
3019        .await;
3020
3021        // Peer should still be backed off
3022        assert!(peers.backed_off_peers.contains_key(&peer_id));
3023        assert!(peers.peers.get(&peer_id).unwrap().backed_off);
3024
3025        // Sleep for the throttle duration
3026        tokio::time::sleep(throttle_duration).await;
3027
3028        // After throttle duration, event! will poll until we get a Connect action
3029        match event!(peers) {
3030            PeerAction::Connect { peer_id: id, .. } => assert_eq!(id, peer_id),
3031            _ => unreachable!(),
3032        }
3033
3034        // After connection is initiated, peer should no longer be backed off
3035        assert!(!peers.backed_off_peers.contains_key(&peer_id));
3036        assert!(!peers.peers.get(&peer_id).unwrap().backed_off);
3037    }
3038
3039    #[tokio::test]
3040    async fn test_backed_off_peer_can_accept_incoming_connection() {
3041        let throttle_duration = Duration::from_millis(100);
3042        let config =
3043            PeersConfig { incoming_ip_throttle_duration: throttle_duration, ..PeersConfig::test() };
3044        let mut peers = PeersManager::new(config);
3045
3046        let peer_id = PeerId::random();
3047        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
3048
3049        // Add as regular peer
3050        peers.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
3051
3052        match event!(peers) {
3053            PeerAction::PeerAdded(id) => assert_eq!(id, peer_id),
3054            _ => unreachable!(),
3055        }
3056
3057        match event!(peers) {
3058            PeerAction::Connect { .. } => {}
3059            _ => unreachable!(),
3060        }
3061
3062        // Simulate outbound connection established
3063        peers.on_active_outgoing_established(peer_id);
3064        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::Out);
3065
3066        // Gracefully close the session - this will back off the peer
3067        peers.on_active_session_gracefully_closed(peer_id);
3068
3069        let peer = peers.peers.get(&peer_id).unwrap();
3070        assert_eq!(peer.state, PeerConnectionState::Idle);
3071        assert!(peer.backed_off);
3072        assert!(peers.backed_off_peers.contains_key(&peer_id));
3073
3074        // Now simulate an incoming connection from the backed-off peer
3075        // First, handle the incoming pending session
3076        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
3077        assert_eq!(peers.connection_info.num_pending_in, 1);
3078
3079        // Establish the incoming session
3080        peers.on_incoming_session_established(peer_id, addr);
3081
3082        // Peer should have been added to incoming connections
3083        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::In);
3084        assert_eq!(peers.connection_info.num_inbound, 1);
3085
3086        // Peer should still be backed off for outbound connections
3087        assert!(peers.backed_off_peers.contains_key(&peer_id));
3088        assert!(peers.peers.get(&peer_id).unwrap().backed_off);
3089
3090        // Verify we don't try to reconnect outbound while peer is backed off
3091        poll_fn(|cx| {
3092            assert!(peers.poll(cx).is_pending());
3093            Poll::Ready(())
3094        })
3095        .await;
3096
3097        // No outbound connection should be attempted while backed off
3098        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::In);
3099
3100        // After throttle duration, the backoff should be cleared
3101        tokio::time::sleep(throttle_duration).await;
3102
3103        poll_fn(|cx| {
3104            let _ = peers.poll(cx);
3105            Poll::Ready(())
3106        })
3107        .await;
3108
3109        // Backoff should be cleared now
3110        assert!(!peers.backed_off_peers.contains_key(&peer_id));
3111        assert!(!peers.peers.get(&peer_id).unwrap().backed_off);
3112
3113        // Peer should still be in incoming state
3114        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::In);
3115    }
3116
3117    #[tokio::test]
3118    async fn test_incoming_outgoing_already_connected() {
3119        let peer_id = PeerId::random();
3120        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
3121        let mut peers = PeersManager::default();
3122
3123        peers.on_incoming_pending_session(addr.ip()).unwrap();
3124        peers.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
3125
3126        match event!(peers) {
3127            PeerAction::PeerAdded(_) => {}
3128            _ => unreachable!(),
3129        }
3130
3131        match event!(peers) {
3132            PeerAction::Connect { .. } => {}
3133            _ => unreachable!(),
3134        }
3135
3136        peers.on_incoming_session_established(peer_id, addr);
3137        peers.on_already_connected(Direction::Outgoing(peer_id));
3138        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::In);
3139        assert_eq!(peers.connection_info.num_inbound, 1);
3140        assert_eq!(peers.connection_info.num_pending_out, 0);
3141        assert_eq!(peers.connection_info.num_pending_in, 0);
3142        assert_eq!(peers.connection_info.num_outbound, 0);
3143    }
3144
3145    #[tokio::test]
3146    async fn test_already_connected_incoming_outgoing_connection_error() {
3147        let peer_id = PeerId::random();
3148        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
3149        let mut peers = PeersManager::default();
3150
3151        peers.on_incoming_pending_session(addr.ip()).unwrap();
3152        peers.add_peer(peer_id, PeerAddr::from_tcp(addr), None);
3153
3154        match event!(peers) {
3155            PeerAction::PeerAdded(_) => {}
3156            _ => unreachable!(),
3157        }
3158
3159        match event!(peers) {
3160            PeerAction::Connect { .. } => {}
3161            _ => unreachable!(),
3162        }
3163
3164        peers.on_incoming_session_established(peer_id, addr);
3165
3166        peers.on_outgoing_connection_failure(
3167            &addr,
3168            &peer_id,
3169            &io::Error::new(io::ErrorKind::ConnectionRefused, ""),
3170        );
3171        assert_eq!(peers.peers.get(&peer_id).unwrap().state, PeerConnectionState::In);
3172        assert_eq!(peers.connection_info.num_inbound, 1);
3173        assert_eq!(peers.connection_info.num_pending_out, 0);
3174        assert_eq!(peers.connection_info.num_pending_in, 0);
3175        assert_eq!(peers.connection_info.num_outbound, 0);
3176    }
3177
3178    #[tokio::test]
3179    async fn test_max_concurrent_dials() {
3180        let config = PeersConfig::default();
3181        let mut peer_manager = PeersManager::new(config);
3182        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
3183        let peer_addr = PeerAddr::from_tcp(SocketAddr::new(ip, 8008));
3184        for _ in 0..peer_manager.connection_info.config.max_concurrent_outbound_dials * 2 {
3185            peer_manager.add_peer(PeerId::random(), peer_addr, None);
3186        }
3187
3188        peer_manager.fill_outbound_slots();
3189        let dials = peer_manager
3190            .queued_actions
3191            .iter()
3192            .filter(|ev| matches!(ev, PeerAction::Connect { .. }))
3193            .count();
3194        assert_eq!(dials, peer_manager.connection_info.config.max_concurrent_outbound_dials);
3195    }
3196
3197    #[tokio::test]
3198    async fn test_max_num_of_pending_dials() {
3199        let config = PeersConfig::default();
3200        let mut peer_manager = PeersManager::new(config);
3201        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2));
3202        let peer_addr = PeerAddr::from_tcp(SocketAddr::new(ip, 8008));
3203
3204        // add more peers than allowed
3205        for _ in 0..peer_manager.connection_info.config.max_concurrent_outbound_dials * 2 {
3206            peer_manager.add_peer(PeerId::random(), peer_addr, None);
3207        }
3208
3209        for _ in 0..peer_manager.connection_info.config.max_concurrent_outbound_dials * 2 {
3210            match event!(peer_manager) {
3211                PeerAction::PeerAdded(_) => {}
3212                _ => unreachable!(),
3213            }
3214        }
3215
3216        for _ in 0..peer_manager.connection_info.config.max_concurrent_outbound_dials {
3217            match event!(peer_manager) {
3218                PeerAction::Connect { .. } => {}
3219                _ => unreachable!(),
3220            }
3221        }
3222
3223        // generate 'Connect' actions
3224        peer_manager.fill_outbound_slots();
3225
3226        // all dialed connections should be in 'PendingOut' state
3227        let dials = peer_manager.connection_info.num_pending_out;
3228        assert_eq!(dials, peer_manager.connection_info.config.max_concurrent_outbound_dials);
3229
3230        let num_pendingout_states = peer_manager
3231            .peers
3232            .iter()
3233            .filter(|(_, peer)| peer.state == PeerConnectionState::PendingOut)
3234            .map(|(peer_id, _)| *peer_id)
3235            .collect::<Vec<PeerId>>();
3236        assert_eq!(
3237            num_pendingout_states.len(),
3238            peer_manager.connection_info.config.max_concurrent_outbound_dials
3239        );
3240
3241        // establish dialed connections
3242        for peer_id in &num_pendingout_states {
3243            peer_manager.on_active_outgoing_established(*peer_id);
3244        }
3245
3246        // all dialed connections should now be in 'Out' state
3247        for peer_id in &num_pendingout_states {
3248            assert_eq!(peer_manager.peers.get(peer_id).unwrap().state, PeerConnectionState::Out);
3249        }
3250
3251        // no more pending outbound connections
3252        assert_eq!(peer_manager.connection_info.num_pending_out, 0);
3253    }
3254
3255    #[tokio::test]
3256    async fn test_connect() {
3257        let peer = PeerId::random();
3258        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
3259        let mut peers = PeersManager::default();
3260        peers.add_and_connect(peer, PeerAddr::from_tcp(socket_addr), None);
3261        assert_eq!(peers.peers.get(&peer).unwrap().state, PeerConnectionState::PendingOut);
3262
3263        match event!(peers) {
3264            PeerAction::Connect { peer_id, remote_addr } => {
3265                assert_eq!(peer_id, peer);
3266                assert_eq!(remote_addr, socket_addr);
3267            }
3268            _ => unreachable!(),
3269        }
3270
3271        let (record, _) = peers.peer_by_id(peer).unwrap();
3272        assert_eq!(record.tcp_addr(), socket_addr);
3273        assert_eq!(record.udp_addr(), socket_addr);
3274
3275        // connect again
3276        peers.add_and_connect(peer, PeerAddr::from_tcp(socket_addr), None);
3277
3278        let (record, _) = peers.peer_by_id(peer).unwrap();
3279        assert_eq!(record.tcp_addr(), socket_addr);
3280        assert_eq!(record.udp_addr(), socket_addr);
3281    }
3282
3283    #[tokio::test]
3284    async fn test_incoming_connection_from_banned() {
3285        let peer = PeerId::random();
3286        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
3287        let config = PeersConfig::test().with_max_inbound(3);
3288        let mut peers = PeersManager::new(config);
3289        peers.add_peer(peer, PeerAddr::from_tcp(socket_addr), None);
3290
3291        match event!(peers) {
3292            PeerAction::PeerAdded(peer_id) => {
3293                assert_eq!(peer_id, peer);
3294            }
3295            _ => unreachable!(),
3296        }
3297        match event!(peers) {
3298            PeerAction::Connect { peer_id, .. } => {
3299                assert_eq!(peer_id, peer);
3300            }
3301            _ => unreachable!(),
3302        }
3303
3304        poll_fn(|cx| {
3305            assert!(peers.poll(cx).is_pending());
3306            Poll::Ready(())
3307        })
3308        .await;
3309
3310        // simulate new connection drops with error
3311        loop {
3312            peers.on_active_session_dropped(
3313                &socket_addr,
3314                &peer,
3315                &EthStreamError::InvalidMessage(reth_eth_wire::message::MessageError::Invalid(
3316                    reth_eth_wire::EthVersion::Eth68,
3317                    reth_eth_wire::EthMessageID::Status,
3318                )),
3319            );
3320
3321            if peers.peers.get(&peer).unwrap().is_banned() {
3322                break;
3323            }
3324
3325            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
3326            peers.on_incoming_session_established(peer, socket_addr);
3327
3328            match event!(peers) {
3329                PeerAction::Connect { peer_id, .. } => {
3330                    assert_eq!(peer_id, peer);
3331                }
3332                _ => unreachable!(),
3333            }
3334        }
3335
3336        assert!(peers.peers.get(&peer).unwrap().is_banned());
3337
3338        // fill all incoming slots
3339        for _ in 0..peers.connection_info.config.max_inbound {
3340            assert!(peers.on_incoming_pending_session(socket_addr.ip()).is_ok());
3341            peers.on_incoming_session_established(peer, socket_addr);
3342
3343            match event!(peers) {
3344                PeerAction::DisconnectBannedIncoming { peer_id } => {
3345                    assert_eq!(peer_id, peer);
3346                }
3347                _ => unreachable!(),
3348            }
3349        }
3350
3351        poll_fn(|cx| {
3352            assert!(peers.poll(cx).is_pending());
3353            Poll::Ready(())
3354        })
3355        .await;
3356
3357        assert_eq!(peers.connection_info.num_inbound, 0);
3358
3359        let new_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 3)), 8008);
3360
3361        // Assert we can still accept new connections
3362        assert!(peers.on_incoming_pending_session(new_addr.ip()).is_ok());
3363        assert_eq!(peers.connection_info.num_pending_in, 1);
3364
3365        // the triggered DisconnectBannedIncoming will result in dropped connections, assert that
3366        // connection info is updated via the peer's state which would be a noop here since the
3367        // banned peer's state is idle
3368        peers.on_active_session_gracefully_closed(peer);
3369        assert_eq!(peers.connection_info.num_inbound, 0);
3370    }
3371
3372    #[tokio::test]
3373    async fn test_add_pending_connect() {
3374        let peer = PeerId::random();
3375        let socket_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
3376        let mut peers = PeersManager::default();
3377        peers.add_and_connect(peer, PeerAddr::from_tcp(socket_addr), None);
3378        assert_eq!(peers.peers.get(&peer).unwrap().state, PeerConnectionState::PendingOut);
3379        assert_eq!(peers.connection_info.num_pending_out, 1);
3380    }
3381
3382    #[tokio::test]
3383    async fn test_dns_updates_peer_address() {
3384        let peer_id = PeerId::random();
3385        let initial_socket = SocketAddr::new("1.1.1.1".parse::<IpAddr>().unwrap(), 8008);
3386        let updated_ip = "2.2.2.2".parse::<IpAddr>().unwrap();
3387
3388        let trusted = TrustedPeer {
3389            host: url::Host::Ipv4("2.2.2.2".parse().unwrap()),
3390            tcp_port: 8008,
3391            udp_port: 8008,
3392            id: peer_id,
3393        };
3394
3395        let config = PeersConfig::test().with_trusted_nodes(vec![trusted.clone()]);
3396        let mut manager = PeersManager::new(config);
3397        manager
3398            .trusted_peers_resolver
3399            .set_interval(tokio::time::interval(Duration::from_millis(1)));
3400
3401        manager.peers.insert(
3402            peer_id,
3403            Peer::trusted(PeerAddr::new_with_ports(initial_socket.ip(), 8008, Some(8008))),
3404        );
3405
3406        for _ in 0..100 {
3407            let _ = event!(manager);
3408            if manager.peers.get(&peer_id).unwrap().addr.tcp().ip() == updated_ip {
3409                break;
3410            }
3411            tokio::time::sleep(Duration::from_millis(10)).await;
3412        }
3413
3414        let updated_peer = manager.peers.get(&peer_id).unwrap();
3415        assert_eq!(updated_peer.addr.tcp().ip(), updated_ip);
3416    }
3417
3418    #[tokio::test]
3419    async fn test_ip_filter_blocks_inbound_connection() {
3420        use reth_net_banlist::IpFilter;
3421        use std::net::IpAddr;
3422
3423        // Create a filter that only allows 192.168.0.0/16
3424        let ip_filter = IpFilter::from_cidr_string("192.168.0.0/16").unwrap();
3425        let config = PeersConfig::test().with_ip_filter(ip_filter);
3426        let mut peers = PeersManager::new(config);
3427
3428        // Try to connect from an allowed IP
3429        let allowed_ip: IpAddr = "192.168.1.100".parse().unwrap();
3430        assert!(peers.on_incoming_pending_session(allowed_ip).is_ok());
3431
3432        // Try to connect from a disallowed IP
3433        let disallowed_ip: IpAddr = "10.0.0.1".parse().unwrap();
3434        assert!(peers.on_incoming_pending_session(disallowed_ip).is_err());
3435    }
3436
3437    #[tokio::test]
3438    async fn test_ip_filter_blocks_outbound_connection() {
3439        use reth_net_banlist::IpFilter;
3440        use std::net::SocketAddr;
3441
3442        // Create a filter that only allows 192.168.0.0/16
3443        let ip_filter = IpFilter::from_cidr_string("192.168.0.0/16").unwrap();
3444        let config = PeersConfig::test().with_ip_filter(ip_filter);
3445        let mut peers = PeersManager::new(config);
3446
3447        let peer_id = PeerId::new([1; 64]);
3448
3449        // Try to add a peer with an allowed IP
3450        let allowed_addr: SocketAddr = "192.168.1.100:30303".parse().unwrap();
3451        peers.add_peer(peer_id, PeerAddr::from_tcp(allowed_addr), None);
3452        assert!(peers.peers.contains_key(&peer_id));
3453
3454        // Try to add a peer with a disallowed IP
3455        let peer_id2 = PeerId::new([2; 64]);
3456        let disallowed_addr: SocketAddr = "10.0.0.1:30303".parse().unwrap();
3457        peers.add_peer(peer_id2, PeerAddr::from_tcp(disallowed_addr), None);
3458        assert!(!peers.peers.contains_key(&peer_id2));
3459    }
3460
3461    #[tokio::test]
3462    async fn test_ip_filter_ipv6() {
3463        use reth_net_banlist::IpFilter;
3464        use std::net::IpAddr;
3465
3466        // Create a filter that only allows IPv6 range 2001:db8::/32
3467        let ip_filter = IpFilter::from_cidr_string("2001:db8::/32").unwrap();
3468        let config = PeersConfig::test().with_ip_filter(ip_filter);
3469        let mut peers = PeersManager::new(config);
3470
3471        // Try to connect from an allowed IPv6 address
3472        let allowed_ip: IpAddr = "2001:db8::1".parse().unwrap();
3473        assert!(peers.on_incoming_pending_session(allowed_ip).is_ok());
3474
3475        // Try to connect from a disallowed IPv6 address
3476        let disallowed_ip: IpAddr = "2001:db9::1".parse().unwrap();
3477        assert!(peers.on_incoming_pending_session(disallowed_ip).is_err());
3478    }
3479
3480    #[tokio::test]
3481    async fn test_ip_filter_multiple_ranges() {
3482        use reth_net_banlist::IpFilter;
3483        use std::net::IpAddr;
3484
3485        // Create a filter that allows multiple ranges
3486        let ip_filter = IpFilter::from_cidr_string("192.168.0.0/16,10.0.0.0/8").unwrap();
3487        let config = PeersConfig::test().with_ip_filter(ip_filter);
3488        let mut peers = PeersManager::new(config);
3489
3490        // Try IPs from both allowed ranges
3491        let ip1: IpAddr = "192.168.1.1".parse().unwrap();
3492        let ip2: IpAddr = "10.5.10.20".parse().unwrap();
3493        assert!(peers.on_incoming_pending_session(ip1).is_ok());
3494        assert!(peers.on_incoming_pending_session(ip2).is_ok());
3495
3496        // Try IP from disallowed range
3497        let disallowed_ip: IpAddr = "172.16.0.1".parse().unwrap();
3498        assert!(peers.on_incoming_pending_session(disallowed_ip).is_err());
3499    }
3500
3501    #[tokio::test]
3502    async fn test_ip_filter_no_restriction() {
3503        use reth_net_banlist::IpFilter;
3504        use std::net::IpAddr;
3505
3506        // Create a filter with no restrictions (allow all)
3507        let ip_filter = IpFilter::allow_all();
3508        let config = PeersConfig::test().with_ip_filter(ip_filter);
3509        let mut peers = PeersManager::new(config);
3510
3511        // All IPs should be allowed
3512        let ip1: IpAddr = "192.168.1.1".parse().unwrap();
3513        let ip2: IpAddr = "10.0.0.1".parse().unwrap();
3514        let ip3: IpAddr = "8.8.8.8".parse().unwrap();
3515        assert!(peers.on_incoming_pending_session(ip1).is_ok());
3516        assert!(peers.on_incoming_pending_session(ip2).is_ok());
3517        assert!(peers.on_incoming_pending_session(ip3).is_ok());
3518    }
3519
3520    #[tokio::test]
3521    async fn test_best_unconnected_prefers_fork_id_as_tiebreaker() {
3522        let mut peers = PeersManager::default();
3523        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8008);
3524
3525        let fork_id = ForkId { hash: ForkHash([0xaa, 0xbb, 0xcc, 0xdd]), next: 0 };
3526
3527        // add two peers with equal reputation, only one has a fork_id
3528        let no_fork = PeerId::random();
3529        peers.add_peer(no_fork, PeerAddr::from_tcp(addr), None);
3530
3531        let with_fork = PeerId::random();
3532        peers.add_peer(with_fork, PeerAddr::from_tcp(addr), None);
3533        peers.peers.get_mut(&with_fork).unwrap().fork_id = Some(Box::new(fork_id));
3534
3535        let (best_id, _) = peers.best_unconnected().unwrap();
3536        assert_eq!(best_id, with_fork, "fork_id should break tie when reputation is equal");
3537    }
3538
3539    #[tokio::test]
3540    async fn test_add_trusted_peer_node_resolves_absent_peer() {
3541        let peer_id = PeerId::random();
3542        let trusted = TrustedPeer {
3543            host: url::Host::Domain("example.invalid".to_string()),
3544            tcp_port: 30303,
3545            udp_port: 30303,
3546            id: peer_id,
3547        };
3548
3549        let mut manager = PeersManager::default();
3550        manager.add_trusted_peer_node(trusted);
3551
3552        assert!(manager.trusted_peer_ids.contains(&peer_id));
3553        assert_eq!(manager.trusted_peers_resolver.trusted_peers.len(), 1);
3554        assert!(!manager.peers.contains_key(&peer_id));
3555
3556        let resolved = NodeRecord {
3557            address: "10.0.0.1".parse::<IpAddr>().unwrap(),
3558            tcp_port: 30303,
3559            udp_port: 30303,
3560            id: peer_id,
3561        };
3562        manager.on_resolved_peer(peer_id, resolved);
3563
3564        let peer = manager.peers.get(&peer_id).expect("peer should be added after resolution");
3565        assert!(peer.kind.is_trusted());
3566        assert_eq!(peer.addr.tcp().ip(), "10.0.0.1".parse::<IpAddr>().unwrap());
3567    }
3568
3569    #[tokio::test]
3570    async fn test_removed_trusted_peer_ignores_late_resolution() {
3571        let peer_id = PeerId::random();
3572        let trusted = TrustedPeer {
3573            host: url::Host::Domain("example.invalid".to_string()),
3574            tcp_port: 30303,
3575            udp_port: 30303,
3576            id: peer_id,
3577        };
3578
3579        let mut manager = PeersManager::default();
3580        manager.add_trusted_peer_node(trusted);
3581        manager.remove_peer_from_trusted_set(peer_id);
3582
3583        let resolved = NodeRecord {
3584            address: "10.0.0.1".parse::<IpAddr>().unwrap(),
3585            tcp_port: 30303,
3586            udp_port: 30303,
3587            id: peer_id,
3588        };
3589        manager.on_resolved_peer(peer_id, resolved);
3590
3591        assert!(!manager.peers.contains_key(&peer_id));
3592        assert!(!manager.trusted_peer_ids.contains(&peer_id));
3593    }
3594
3595    #[tokio::test]
3596    async fn test_rotation_disconnects_eligible_peer() {
3597        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
3598        let peer = PeerId::random();
3599        let mut peers = PeersManager::new(
3600            PeersConfig::test()
3601                .with_max_outbound(1)
3602                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3603        );
3604
3605        peers.add_peer(peer, PeerAddr::from_tcp(addr), None);
3606        match event!(peers) {
3607            PeerAction::PeerAdded(_) => {}
3608            _ => unreachable!(),
3609        }
3610        match event!(peers) {
3611            PeerAction::Connect { .. } => {}
3612            _ => unreachable!(),
3613        }
3614        peers.on_active_outgoing_established(peer);
3615
3616        set_connected_at(
3617            &mut peers,
3618            peer,
3619            std::time::Instant::now() - Duration::from_secs(11 * 60),
3620        );
3621
3622        tokio::time::sleep(Duration::from_millis(100)).await;
3623
3624        match event!(peers) {
3625            PeerAction::Disconnect { peer_id, reason } => {
3626                assert_eq!(peer_id, peer);
3627                assert_eq!(reason, Some(DisconnectReason::UselessPeer));
3628            }
3629            _ => unreachable!(),
3630        }
3631    }
3632
3633    #[tokio::test]
3634    async fn test_rotation_then_remove_uses_active_session_removal() {
3635        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 7)), 8008);
3636        let peer = PeerId::random();
3637        let mut peers = PeersManager::new(
3638            PeersConfig::test()
3639                .with_max_outbound(1)
3640                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3641        );
3642
3643        peers.add_peer(peer, PeerAddr::from_tcp(addr), None);
3644        match event!(peers) {
3645            PeerAction::PeerAdded(_) => {}
3646            _ => unreachable!(),
3647        }
3648        match event!(peers) {
3649            PeerAction::Connect { .. } => {}
3650            _ => unreachable!(),
3651        }
3652        peers.on_active_outgoing_established(peer);
3653        set_connected_at(
3654            &mut peers,
3655            peer,
3656            std::time::Instant::now() - Duration::from_secs(11 * 60),
3657        );
3658
3659        tokio::time::sleep(Duration::from_millis(100)).await;
3660
3661        match event!(peers) {
3662            PeerAction::Disconnect { peer_id, reason } => {
3663                assert_eq!(peer_id, peer);
3664                assert_eq!(reason, Some(DisconnectReason::UselessPeer));
3665            }
3666            _ => unreachable!(),
3667        }
3668        assert_eq!(peers.peers.get(&peer).unwrap().state, PeerConnectionState::Out);
3669        assert_eq!(peers.connection_info.num_outbound, 1);
3670
3671        peers.remove_peer(peer);
3672        match event!(peers) {
3673            PeerAction::PeerRemoved(peer_id) => assert_eq!(peer_id, peer),
3674            _ => unreachable!(),
3675        }
3676        match event!(peers) {
3677            PeerAction::Disconnect { peer_id, reason } => {
3678                assert_eq!(peer_id, peer);
3679                assert_eq!(reason, Some(DisconnectReason::DisconnectRequested));
3680            }
3681            _ => unreachable!(),
3682        }
3683
3684        let p = peers.peers.get(&peer).unwrap();
3685        assert_eq!(p.state, PeerConnectionState::DisconnectingOut);
3686        assert!(p.remove_after_disconnect);
3687        assert_eq!(peers.connection_info.num_outbound, 1);
3688
3689        peers.on_active_session_gracefully_closed(peer);
3690        assert_eq!(peers.connection_info.num_outbound, 0);
3691        assert!(!peers.peers.contains_key(&peer));
3692    }
3693
3694    #[tokio::test]
3695    async fn test_rotation_skips_trusted_peers() {
3696        let _addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 3)), 8008);
3697        let peer = PeerId::random();
3698        let trusted = TrustedPeer {
3699            host: Host::Ipv4(Ipv4Addr::new(127, 0, 1, 3)),
3700            tcp_port: 8008,
3701            udp_port: 8008,
3702            id: peer,
3703        };
3704        let mut peers = PeersManager::new(
3705            PeersConfig::test()
3706                .with_max_outbound(1)
3707                .with_trusted_nodes(vec![trusted])
3708                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3709        );
3710
3711        match event!(peers) {
3712            PeerAction::Connect { .. } => {}
3713            _ => unreachable!(),
3714        }
3715        peers.on_active_outgoing_established(peer);
3716        set_connected_at(
3717            &mut peers,
3718            peer,
3719            std::time::Instant::now() - Duration::from_secs(11 * 60),
3720        );
3721
3722        tokio::time::sleep(Duration::from_millis(100)).await;
3723
3724        poll_fn(|cx| {
3725            assert!(peers.poll(cx).is_pending(), "trusted peer must not be rotated");
3726            Poll::Ready(())
3727        })
3728        .await;
3729    }
3730
3731    #[tokio::test]
3732    async fn test_rotation_skips_recently_connected_peers() {
3733        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 4)), 8008);
3734        let peer = PeerId::random();
3735        let mut peers = PeersManager::new(
3736            PeersConfig::test()
3737                .with_max_outbound(1)
3738                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3739        );
3740
3741        peers.add_peer(peer, PeerAddr::from_tcp(addr), None);
3742        match event!(peers) {
3743            PeerAction::PeerAdded(_) => {}
3744            _ => unreachable!(),
3745        }
3746        match event!(peers) {
3747            PeerAction::Connect { .. } => {}
3748            _ => unreachable!(),
3749        }
3750        peers.on_active_outgoing_established(peer);
3751
3752        tokio::time::sleep(Duration::from_millis(100)).await;
3753
3754        poll_fn(|cx| {
3755            assert!(peers.poll(cx).is_pending(), "recently connected peer must not be rotated");
3756            Poll::Ready(())
3757        })
3758        .await;
3759    }
3760
3761    #[tokio::test]
3762    async fn test_rotation_skips_when_slots_available() {
3763        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 5)), 8008);
3764        let peer = PeerId::random();
3765        let mut peers = PeersManager::new(
3766            PeersConfig::test()
3767                .with_max_outbound(2)
3768                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3769        );
3770
3771        peers.add_peer(peer, PeerAddr::from_tcp(addr), None);
3772        match event!(peers) {
3773            PeerAction::PeerAdded(_) => {}
3774            _ => unreachable!(),
3775        }
3776        match event!(peers) {
3777            PeerAction::Connect { .. } => {}
3778            _ => unreachable!(),
3779        }
3780        peers.on_active_outgoing_established(peer);
3781        set_connected_at(
3782            &mut peers,
3783            peer,
3784            std::time::Instant::now() - Duration::from_secs(11 * 60),
3785        );
3786
3787        tokio::time::sleep(Duration::from_millis(100)).await;
3788
3789        poll_fn(|cx| {
3790            assert!(
3791                peers.poll(cx).is_pending(),
3792                "rotation must not fire when outbound slots are available"
3793            );
3794            Poll::Ready(())
3795        })
3796        .await;
3797    }
3798
3799    #[tokio::test]
3800    async fn test_rotation_disconnects_inbound_peer() {
3801        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 6)), 8008);
3802        let peer = PeerId::random();
3803        let mut peers = PeersManager::new(
3804            PeersConfig::test()
3805                .with_max_inbound(1)
3806                .with_peer_rotation_interval(Some(Duration::from_millis(50))),
3807        );
3808
3809        assert!(peers.on_incoming_pending_session(addr.ip()).is_ok());
3810        peers.on_incoming_session_established(peer, addr);
3811
3812        match event!(peers) {
3813            PeerAction::PeerAdded(_) => {}
3814            _ => unreachable!(),
3815        }
3816
3817        set_connected_at(
3818            &mut peers,
3819            peer,
3820            std::time::Instant::now() - Duration::from_secs(11 * 60),
3821        );
3822
3823        tokio::time::sleep(Duration::from_millis(100)).await;
3824
3825        match event!(peers) {
3826            PeerAction::Disconnect { peer_id, reason } => {
3827                assert_eq!(peer_id, peer);
3828                assert_eq!(reason, Some(DisconnectReason::UselessPeer));
3829            }
3830            _ => unreachable!(),
3831        }
3832    }
3833}