reth_network/session/
mod.rs

1//! Support for handling peer sessions.
2
3mod active;
4mod conn;
5mod counter;
6mod handle;
7mod types;
8pub use types::BlockRangeInfo;
9
10use crate::{
11    message::PeerMessage,
12    metrics::SessionManagerMetrics,
13    protocol::{IntoRlpxSubProtocol, OnNotSupported, RlpxSubProtocolHandlers, RlpxSubProtocols},
14    session::active::ActiveSession,
15};
16use active::QueuedOutgoingMessages;
17use counter::SessionCounter;
18use futures::{future::Either, io, FutureExt, StreamExt};
19use reth_ecies::{stream::ECIESStream, ECIESError};
20use reth_eth_wire::{
21    errors::EthStreamError, handshake::EthRlpxHandshake, multiplex::RlpxProtocolMultiplexer,
22    BlockRangeUpdate, Capabilities, DisconnectReason, EthStream, EthVersion,
23    HelloMessageWithProtocols, NetworkPrimitives, UnauthedP2PStream, UnifiedStatus,
24    HANDSHAKE_TIMEOUT,
25};
26use reth_ethereum_forks::{ForkFilter, ForkId, ForkTransition, Head};
27use reth_metrics::common::mpsc::MeteredPollSender;
28use reth_network_api::{PeerRequest, PeerRequestSender};
29use reth_network_peers::PeerId;
30use reth_network_types::SessionsConfig;
31use reth_tasks::TaskSpawner;
32use rustc_hash::FxHashMap;
33use secp256k1::SecretKey;
34use std::{
35    collections::HashMap,
36    future::Future,
37    net::SocketAddr,
38    sync::{atomic::AtomicU64, Arc},
39    task::{Context, Poll},
40    time::{Duration, Instant},
41};
42use tokio::{
43    io::{AsyncRead, AsyncWrite},
44    net::TcpStream,
45    sync::{mpsc, mpsc::error::TrySendError, oneshot},
46};
47use tokio_stream::wrappers::ReceiverStream;
48use tokio_util::sync::PollSender;
49use tracing::{debug, instrument, trace};
50
51use crate::session::active::RANGE_UPDATE_INTERVAL;
52pub use conn::EthRlpxConnection;
53pub use handle::{
54    ActiveSessionHandle, ActiveSessionMessage, PendingSessionEvent, PendingSessionHandle,
55    SessionCommand,
56};
57pub use reth_network_api::{Direction, PeerInfo};
58
59/// Internal identifier for active sessions.
60#[derive(Debug, Clone, Copy, PartialOrd, PartialEq, Eq, Hash)]
61pub struct SessionId(usize);
62
63/// Manages a set of sessions.
64#[must_use = "Session Manager must be polled to process session events."]
65#[derive(Debug)]
66pub struct SessionManager<N: NetworkPrimitives> {
67    /// Tracks the identifier for the next session.
68    next_id: usize,
69    /// Keeps track of all sessions
70    counter: SessionCounter,
71    ///  The maximum initial time an [`ActiveSession`] waits for a response from the peer before it
72    /// responds to an _internal_ request with a `TimeoutError`
73    initial_internal_request_timeout: Duration,
74    /// If an [`ActiveSession`] does not receive a response at all within this duration then it is
75    /// considered a protocol violation and the session will initiate a drop.
76    protocol_breach_request_timeout: Duration,
77    /// The timeout after which a pending session attempt is considered failed.
78    pending_session_timeout: Duration,
79    /// The secret key used for authenticating sessions.
80    secret_key: SecretKey,
81    /// The `Status` message to send to peers.
82    status: UnifiedStatus,
83    /// The `HelloMessage` message to send to peers.
84    hello_message: HelloMessageWithProtocols,
85    /// The [`ForkFilter`] used to validate the peer's `Status` message.
86    fork_filter: ForkFilter,
87    /// Size of the command buffer per session.
88    session_command_buffer: usize,
89    /// The executor for spawned tasks.
90    executor: Box<dyn TaskSpawner>,
91    /// All pending session that are currently handshaking, exchanging `Hello`s.
92    ///
93    /// Events produced during the authentication phase are reported to this manager. Once the
94    /// session is authenticated, it can be moved to the `active_session` set.
95    pending_sessions: FxHashMap<SessionId, PendingSessionHandle>,
96    /// All active sessions that are ready to exchange messages.
97    active_sessions: HashMap<PeerId, ActiveSessionHandle<N>>,
98    /// The original Sender half of the [`PendingSessionEvent`] channel.
99    ///
100    /// When a new (pending) session is created, the corresponding [`PendingSessionHandle`] will
101    /// get a clone of this sender half.
102    pending_sessions_tx: mpsc::Sender<PendingSessionEvent<N>>,
103    /// Receiver half that listens for [`PendingSessionEvent`] produced by pending sessions.
104    pending_session_rx: ReceiverStream<PendingSessionEvent<N>>,
105    /// The original Sender half of the [`ActiveSessionMessage`] channel.
106    ///
107    /// When active session state is reached, the corresponding [`ActiveSessionHandle`] will get a
108    /// clone of this sender half.
109    active_session_tx: MeteredPollSender<ActiveSessionMessage<N>>,
110    /// Receiver half that listens for [`ActiveSessionMessage`] produced by pending sessions.
111    active_session_rx: ReceiverStream<ActiveSessionMessage<N>>,
112    /// Additional `RLPx` sub-protocols to be used by the session manager.
113    extra_protocols: RlpxSubProtocols,
114    /// Tracks the ongoing graceful disconnections attempts for incoming connections.
115    disconnections_counter: DisconnectionsCounter,
116    /// Metrics for the session manager.
117    metrics: SessionManagerMetrics,
118    /// The [`EthRlpxHandshake`] is used to perform the initial handshake with the peer.
119    handshake: Arc<dyn EthRlpxHandshake>,
120    /// Shared local range information that gets propagated to active sessions.
121    /// This represents the range of blocks that this node can serve to other peers.
122    local_range_info: BlockRangeInfo,
123}
124
125// === impl SessionManager ===
126
127impl<N: NetworkPrimitives> SessionManager<N> {
128    /// Creates a new empty [`SessionManager`].
129    #[expect(clippy::too_many_arguments)]
130    pub fn new(
131        secret_key: SecretKey,
132        config: SessionsConfig,
133        executor: Box<dyn TaskSpawner>,
134        status: UnifiedStatus,
135        hello_message: HelloMessageWithProtocols,
136        fork_filter: ForkFilter,
137        extra_protocols: RlpxSubProtocols,
138        handshake: Arc<dyn EthRlpxHandshake>,
139    ) -> Self {
140        let (pending_sessions_tx, pending_sessions_rx) = mpsc::channel(config.session_event_buffer);
141        let (active_session_tx, active_session_rx) = mpsc::channel(config.session_event_buffer);
142        let active_session_tx = PollSender::new(active_session_tx);
143
144        // Initialize local range info from the status
145        let local_range_info = BlockRangeInfo::new(
146            status.earliest_block.unwrap_or_default(),
147            status.latest_block.unwrap_or_default(),
148            status.blockhash,
149        );
150
151        Self {
152            next_id: 0,
153            counter: SessionCounter::new(config.limits),
154            initial_internal_request_timeout: config.initial_internal_request_timeout,
155            protocol_breach_request_timeout: config.protocol_breach_request_timeout,
156            pending_session_timeout: config.pending_session_timeout,
157            secret_key,
158            status,
159            hello_message,
160            fork_filter,
161            session_command_buffer: config.session_command_buffer,
162            executor,
163            pending_sessions: Default::default(),
164            active_sessions: Default::default(),
165            pending_sessions_tx,
166            pending_session_rx: ReceiverStream::new(pending_sessions_rx),
167            active_session_tx: MeteredPollSender::new(active_session_tx, "network_active_session"),
168            active_session_rx: ReceiverStream::new(active_session_rx),
169            extra_protocols,
170            disconnections_counter: Default::default(),
171            metrics: Default::default(),
172            handshake,
173            local_range_info,
174        }
175    }
176
177    /// Returns the currently tracked [`ForkId`].
178    pub(crate) const fn fork_id(&self) -> ForkId {
179        self.fork_filter.current()
180    }
181
182    /// Check whether the provided [`ForkId`] is compatible based on the validation rules in
183    /// `EIP-2124`.
184    pub fn is_valid_fork_id(&self, fork_id: ForkId) -> bool {
185        self.fork_filter.validate(fork_id).is_ok()
186    }
187
188    /// Returns the next unique [`SessionId`].
189    const fn next_id(&mut self) -> SessionId {
190        let id = self.next_id;
191        self.next_id += 1;
192        SessionId(id)
193    }
194
195    /// Returns the current status of the session.
196    pub const fn status(&self) -> UnifiedStatus {
197        self.status
198    }
199
200    /// Returns the secret key used for authenticating sessions.
201    pub const fn secret_key(&self) -> SecretKey {
202        self.secret_key
203    }
204
205    /// Returns a borrowed reference to the active sessions.
206    pub const fn active_sessions(&self) -> &HashMap<PeerId, ActiveSessionHandle<N>> {
207        &self.active_sessions
208    }
209
210    /// Returns the session hello message.
211    pub fn hello_message(&self) -> HelloMessageWithProtocols {
212        self.hello_message.clone()
213    }
214
215    /// Adds an additional protocol handler to the `RLPx` sub-protocol list.
216    pub(crate) fn add_rlpx_sub_protocol(&mut self, protocol: impl IntoRlpxSubProtocol) {
217        self.extra_protocols.push(protocol)
218    }
219
220    /// Returns the number of currently pending connections.
221    #[inline]
222    pub(crate) fn num_pending_connections(&self) -> usize {
223        self.pending_sessions.len()
224    }
225
226    /// Spawns the given future onto a new task that is tracked in the `spawned_tasks`
227    /// [`JoinSet`](tokio::task::JoinSet).
228    fn spawn<F>(&self, f: F)
229    where
230        F: Future<Output = ()> + Send + 'static,
231    {
232        self.executor.spawn(f.boxed());
233    }
234
235    /// Invoked on a received status update.
236    ///
237    /// If the updated activated another fork, this will return a [`ForkTransition`] and updates the
238    /// active [`ForkId`]. See also [`ForkFilter::set_head`].
239    pub(crate) fn on_status_update(&mut self, head: Head) -> Option<ForkTransition> {
240        self.status.blockhash = head.hash;
241        self.status.total_difficulty = Some(head.total_difficulty);
242        let transition = self.fork_filter.set_head(head);
243        self.status.forkid = self.fork_filter.current();
244        self.status.latest_block = Some(head.number);
245
246        transition
247    }
248
249    /// An incoming TCP connection was received. This starts the authentication process to turn this
250    /// stream into an active peer session.
251    ///
252    /// Returns an error if the configured limit has been reached.
253    pub(crate) fn on_incoming(
254        &mut self,
255        stream: TcpStream,
256        remote_addr: SocketAddr,
257    ) -> Result<SessionId, ExceedsSessionLimit> {
258        self.counter.ensure_pending_inbound()?;
259
260        let session_id = self.next_id();
261
262        trace!(
263            target: "net::session",
264            ?remote_addr,
265            ?session_id,
266            "new pending incoming session"
267        );
268
269        let (disconnect_tx, disconnect_rx) = oneshot::channel();
270        let pending_events = self.pending_sessions_tx.clone();
271        let secret_key = self.secret_key;
272        let hello_message = self.hello_message.clone();
273        let status = self.status;
274        let fork_filter = self.fork_filter.clone();
275        let extra_handlers = self.extra_protocols.on_incoming(remote_addr);
276        self.spawn(pending_session_with_timeout(
277            self.pending_session_timeout,
278            session_id,
279            remote_addr,
280            Direction::Incoming,
281            pending_events.clone(),
282            start_pending_incoming_session(
283                self.handshake.clone(),
284                disconnect_rx,
285                session_id,
286                stream,
287                pending_events,
288                remote_addr,
289                secret_key,
290                hello_message,
291                status,
292                fork_filter,
293                extra_handlers,
294            ),
295        ));
296
297        let handle = PendingSessionHandle {
298            disconnect_tx: Some(disconnect_tx),
299            direction: Direction::Incoming,
300        };
301        self.pending_sessions.insert(session_id, handle);
302        self.counter.inc_pending_inbound();
303        Ok(session_id)
304    }
305
306    /// Starts a new pending session from the local node to the given remote node.
307    pub fn dial_outbound(&mut self, remote_addr: SocketAddr, remote_peer_id: PeerId) {
308        // The error can be dropped because no dial will be made if it would exceed the limit
309        if self.counter.ensure_pending_outbound().is_ok() {
310            let session_id = self.next_id();
311            let (disconnect_tx, disconnect_rx) = oneshot::channel();
312            let pending_events = self.pending_sessions_tx.clone();
313            let secret_key = self.secret_key;
314            let hello_message = self.hello_message.clone();
315            let fork_filter = self.fork_filter.clone();
316            let status = self.status;
317            let extra_handlers = self.extra_protocols.on_outgoing(remote_addr, remote_peer_id);
318            self.spawn(pending_session_with_timeout(
319                self.pending_session_timeout,
320                session_id,
321                remote_addr,
322                Direction::Outgoing(remote_peer_id),
323                pending_events.clone(),
324                start_pending_outbound_session(
325                    self.handshake.clone(),
326                    disconnect_rx,
327                    pending_events,
328                    session_id,
329                    remote_addr,
330                    remote_peer_id,
331                    secret_key,
332                    hello_message,
333                    status,
334                    fork_filter,
335                    extra_handlers,
336                ),
337            ));
338
339            let handle = PendingSessionHandle {
340                disconnect_tx: Some(disconnect_tx),
341                direction: Direction::Outgoing(remote_peer_id),
342            };
343            self.pending_sessions.insert(session_id, handle);
344            self.counter.inc_pending_outbound();
345        }
346    }
347
348    /// Initiates a shutdown of the channel.
349    ///
350    /// This will trigger the disconnect on the session task to gracefully terminate. The result
351    /// will be picked up by the receiver.
352    pub fn disconnect(&self, node: PeerId, reason: Option<DisconnectReason>) {
353        if let Some(session) = self.active_sessions.get(&node) {
354            session.disconnect(reason);
355        }
356    }
357
358    /// Initiates a shutdown of all sessions.
359    ///
360    /// It will trigger the disconnect on all the session tasks to gracefully terminate. The result
361    /// will be picked by the receiver.
362    pub fn disconnect_all(&self, reason: Option<DisconnectReason>) {
363        for session in self.active_sessions.values() {
364            session.disconnect(reason);
365        }
366    }
367
368    /// Disconnects all pending sessions.
369    pub fn disconnect_all_pending(&mut self) {
370        for session in self.pending_sessions.values_mut() {
371            session.disconnect();
372        }
373    }
374
375    /// Sends a message to the peer's session
376    pub fn send_message(&self, peer_id: &PeerId, msg: PeerMessage<N>) {
377        if let Some(session) = self.active_sessions.get(peer_id) {
378            let _ = session.commands_to_session.try_send(SessionCommand::Message(msg)).inspect_err(
379                |e| {
380                    if let TrySendError::Full(_) = e {
381                        debug!(
382                            target: "net::session",
383                            ?peer_id,
384                            "session command buffer full, dropping message"
385                        );
386                        self.metrics.total_outgoing_peer_messages_dropped.increment(1);
387                    }
388                },
389            );
390        }
391    }
392
393    /// Removes the [`PendingSessionHandle`] if it exists.
394    fn remove_pending_session(&mut self, id: &SessionId) -> Option<PendingSessionHandle> {
395        let session = self.pending_sessions.remove(id)?;
396        self.counter.dec_pending(&session.direction);
397        Some(session)
398    }
399
400    /// Removes the [`PendingSessionHandle`] if it exists.
401    fn remove_active_session(&mut self, id: &PeerId) -> Option<ActiveSessionHandle<N>> {
402        let session = self.active_sessions.remove(id)?;
403        self.counter.dec_active(&session.direction);
404        Some(session)
405    }
406
407    /// Try to gracefully disconnect an incoming connection by initiating a ECIES connection and
408    /// sending a disconnect. If [`SessionManager`] is at capacity for ongoing disconnections, will
409    /// simply drop the incoming connection.
410    pub(crate) fn try_disconnect_incoming_connection(
411        &self,
412        stream: TcpStream,
413        reason: DisconnectReason,
414    ) {
415        if !self.disconnections_counter.has_capacity() {
416            // drop the connection if we don't have capacity for gracefully disconnecting
417            return
418        }
419
420        let guard = self.disconnections_counter.clone();
421        let secret_key = self.secret_key;
422
423        self.spawn(async move {
424            trace!(
425                target: "net::session",
426                "gracefully disconnecting incoming connection"
427            );
428            if let Ok(stream) = get_ecies_stream(stream, secret_key, Direction::Incoming).await {
429                let mut unauth = UnauthedP2PStream::new(stream);
430                let _ = unauth.send_disconnect(reason).await;
431                drop(guard);
432            }
433        });
434    }
435
436    /// This polls all the session handles and returns [`SessionEvent`].
437    ///
438    /// Active sessions are prioritized.
439    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<SessionEvent<N>> {
440        // Poll events from active sessions
441        match self.active_session_rx.poll_next_unpin(cx) {
442            Poll::Pending => {}
443            Poll::Ready(None) => {
444                unreachable!("Manager holds both channel halves.")
445            }
446            Poll::Ready(Some(event)) => {
447                return match event {
448                    ActiveSessionMessage::Disconnected { peer_id, remote_addr } => {
449                        trace!(
450                            target: "net::session",
451                            ?peer_id,
452                            "gracefully disconnected active session."
453                        );
454                        self.remove_active_session(&peer_id);
455                        Poll::Ready(SessionEvent::Disconnected { peer_id, remote_addr })
456                    }
457                    ActiveSessionMessage::ClosedOnConnectionError {
458                        peer_id,
459                        remote_addr,
460                        error,
461                    } => {
462                        trace!(target: "net::session", ?peer_id, %error,"closed session.");
463                        self.remove_active_session(&peer_id);
464                        Poll::Ready(SessionEvent::SessionClosedOnConnectionError {
465                            remote_addr,
466                            peer_id,
467                            error,
468                        })
469                    }
470                    ActiveSessionMessage::ValidMessage { peer_id, message } => {
471                        Poll::Ready(SessionEvent::ValidMessage { peer_id, message })
472                    }
473                    ActiveSessionMessage::BadMessage { peer_id } => {
474                        Poll::Ready(SessionEvent::BadMessage { peer_id })
475                    }
476                    ActiveSessionMessage::ProtocolBreach { peer_id } => {
477                        Poll::Ready(SessionEvent::ProtocolBreach { peer_id })
478                    }
479                }
480            }
481        }
482
483        // Poll the pending session event stream
484        let event = match self.pending_session_rx.poll_next_unpin(cx) {
485            Poll::Pending => return Poll::Pending,
486            Poll::Ready(None) => unreachable!("Manager holds both channel halves."),
487            Poll::Ready(Some(event)) => event,
488        };
489        match event {
490            PendingSessionEvent::Established {
491                session_id,
492                remote_addr,
493                local_addr,
494                peer_id,
495                capabilities,
496                conn,
497                status,
498                direction,
499                client_id,
500            } => {
501                // move from pending to established.
502                self.remove_pending_session(&session_id);
503
504                // If there's already a session to the peer then we disconnect right away
505                if self.active_sessions.contains_key(&peer_id) {
506                    trace!(
507                        target: "net::session",
508                        ?session_id,
509                        ?remote_addr,
510                        ?peer_id,
511                        ?direction,
512                        "already connected"
513                    );
514
515                    self.spawn(async move {
516                        // send a disconnect message
517                        let _ =
518                            conn.into_inner().disconnect(DisconnectReason::AlreadyConnected).await;
519                    });
520
521                    return Poll::Ready(SessionEvent::AlreadyConnected {
522                        peer_id,
523                        remote_addr,
524                        direction,
525                    })
526                }
527
528                let (commands_to_session, commands_rx) = mpsc::channel(self.session_command_buffer);
529
530                let (to_session_tx, messages_rx) = mpsc::channel(self.session_command_buffer);
531
532                let messages = PeerRequestSender::new(peer_id, to_session_tx);
533
534                let timeout = Arc::new(AtomicU64::new(
535                    self.initial_internal_request_timeout.as_millis() as u64,
536                ));
537
538                // negotiated version
539                let version = conn.version();
540
541                // Configure the interval at which the range information is updated, starting with
542                // ETH69
543                let range_update_interval = (conn.version() >= EthVersion::Eth69).then(|| {
544                    let mut interval = tokio::time::interval(RANGE_UPDATE_INTERVAL);
545                    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
546                    interval
547                });
548
549                let session = ActiveSession {
550                    next_id: 0,
551                    remote_peer_id: peer_id,
552                    remote_addr,
553                    remote_capabilities: Arc::clone(&capabilities),
554                    session_id,
555                    commands_rx: ReceiverStream::new(commands_rx),
556                    to_session_manager: self.active_session_tx.clone(),
557                    pending_message_to_session: None,
558                    internal_request_rx: ReceiverStream::new(messages_rx).fuse(),
559                    inflight_requests: Default::default(),
560                    conn,
561                    queued_outgoing: QueuedOutgoingMessages::new(
562                        self.metrics.queued_outgoing_messages.clone(),
563                    ),
564                    received_requests_from_remote: Default::default(),
565                    internal_request_timeout_interval: tokio::time::interval(
566                        self.initial_internal_request_timeout,
567                    ),
568                    internal_request_timeout: Arc::clone(&timeout),
569                    protocol_breach_request_timeout: self.protocol_breach_request_timeout,
570                    terminate_message: None,
571                    range_info: None,
572                    local_range_info: self.local_range_info.clone(),
573                    range_update_interval,
574                };
575
576                self.spawn(session);
577
578                let client_version = client_id.into();
579                let handle = ActiveSessionHandle {
580                    status: status.clone(),
581                    direction,
582                    session_id,
583                    remote_id: peer_id,
584                    version,
585                    established: Instant::now(),
586                    capabilities: Arc::clone(&capabilities),
587                    commands_to_session,
588                    client_version: Arc::clone(&client_version),
589                    remote_addr,
590                    local_addr,
591                };
592
593                self.active_sessions.insert(peer_id, handle);
594                self.counter.inc_active(&direction);
595
596                if direction.is_outgoing() {
597                    self.metrics.total_dial_successes.increment(1);
598                }
599
600                Poll::Ready(SessionEvent::SessionEstablished {
601                    peer_id,
602                    remote_addr,
603                    client_version,
604                    version,
605                    capabilities,
606                    status,
607                    messages,
608                    direction,
609                    timeout,
610                    range_info: None,
611                })
612            }
613            PendingSessionEvent::Disconnected { remote_addr, session_id, direction, error } => {
614                trace!(
615                    target: "net::session",
616                    ?session_id,
617                    ?remote_addr,
618                    ?error,
619                    "disconnected pending session"
620                );
621                self.remove_pending_session(&session_id);
622                match direction {
623                    Direction::Incoming => {
624                        Poll::Ready(SessionEvent::IncomingPendingSessionClosed {
625                            remote_addr,
626                            error,
627                        })
628                    }
629                    Direction::Outgoing(peer_id) => {
630                        Poll::Ready(SessionEvent::OutgoingPendingSessionClosed {
631                            remote_addr,
632                            peer_id,
633                            error,
634                        })
635                    }
636                }
637            }
638            PendingSessionEvent::OutgoingConnectionError {
639                remote_addr,
640                session_id,
641                peer_id,
642                error,
643            } => {
644                trace!(
645                    target: "net::session",
646                    %error,
647                    ?session_id,
648                    ?remote_addr,
649                    ?peer_id,
650                    "connection refused"
651                );
652                self.remove_pending_session(&session_id);
653                Poll::Ready(SessionEvent::OutgoingConnectionError { remote_addr, peer_id, error })
654            }
655            PendingSessionEvent::EciesAuthError { remote_addr, session_id, error, direction } => {
656                trace!(
657                    target: "net::session",
658                    %error,
659                    ?session_id,
660                    ?remote_addr,
661                    "ecies auth failed"
662                );
663                self.remove_pending_session(&session_id);
664                match direction {
665                    Direction::Incoming => {
666                        Poll::Ready(SessionEvent::IncomingPendingSessionClosed {
667                            remote_addr,
668                            error: Some(PendingSessionHandshakeError::Ecies(error)),
669                        })
670                    }
671                    Direction::Outgoing(peer_id) => {
672                        Poll::Ready(SessionEvent::OutgoingPendingSessionClosed {
673                            remote_addr,
674                            peer_id,
675                            error: Some(PendingSessionHandshakeError::Ecies(error)),
676                        })
677                    }
678                }
679            }
680        }
681    }
682
683    /// Updates the advertised block range that this node can serve to other peers starting with
684    /// Eth69.
685    ///
686    /// This method updates both the local status message that gets sent to peers during handshake
687    /// and the shared local range information that gets propagated to active sessions (Eth69).
688    /// The range information is used in ETH69 protocol where peers announce the range of blocks
689    /// they can serve to optimize data synchronization.
690    pub(crate) fn update_advertised_block_range(&mut self, block_range_update: BlockRangeUpdate) {
691        self.status.earliest_block = Some(block_range_update.earliest);
692        self.status.latest_block = Some(block_range_update.latest);
693        self.status.blockhash = block_range_update.latest_hash;
694
695        // Update the shared local range info that gets propagated to active sessions
696        self.local_range_info.update(
697            block_range_update.earliest,
698            block_range_update.latest,
699            block_range_update.latest_hash,
700        );
701    }
702}
703
704/// A counter for ongoing graceful disconnections attempts.
705#[derive(Default, Debug, Clone)]
706struct DisconnectionsCounter(Arc<()>);
707
708impl DisconnectionsCounter {
709    const MAX_CONCURRENT_GRACEFUL_DISCONNECTIONS: usize = 15;
710
711    /// Returns true if the [`DisconnectionsCounter`] still has capacity
712    /// for an additional graceful disconnection.
713    fn has_capacity(&self) -> bool {
714        Arc::strong_count(&self.0) <= Self::MAX_CONCURRENT_GRACEFUL_DISCONNECTIONS
715    }
716}
717
718/// Events produced by the [`SessionManager`]
719#[derive(Debug)]
720pub enum SessionEvent<N: NetworkPrimitives> {
721    /// A new session was successfully authenticated.
722    ///
723    /// This session is now able to exchange data.
724    SessionEstablished {
725        /// The remote node's public key
726        peer_id: PeerId,
727        /// The remote node's socket address
728        remote_addr: SocketAddr,
729        /// The user agent of the remote node, usually containing the client name and version
730        client_version: Arc<str>,
731        /// The capabilities the remote node has announced
732        capabilities: Arc<Capabilities>,
733        /// negotiated eth version
734        version: EthVersion,
735        /// The Status message the peer sent during the `eth` handshake
736        status: Arc<UnifiedStatus>,
737        /// The channel for sending messages to the peer with the session
738        messages: PeerRequestSender<PeerRequest<N>>,
739        /// The direction of the session, either `Inbound` or `Outgoing`
740        direction: Direction,
741        /// The maximum time that the session waits for a response from the peer before timing out
742        /// the connection
743        timeout: Arc<AtomicU64>,
744        /// The range info for the peer.
745        range_info: Option<BlockRangeInfo>,
746    },
747    /// The peer was already connected with another session.
748    AlreadyConnected {
749        /// The remote node's public key
750        peer_id: PeerId,
751        /// The remote node's socket address
752        remote_addr: SocketAddr,
753        /// The direction of the session, either `Inbound` or `Outgoing`
754        direction: Direction,
755    },
756    /// A session received a valid message via `RLPx`.
757    ValidMessage {
758        /// The remote node's public key
759        peer_id: PeerId,
760        /// Message received from the peer.
761        message: PeerMessage<N>,
762    },
763    /// Received a bad message from the peer.
764    BadMessage {
765        /// Identifier of the remote peer.
766        peer_id: PeerId,
767    },
768    /// Remote peer is considered in protocol violation
769    ProtocolBreach {
770        /// Identifier of the remote peer.
771        peer_id: PeerId,
772    },
773    /// Closed an incoming pending session during handshaking.
774    IncomingPendingSessionClosed {
775        /// The remote node's socket address
776        remote_addr: SocketAddr,
777        /// The pending handshake session error that caused the session to close
778        error: Option<PendingSessionHandshakeError>,
779    },
780    /// Closed an outgoing pending session during handshaking.
781    OutgoingPendingSessionClosed {
782        /// The remote node's socket address
783        remote_addr: SocketAddr,
784        /// The remote node's public key
785        peer_id: PeerId,
786        /// The pending handshake session error that caused the session to close
787        error: Option<PendingSessionHandshakeError>,
788    },
789    /// Failed to establish a tcp stream
790    OutgoingConnectionError {
791        /// The remote node's socket address
792        remote_addr: SocketAddr,
793        /// The remote node's public key
794        peer_id: PeerId,
795        /// The error that caused the outgoing connection to fail
796        error: io::Error,
797    },
798    /// Session was closed due to an error
799    SessionClosedOnConnectionError {
800        /// The id of the remote peer.
801        peer_id: PeerId,
802        /// The socket we were connected to.
803        remote_addr: SocketAddr,
804        /// The error that caused the session to close
805        error: EthStreamError,
806    },
807    /// Active session was gracefully disconnected.
808    Disconnected {
809        /// The remote node's public key
810        peer_id: PeerId,
811        /// The remote node's socket address that we were connected to
812        remote_addr: SocketAddr,
813    },
814}
815
816/// Errors that can occur during handshaking/authenticating the underlying streams.
817#[derive(Debug, thiserror::Error)]
818pub enum PendingSessionHandshakeError {
819    /// The pending session failed due to an error while establishing the `eth` stream
820    #[error(transparent)]
821    Eth(EthStreamError),
822    /// The pending session failed due to an error while establishing the ECIES stream
823    #[error(transparent)]
824    Ecies(ECIESError),
825    /// Thrown when the authentication timed out
826    #[error("authentication timed out")]
827    Timeout,
828    /// Thrown when the remote lacks the required capability
829    #[error("Mandatory extra capability unsupported")]
830    UnsupportedExtraCapability,
831}
832
833impl PendingSessionHandshakeError {
834    /// Returns the [`DisconnectReason`] if the error is a disconnect message
835    pub const fn as_disconnected(&self) -> Option<DisconnectReason> {
836        match self {
837            Self::Eth(eth_err) => eth_err.as_disconnected(),
838            _ => None,
839        }
840    }
841}
842
843/// The error thrown when the max configured limit has been reached and no more connections are
844/// accepted.
845#[derive(Debug, Clone, thiserror::Error)]
846#[error("session limit reached {0}")]
847pub struct ExceedsSessionLimit(pub(crate) u32);
848
849/// Starts a pending session authentication with a timeout.
850pub(crate) async fn pending_session_with_timeout<F, N: NetworkPrimitives>(
851    timeout: Duration,
852    session_id: SessionId,
853    remote_addr: SocketAddr,
854    direction: Direction,
855    events: mpsc::Sender<PendingSessionEvent<N>>,
856    f: F,
857) where
858    F: Future<Output = ()>,
859{
860    if tokio::time::timeout(timeout, f).await.is_err() {
861        trace!(target: "net::session", ?remote_addr, ?direction, "pending session timed out");
862        let event = PendingSessionEvent::Disconnected {
863            remote_addr,
864            session_id,
865            direction,
866            error: Some(PendingSessionHandshakeError::Timeout),
867        };
868        let _ = events.send(event).await;
869    }
870}
871
872/// Starts the authentication process for a connection initiated by a remote peer.
873///
874/// This will wait for the _incoming_ handshake request and answer it.
875#[expect(clippy::too_many_arguments)]
876pub(crate) async fn start_pending_incoming_session<N: NetworkPrimitives>(
877    handshake: Arc<dyn EthRlpxHandshake>,
878    disconnect_rx: oneshot::Receiver<()>,
879    session_id: SessionId,
880    stream: TcpStream,
881    events: mpsc::Sender<PendingSessionEvent<N>>,
882    remote_addr: SocketAddr,
883    secret_key: SecretKey,
884    hello: HelloMessageWithProtocols,
885    status: UnifiedStatus,
886    fork_filter: ForkFilter,
887    extra_handlers: RlpxSubProtocolHandlers,
888) {
889    authenticate(
890        handshake,
891        disconnect_rx,
892        events,
893        stream,
894        session_id,
895        remote_addr,
896        secret_key,
897        Direction::Incoming,
898        hello,
899        status,
900        fork_filter,
901        extra_handlers,
902    )
903    .await
904}
905
906/// Starts the authentication process for a connection initiated by a remote peer.
907#[instrument(skip_all, fields(%remote_addr, peer_id), target = "net")]
908#[expect(clippy::too_many_arguments)]
909async fn start_pending_outbound_session<N: NetworkPrimitives>(
910    handshake: Arc<dyn EthRlpxHandshake>,
911    disconnect_rx: oneshot::Receiver<()>,
912    events: mpsc::Sender<PendingSessionEvent<N>>,
913    session_id: SessionId,
914    remote_addr: SocketAddr,
915    remote_peer_id: PeerId,
916    secret_key: SecretKey,
917    hello: HelloMessageWithProtocols,
918    status: UnifiedStatus,
919    fork_filter: ForkFilter,
920    extra_handlers: RlpxSubProtocolHandlers,
921) {
922    let stream = match TcpStream::connect(remote_addr).await {
923        Ok(stream) => {
924            if let Err(err) = stream.set_nodelay(true) {
925                tracing::warn!(target: "net::session", "set nodelay failed: {:?}", err);
926            }
927            stream
928        }
929        Err(error) => {
930            let _ = events
931                .send(PendingSessionEvent::OutgoingConnectionError {
932                    remote_addr,
933                    session_id,
934                    peer_id: remote_peer_id,
935                    error,
936                })
937                .await;
938            return
939        }
940    };
941    authenticate(
942        handshake,
943        disconnect_rx,
944        events,
945        stream,
946        session_id,
947        remote_addr,
948        secret_key,
949        Direction::Outgoing(remote_peer_id),
950        hello,
951        status,
952        fork_filter,
953        extra_handlers,
954    )
955    .await
956}
957
958/// Authenticates a session
959#[expect(clippy::too_many_arguments)]
960async fn authenticate<N: NetworkPrimitives>(
961    handshake: Arc<dyn EthRlpxHandshake>,
962    disconnect_rx: oneshot::Receiver<()>,
963    events: mpsc::Sender<PendingSessionEvent<N>>,
964    stream: TcpStream,
965    session_id: SessionId,
966    remote_addr: SocketAddr,
967    secret_key: SecretKey,
968    direction: Direction,
969    hello: HelloMessageWithProtocols,
970    status: UnifiedStatus,
971    fork_filter: ForkFilter,
972    extra_handlers: RlpxSubProtocolHandlers,
973) {
974    let local_addr = stream.local_addr().ok();
975    let stream = match get_ecies_stream(stream, secret_key, direction).await {
976        Ok(stream) => stream,
977        Err(error) => {
978            let _ = events
979                .send(PendingSessionEvent::EciesAuthError {
980                    remote_addr,
981                    session_id,
982                    error,
983                    direction,
984                })
985                .await;
986            return
987        }
988    };
989
990    let unauthed = UnauthedP2PStream::new(stream);
991
992    let auth = authenticate_stream(
993        handshake,
994        unauthed,
995        session_id,
996        remote_addr,
997        local_addr,
998        direction,
999        hello,
1000        status,
1001        fork_filter,
1002        extra_handlers,
1003    )
1004    .boxed();
1005
1006    match futures::future::select(disconnect_rx, auth).await {
1007        Either::Left((_, _)) => {
1008            let _ = events
1009                .send(PendingSessionEvent::Disconnected {
1010                    remote_addr,
1011                    session_id,
1012                    direction,
1013                    error: None,
1014                })
1015                .await;
1016        }
1017        Either::Right((res, _)) => {
1018            let _ = events.send(res).await;
1019        }
1020    }
1021}
1022
1023/// Returns an [`ECIESStream`] if it can be built. If not, send a
1024/// [`PendingSessionEvent::EciesAuthError`] and returns `None`
1025async fn get_ecies_stream<Io: AsyncRead + AsyncWrite + Unpin>(
1026    stream: Io,
1027    secret_key: SecretKey,
1028    direction: Direction,
1029) -> Result<ECIESStream<Io>, ECIESError> {
1030    match direction {
1031        Direction::Incoming => ECIESStream::incoming(stream, secret_key).await,
1032        Direction::Outgoing(remote_peer_id) => {
1033            ECIESStream::connect(stream, secret_key, remote_peer_id).await
1034        }
1035    }
1036}
1037
1038/// Authenticate the stream via handshake
1039///
1040/// On Success return the authenticated stream as [`PendingSessionEvent`].
1041///
1042/// If additional [`RlpxSubProtocolHandlers`] are provided, the hello message will be updated to
1043/// also negotiate the additional protocols.
1044#[expect(clippy::too_many_arguments)]
1045async fn authenticate_stream<N: NetworkPrimitives>(
1046    handshake: Arc<dyn EthRlpxHandshake>,
1047    stream: UnauthedP2PStream<ECIESStream<TcpStream>>,
1048    session_id: SessionId,
1049    remote_addr: SocketAddr,
1050    local_addr: Option<SocketAddr>,
1051    direction: Direction,
1052    mut hello: HelloMessageWithProtocols,
1053    mut status: UnifiedStatus,
1054    fork_filter: ForkFilter,
1055    mut extra_handlers: RlpxSubProtocolHandlers,
1056) -> PendingSessionEvent<N> {
1057    // Add extra protocols to the hello message
1058    extra_handlers.retain(|handler| hello.try_add_protocol(handler.protocol()).is_ok());
1059
1060    // conduct the p2p rlpx handshake and return the rlpx authenticated stream
1061    let (mut p2p_stream, their_hello) = match stream.handshake(hello).await {
1062        Ok(stream_res) => stream_res,
1063        Err(err) => {
1064            return PendingSessionEvent::Disconnected {
1065                remote_addr,
1066                session_id,
1067                direction,
1068                error: Some(PendingSessionHandshakeError::Eth(err.into())),
1069            }
1070        }
1071    };
1072
1073    // if we have extra handlers, check if it must be supported by the remote
1074    if !extra_handlers.is_empty() {
1075        // ensure that no extra handlers that aren't supported are not mandatory
1076        while let Some(pos) = extra_handlers.iter().position(|handler| {
1077            p2p_stream
1078                .shared_capabilities()
1079                .ensure_matching_capability(&handler.protocol().cap)
1080                .is_err()
1081        }) {
1082            let handler = extra_handlers.remove(pos);
1083            if handler.on_unsupported_by_peer(
1084                p2p_stream.shared_capabilities(),
1085                direction,
1086                their_hello.id,
1087            ) == OnNotSupported::Disconnect
1088            {
1089                return PendingSessionEvent::Disconnected {
1090                    remote_addr,
1091                    session_id,
1092                    direction,
1093                    error: Some(PendingSessionHandshakeError::UnsupportedExtraCapability),
1094                };
1095            }
1096        }
1097    }
1098
1099    // Ensure we negotiated mandatory eth protocol
1100    let eth_version = match p2p_stream.shared_capabilities().eth_version() {
1101        Ok(version) => version,
1102        Err(err) => {
1103            return PendingSessionEvent::Disconnected {
1104                remote_addr,
1105                session_id,
1106                direction,
1107                error: Some(PendingSessionHandshakeError::Eth(err.into())),
1108            }
1109        }
1110    };
1111
1112    // Before trying status handshake, set up the version to negotiated shared version
1113    status.set_eth_version(eth_version);
1114
1115    let (conn, their_status) = if p2p_stream.shared_capabilities().len() == 1 {
1116        // if the shared caps are 1, we know both support the eth version
1117        // if the hello handshake was successful we can try status handshake
1118
1119        // perform the eth protocol handshake
1120        match handshake
1121            .handshake(&mut p2p_stream, status, fork_filter.clone(), HANDSHAKE_TIMEOUT)
1122            .await
1123        {
1124            Ok(their_status) => {
1125                let eth_stream = EthStream::new(eth_version, p2p_stream);
1126                (eth_stream.into(), their_status)
1127            }
1128            Err(err) => {
1129                return PendingSessionEvent::Disconnected {
1130                    remote_addr,
1131                    session_id,
1132                    direction,
1133                    error: Some(PendingSessionHandshakeError::Eth(err)),
1134                }
1135            }
1136        }
1137    } else {
1138        // Multiplex the stream with the extra protocols
1139        let mut multiplex_stream = RlpxProtocolMultiplexer::new(p2p_stream);
1140
1141        // install additional handlers
1142        for handler in extra_handlers.into_iter() {
1143            let cap = handler.protocol().cap;
1144            let remote_peer_id = their_hello.id;
1145
1146            multiplex_stream
1147                .install_protocol(&cap, move |conn| {
1148                    handler.into_connection(direction, remote_peer_id, conn)
1149                })
1150                .ok();
1151        }
1152
1153        let (multiplex_stream, their_status) = match multiplex_stream
1154            .into_eth_satellite_stream(status, fork_filter, handshake)
1155            .await
1156        {
1157            Ok((multiplex_stream, their_status)) => (multiplex_stream, their_status),
1158            Err(err) => {
1159                return PendingSessionEvent::Disconnected {
1160                    remote_addr,
1161                    session_id,
1162                    direction,
1163                    error: Some(PendingSessionHandshakeError::Eth(err)),
1164                }
1165            }
1166        };
1167
1168        (multiplex_stream.into(), their_status)
1169    };
1170
1171    PendingSessionEvent::Established {
1172        session_id,
1173        remote_addr,
1174        local_addr,
1175        peer_id: their_hello.id,
1176        capabilities: Arc::new(Capabilities::from(their_hello.capabilities)),
1177        status: Arc::new(their_status),
1178        conn,
1179        direction,
1180        client_id: their_hello.client_version,
1181    }
1182}