Skip to content

Commit

Permalink
refactor(connection-limits): make check_limit a free-function
Browse files Browse the repository at this point in the history
Pull-Request: #4958.
  • Loading branch information
zhiqiangxu authored Nov 30, 2023
1 parent 7b7cf5f commit ec2258e
Showing 1 changed file with 16 additions and 21 deletions.
37 changes: 16 additions & 21 deletions misc/connection-limits/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -79,22 +79,17 @@ impl Behaviour {
established_per_peer: Default::default(),
}
}
}

fn check_limit(
&mut self,
limit: Option<u32>,
current: usize,
kind: Kind,
) -> Result<(), ConnectionDenied> {
let limit = limit.unwrap_or(u32::MAX);
let current = current as u32;
fn check_limit(limit: Option<u32>, current: usize, kind: Kind) -> Result<(), ConnectionDenied> {
let limit = limit.unwrap_or(u32::MAX);
let current = current as u32;

if current >= limit {
return Err(ConnectionDenied::new(Exceeded { limit, kind }));
}

Ok(())
if current >= limit {
return Err(ConnectionDenied::new(Exceeded { limit, kind }));
}

Ok(())
}

/// A connection limit has been exceeded.
Expand Down Expand Up @@ -210,7 +205,7 @@ impl NetworkBehaviour for Behaviour {
_: &Multiaddr,
_: &Multiaddr,
) -> Result<(), ConnectionDenied> {
self.check_limit(
check_limit(
self.limits.max_pending_incoming,
self.pending_inbound_connections.len(),
Kind::PendingIncoming,
Expand All @@ -230,20 +225,20 @@ impl NetworkBehaviour for Behaviour {
) -> Result<THandler<Self>, ConnectionDenied> {
self.pending_inbound_connections.remove(&connection_id);

self.check_limit(
check_limit(
self.limits.max_established_incoming,
self.established_inbound_connections.len(),
Kind::EstablishedIncoming,
)?;
self.check_limit(
check_limit(
self.limits.max_established_per_peer,
self.established_per_peer
.get(&peer)
.map(|connections| connections.len())
.unwrap_or(0),
Kind::EstablishedPerPeer,
)?;
self.check_limit(
check_limit(
self.limits.max_established_total,
self.established_inbound_connections.len()
+ self.established_outbound_connections.len(),
Expand All @@ -260,7 +255,7 @@ impl NetworkBehaviour for Behaviour {
_: &[Multiaddr],
_: Endpoint,
) -> Result<Vec<Multiaddr>, ConnectionDenied> {
self.check_limit(
check_limit(
self.limits.max_pending_outgoing,
self.pending_outbound_connections.len(),
Kind::PendingOutgoing,
Expand All @@ -280,20 +275,20 @@ impl NetworkBehaviour for Behaviour {
) -> Result<THandler<Self>, ConnectionDenied> {
self.pending_outbound_connections.remove(&connection_id);

self.check_limit(
check_limit(
self.limits.max_established_outgoing,
self.established_outbound_connections.len(),
Kind::EstablishedOutgoing,
)?;
self.check_limit(
check_limit(
self.limits.max_established_per_peer,
self.established_per_peer
.get(&peer)
.map(|connections| connections.len())
.unwrap_or(0),
Kind::EstablishedPerPeer,
)?;
self.check_limit(
check_limit(
self.limits.max_established_total,
self.established_inbound_connections.len()
+ self.established_outbound_connections.len(),
Expand Down

0 comments on commit ec2258e

Please sign in to comment.