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 and dials it if there's an
322    /// outbound slot.
323    pub(crate) fn add_peer_kind(
324        &mut self,
325        peer_id: PeerId,
326        kind: Option<PeerKind>,
327        addr: PeerAddr,
328    ) {
329        self.peers_manager.add_requested_peer(peer_id, kind, addr)
330    }
331
332    /// Connects a peer and its address with the given kind
333    pub(crate) fn add_and_connect(&mut self, peer_id: PeerId, kind: PeerKind, addr: PeerAddr) {
334        self.peers_manager.add_and_connect_kind(peer_id, kind, addr, None)
335    }
336
337    /// Removes a peer and its address with the given kind from the peerset.
338    pub(crate) fn remove_peer_kind(&mut self, peer_id: PeerId, kind: PeerKind) {
339        match kind {
340            PeerKind::Basic | PeerKind::Static => self.peers_manager.remove_peer(peer_id),
341            PeerKind::Trusted => self.peers_manager.remove_peer_from_trusted_set(peer_id),
342        }
343    }
344
345    /// Event hook for events received from the discovery service.
346    fn on_discovery_event(&mut self, event: DiscoveryEvent) {
347        match event {
348            DiscoveryEvent::NewNode(DiscoveredEvent::EventQueued { peer_id, addr, fork_id }) => {
349                self.queued_messages.push_back(StateAction::DiscoveredNode {
350                    peer_id,
351                    addr,
352                    fork_id,
353                });
354            }
355            DiscoveryEvent::EnrForkId(record, fork_id) => {
356                let peer_id = record.id;
357                let tcp_addr = record.tcp_addr();
358                if tcp_addr.port() == 0 {
359                    return
360                }
361                let udp_addr = record.udp_addr();
362                let addr = PeerAddr::new(tcp_addr, Some(udp_addr));
363                self.queued_messages.push_back(StateAction::DiscoveredEnrForkId {
364                    peer_id,
365                    addr,
366                    fork_id,
367                });
368            }
369        }
370    }
371
372    /// Event hook for new actions derived from the peer management set.
373    fn on_peer_action(&mut self, action: PeerAction) {
374        match action {
375            PeerAction::Connect { peer_id, remote_addr } => {
376                self.queued_messages.push_back(StateAction::Connect { peer_id, remote_addr });
377            }
378            PeerAction::Disconnect { peer_id, reason } => {
379                self.state_fetcher.on_pending_disconnect(&peer_id);
380                self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason });
381            }
382            PeerAction::DisconnectBannedIncoming { peer_id } |
383            PeerAction::DisconnectUntrustedIncoming { peer_id } => {
384                self.state_fetcher.on_pending_disconnect(&peer_id);
385                self.queued_messages.push_back(StateAction::Disconnect { peer_id, reason: None });
386            }
387            PeerAction::DiscoveryBanPeerId { peer_id, ip_addr } => {
388                self.ban_discovery(peer_id, ip_addr)
389            }
390            PeerAction::DiscoveryBanIp { ip_addr } => self.ban_ip_discovery(ip_addr),
391            PeerAction::PeerAdded(peer_id) => {
392                self.queued_messages.push_back(StateAction::PeerAdded(peer_id))
393            }
394            PeerAction::PeerRemoved(peer_id) => {
395                self.queued_messages.push_back(StateAction::PeerRemoved(peer_id))
396            }
397            PeerAction::BanPeer { .. } | PeerAction::UnBanPeer { .. } => {}
398        }
399    }
400
401    /// Sends The message to the peer's session and queues in a response.
402    ///
403    /// Caution: this will replace an already pending response. It's the responsibility of the
404    /// caller to select the peer.
405    fn handle_block_request(&mut self, peer_id: PeerId, request: BlockRequest) {
406        if let Some(ref mut peer) = self.active_peers.get_mut(&peer_id) {
407            let (request, response) = match request {
408                BlockRequest::GetBlockHeaders(request) => {
409                    let (response, rx) = oneshot::channel();
410                    let request = PeerRequest::GetBlockHeaders { request, response };
411                    let response = PeerResponse::BlockHeaders { response: rx };
412                    (request, response)
413                }
414                BlockRequest::GetBlockBodies(request) => {
415                    let (response, rx) = oneshot::channel();
416                    let request = PeerRequest::GetBlockBodies { request, response };
417                    let response = PeerResponse::BlockBodies { response: rx };
418                    (request, response)
419                }
420                BlockRequest::GetBlockAccessLists(request) => {
421                    let (response, rx) = oneshot::channel();
422                    let request = PeerRequest::GetBlockAccessLists { request, response };
423                    let response = PeerResponse::BlockAccessLists { response: rx };
424                    (request, response)
425                }
426                BlockRequest::GetReceipts(request) => {
427                    if peer.capabilities.supports_eth_v70() {
428                        let (response, rx) = oneshot::channel();
429                        let request = PeerRequest::GetReceipts70 {
430                            request: GetReceipts70 {
431                                first_block_receipt_index: 0,
432                                block_hashes: request.0,
433                            },
434                            response,
435                        };
436                        let response = PeerResponse::Receipts70 { response: rx };
437                        (request, response)
438                    } else if peer.capabilities.supports_eth_v69() {
439                        let (response, rx) = oneshot::channel();
440                        let request = PeerRequest::GetReceipts69 { request, response };
441                        let response = PeerResponse::Receipts69 { response: rx };
442                        (request, response)
443                    } else {
444                        let (response, rx) = oneshot::channel();
445                        let request = PeerRequest::GetReceipts { request, response };
446                        let response = PeerResponse::Receipts { response: rx };
447                        (request, response)
448                    }
449                }
450                BlockRequest::GetSnap(request) => {
451                    let (response, rx) = oneshot::channel();
452                    let request = PeerRequest::GetSnap { request: *request, response };
453                    let response = PeerResponse::Snap { response: rx };
454                    (request, response)
455                }
456            };
457            let _ = peer.request_tx.to_session_tx.try_send(request);
458            peer.pending_response = Some(response);
459        }
460    }
461
462    /// Handle the outcome of processed response, for example directly queue another request.
463    fn on_block_response_outcome(&mut self, outcome: BlockResponseOutcome) {
464        match outcome {
465            BlockResponseOutcome::Request(peer, request) => {
466                self.handle_block_request(peer, request);
467            }
468            BlockResponseOutcome::BadResponse(peer, reputation_change) => {
469                self.peers_manager.apply_reputation_change(&peer, reputation_change);
470            }
471        }
472    }
473
474    /// Invoked when received a response from a connected peer.
475    ///
476    /// Delegates the response result to the fetcher which may return an outcome specific
477    /// instruction that needs to be handled in [`Self::on_block_response_outcome`]. This could be
478    /// a follow-up request or an instruction to slash the peer's reputation.
479    fn on_eth_response(&mut self, peer: PeerId, resp: PeerResponseResult<N>) {
480        let outcome = match resp {
481            PeerResponseResult::BlockHeaders(res) => {
482                self.state_fetcher.on_block_headers_response(peer, res)
483            }
484            PeerResponseResult::BlockBodies(res) => {
485                self.state_fetcher.on_block_bodies_response(peer, res)
486            }
487            PeerResponseResult::Receipts(res) => {
488                // Legacy eth/66-68: strip bloom filters and wrap in ReceiptsResponse
489                let normalized = res.map(|blocks| {
490                    let receipts = blocks
491                        .into_iter()
492                        .map(|block_receipts| {
493                            block_receipts.into_iter().map(|rwb| rwb.receipt).collect()
494                        })
495                        .collect();
496                    ReceiptsResponse::new(receipts)
497                });
498                self.state_fetcher.on_receipts_response(peer, normalized)
499            }
500            PeerResponseResult::Receipts69(res) => {
501                let normalized = res.map(ReceiptsResponse::new);
502                self.state_fetcher.on_receipts_response(peer, normalized)
503            }
504            PeerResponseResult::Receipts70(res) => {
505                let normalized = res.map(ReceiptsResponse::from);
506                self.state_fetcher.on_receipts_response(peer, normalized)
507            }
508            PeerResponseResult::BlockAccessLists(res) => {
509                self.state_fetcher.on_block_access_lists_response(peer, res)
510            }
511            PeerResponseResult::Snap(res) => self.state_fetcher.on_snap_response(peer, res),
512            _ => None,
513        };
514
515        if let Some(outcome) = outcome {
516            self.on_block_response_outcome(outcome);
517        }
518    }
519
520    /// Advances the state
521    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<StateAction<N>> {
522        loop {
523            // drain buffered messages
524            if let Some(message) = self.queued_messages.pop_front() {
525                return Poll::Ready(message)
526            }
527
528            while let Poll::Ready(discovery) = self.discovery.poll(cx) {
529                self.on_discovery_event(discovery);
530            }
531
532            while let Poll::Ready(action) = self.state_fetcher.poll(cx) {
533                match action {
534                    FetchAction::BlockRequest { peer_id, request } => {
535                        self.handle_block_request(peer_id, request)
536                    }
537                }
538            }
539
540            loop {
541                // need to buffer results here to make borrow checker happy
542                let mut closed_sessions = Vec::new();
543                let mut received_responses = Vec::new();
544
545                // poll all connected peers for responses
546                for (id, peer) in &mut self.active_peers {
547                    let Some(mut response) = peer.pending_response.take() else { continue };
548                    match response.poll(cx) {
549                        Poll::Ready(res) => {
550                            // check if the error is due to a closed channel to the session
551                            if res.err().is_some_and(|err| err.is_channel_closed()) {
552                                debug!(
553                                    target: "net",
554                                    ?id,
555                                    "Request canceled, response channel from session closed."
556                                );
557                                // if the channel is closed, this means the peer session is also
558                                // closed, in which case we can invoke the
559                                // [Self::on_closed_session]
560                                // immediately, preventing followup requests and propagate the
561                                // connection dropped error
562                                closed_sessions.push(*id);
563                            } else {
564                                received_responses.push((*id, res));
565                            }
566                        }
567                        Poll::Pending => {
568                            // not ready yet, store again.
569                            peer.pending_response = Some(response);
570                        }
571                    };
572                }
573
574                for peer in closed_sessions {
575                    self.on_session_closed(peer)
576                }
577
578                if received_responses.is_empty() {
579                    break;
580                }
581
582                for (peer_id, resp) in received_responses {
583                    self.on_eth_response(peer_id, resp);
584                }
585            }
586
587            // poll peer manager
588            while let Poll::Ready(action) = self.peers_manager.poll(cx) {
589                self.on_peer_action(action);
590            }
591
592            // We need to poll again in case we have received any responses because they may have
593            // triggered follow-up requests.
594            if self.queued_messages.is_empty() {
595                return Poll::Pending
596            }
597        }
598    }
599}
600
601/// Tracks the state of a Peer with an active Session.
602///
603/// For example known blocks,so we can decide what to announce.
604#[derive(Debug)]
605pub(crate) struct ActivePeer<N: NetworkPrimitives> {
606    /// Best block of the peer.
607    pub(crate) best_hash: B256,
608    /// The capabilities of the remote peer.
609    pub(crate) capabilities: Arc<Capabilities>,
610    /// A communication channel directly to the session task.
611    pub(crate) request_tx: PeerRequestSender<PeerRequest<N>>,
612    /// The response receiver for a currently active request to that peer.
613    pub(crate) pending_response: Option<PeerResponse<N>>,
614    /// Blocks we know the peer has.
615    pub(crate) blocks: LruCache<B256>,
616}
617
618/// Everything [`NetworkState::on_session_activated`] needs to register a newly established
619/// session.
620pub(crate) struct SessionActivation<N: NetworkPrimitives> {
621    /// The remote peer's identifier.
622    pub(crate) peer: PeerId,
623    /// The capabilities the peer announced.
624    pub(crate) capabilities: Arc<Capabilities>,
625    /// The `Status` message the peer sent during the `eth` handshake.
626    pub(crate) status: Arc<UnifiedStatus>,
627    /// A communication channel directly to the session task.
628    pub(crate) request_tx: PeerRequestSender<PeerRequest<N>>,
629    /// The maximum time the session waits for a response from the peer.
630    pub(crate) timeout: Arc<AtomicU64>,
631    /// The range info for the peer.
632    pub(crate) range_info: Option<BlockRangeInfo>,
633    /// Whether the connection negotiated `snap/2` and can serve [`PeerRequest::GetSnap`].
634    pub(crate) supports_snap: bool,
635}
636
637/// Message variants triggered by the [`NetworkState`]
638#[derive(Debug)]
639pub(crate) enum StateAction<N: NetworkPrimitives> {
640    /// Dispatch a `NewBlock` message to the peer
641    NewBlock {
642        /// Target of the message
643        peer_id: PeerId,
644        /// The `NewBlock` message
645        block: NewBlockMessage<N::NewBlockPayload>,
646    },
647    NewBlockHashes {
648        /// Target of the message
649        peer_id: PeerId,
650        /// `NewBlockHashes` message to send to the peer.
651        hashes: NewBlockHashes,
652    },
653    /// Create a new connection to the given node.
654    Connect { remote_addr: SocketAddr, peer_id: PeerId },
655    /// Disconnect an existing connection
656    Disconnect {
657        peer_id: PeerId,
658        /// Why the disconnect was initiated
659        reason: Option<DisconnectReason>,
660    },
661    /// Retrieved a [`ForkId`] from the peer via ENR request, See <https://eips.ethereum.org/EIPS/eip-868>
662    DiscoveredEnrForkId {
663        peer_id: PeerId,
664        /// The address of the peer.
665        addr: PeerAddr,
666        /// The reported [`ForkId`] by this peer.
667        fork_id: ForkId,
668    },
669    /// A new node was found through the discovery, possibly with a `ForkId`
670    DiscoveredNode { peer_id: PeerId, addr: PeerAddr, fork_id: Option<ForkId> },
671    /// A peer was added
672    PeerAdded(PeerId),
673    /// A peer was dropped
674    PeerRemoved(PeerId),
675}
676
677#[cfg(test)]
678mod tests {
679    use crate::{
680        discovery::Discovery,
681        fetch::StateFetcher,
682        peers::PeersManager,
683        state::{BlockNumReader, NetworkState, SessionActivation},
684        PeerRequest,
685    };
686    use alloy_consensus::Header;
687    use alloy_primitives::B256;
688    use reth_eth_wire::{BlockBodies, Capabilities, Capability, EthNetworkPrimitives, EthVersion};
689    use reth_ethereum_primitives::BlockBody;
690    use reth_network_api::PeerRequestSender;
691    use reth_network_p2p::{bodies::client::BodiesClient, error::RequestError};
692    use reth_network_peers::PeerId;
693    use reth_storage_api::noop::NoopProvider;
694    use std::{
695        future::poll_fn,
696        sync::{atomic::AtomicU64, Arc},
697    };
698    use tokio::sync::mpsc;
699    use tokio_stream::{wrappers::ReceiverStream, StreamExt};
700
701    /// Returns a testing instance of the [`NetworkState`].
702    fn state() -> NetworkState<EthNetworkPrimitives> {
703        let peers = PeersManager::default();
704        let handle = peers.handle();
705        NetworkState {
706            active_peers: Default::default(),
707            peers_manager: Default::default(),
708            queued_messages: Default::default(),
709            client: BlockNumReader(Box::new(NoopProvider::default())),
710            discovery: Discovery::noop(),
711            state_fetcher: StateFetcher::new(handle, Default::default()),
712        }
713    }
714
715    fn capabilities() -> Arc<Capabilities> {
716        Arc::new(vec![Capability::from(EthVersion::Eth67)].into())
717    }
718
719    // tests that ongoing requests are answered with connection dropped if the session that received
720    // that request is drops the request object.
721    #[tokio::test(flavor = "multi_thread")]
722    async fn test_dropped_active_session() {
723        let mut state = state();
724        let client = state.fetch_client();
725
726        let peer_id = PeerId::random();
727        let (tx, session_rx) = mpsc::channel(1);
728        let peer_tx = PeerRequestSender::new(peer_id, tx);
729
730        state.on_session_activated(SessionActivation {
731            peer: peer_id,
732            capabilities: capabilities(),
733            status: Arc::default(),
734            request_tx: peer_tx,
735            timeout: Arc::new(AtomicU64::new(1)),
736            range_info: None,
737            supports_snap: false,
738        });
739
740        assert!(state.active_peers.contains_key(&peer_id));
741
742        let body = BlockBody { ommers: vec![Header::default()], ..Default::default() };
743
744        let body_response = body.clone();
745
746        // this mimics an active session that receives the requests from the state
747        tokio::task::spawn(async move {
748            let mut stream = ReceiverStream::new(session_rx);
749            let resp = stream.next().await.unwrap();
750            match resp {
751                PeerRequest::GetBlockBodies { response, .. } => {
752                    response.send(Ok(BlockBodies(vec![body_response]))).unwrap();
753                }
754                _ => unreachable!(),
755            }
756
757            // wait for the next request, then drop
758            let _resp = stream.next().await.unwrap();
759        });
760
761        // spawn the state as future
762        tokio::task::spawn(async move {
763            loop {
764                poll_fn(|cx| state.poll(cx)).await;
765            }
766        });
767
768        // send requests to the state via the client
769        let (peer, bodies) = client.get_block_bodies(vec![B256::random()]).await.unwrap().split();
770        assert_eq!(peer, peer_id);
771        assert_eq!(bodies, vec![body]);
772
773        let resp = client.get_block_bodies(vec![B256::random()]).await;
774        assert!(resp.is_err());
775        assert_eq!(resp.unwrap_err(), RequestError::ConnectionDropped);
776    }
777}
reth_network/state.rs

reth_network/
state.rs
