-
Notifications
You must be signed in to change notification settings - Fork 999
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat(connection-limit): set bypass rules for connections #5720
base: master
Are you sure you want to change the base?
Changes from 1 commit
d78a671
05c757f
cbf009c
399fd34
639c3d2
227b062
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -67,6 +67,7 @@ use libp2p_swarm::{ | |
/// ``` | ||
pub struct Behaviour { | ||
limits: ConnectionLimits, | ||
bypass_rules: BypassRules, | ||
|
||
pending_inbound_connections: HashSet<ConnectionId>, | ||
pending_outbound_connections: HashSet<ConnectionId>, | ||
|
@@ -76,9 +77,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 +94,10 @@ impl Behaviour { | |
pub fn limits_mut(&mut self) -> &mut ConnectionLimits { | ||
&mut self.limits | ||
} | ||
|
||
pub fn bypass_rules_mut(&mut self) -> &mut BypassRules { | ||
&mut self.bypass_rules | ||
} | ||
} | ||
|
||
fn check_limit(limit: Option<u32>, current: usize, kind: Kind) -> Result<(), ConnectionDenied> { | ||
|
@@ -208,22 +214,57 @@ impl ConnectionLimits { | |
} | ||
} | ||
|
||
#[derive(Debug, Clone, Default)] | ||
pub struct BypassRules { | ||
by_peer_id: HashSet<PeerId>, | ||
by_multiaddr: HashSet<Multiaddr>, | ||
} | ||
impl BypassRules { | ||
pub fn new(peer_ids: HashSet<PeerId>, remote_multiaddrs: HashSet<Multiaddr>) -> Self { | ||
Self { | ||
by_peer_id: peer_ids, | ||
by_multiaddr: remote_multiaddrs, | ||
} | ||
} | ||
pub fn bypass_peer_id(&mut self, peer_id: &PeerId) { | ||
self.by_peer_id.insert(*peer_id); | ||
} | ||
pub fn remove_peer_id(&mut self, peer_id: &PeerId) { | ||
self.by_peer_id.remove(peer_id); | ||
} | ||
pub fn bypass_multiaddr(&mut self, multiaddr: Multiaddr) { | ||
self.by_multiaddr.insert(multiaddr); | ||
} | ||
pub fn remove_multiaddr(&mut self, multiaddr: &Multiaddr) { | ||
self.by_multiaddr.remove(multiaddr); | ||
} | ||
pub fn is_peer_bypassed(&self, peer: &PeerId) -> bool { | ||
self.by_peer_id.contains(peer) | ||
} | ||
pub fn is_addr_bypassed(&self, addr: &Multiaddr) -> bool { | ||
self.by_multiaddr.contains(addr) | ||
} | ||
} | ||
|
||
impl NetworkBehaviour for Behaviour { | ||
type ConnectionHandler = dummy::ConnectionHandler; | ||
type ToSwarm = Infallible; | ||
|
||
fn handle_pending_inbound_connection( | ||
&mut self, | ||
connection_id: ConnectionId, | ||
_: &Multiaddr, | ||
_: &Multiaddr, | ||
local_addr: &Multiaddr, | ||
remote_addr: &Multiaddr, | ||
) -> Result<(), ConnectionDenied> { | ||
check_limit( | ||
self.limits.max_pending_incoming, | ||
self.pending_inbound_connections.len(), | ||
Kind::PendingIncoming, | ||
)?; | ||
|
||
if !(self.bypass_rules.is_addr_bypassed(local_addr) | ||
|| self.bypass_rules.is_addr_bypassed(remote_addr)) | ||
{ | ||
check_limit( | ||
self.limits.max_pending_incoming, | ||
self.pending_inbound_connections.len(), | ||
Kind::PendingIncoming, | ||
)?; | ||
} | ||
self.pending_inbound_connections.insert(connection_id); | ||
|
||
Ok(()) | ||
|
@@ -233,46 +274,60 @@ impl NetworkBehaviour for Behaviour { | |
&mut self, | ||
connection_id: ConnectionId, | ||
peer: PeerId, | ||
_: &Multiaddr, | ||
_: &Multiaddr, | ||
local_addr: &Multiaddr, | ||
remote_addr: &Multiaddr, | ||
) -> Result<THandler<Self>, ConnectionDenied> { | ||
self.pending_inbound_connections.remove(&connection_id); | ||
|
||
check_limit( | ||
self.limits.max_established_incoming, | ||
self.established_inbound_connections.len(), | ||
Kind::EstablishedIncoming, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_per_peer, | ||
self.established_per_peer | ||
.get(&peer) | ||
.map(|connections| connections.len()) | ||
.unwrap_or(0), | ||
Kind::EstablishedPerPeer, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_total, | ||
self.established_inbound_connections.len() | ||
+ self.established_outbound_connections.len(), | ||
Kind::EstablishedTotal, | ||
)?; | ||
|
||
if !(self.bypass_rules.is_addr_bypassed(local_addr) | ||
|| self.bypass_rules.is_addr_bypassed(remote_addr) | ||
|| self.bypass_rules.is_peer_bypassed(&peer)) | ||
{ | ||
check_limit( | ||
self.limits.max_established_incoming, | ||
self.established_inbound_connections.len(), | ||
Kind::EstablishedIncoming, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_per_peer, | ||
self.established_per_peer | ||
.get(&peer) | ||
.map(|connections| connections.len()) | ||
.unwrap_or(0), | ||
Kind::EstablishedPerPeer, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_total, | ||
self.established_inbound_connections.len() | ||
+ self.established_outbound_connections.len(), | ||
Kind::EstablishedTotal, | ||
)?; | ||
} | ||
Ok(dummy::ConnectionHandler) | ||
} | ||
|
||
fn handle_pending_outbound_connection( | ||
&mut self, | ||
connection_id: ConnectionId, | ||
_: Option<PeerId>, | ||
_: &[Multiaddr], | ||
maybe_peer: Option<PeerId>, | ||
addresses: &[Multiaddr], | ||
_: Endpoint, | ||
) -> Result<Vec<Multiaddr>, ConnectionDenied> { | ||
check_limit( | ||
self.limits.max_pending_outgoing, | ||
self.pending_outbound_connections.len(), | ||
Kind::PendingOutgoing, | ||
)?; | ||
let mut is_bypassed = false; | ||
if let Some(peer) = maybe_peer { | ||
is_bypassed = self.bypass_rules.is_peer_bypassed(&peer) | ||
} | ||
is_bypassed = is_bypassed | ||
|| addresses | ||
.iter() | ||
.any(|addr| self.bypass_rules.is_addr_bypassed(addr)); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It can still happen that if there are multiple addresses but not all are in the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm not sure because the implementation will get complex very soon if we want to implement AND/OR logic. If we want to make it as simple as possible we should only check against PeerID and use a closure for more fine-grain control, what do you think? |
||
if !is_bypassed { | ||
check_limit( | ||
self.limits.max_pending_outgoing, | ||
self.pending_outbound_connections.len(), | ||
Kind::PendingOutgoing, | ||
)?; | ||
} | ||
|
||
self.pending_outbound_connections.insert(connection_id); | ||
|
||
|
@@ -283,31 +338,33 @@ impl NetworkBehaviour for Behaviour { | |
&mut self, | ||
connection_id: ConnectionId, | ||
peer: PeerId, | ||
_: &Multiaddr, | ||
addr: &Multiaddr, | ||
_: Endpoint, | ||
_: PortUse, | ||
) -> Result<THandler<Self>, ConnectionDenied> { | ||
self.pending_outbound_connections.remove(&connection_id); | ||
|
||
check_limit( | ||
self.limits.max_established_outgoing, | ||
self.established_outbound_connections.len(), | ||
Kind::EstablishedOutgoing, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_per_peer, | ||
self.established_per_peer | ||
.get(&peer) | ||
.map(|connections| connections.len()) | ||
.unwrap_or(0), | ||
Kind::EstablishedPerPeer, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_total, | ||
self.established_inbound_connections.len() | ||
+ self.established_outbound_connections.len(), | ||
Kind::EstablishedTotal, | ||
)?; | ||
if !(self.bypass_rules.is_peer_bypassed(&peer) || self.bypass_rules.is_addr_bypassed(addr)) | ||
{ | ||
check_limit( | ||
self.limits.max_established_outgoing, | ||
self.established_outbound_connections.len(), | ||
Kind::EstablishedOutgoing, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_per_peer, | ||
self.established_per_peer | ||
.get(&peer) | ||
.map(|connections| connections.len()) | ||
.unwrap_or(0), | ||
Kind::EstablishedPerPeer, | ||
)?; | ||
check_limit( | ||
self.limits.max_established_total, | ||
self.established_inbound_connections.len() | ||
+ self.established_outbound_connections.len(), | ||
Kind::EstablishedTotal, | ||
)?; | ||
} | ||
|
||
Ok(dummy::ConnectionHandler) | ||
} | ||
|
@@ -544,13 +601,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(), | ||
} | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
local_addr
is our own address at which we accept the connection. I don't think any bypass rules should apply to it?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's what "allow connections from specific listeners" mean, at least that's what I believe it will do. For example you set a limit for connections from the Internet but not for intranet.