diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md index a84ff4deb32..9ebed7a5a4b 100644 --- a/misc/connection-limits/CHANGELOG.md +++ b/misc/connection-limits/CHANGELOG.md @@ -1,3 +1,9 @@ +## 0.5.1 + +- Allow setting Peer IDs and custom filter for bypassing limit check. + Connections to specific peers OR matching the custom filter won't be counted toward limits. + See [PR 5720](https://github.com/libp2p/rust-libp2p/pull/5720). + ## 0.5.0 - Deprecate `void` crate. diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs index c8df5be5653..47c0eb61a59 100644 --- a/misc/connection-limits/src/lib.rs +++ b/misc/connection-limits/src/lib.rs @@ -33,6 +33,12 @@ use libp2p_swarm::{ THandlerInEvent, THandlerOutEvent, ToSwarm, }; +type BypassFn = fn( + remote_peer: Option<&PeerId>, + remote_addresses: &[Multiaddr], + local_addresses: Option<&Multiaddr>, +) -> bool; + /// A [`NetworkBehaviour`] that enforces a set of [`ConnectionLimits`]. /// /// For these limits to take effect, this needs to be composed @@ -46,6 +52,9 @@ use libp2p_swarm::{ /// contain a [`ConnectionDenied`] type that can be downcast to [`Exceeded`] error if (and only if) /// **this** behaviour denied the connection. /// +/// You can also set Peer IDs and Multiaddrs that bypass the said limit. Connections that +/// match the bypass rules will not be checked against limits. +/// /// If you employ multiple [`NetworkBehaviour`]s that manage connections, /// it may also be a different error. /// @@ -67,6 +76,7 @@ use libp2p_swarm::{ /// ``` pub struct Behaviour { limits: ConnectionLimits, + bypass_rules: BypassRules, pending_inbound_connections: HashSet, pending_outbound_connections: HashSet, @@ -76,9 +86,10 @@ pub struct Behaviour { } impl Behaviour { - pub fn new(limits: ConnectionLimits) -> Self { + pub fn new(limits: ConnectionLimits, bypass_rules: BypassRules) -> Self { Self { limits, + bypass_rules, pending_inbound_connections: Default::default(), pending_outbound_connections: Default::default(), established_inbound_connections: Default::default(), @@ -92,6 +103,11 @@ impl Behaviour { pub fn limits_mut(&mut self) -> &mut ConnectionLimits { &mut self.limits } + + /// Returns a mutable reference to [`BypassRules`]. + pub fn bypass_rules_mut(&mut self) -> &mut BypassRules { + &mut self.bypass_rules + } } fn check_limit(limit: Option, current: usize, kind: Kind) -> Result<(), ConnectionDenied> { @@ -208,6 +224,62 @@ impl ConnectionLimits { } } +/// A set of rules that allows bypass of limits. +/// Connections that match the rules won't be counted toward limits. +#[derive(Debug, Clone, Default)] +pub struct BypassRules { + /// Peer IDs that bypass limit check, regardless of inbound or outbound. + by_peer_id: HashSet, + /// The custom bypass rule. Note that the rules are applied in OR logic, + /// allow connections that matches ANY of the rules. + /// Cannot be mutated at runtime. + custom: Option, +} +impl BypassRules { + pub fn new(peer_ids: HashSet) -> Self { + Self { + by_peer_id: peer_ids, + custom: None, + } + } + /// Add the peer to bypass list. + pub fn bypass_peer_id(&mut self, peer_id: &PeerId) { + self.by_peer_id.insert(*peer_id); + } + /// Remove the peer from bypass list. + pub fn remove_peer_id(&mut self, peer_id: &PeerId) { + self.by_peer_id.remove(peer_id); + } + /// Set the custom bypass rule. Note that the rules are applied in OR logic, + /// allow connections that matches ANY of the rules. + /// This rule cannot be mutated at runtime, or capture any environment. + pub fn set_custom_bypass(mut self, bypass_fn: BypassFn) -> Self { + let _ = self.custom.insert(bypass_fn); + self + } + /// Remove the address to bypass list. + pub fn remove_custom_bypass(&mut self) { + self.custom.take(); + } + /// Whether the peer is in the bypass list. + pub fn is_peer_bypassed(&self, peer: &PeerId) -> bool { + self.by_peer_id.contains(peer) + } + /// Whether the connection is bypassed, the rules are applied in OR logic. + pub fn is_bypassed( + &self, + remote_peer: Option<&PeerId>, + remote_addresses: &[Multiaddr], + local_addresses: Option<&Multiaddr>, + ) -> bool { + let is_peer_id_bypassed = remote_peer.is_some_and(|peer| self.by_peer_id.contains(peer)); + let is_connection_bypassed = self + .custom + .is_some_and(|f| f(remote_peer, remote_addresses, local_addresses)); + is_peer_id_bypassed || is_connection_bypassed + } +} + impl NetworkBehaviour for Behaviour { type ConnectionHandler = dummy::ConnectionHandler; type ToSwarm = Infallible; @@ -215,9 +287,16 @@ impl NetworkBehaviour for Behaviour { fn handle_pending_inbound_connection( &mut self, connection_id: ConnectionId, - _: &Multiaddr, - _: &Multiaddr, + local_addr: &Multiaddr, + remote_addr: &Multiaddr, ) -> Result<(), ConnectionDenied> { + if self + .bypass_rules + .is_bypassed(None, std::slice::from_ref(remote_addr), Some(local_addr)) + { + return Ok(()); + } + check_limit( self.limits.max_pending_incoming, self.pending_inbound_connections.len(), @@ -233,11 +312,19 @@ impl NetworkBehaviour for Behaviour { &mut self, connection_id: ConnectionId, peer: PeerId, - _: &Multiaddr, - _: &Multiaddr, + local_addr: &Multiaddr, + remote_addr: &Multiaddr, ) -> Result, ConnectionDenied> { self.pending_inbound_connections.remove(&connection_id); + if self.bypass_rules.is_bypassed( + Some(&peer), + std::slice::from_ref(remote_addr), + Some(local_addr), + ) { + return Ok(dummy::ConnectionHandler); + } + check_limit( self.limits.max_established_incoming, self.established_inbound_connections.len(), @@ -264,10 +351,17 @@ impl NetworkBehaviour for Behaviour { fn handle_pending_outbound_connection( &mut self, connection_id: ConnectionId, - _: Option, - _: &[Multiaddr], + maybe_peer: Option, + addresses: &[Multiaddr], _: Endpoint, ) -> Result, ConnectionDenied> { + if self + .bypass_rules + .is_bypassed(maybe_peer.as_ref(), addresses, None) + { + return Ok(vec![]); + } + check_limit( self.limits.max_pending_outgoing, self.pending_outbound_connections.len(), @@ -283,11 +377,17 @@ impl NetworkBehaviour for Behaviour { &mut self, connection_id: ConnectionId, peer: PeerId, - _: &Multiaddr, + addr: &Multiaddr, _: Endpoint, _: PortUse, ) -> Result, ConnectionDenied> { self.pending_outbound_connections.remove(&connection_id); + if self + .bypass_rules + .is_bypassed(Some(&peer), std::slice::from_ref(addr), None) + { + return Ok(dummy::ConnectionHandler); + } check_limit( self.limits.max_established_outgoing, @@ -544,13 +644,13 @@ mod tests { impl Behaviour { fn new(limits: ConnectionLimits) -> Self { Self { - limits: super::Behaviour::new(limits), + limits: super::Behaviour::new(limits, Default::default()), connection_denier: None.into(), } } fn new_with_connection_denier(limits: ConnectionLimits) -> Self { Self { - limits: super::Behaviour::new(limits), + limits: super::Behaviour::new(limits, Default::default()), connection_denier: Some(ConnectionDenier {}).into(), } }