reth_network/
state.rs

1//! Keeps track of the state of the network.
2
3use crate::{
4    cache::LruCache,
5    discovery::Discovery,
6    fetch::{BlockResponseOutcome, FetchAction, NewPeerInfo, StateFetcher},
7    message::{BlockRequest, NewBlockMessage, PeerResponse, PeerResponseResult},
8    peers::{PeerAction, PeersManager},
9    session::BlockRangeInfo,
10    FetchClient,
11};
12use alloy_consensus::BlockHeader;
13use alloy_primitives::{
14    map::{FbBuildHasher, HashMap},
15    B256,
16};
17use rand::seq::SliceRandom;
18use reth_eth_wire::{
19    BlockHashNumber, Capabilities, DisconnectReason, EthNetworkPrimitives, GetReceipts70,
20    NetworkPrimitives, NewBlockHashes, NewBlockPayload, UnifiedStatus,
21};
22use reth_ethereum_forks::ForkId;
23use reth_network_api::{DiscoveredEvent, DiscoveryEvent, PeerRequest, PeerRequestSender};
24use reth_network_p2p::receipts::client::ReceiptsResponse;
25use reth_network_peers::PeerId;
26use reth_network_types::{PeerAddr, PeerKind};
27use reth_primitives_traits::Block;
28use std::{
29    collections::VecDeque,
30    fmt,
31    net::{IpAddr, SocketAddr},
32    ops::Deref,
33    sync::{
34        atomic::{AtomicU64, AtomicUsize},
35        Arc,
36    },
37    task::{Context, Poll},
38};
39use tokio::sync::oneshot;
40use tracing::{debug, trace};
41
42/// Cache limit of blocks to keep track of for a single peer.
43const PEER_BLOCK_CACHE_LIMIT: u32 = 512;
44
45/// Wrapper type for the [`BlockNumReader`] trait.
46pub(crate) struct BlockNumReader(Box<dyn reth_storage_api::BlockNumReader>);
47
48impl BlockNumReader {
49    /// Create a new instance with the given reader.
50    pub fn new(reader: impl reth_storage_api::BlockNumReader + 'static) -> Self {
51        Self(Box::new(reader))
52    }
53}
54
55impl fmt::Debug for BlockNumReader {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        f.debug_struct("BlockNumReader").field("inner", &"<dyn BlockNumReader>").finish()
58    }
59}
60
61impl Deref for BlockNumReader {
62    type Target = Box<dyn reth_storage_api::BlockNumReader>;
63
64    fn deref(&self) -> &Self::Target {
65        &self.0
66    }
67}
68
69/// The [`NetworkState`] keeps track of the state of all peers in the network.
70///
71/// This includes:
72///   - [`Discovery`]: manages the discovery protocol, essentially a stream of discovery updates
73///   - [`PeersManager`]: keeps track of connected peers and issues new outgoing connections
74///     depending on the configured capacity.
75///   - [`StateFetcher`]: streams download request (received from outside via channel) which are
76///     then send to the session of the peer.
77///
78/// This type is also responsible for responding for received request.
79#[derive(Debug)]
80pub struct NetworkState<N: NetworkPrimitives = EthNetworkPrimitives> {
81    /// All active peers and their state.
82    active_peers: HashMap<PeerId, ActivePeer<N>, FbBuildHasher<64>>,
83    /// Manages connections to peers.
84    peers_manager: PeersManager,
85    /// Buffered messages until polled.
86    queued_messages: VecDeque<StateAction<N>>,
87    /// The client type that can interact with the chain.
88    ///
89    /// This type is used to fetch the block number after we established a session and received the
90    /// [`UnifiedStatus`] block hash.
91    client: BlockNumReader,
92    /// Network discovery.
93    discovery: Discovery,
94    /// The type that handles requests.
95    ///
96    /// The fetcher streams `RLPx` related requests on a per-peer basis to this type. This type
97    /// will then queue in the request and notify the fetcher once the result has been
98    /// received.
99    state_fetcher: StateFetcher<N>,
100}
101
102impl<N: NetworkPrimitives> NetworkState<N> {
103    /// Create a new state instance with the given params
104    pub(crate) fn new(
105        client: BlockNumReader,
106        discovery: Discovery,
107        peers_manager: PeersManager,
108        num_active_peers: Arc<AtomicUsize>,
109    ) -> Self {
110        let state_fetcher = StateFetcher::new(peers_manager.handle(), num_active_peers);
111        Self {
112            active_peers: Default::default(),
113            peers_manager,
114            queued_messages: Default::default(),
115            client,
116            discovery,
117            state_fetcher,
118        }
119    }
120
121    /// Returns mutable access to the [`PeersManager`]
122    pub(crate) const fn peers_mut(&mut self) -> &mut PeersManager {
123        &mut self.peers_manager
124    }
125
126    /// Returns mutable access to the [`Discovery`]
127    pub(crate) const fn discovery_mut(&mut self) -> &mut Discovery {
128        &mut self.discovery
129    }
130
131    /// Returns access to the [`PeersManager`]
132    pub(crate) const fn peers(&self) -> &PeersManager {
133        &self.peers_manager
134    }
135
136    /// Returns a new [`FetchClient`]
137    pub(crate) fn fetch_client(&self) -> FetchClient<N> {
138        self.state_fetcher.client()
139    }
140
141    /// How many peers we're currently connected to.
142    pub fn num_active_peers(&self) -> usize {
143        self.active_peers.len()
144    }
145
146    /// Event hook for an activated session for the peer.
147    ///
148    /// Returns `Ok` if the session is valid, returns an `Err` if the session is not accepted and
149    /// should be rejected.
150    pub(crate) fn on_session_activated(&mut self, activation: SessionActivation<N>) {
151        let SessionActivation {
152            peer,
153            capabilities,
154            status,
155            request_tx,
156            timeout,
157            range_info,
158            supports_snap,
159        } = activation;
160
161        debug_assert!(!self.active_peers.contains_key(&peer), "Already connected; not possible");
162
163        // Use the block number from the peer's status (eth/69+) if available,
164        // otherwise fall back to a local lookup by hash.
165        let block_number = status.latest_block.unwrap_or_else(|| {
166            self.client.block_number(status.blockhash).ok().flatten().unwrap_or_default()
167        });
168        self.state_fetcher.new_active_peer(NewPeerInfo {
169            peer_id: peer,
170            best_hash: status.blockhash,
171            best_number: block_number,
172            capabilities: Arc::clone(&capabilities),
173            timeout,
174            range_info,
175            supports_snap,
176        });
177
178        self.active_peers.insert(
179            peer,
180            ActivePeer {
181                best_hash: status.blockhash,
182                capabilities,
183                request_tx,
184                pending_response: None,
185                blocks: LruCache::new(PEER_BLOCK_CACHE_LIMIT),
186            },
187        );
188    }
189
190    /// Event hook for a disconnected session for the given peer.
191    ///
192    /// This will remove the peer from the available set of peers and close all inflight requests.
193    pub(crate) fn on_session_closed(&mut self, peer: PeerId) {
194        self.active_peers.remove(&peer);
195        self.state_fetcher.on_session_closed(&peer);
196    }
197
198    /// Starts propagating the new block to peers that haven't reported the block yet.
199    ///
200    /// This is supposed to be invoked after the block was validated.
201    ///
202    /// > It then sends the block to a small fraction of connected peers (usually the square root of
203    /// > the total number of peers) using the `NewBlock` message.
204    ///
205    /// See also <https://github.com/ethereum/devp2p/blob/master/caps/eth.md>
206    pub(crate) fn announce_new_block(&mut self, msg: NewBlockMessage<N::NewBlockPayload>) {
207        // send a `NewBlock` message to a fraction of the connected peers (square root of the total
208        // number of peers)
209        let num_propagate = (self.active_peers.len() as f64).sqrt() as u64 + 1;
210
211        let number = msg.block.block().header().number();
212        let mut count = 0;
213
214        // Shuffle to propagate to a random sample of peers on every block announcement
215        let mut peers: Vec<_> = self.active_peers.iter_mut().collect();
216        peers.shuffle(&mut rand::rng());
217
218        for (peer_id, peer) in peers {
219            if peer.blocks.contains(&msg.hash) {
220                // skip peers which already reported the block
221                continue
222            }
223
224            // Queue a `NewBlock` message for the peer
225            if count < num_propagate {
226                self.queued_messages
227                    .push_back(StateAction::NewBlock { peer_id: *peer_id, block: msg.clone() });
228
229                // update peer block info
230                if self.state_fetcher.update_peer_block(peer_id, msg.hash, number) {
231                    peer.best_hash = msg.hash;
232                }
233
234                // mark the block as seen by the peer
235                peer.blocks.insert(msg.hash);
236
237                count += 1;
238            }
239
240            if count >= num_propagate {
241                break
242            }
243        }
244    }
245
246    /// Completes the block propagation process started in [`NetworkState::announce_new_block()`]
247    /// but sending `NewBlockHash` broadcast to all peers that haven't seen it yet.
248    pub(crate) fn announce_new_block_hash(&mut self, msg: NewBlockMessage<N::NewBlockPayload>) {
249        let number = msg.block.block().header().number();
250        let hashes = NewBlockHashes(vec![BlockHashNumber { hash: msg.hash, number }]);
251        for (peer_id, peer) in &mut self.active_peers {
252            if peer.blocks.contains(&msg.hash) {
253                // skip peers which already reported the block
254                continue
255            }
256
257            if self.state_fetcher.update_peer_block(peer_id, msg.hash, number) {
258                peer.best_hash = msg.hash;
259            }
260
261            self.queued_messages.push_back(StateAction::NewBlockHashes {
262                peer_id: *peer_id,
263                hashes: hashes.clone(),
264            });
265        }
266    }
267
268    /// Updates the block information for the peer.
269    pub(crate) fn update_peer_block(&mut self, peer_id: &PeerId, hash: B256, number: u64) {
270        if let Some(peer) = self.active_peers.get_mut(peer_id) {
271            peer.best_hash = hash;
272        }
273        self.state_fetcher.update_peer_block(peer_id, hash, number);
274    }
275
276    /// Invoked when a new [`ForkId`] is activated.
277    pub(crate) fn update_fork_id(&self, fork_id: ForkId) {
278        self.discovery.update_fork_id(fork_id)
279    }
280
281    /// Invoked after a `NewBlock` message was received by the peer.
282    ///
283    /// This will keep track of blocks we know a peer has
284    pub(crate) fn on_new_block(&mut self, peer_id: PeerId, hash: B256) {
285        // Mark the blocks as seen
286        if let Some(peer) = self.active_peers.get_mut(&peer_id) {
287            peer.blocks.insert(hash);
288        }
289    }
290
291    /// Invoked for a `NewBlockHashes` broadcast message.
292    pub(crate) fn on_new_block_hashes(&mut self, peer_id: PeerId, hashes: Vec<BlockHashNumber>) {
293        // Mark the blocks as seen
294        if let Some(peer) = self.active_peers.get_mut(&peer_id) {
295            peer.blocks.extend(hashes.into_iter().map(|b| b.hash));
296        }
297    }
298
299    /// Bans the [`IpAddr`] in the discovery service.
300    pub(crate) fn ban_ip_discovery(&self, ip: IpAddr) {
301        trace!(target: "net", ?ip, "Banning discovery");
302        self.discovery.ban_ip(ip)
303    }
304
305    /// Bans the [`PeerId`] and [`IpAddr`] in the discovery service.
306    pub(crate) fn ban_discovery(&self, peer_id: PeerId, ip: IpAddr) {
307        trace!(target: "net", ?peer_id, ?ip, "Banning discovery");
308        self.discovery.ban(peer_id, ip)
309    }
310
311    /// Marks the given peer as trusted.
312    pub(crate) fn add_trusted_peer_id(&mut self, peer_id: PeerId) {
313        self.peers_manager.add_trusted_peer_id(peer_id)
314    }
315
316    /// Adds a trusted peer that may use a hostname, with periodic DNS re-resolution.
317    pub(crate) fn add_trusted_peer_node(&mut self, trusted: reth_network_peers::TrustedPeer) {
318        self.peers_manager.add_trusted_peer_node(trusted)
319    }
320
321    /// Adds a peer and its address with the given kind to the peerset.
322    pub(crate) fn add_peer_kind(
323        &mut self,
324        peer_id: PeerId,
325        kind: Option<PeerKind>,
326        addr: PeerAddr,
327    ) {
328        self.peers_manager.add_peer_kind(peer_id, kind, addr, None)
329    }
330
331    /// Connects a peer and its address with the given kind
332    pub(crate) fn add_and_connect(&mut self, peer_id: PeerId, kind: PeerKind, addr: PeerAddr) {
333        self.peers_manager.add_and_connect_kind(peer_id, kind, addr, None)
334    }
335
336    /// Removes a peer and its address with the given kind from the peerset.
337    pub(crate) fn remove_peer_kind(&mut self, peer_id: PeerId, kind: PeerKind) {
338        match kind {
339            PeerKind::Basic | PeerKind::Static => self.peers_manager.remove_peer(peer_id),
340            PeerKind::Trusted => self.peers_manager.remove_peer_from_trusted_set(peer_id),
341        }
342    }
343
344    /// Event hook for events received from the discovery service.
345    fn on_discovery_event(&mut self, event: DiscoveryEvent) {
346        match event {
347            DiscoveryEvent::NewNode(DiscoveredEvent::EventQueued { peer_id, addr, fork_id }) => {
348                self.queued_messages.push_back(StateAction::DiscoveredNode {
349                    peer_id,
350                    addr,
351                    fork_id,
352                });
353            }
354            DiscoveryEvent::EnrForkId(record, fork_id) => {
355                let peer_id = record.id;
356                let tcp_addr = record.tcp_addr();
357                if tcp_addr.port() == 0 {
358                    return
359                }
360                let udp_addr = record.udp_addr();
361                let addr = PeerAddr::new(tcp_addr, Some(udp_addr));
362                self.queued_messages.push_back(StateAction::DiscoveredEnrForkId {
363                    peer_id,
364                    addr,
365                    fork_id,
366                });
367            }
368        }
369    }
370
371    /// Event hook for new actions derived from the peer management set.
372    fn on_peer_action(&mut self, action: PeerAction) {
373        match action {
374            PeerAction::Connect { peer_id, remote_addr } => {
375                self.queued_messages.push_back(StateAction::Connect { peer_id, remote_addr });
376            }
377            PeerAction::Disconnect { peer_id, reason } => {
378                self.state_fetcher.on_pending_disconnect(&peer_id);
379                self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason });
380            }
381            PeerAction::DisconnectBannedIncoming { peer_id } |
382            PeerAction::DisconnectUntrustedIncoming { peer_id } => {
383                self.state_fetcher.on_pending_disconnect(&peer_id);
384                self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason: None });
385            }
386            PeerAction::DiscoveryBanPeerId { peer_id, ip_addr } => {
387                self.ban_discovery(peer_id, ip_addr)
388            }
389            PeerAction::DiscoveryBanIp { ip_addr } => self.ban_ip_discovery(ip_addr),
390            PeerAction::PeerAdded(peer_id) => {
391                self.queued_messages.push_back(StateAction::PeerAdded(peer_id))
392            }
393            PeerAction::PeerRemoved(peer_id) => {
394                self.queued_messages.push_back(StateAction::PeerRemoved(peer_id))
395            }
396            PeerAction::BanPeer { .. } | PeerAction::UnBanPeer { .. } => {}
397        }
398    }
399
400    /// Sends The message to the peer's session and queues in a response.
401    ///
402    /// Caution: this will replace an already pending response. It's the responsibility of the
403    /// caller to select the peer.
404    fn handle_block_request(&mut self, peer_id: PeerId, request: BlockRequest) {
405        if let Some(ref mut peer) = self.active_peers.get_mut(&peer_id) {
406            let (request, response) = match request {
407                BlockRequest::GetBlockHeaders(request) => {
408                    let (response, rx) = oneshot::channel();
409                    let request = PeerRequest::GetBlockHeaders { request, response };
410                    let response = PeerResponse::BlockHeaders { response: rx };
411                    (request, response)
412                }
413                BlockRequest::GetBlockBodies(request) => {
414                    let (response, rx) = oneshot::channel();
415                    let request = PeerRequest::GetBlockBodies { request, response };
416                    let response = PeerResponse::BlockBodies { response: rx };
417                    (request, response)
418                }
419                BlockRequest::GetBlockAccessLists(request) => {
420                    let (response, rx) = oneshot::channel();
421                    let request = PeerRequest::GetBlockAccessLists { request, response };
422                    let response = PeerResponse::BlockAccessLists { response: rx };
423                    (request, response)
424                }
425                BlockRequest::GetReceipts(request) => {
426                    if peer.capabilities.supports_eth_v70() {
427                        let (response, rx) = oneshot::channel();
428                        let request = PeerRequest::GetReceipts70 {
429                            request: GetReceipts70 {
430                                first_block_receipt_index: 0,
431                                block_hashes: request.0,
432                            },
433                            response,
434                        };
435                        let response = PeerResponse::Receipts70 { response: rx };
436                        (request, response)
437                    } else if peer.capabilities.supports_eth_v69() {
438                        let (response, rx) = oneshot::channel();
439                        let request = PeerRequest::GetReceipts69 { request, response };
440                        let response = PeerResponse::Receipts69 { response: rx };
441                        (request, response)
442                    } else {
443                        let (response, rx) = oneshot::channel();
444                        let request = PeerRequest::GetReceipts { request, response };
445                        let response = PeerResponse::Receipts { response: rx };
446                        (request, response)
447                    }
448                }
449                BlockRequest::GetSnap(request) => {
450                    let (response, rx) = oneshot::channel();
451                    let request = PeerRequest::GetSnap { request, response };
452                    let response = PeerResponse::Snap { response: rx };
453                    (request, response)
454                }
455            };
456            let _ = peer.request_tx.to_session_tx.try_send(request);
457            peer.pending_response = Some(response);
458        }
459    }
460
461    /// Handle the outcome of processed response, for example directly queue another request.
462    fn on_block_response_outcome(&mut self, outcome: BlockResponseOutcome) {
463        match outcome {
464            BlockResponseOutcome::Request(peer, request) => {
465                self.handle_block_request(peer, request);
466            }
467            BlockResponseOutcome::BadResponse(peer, reputation_change) => {
468                self.peers_manager.apply_reputation_change(&peer, reputation_change);
469            }
470        }
471    }
472
473    /// Invoked when received a response from a connected peer.
474    ///
475    /// Delegates the response result to the fetcher which may return an outcome specific
476    /// instruction that needs to be handled in [`Self::on_block_response_outcome`]. This could be
477    /// a follow-up request or an instruction to slash the peer's reputation.
478    fn on_eth_response(&mut self, peer: PeerId, resp: PeerResponseResult<N>) {
479        let outcome = match resp {
480            PeerResponseResult::BlockHeaders(res) => {
481                self.state_fetcher.on_block_headers_response(peer, res)
482            }
483            PeerResponseResult::BlockBodies(res) => {
484                self.state_fetcher.on_block_bodies_response(peer, res)
485            }
486            PeerResponseResult::Receipts(res) => {
487                // Legacy eth/66-68: strip bloom filters and wrap in ReceiptsResponse
488                let normalized = res.map(|blocks| {
489                    let receipts = blocks
490                        .into_iter()
491                        .map(|block_receipts| {
492                            block_receipts.into_iter().map(|rwb| rwb.receipt).collect()
493                        })
494                        .collect();
495                    ReceiptsResponse::new(receipts)
496                });
497                self.state_fetcher.on_receipts_response(peer, normalized)
498            }
499            PeerResponseResult::Receipts69(res) => {
500                let normalized = res.map(ReceiptsResponse::new);
501                self.state_fetcher.on_receipts_response(peer, normalized)
502            }
503            PeerResponseResult::Receipts70(res) => {
504                let normalized = res.map(ReceiptsResponse::from);
505                self.state_fetcher.on_receipts_response(peer, normalized)
506            }
507            PeerResponseResult::BlockAccessLists(res) => {
508                self.state_fetcher.on_block_access_lists_response(peer, res)
509            }
510            PeerResponseResult::Snap(res) => self.state_fetcher.on_snap_response(peer, res),
511            _ => None,
512        };
513
514        if let Some(outcome) = outcome {
515            self.on_block_response_outcome(outcome);
516        }
517    }
518
519    /// Advances the state
520    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<StateAction<N>> {
521        loop {
522            // drain buffered messages
523            if let Some(message) = self.queued_messages.pop_front() {
524                return Poll::Ready(message)
525            }
526
527            while let Poll::Ready(discovery) = self.discovery.poll(cx) {
528                self.on_discovery_event(discovery);
529            }
530
531            while let Poll::Ready(action) = self.state_fetcher.poll(cx) {
532                match action {
533                    FetchAction::BlockRequest { peer_id, request } => {
534                        self.handle_block_request(peer_id, request)
535                    }
536                }
537            }
538
539            loop {
540                // need to buffer results here to make borrow checker happy
541                let mut closed_sessions = Vec::new();
542                let mut received_responses = Vec::new();
543
544                // poll all connected peers for responses
545                for (id, peer) in &mut self.active_peers {
546                    let Some(mut response) = peer.pending_response.take() else { continue };
547                    match response.poll(cx) {
548                        Poll::Ready(res) => {
549                            // check if the error is due to a closed channel to the session
550                            if res.err().is_some_and(|err| err.is_channel_closed()) {
551                                debug!(
552                                    target: "net",
553                                    ?id,
554                                    "Request canceled, response channel from session closed."
555                                );
556                                // if the channel is closed, this means the peer session is also
557                                // closed, in which case we can invoke the
558                                // [Self::on_closed_session]
559                                // immediately, preventing followup requests and propagate the
560                                // connection dropped error
561                                closed_sessions.push(*id);
562                            } else {
563                                received_responses.push((*id, res));
564                            }
565                        }
566                        Poll::Pending => {
567                            // not ready yet, store again.
568                            peer.pending_response = Some(response);
569                        }
570                    };
571                }
572
573                for peer in closed_sessions {
574                    self.on_session_closed(peer)
575                }
576
577                if received_responses.is_empty() {
578                    break;
579                }
580
581                for (peer_id, resp) in received_responses {
582                    self.on_eth_response(peer_id, resp);
583                }
584            }
585
586            // poll peer manager
587            while let Poll::Ready(action) = self.peers_manager.poll(cx) {
588                self.on_peer_action(action);
589            }
590
591            // We need to poll again in case we have received any responses because they may have
592            // triggered follow-up requests.
593            if self.queued_messages.is_empty() {
594                return Poll::Pending
595            }
596        }
597    }
598}
599
600/// Tracks the state of a Peer with an active Session.
601///
602/// For example known blocks,so we can decide what to announce.
603#[derive(Debug)]
604pub(crate) struct ActivePeer<N: NetworkPrimitives> {
605    /// Best block of the peer.
606    pub(crate) best_hash: B256,
607    /// The capabilities of the remote peer.
608    pub(crate) capabilities: Arc<Capabilities>,
609    /// A communication channel directly to the session task.
610    pub(crate) request_tx: PeerRequestSender<PeerRequest<N>>,
611    /// The response receiver for a currently active request to that peer.
612    pub(crate) pending_response: Option<PeerResponse<N>>,
613    /// Blocks we know the peer has.
614    pub(crate) blocks: LruCache<B256>,
615}
616
617/// Everything [`NetworkState::on_session_activated`] needs to register a newly established
618/// session.
619pub(crate) struct SessionActivation<N: NetworkPrimitives> {
620    /// The remote peer's identifier.
621    pub(crate) peer: PeerId,
622    /// The capabilities the peer announced.
623    pub(crate) capabilities: Arc<Capabilities>,
624    /// The `Status` message the peer sent during the `eth` handshake.
625    pub(crate) status: Arc<UnifiedStatus>,
626    /// A communication channel directly to the session task.
627    pub(crate) request_tx: PeerRequestSender<PeerRequest<N>>,
628    /// The maximum time the session waits for a response from the peer.
629    pub(crate) timeout: Arc<AtomicU64>,
630    /// The range info for the peer.
631    pub(crate) range_info: Option<BlockRangeInfo>,
632    /// Whether the connection negotiated `snap/2` and can serve [`PeerRequest::GetSnap`].
633    pub(crate) supports_snap: bool,
634}
635
636/// Message variants triggered by the [`NetworkState`]
637#[derive(Debug)]
638pub(crate) enum StateAction<N: NetworkPrimitives> {
639    /// Dispatch a `NewBlock` message to the peer
640    NewBlock {
641        /// Target of the message
642        peer_id: PeerId,
643        /// The `NewBlock` message
644        block: NewBlockMessage<N::NewBlockPayload>,
645    },
646    NewBlockHashes {
647        /// Target of the message
648        peer_id: PeerId,
649        /// `NewBlockHashes` message to send to the peer.
650        hashes: NewBlockHashes,
651    },
652    /// Create a new connection to the given node.
653    Connect { remote_addr: SocketAddr, peer_id: PeerId },
654    /// Disconnect an existing connection
655    Disconnect {
656        peer_id: PeerId,
657        /// Why the disconnect was initiated
658        reason: Option<DisconnectReason>,
659    },
660    /// Retrieved a [`ForkId`] from the peer via ENR request, See <https://eips.ethereum.org/EIPS/eip-868>
661    DiscoveredEnrForkId {
662        peer_id: PeerId,
663        /// The address of the peer.
664        addr: PeerAddr,
665        /// The reported [`ForkId`] by this peer.
666        fork_id: ForkId,
667    },
668    /// A new node was found through the discovery, possibly with a `ForkId`
669    DiscoveredNode { peer_id: PeerId, addr: PeerAddr, fork_id: Option<ForkId> },
670    /// A peer was added
671    PeerAdded(PeerId),
672    /// A peer was dropped
673    PeerRemoved(PeerId),
674}
675
676#[cfg(test)]
677mod tests {
678    use crate::{
679        discovery::Discovery,
680        fetch::StateFetcher,
681        peers::PeersManager,
682        state::{BlockNumReader, NetworkState, SessionActivation},
683        PeerRequest,
684    };
685    use alloy_consensus::Header;
686    use alloy_primitives::B256;
687    use reth_eth_wire::{BlockBodies, Capabilities, Capability, EthNetworkPrimitives, EthVersion};
688    use reth_ethereum_primitives::BlockBody;
689    use reth_network_api::PeerRequestSender;
690    use reth_network_p2p::{bodies::client::BodiesClient, error::RequestError};
691    use reth_network_peers::PeerId;
692    use reth_storage_api::noop::NoopProvider;
693    use std::{
694        future::poll_fn,
695        sync::{atomic::AtomicU64, Arc},
696    };
697    use tokio::sync::mpsc;
698    use tokio_stream::{wrappers::ReceiverStream, StreamExt};
699
700    /// Returns a testing instance of the [`NetworkState`].
701    fn state() -> NetworkState<EthNetworkPrimitives> {
702        let peers = PeersManager::default();
703        let handle = peers.handle();
704        NetworkState {
705            active_peers: Default::default(),
706            peers_manager: Default::default(),
707            queued_messages: Default::default(),
708            client: BlockNumReader(Box::new(NoopProvider::default())),
709            discovery: Discovery::noop(),
710            state_fetcher: StateFetcher::new(handle, Default::default()),
711        }
712    }
713
714    fn capabilities() -> Arc<Capabilities> {
715        Arc::new(vec![Capability::from(EthVersion::Eth67)].into())
716    }
717
718    // tests that ongoing requests are answered with connection dropped if the session that received
719    // that request is drops the request object.
720    #[tokio::test(flavor = "multi_thread")]
721    async fn test_dropped_active_session() {
722        let mut state = state();
723        let client = state.fetch_client();
724
725        let peer_id = PeerId::random();
726        let (tx, session_rx) = mpsc::channel(1);
727        let peer_tx = PeerRequestSender::new(peer_id, tx);
728
729        state.on_session_activated(SessionActivation {
730            peer: peer_id,
731            capabilities: capabilities(),
732            status: Arc::default(),
733            request_tx: peer_tx,
734            timeout: Arc::new(AtomicU64::new(1)),
735            range_info: None,
736            supports_snap: false,
737        });
738
739        assert!(state.active_peers.contains_key(&peer_id));
740
741        let body = BlockBody { ommers: vec![Header::default()], ..Default::default() };
742
743        let body_response = body.clone();
744
745        // this mimics an active session that receives the requests from the state
746        tokio::task::spawn(async move {
747            let mut stream = ReceiverStream::new(session_rx);
748            let resp = stream.next().await.unwrap();
749            match resp {
750                PeerRequest::GetBlockBodies { response, .. } => {
751                    response.send(Ok(BlockBodies(vec![body_response]))).unwrap();
752                }
753                _ => unreachable!(),
754            }
755
756            // wait for the next request, then drop
757            let _resp = stream.next().await.unwrap();
758        });
759
760        // spawn the state as future
761        tokio::task::spawn(async move {
762            loop {
763                poll_fn(|cx| state.poll(cx)).await;
764            }
765        });
766
767        // send requests to the state via the client
768        let (peer, bodies) = client.get_block_bodies(vec![B256::random()]).await.unwrap().split();
769        assert_eq!(peer, peer_id);
770        assert_eq!(bodies, vec![body]);
771
772        let resp = client.get_block_bodies(vec![B256::random()]).await;
773        assert!(resp.is_err());
774        assert_eq!(resp.unwrap_err(), RequestError::ConnectionDropped);
775    }
776}
reth_network/state.rs

reth_network/
state.rs
