From a3e56eaffa89b7be8ac3eb2c8482e8025ef9418f Mon Sep 17 00:00:00 2001
From: Darex <1010adigupta@gmail.com>
Date: Thu, 18 Jan 2024 23:50:48 +0530
Subject: [PATCH] Cap max graceful disconnects (#6117)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/net/network/src/session/mod.rs | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/crates/net/network/src/session/mod.rs b/crates/net/network/src/session/mod.rs
index 9762a1d6a3..3ee92f0a4b 100644
--- a/crates/net/network/src/session/mod.rs
+++ b/crates/net/network/src/session/mod.rs
@@ -53,6 +53,9 @@ pub use handle::{
 use reth_eth_wire::multiplex::RlpxProtocolMultiplexer;
 pub use reth_network_api::{Direction, PeerInfo};
 
+/// Maximum allowed graceful disconnects at a time.
+const MAX_GRACEFUL_DISCONNECTS: usize = 15;
+
 /// Internal identifier for active sessions.
 #[derive(Debug, Clone, Copy, PartialOrd, PartialEq, Eq, Hash)]
 pub struct SessionId(usize);
@@ -110,6 +113,8 @@ pub struct SessionManager {
     bandwidth_meter: BandwidthMeter,
     /// Metrics for the session manager.
     metrics: SessionManagerMetrics,
+    /// Tracks the number of active graceful disconnects for incoming connections.
+    graceful_disconnects_counter: GracefulDisconnects,
 }
 
 // === impl SessionManager ===
@@ -151,6 +156,7 @@ impl SessionManager {
             bandwidth_meter,
             extra_protocols,
             metrics: Default::default(),
+            graceful_disconnects_counter: Default::default(),
         }
     }
 
@@ -310,12 +316,18 @@ impl SessionManager {
         stream: TcpStream,
         reason: DisconnectReason,
     ) {
+        let counter = self.graceful_disconnects_counter.clone();
+        if counter.exceeds_limit() {
+            // simply drop the connection if there are too many active disconnects already
+            return
+        }
         let secret_key = self.secret_key;
 
         self.spawn(async move {
             if let Ok(stream) = get_eciess_stream(stream, secret_key, Direction::Incoming).await {
                 let _ = UnauthedP2PStream::new(stream).send_disconnect(reason).await;
             }
+            drop(counter)
         });
     }
 
@@ -624,6 +636,18 @@ impl SessionManager {
     }
 }
 
+/// Keep track of graceful disconnects for incoming connections.
+#[derive(Debug, Clone, Default)]
+struct GracefulDisconnects(Arc<()>);
+
+impl GracefulDisconnects {
+    /// Returns true if the number of graceful disconnects exceeds the limit
+    /// [MAX_GRACEFUL_DISCONNECTS]
+    fn exceeds_limit(&self) -> bool {
+        Arc::strong_count(&self.0) > MAX_GRACEFUL_DISCONNECTS
+    }
+}
+
 /// Events produced by the [`SessionManager`]
 #[derive(Debug)]
 pub enum SessionEvent {