From 76f1fcbee3442386e04817a0c8371c368df1ccfc Mon Sep 17 00:00:00 2001
From: Age Manning <Age@AgeManning.com>
Date: Sat, 31 Jul 2021 03:22:53 +1000
Subject: [PATCH] core: Add a total established connection limit (#2137)

Co-authored-by: Max Inden <mail@max-inden.de>
---
 core/CHANGELOG.md           |  3 +++
 core/src/connection/pool.rs | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/core/CHANGELOG.md b/core/CHANGELOG.md
index 10e0066f..cc19a911 100644
--- a/core/CHANGELOG.md
+++ b/core/CHANGELOG.md
@@ -1,5 +1,7 @@
 # 0.30.0 [unreleased]
 
+- Add `ConnectionLimit::with_max_established` (see [PR 2137]).
+
 - Add `Keypair::to_protobuf_encoding` (see [PR 2142]).
 
 - Change `PublicKey::into_protobuf_encoding` to `PublicKey::to_protobuf_encoding` (see [PR 2145]).
@@ -12,6 +14,7 @@
 
 [PR 2145]: https://github.com/libp2p/rust-libp2p/pull/2145
 [PR 2142]: https://github.com/libp2p/rust-libp2p/pull/2142
+[PR 2137]: https://github.com/libp2p/rust-libp2p/pull/2137/
 
 # 0.29.0 [2021-07-12]
 
diff --git a/core/src/connection/pool.rs b/core/src/connection/pool.rs
index 6615219e..225a3063 100644
--- a/core/src/connection/pool.rs
+++ b/core/src/connection/pool.rs
@@ -987,6 +987,9 @@ impl ConnectionCounters {
     fn check_max_established(&self, endpoint: &ConnectedPoint)
         -> Result<(), ConnectionLimit>
     {
+        // Check total connection limit.
+        Self::check(self.num_established(), self.limits.max_established_total)?;
+        // Check incoming/outgoing connection limits
         match endpoint {
             ConnectedPoint::Dialer { .. } =>
                 Self::check(self.established_outgoing, self.limits.max_established_outgoing),
@@ -1031,6 +1034,7 @@ pub struct ConnectionLimits {
     max_established_incoming: Option<u32>,
     max_established_outgoing: Option<u32>,
     max_established_per_peer: Option<u32>,
+    max_established_total: Option<u32>,
 }
 
 impl ConnectionLimits {
@@ -1058,6 +1062,17 @@ impl ConnectionLimits {
         self
     }
 
+    /// Configures the maximum number of concurrent established connections (both
+    /// inbound and outbound).
+    ///
+    /// Note: This should be used in conjunction with
+    /// [`ConnectionLimits::with_max_established_incoming`] to prevent possible
+    /// eclipse attacks (all connections being inbound).
+    pub fn with_max_established(mut self, limit: Option<u32>) -> Self {
+        self.max_established_total = limit;
+        self
+    }
+
     /// Configures the maximum number of concurrent established connections per peer,
     /// regardless of direction (incoming or outgoing).
     pub fn with_max_established_per_peer(mut self, limit: Option<u32>) -> Self {