Cap max graceful disconnects (#6117)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
This commit is contained in:
Darex
2024-01-18 23:50:48 +05:30
committed by GitHub
parent 63adf39290
commit a3e56eaffa

View File

@@ -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 {