From d78a671803c67b7018f98b70772a6749cefa99e4 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Fri, 6 Dec 2024 15:56:12 +0800
Subject: [PATCH 01/12] set bypass rules

---
 misc/connection-limits/src/lib.rs | 179 ++++++++++++++++++++----------
 1 file changed, 118 insertions(+), 61 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index c8df5be5653..59afe9be58d 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -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,6 +214,38 @@ 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;
@@ -215,15 +253,18 @@ impl NetworkBehaviour for Behaviour {
     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));
+        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(),
             }
         }

From 639c3d21bf99d8445c10162f3021fc2c38fa88f7 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Fri, 3 Jan 2025 14:54:22 +0800
Subject: [PATCH 02/12] changelog and documentation

---
 misc/connection-limits/CHANGELOG.md |  3 +++
 misc/connection-limits/src/lib.rs   | 17 +++++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md
index f2722b3745a..f8a21fa2e80 100644
--- a/misc/connection-limits/CHANGELOG.md
+++ b/misc/connection-limits/CHANGELOG.md
@@ -3,6 +3,9 @@
 - Deprecate `void` crate.
   See [PR 5676](https://github.com/libp2p/rust-libp2p/pull/5676).
 
+- Allow setting Peer IDs and Multiaddrs that bypass limit check.
+  See [PR 5720](https://github.com/libp2p/rust-libp2p/pull/5720).
+
 ## 0.4.0
 
 <!-- Update to libp2p-swarm v0.45.0 -->
diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index 59afe9be58d..58b53a86d76 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -46,6 +46,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.
 ///
@@ -95,6 +98,7 @@ impl Behaviour {
         &mut self.limits
     }
 
+    /// Returns a mutable reference to [`BypassRules`].
     pub fn bypass_rules_mut(&mut self) -> &mut BypassRules {
         &mut self.bypass_rules
     }
@@ -214,9 +218,12 @@ impl ConnectionLimits {
     }
 }
 
+/// A set of rules that allows bypass of limits.
 #[derive(Debug, Clone, Default)]
 pub struct BypassRules {
+    /// Peer IDs that bypass limit check, regardless of inbound or outbound.
     by_peer_id: HashSet<PeerId>,
+    /// Addresses that bypass limit check, regardless of inbound or outbound.
     by_multiaddr: HashSet<Multiaddr>,
 }
 impl BypassRules {
@@ -226,21 +233,27 @@ impl BypassRules {
             by_multiaddr: remote_multiaddrs,
         }
     }
+    /// 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);
     }
-    pub fn bypass_multiaddr(&mut self, multiaddr: Multiaddr) {
-        self.by_multiaddr.insert(multiaddr);
+    /// Add the address to bypass list.
+    pub fn bypass_multiaddr(&mut self, multiaddr: &Multiaddr) {
+        self.by_multiaddr.insert(multiaddr.clone());
     }
+    /// Remove the address to bypass list.
     pub fn remove_multiaddr(&mut self, multiaddr: &Multiaddr) {
         self.by_multiaddr.remove(multiaddr);
     }
+    /// 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 address is in the bypass list.
     pub fn is_addr_bypassed(&self, addr: &Multiaddr) -> bool {
         self.by_multiaddr.contains(addr)
     }

From 227b0620b9ec31215ed43c4946343c05757d00c5 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Fri, 10 Jan 2025 12:13:09 +0800
Subject: [PATCH 03/12] replace multiaddr bypass with filter

---
 misc/connection-limits/CHANGELOG.md |  3 +-
 misc/connection-limits/src/lib.rs   | 76 ++++++++++++++++++-----------
 2 files changed, 50 insertions(+), 29 deletions(-)

diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md
index f8a21fa2e80..b5bb33d978e 100644
--- a/misc/connection-limits/CHANGELOG.md
+++ b/misc/connection-limits/CHANGELOG.md
@@ -3,7 +3,8 @@
 - Deprecate `void` crate.
   See [PR 5676](https://github.com/libp2p/rust-libp2p/pull/5676).
 
-- Allow setting Peer IDs and Multiaddrs that bypass limit check.
+- 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.4.0
diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index 58b53a86d76..aedfca10ca3 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
@@ -219,18 +225,21 @@ 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<PeerId>,
-    /// Addresses that bypass limit check, regardless of inbound or outbound.
-    by_multiaddr: HashSet<Multiaddr>,
+    /// 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<BypassFn>,
 }
 impl BypassRules {
-    pub fn new(peer_ids: HashSet<PeerId>, remote_multiaddrs: HashSet<Multiaddr>) -> Self {
+    pub fn new(peer_ids: HashSet<PeerId>) -> Self {
         Self {
             by_peer_id: peer_ids,
-            by_multiaddr: remote_multiaddrs,
+            custom: None,
         }
     }
     /// Add the peer to bypass list.
@@ -241,21 +250,33 @@ impl BypassRules {
     pub fn remove_peer_id(&mut self, peer_id: &PeerId) {
         self.by_peer_id.remove(peer_id);
     }
-    /// Add the address to bypass list.
-    pub fn bypass_multiaddr(&mut self, multiaddr: &Multiaddr) {
-        self.by_multiaddr.insert(multiaddr.clone());
+    /// 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_multiaddr(&mut self, multiaddr: &Multiaddr) {
-        self.by_multiaddr.remove(multiaddr);
+    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 address is in the bypass list.
-    pub fn is_addr_bypassed(&self, addr: &Multiaddr) -> bool {
-        self.by_multiaddr.contains(addr)
+    /// 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
     }
 }
 
@@ -269,8 +290,9 @@ impl NetworkBehaviour for Behaviour {
         local_addr: &Multiaddr,
         remote_addr: &Multiaddr,
     ) -> Result<(), ConnectionDenied> {
-        if !(self.bypass_rules.is_addr_bypassed(local_addr)
-            || self.bypass_rules.is_addr_bypassed(remote_addr))
+        if !self
+            .bypass_rules
+            .is_bypassed(None, std::slice::from_ref(remote_addr), Some(local_addr))
         {
             check_limit(
                 self.limits.max_pending_incoming,
@@ -292,10 +314,11 @@ impl NetworkBehaviour for Behaviour {
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_inbound_connections.remove(&connection_id);
 
-        if !(self.bypass_rules.is_addr_bypassed(local_addr)
-            || self.bypass_rules.is_addr_bypassed(remote_addr)
-            || self.bypass_rules.is_peer_bypassed(&peer))
-        {
+        if !self.bypass_rules.is_bypassed(
+            Some(&peer),
+            std::slice::from_ref(remote_addr),
+            Some(local_addr),
+        ) {
             check_limit(
                 self.limits.max_established_incoming,
                 self.established_inbound_connections.len(),
@@ -326,15 +349,10 @@ impl NetworkBehaviour for Behaviour {
         addresses: &[Multiaddr],
         _: Endpoint,
     ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
-        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));
-        if !is_bypassed {
+        if !self
+            .bypass_rules
+            .is_bypassed(maybe_peer.as_ref(), addresses, None)
+        {
             check_limit(
                 self.limits.max_pending_outgoing,
                 self.pending_outbound_connections.len(),
@@ -356,7 +374,9 @@ impl NetworkBehaviour for Behaviour {
         _: PortUse,
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_outbound_connections.remove(&connection_id);
-        if !(self.bypass_rules.is_peer_bypassed(&peer) || self.bypass_rules.is_addr_bypassed(addr))
+        if !self
+            .bypass_rules
+            .is_bypassed(Some(&peer), std::slice::from_ref(addr), None)
         {
             check_limit(
                 self.limits.max_established_outgoing,

From efa83b125be83c63e730c8411230c5f7231f45d8 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Thu, 23 Jan 2025 20:42:48 +0800
Subject: [PATCH 04/12] reduce diff

fix: bypassed pending connections still take up slots
---
 misc/connection-limits/src/lib.rs | 108 ++++++++++++++++--------------
 1 file changed, 59 insertions(+), 49 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index aedfca10ca3..ae73a9d26a0 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -294,12 +294,15 @@ impl NetworkBehaviour for Behaviour {
             .bypass_rules
             .is_bypassed(None, std::slice::from_ref(remote_addr), Some(local_addr))
         {
-            check_limit(
-                self.limits.max_pending_incoming,
-                self.pending_inbound_connections.len(),
-                Kind::PendingIncoming,
-            )?;
+            return Ok(());
         }
+
+        check_limit(
+            self.limits.max_pending_incoming,
+            self.pending_inbound_connections.len(),
+            Kind::PendingIncoming,
+        )?;
+
         self.pending_inbound_connections.insert(connection_id);
 
         Ok(())
@@ -314,31 +317,34 @@ impl NetworkBehaviour for Behaviour {
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_inbound_connections.remove(&connection_id);
 
-        if !self.bypass_rules.is_bypassed(
+        if self.bypass_rules.is_bypassed(
             Some(&peer),
             std::slice::from_ref(remote_addr),
             Some(local_addr),
         ) {
-            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,
-            )?;
+            return Ok(dummy::ConnectionHandler);
         }
+
+        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)
     }
 
@@ -353,13 +359,15 @@ impl NetworkBehaviour for Behaviour {
             .bypass_rules
             .is_bypassed(maybe_peer.as_ref(), addresses, None)
         {
-            check_limit(
-                self.limits.max_pending_outgoing,
-                self.pending_outbound_connections.len(),
-                Kind::PendingOutgoing,
-            )?;
+            return Ok(vec![]);
         }
 
+        check_limit(
+            self.limits.max_pending_outgoing,
+            self.pending_outbound_connections.len(),
+            Kind::PendingOutgoing,
+        )?;
+
         self.pending_outbound_connections.insert(connection_id);
 
         Ok(vec![])
@@ -378,27 +386,29 @@ impl NetworkBehaviour for Behaviour {
             .bypass_rules
             .is_bypassed(Some(&peer), std::slice::from_ref(addr), None)
         {
-            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,
-            )?;
+            return Ok(dummy::ConnectionHandler);
         }
 
+        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)
     }
 

From 47b81c3f4326b5a08a4ab5895df569cd247938f9 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Thu, 23 Jan 2025 20:49:54 +0800
Subject: [PATCH 05/12] fix incorrect condition

---
 misc/connection-limits/src/lib.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index ae73a9d26a0..47c0eb61a59 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -290,7 +290,7 @@ impl NetworkBehaviour for Behaviour {
         local_addr: &Multiaddr,
         remote_addr: &Multiaddr,
     ) -> Result<(), ConnectionDenied> {
-        if !self
+        if self
             .bypass_rules
             .is_bypassed(None, std::slice::from_ref(remote_addr), Some(local_addr))
         {
@@ -355,7 +355,7 @@ impl NetworkBehaviour for Behaviour {
         addresses: &[Multiaddr],
         _: Endpoint,
     ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
-        if !self
+        if self
             .bypass_rules
             .is_bypassed(maybe_peer.as_ref(), addresses, None)
         {
@@ -382,7 +382,7 @@ impl NetworkBehaviour for Behaviour {
         _: PortUse,
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_outbound_connections.remove(&connection_id);
-        if !self
+        if self
             .bypass_rules
             .is_bypassed(Some(&peer), std::slice::from_ref(addr), None)
         {

From aa01b26cb7349b3cf603656428ad086ce1f27874 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Fri, 31 Jan 2025 19:28:55 +0800
Subject: [PATCH 06/12] bypass by PeerId only

---
 misc/connection-limits/CHANGELOG.md |  4 +-
 misc/connection-limits/src/lib.rs   | 82 ++++++-----------------------
 2 files changed, 17 insertions(+), 69 deletions(-)

diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md
index 9ebed7a5a4b..f6ffd00741e 100644
--- a/misc/connection-limits/CHANGELOG.md
+++ b/misc/connection-limits/CHANGELOG.md
@@ -1,7 +1,7 @@
 ## 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.
+- Allow setting Peer IDs for bypassing limit check. 
+  Connections to the specified peers won't be counted toward limits.
   See [PR 5720](https://github.com/libp2p/rust-libp2p/pull/5720).
 
 ## 0.5.0
diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index 47c0eb61a59..f375da9f552 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -33,12 +33,6 @@ 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
@@ -52,8 +46,8 @@ type BypassFn = fn(
 /// 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.
+/// You can also set Peer IDs that bypass the said limit. Connections that
+/// match the bypass rules will not be checked against or counted for limits.
 ///
 /// If you employ multiple [`NetworkBehaviour`]s that manage connections,
 /// it may also be a different error.
@@ -226,20 +220,15 @@ 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)]
+#[derive(Default)]
 pub struct BypassRules {
     /// Peer IDs that bypass limit check, regardless of inbound or outbound.
     by_peer_id: HashSet<PeerId>,
-    /// 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<BypassFn>,
 }
 impl BypassRules {
     pub fn new(peer_ids: HashSet<PeerId>) -> Self {
         Self {
             by_peer_id: peer_ids,
-            custom: None,
         }
     }
     /// Add the peer to bypass list.
@@ -250,33 +239,9 @@ impl BypassRules {
     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
+    /// Whether the connection is bypassed.
+    pub fn is_bypassed(&mut self, remote_peer: &PeerId) -> bool {
+        self.by_peer_id.contains(remote_peer)
     }
 }
 
@@ -287,16 +252,9 @@ impl NetworkBehaviour for Behaviour {
     fn handle_pending_inbound_connection(
         &mut self,
         connection_id: ConnectionId,
-        local_addr: &Multiaddr,
-        remote_addr: &Multiaddr,
+        _: &Multiaddr,
+        _: &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(),
@@ -312,16 +270,12 @@ impl NetworkBehaviour for Behaviour {
         &mut self,
         connection_id: ConnectionId,
         peer: PeerId,
-        local_addr: &Multiaddr,
-        remote_addr: &Multiaddr,
+        _: &Multiaddr,
+        _: &Multiaddr,
     ) -> Result<THandler<Self>, 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),
-        ) {
+        if self.bypass_rules.is_bypassed(&peer) {
             return Ok(dummy::ConnectionHandler);
         }
 
@@ -352,13 +306,10 @@ impl NetworkBehaviour for Behaviour {
         &mut self,
         connection_id: ConnectionId,
         maybe_peer: Option<PeerId>,
-        addresses: &[Multiaddr],
+        _: &[Multiaddr],
         _: Endpoint,
     ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
-        if self
-            .bypass_rules
-            .is_bypassed(maybe_peer.as_ref(), addresses, None)
-        {
+        if maybe_peer.is_some_and(|peer| self.bypass_rules.is_bypassed(&peer)) {
             return Ok(vec![]);
         }
 
@@ -377,15 +328,12 @@ impl NetworkBehaviour for Behaviour {
         &mut self,
         connection_id: ConnectionId,
         peer: PeerId,
-        addr: &Multiaddr,
+        _: &Multiaddr,
         _: Endpoint,
         _: PortUse,
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_outbound_connections.remove(&connection_id);
-        if self
-            .bypass_rules
-            .is_bypassed(Some(&peer), std::slice::from_ref(addr), None)
-        {
+        if self.bypass_rules.is_bypassed(&peer) {
             return Ok(dummy::ConnectionHandler);
         }
 

From 092f3136c7e53a006ac025063341b81c5d5d8e8b Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Sat, 1 Feb 2025 15:44:59 +0800
Subject: [PATCH 07/12] move bypass list into Behaviour

---
 misc/connection-limits/src/lib.rs | 60 +++++++++++--------------------
 1 file changed, 20 insertions(+), 40 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index f375da9f552..26a4f4944c7 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -70,7 +70,8 @@ use libp2p_swarm::{
 /// ```
 pub struct Behaviour {
     limits: ConnectionLimits,
-    bypass_rules: BypassRules,
+    /// Peer IDs that bypass limit check, regardless of inbound or outbound.
+    bypass_peer_id: HashSet<PeerId>,
 
     pending_inbound_connections: HashSet<ConnectionId>,
     pending_outbound_connections: HashSet<ConnectionId>,
@@ -80,10 +81,10 @@ pub struct Behaviour {
 }
 
 impl Behaviour {
-    pub fn new(limits: ConnectionLimits, bypass_rules: BypassRules) -> Self {
+    pub fn new(limits: ConnectionLimits) -> Self {
         Self {
             limits,
-            bypass_rules,
+            bypass_peer_id: Default::default(),
             pending_inbound_connections: Default::default(),
             pending_outbound_connections: Default::default(),
             established_inbound_connections: Default::default(),
@@ -98,9 +99,17 @@ impl Behaviour {
         &mut self.limits
     }
 
-    /// Returns a mutable reference to [`BypassRules`].
-    pub fn bypass_rules_mut(&mut self) -> &mut BypassRules {
-        &mut self.bypass_rules
+    /// Add the peer to bypass list.
+    pub fn bypass_peer_id(&mut self, peer_id: &PeerId) {
+        self.bypass_peer_id.insert(*peer_id);
+    }
+    /// Remove the peer from bypass list.
+    pub fn remove_peer_id(&mut self, peer_id: &PeerId) {
+        self.bypass_peer_id.remove(peer_id);
+    }
+    /// Whether the connection is bypassed.
+    pub fn is_bypassed(&mut self, remote_peer: &PeerId) -> bool {
+        self.bypass_peer_id.contains(remote_peer)
     }
 }
 
@@ -218,33 +227,6 @@ impl ConnectionLimits {
     }
 }
 
-/// A set of rules that allows bypass of limits.
-/// Connections that match the rules won't be counted toward limits.
-#[derive(Default)]
-pub struct BypassRules {
-    /// Peer IDs that bypass limit check, regardless of inbound or outbound.
-    by_peer_id: HashSet<PeerId>,
-}
-impl BypassRules {
-    pub fn new(peer_ids: HashSet<PeerId>) -> Self {
-        Self {
-            by_peer_id: peer_ids,
-        }
-    }
-    /// 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);
-    }
-    /// Whether the connection is bypassed.
-    pub fn is_bypassed(&mut self, remote_peer: &PeerId) -> bool {
-        self.by_peer_id.contains(remote_peer)
-    }
-}
-
 impl NetworkBehaviour for Behaviour {
     type ConnectionHandler = dummy::ConnectionHandler;
     type ToSwarm = Infallible;
@@ -275,10 +257,9 @@ impl NetworkBehaviour for Behaviour {
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_inbound_connections.remove(&connection_id);
 
-        if self.bypass_rules.is_bypassed(&peer) {
+        if self.is_bypassed(&peer) {
             return Ok(dummy::ConnectionHandler);
         }
-
         check_limit(
             self.limits.max_established_incoming,
             self.established_inbound_connections.len(),
@@ -309,10 +290,9 @@ impl NetworkBehaviour for Behaviour {
         _: &[Multiaddr],
         _: Endpoint,
     ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
-        if maybe_peer.is_some_and(|peer| self.bypass_rules.is_bypassed(&peer)) {
+        if maybe_peer.is_some_and(|peer| self.is_bypassed(&peer)) {
             return Ok(vec![]);
         }
-
         check_limit(
             self.limits.max_pending_outgoing,
             self.pending_outbound_connections.len(),
@@ -333,7 +313,7 @@ impl NetworkBehaviour for Behaviour {
         _: PortUse,
     ) -> Result<THandler<Self>, ConnectionDenied> {
         self.pending_outbound_connections.remove(&connection_id);
-        if self.bypass_rules.is_bypassed(&peer) {
+        if self.is_bypassed(&peer) {
             return Ok(dummy::ConnectionHandler);
         }
 
@@ -592,13 +572,13 @@ mod tests {
     impl Behaviour {
         fn new(limits: ConnectionLimits) -> Self {
             Self {
-                limits: super::Behaviour::new(limits, Default::default()),
+                limits: super::Behaviour::new(limits),
                 connection_denier: None.into(),
             }
         }
         fn new_with_connection_denier(limits: ConnectionLimits) -> Self {
             Self {
-                limits: super::Behaviour::new(limits, Default::default()),
+                limits: super::Behaviour::new(limits),
                 connection_denier: Some(ConnectionDenier {}).into(),
             }
         }

From 28362123d665532bcaefd8c30e77b2887f37ef30 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Sat, 1 Feb 2025 15:46:57 +0800
Subject: [PATCH 08/12] minor version bump

---
 misc/connection-limits/CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md
index f6ffd00741e..cbc19f9fb0a 100644
--- a/misc/connection-limits/CHANGELOG.md
+++ b/misc/connection-limits/CHANGELOG.md
@@ -1,4 +1,4 @@
-## 0.5.1
+## 0.6.0
 
 - Allow setting Peer IDs for bypassing limit check. 
   Connections to the specified peers won't be counted toward limits.

From 0421f40654d119168681315b860fbd634c8425b2 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Tue, 4 Feb 2025 20:02:43 +0800
Subject: [PATCH 09/12] tests and version bump

---
 Cargo.lock                          |   2 +-
 Cargo.toml                          |   2 +-
 misc/connection-limits/CHANGELOG.md |   2 +-
 misc/connection-limits/Cargo.toml   |   2 +-
 misc/connection-limits/src/lib.rs   | 130 +++++++++++++++++++++++++---
 5 files changed, 124 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c9903798524..1b961bbba82 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2744,7 +2744,7 @@ dependencies = [
 
 [[package]]
 name = "libp2p-connection-limits"
-version = "0.5.0"
+version = "0.5.1"
 dependencies = [
  "libp2p-core",
  "libp2p-identify",
diff --git a/Cargo.toml b/Cargo.toml
index 2a8854cc4bc..24f42e9f824 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -75,7 +75,7 @@ rust-version = "1.83.0"
 libp2p = { version = "0.55.1", path = "libp2p" }
 libp2p-allow-block-list = { version = "0.5.0", path = "misc/allow-block-list" }
 libp2p-autonat = { version = "0.14.0", path = "protocols/autonat" }
-libp2p-connection-limits = { version = "0.5.0", path = "misc/connection-limits" }
+libp2p-connection-limits = { version = "0.5.1", path = "misc/connection-limits" }
 libp2p-core = { version = "0.43.0", path = "core" }
 libp2p-dcutr = { version = "0.13.0", path = "protocols/dcutr" }
 libp2p-dns = { version = "0.43.0", path = "transports/dns" }
diff --git a/misc/connection-limits/CHANGELOG.md b/misc/connection-limits/CHANGELOG.md
index cbc19f9fb0a..f6ffd00741e 100644
--- a/misc/connection-limits/CHANGELOG.md
+++ b/misc/connection-limits/CHANGELOG.md
@@ -1,4 +1,4 @@
-## 0.6.0
+## 0.5.1
 
 - Allow setting Peer IDs for bypassing limit check. 
   Connections to the specified peers won't be counted toward limits.
diff --git a/misc/connection-limits/Cargo.toml b/misc/connection-limits/Cargo.toml
index 77a9bac46ef..8b1534ce352 100644
--- a/misc/connection-limits/Cargo.toml
+++ b/misc/connection-limits/Cargo.toml
@@ -3,7 +3,7 @@ name = "libp2p-connection-limits"
 edition = "2021"
 rust-version = { workspace = true }
 description = "Connection limits for libp2p."
-version = "0.5.0"
+version = "0.5.1"
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"
 keywords = ["peer-to-peer", "libp2p", "networking"]
diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index d1b6417a662..ae1275b6bae 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -108,7 +108,7 @@ impl Behaviour {
         self.bypass_peer_id.remove(peer_id);
     }
     /// Whether the connection is bypassed.
-    pub fn is_bypassed(&mut self, remote_peer: &PeerId) -> bool {
+    pub fn is_bypassed(&self, remote_peer: &PeerId) -> bool {
         self.bypass_peer_id.contains(remote_peer)
     }
 }
@@ -413,8 +413,7 @@ mod tests {
 
     use super::*;
 
-    #[test]
-    fn max_outgoing() {
+    fn fill_outgoing() -> (Swarm<Behaviour>, Multiaddr, u32) {
         use rand::Rng;
 
         let outgoing_limit = rand::thread_rng().gen_range(1..10);
@@ -439,10 +438,15 @@ mod tests {
                 )
                 .expect("Unexpected connection limit.");
         }
+        (network, addr, outgoing_limit)
+    }
 
+    #[test]
+    fn max_outgoing() {
+        let (mut network, addr, outgoing_limit) = fill_outgoing();
         match network
             .dial(
-                DialOpts::peer_id(target)
+                DialOpts::peer_id(PeerId::random())
                     .condition(PeerCondition::Always)
                     .addresses(vec![addr])
                     .build(),
@@ -467,6 +471,52 @@ mod tests {
         );
     }
 
+    #[test]
+    fn outgoing_limit_bypass() {
+        let (mut network, addr, _) = fill_outgoing();
+        let bypassed_peer = PeerId::random();
+        network
+            .behaviour_mut()
+            .limits
+            .bypass_peer_id(&bypassed_peer);
+        assert!(network.behaviour().limits.is_bypassed(&bypassed_peer));
+        if let Err(e) = network.dial(
+            DialOpts::peer_id(bypassed_peer)
+                .addresses(vec![addr.clone()])
+                .build(),
+        ) {
+            match e {
+                DialError::Denied { cause } => {
+                    cause
+                        .downcast::<Exceeded>()
+                        .expect_err("Unexpected connection denied because of limit");
+                }
+                _ => {}
+            }
+        }
+        let not_bypassed_peer = loop {
+            let new_peer = PeerId::random();
+            if new_peer != bypassed_peer {
+                break new_peer;
+            }
+        };
+        match network
+            .dial(
+                DialOpts::peer_id(not_bypassed_peer)
+                    .addresses(vec![addr])
+                    .build(),
+            )
+            .expect_err("Unexpected dialing success.")
+        {
+            DialError::Denied { cause } => {
+                cause
+                    .downcast::<Exceeded>()
+                    .expect("connection denied because of limit");
+            }
+            e => panic!("Unexpected error: {e:?}"),
+        }
+    }
+
     #[test]
     fn max_established_incoming() {
         fn prop(Limit(limit): Limit) {
@@ -507,13 +557,64 @@ mod tests {
             });
         }
 
-        #[derive(Debug, Clone)]
-        struct Limit(u32);
+        quickcheck(prop as fn(_));
+    }
 
-        impl Arbitrary for Limit {
-            fn arbitrary(g: &mut Gen) -> Self {
-                Self(g.gen_range(1..10))
-            }
+    #[test]
+    fn bypass_established_incoming() {
+        fn prop(Limit(limit): Limit) {
+            let mut swarm1 = Swarm::new_ephemeral(|_| {
+                Behaviour::new(
+                    ConnectionLimits::default().with_max_established_incoming(Some(limit)),
+                )
+            });
+            let mut swarm2 = Swarm::new_ephemeral(|_| {
+                Behaviour::new(
+                    ConnectionLimits::default().with_max_established_incoming(Some(limit)),
+                )
+            });
+            let mut swarm3 = Swarm::new_ephemeral(|_| {
+                Behaviour::new(
+                    ConnectionLimits::default().with_max_established_incoming(Some(limit)),
+                )
+            });
+
+            let bypassed_peer_id = *swarm3.local_peer_id();
+            swarm1
+                .behaviour_mut()
+                .limits
+                .bypass_peer_id(&bypassed_peer_id);
+
+            async_std::task::block_on(async {
+                let (listen_addr, _) = swarm1.listen().with_memory_addr_external().await;
+
+                for _ in 0..limit {
+                    swarm2.connect(&mut swarm1).await;
+                }
+
+                swarm3.dial(listen_addr.clone()).unwrap();
+
+                async_std::task::spawn(swarm2.loop_on_next());
+                async_std::task::spawn(swarm3.loop_on_next());
+
+                swarm1
+                    .wait(|event| match event {
+                        SwarmEvent::ConnectionEstablished { peer_id, .. } => {
+                            (peer_id == bypassed_peer_id).then_some(())
+                        }
+                        SwarmEvent::IncomingConnectionError {
+                            error: ListenError::Denied { cause },
+                            ..
+                        } => {
+                            cause
+                                .downcast::<Exceeded>()
+                                .expect_err("Unexpected connection denied because of limit");
+                            None
+                        }
+                        _ => None,
+                    })
+                    .await;
+            });
         }
 
         quickcheck(prop as fn(_));
@@ -637,4 +738,13 @@ mod tests {
             Poll::Pending
         }
     }
+
+    #[derive(Debug, Clone)]
+    struct Limit(u32);
+
+    impl Arbitrary for Limit {
+        fn arbitrary(g: &mut Gen) -> Self {
+            Self(g.gen_range(1..10))
+        }
+    }
 }

From 7817ae0df9c2fd81de5c6ee8e564a58140d69a53 Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Tue, 4 Feb 2025 21:39:41 +0800
Subject: [PATCH 10/12] migrate test to tokio

---
 misc/connection-limits/src/lib.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index ae1275b6bae..faf7055c0d0 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -579,13 +579,14 @@ mod tests {
                 )
             });
 
+            let rt = Runtime::new().unwrap();
             let bypassed_peer_id = *swarm3.local_peer_id();
             swarm1
                 .behaviour_mut()
                 .limits
                 .bypass_peer_id(&bypassed_peer_id);
 
-            async_std::task::block_on(async {
+            rt.block_on(async {
                 let (listen_addr, _) = swarm1.listen().with_memory_addr_external().await;
 
                 for _ in 0..limit {
@@ -594,8 +595,8 @@ mod tests {
 
                 swarm3.dial(listen_addr.clone()).unwrap();
 
-                async_std::task::spawn(swarm2.loop_on_next());
-                async_std::task::spawn(swarm3.loop_on_next());
+                tokio::spawn(swarm2.loop_on_next());
+                tokio::spawn(swarm3.loop_on_next());
 
                 swarm1
                     .wait(|event| match event {

From 3d8a93acaeb060718894a025b0a7897c8c9b198b Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Wed, 5 Feb 2025 10:42:59 +0800
Subject: [PATCH 11/12] clippy lint

---
 misc/connection-limits/src/lib.rs | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index faf7055c0d0..544ac6240ac 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -485,13 +485,10 @@ mod tests {
                 .addresses(vec![addr.clone()])
                 .build(),
         ) {
-            match e {
-                DialError::Denied { cause } => {
-                    cause
-                        .downcast::<Exceeded>()
-                        .expect_err("Unexpected connection denied because of limit");
-                }
-                _ => {}
+            if let DialError::Denied { cause } = e {
+                cause
+                    .downcast::<Exceeded>()
+                    .expect_err("Unexpected connection denied because of limit");
             }
         }
         let not_bypassed_peer = loop {

From 9bb4dff3d23f5331bae3f2fad08213d5c828803f Mon Sep 17 00:00:00 2001
From: drHuangMHT <DrHuangMHT@outlook.com>
Date: Wed, 5 Feb 2025 18:37:44 +0800
Subject: [PATCH 12/12] clippy lint

---
 misc/connection-limits/src/lib.rs | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/misc/connection-limits/src/lib.rs b/misc/connection-limits/src/lib.rs
index 544ac6240ac..bd95dd6760c 100644
--- a/misc/connection-limits/src/lib.rs
+++ b/misc/connection-limits/src/lib.rs
@@ -480,16 +480,14 @@ mod tests {
             .limits
             .bypass_peer_id(&bypassed_peer);
         assert!(network.behaviour().limits.is_bypassed(&bypassed_peer));
-        if let Err(e) = network.dial(
+        if let Err(DialError::Denied { cause }) = network.dial(
             DialOpts::peer_id(bypassed_peer)
                 .addresses(vec![addr.clone()])
                 .build(),
         ) {
-            if let DialError::Denied { cause } = e {
-                cause
-                    .downcast::<Exceeded>()
-                    .expect_err("Unexpected connection denied because of limit");
-            }
+            cause
+                .downcast::<Exceeded>()
+                .expect_err("Unexpected connection denied because of limit");
         }
         let not_bypassed_peer = loop {
             let new_peer = PeerId::random();