From dda6e3e899acb7c94f3bb64363a764139aa044a2 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Thu, 9 Jan 2025 00:06:51 -0500
Subject: [PATCH] Limit each peer to one connection

Prevents dialing the same peer multiple times (successfully).
---
 Cargo.lock                                 |  1 +
 coordinator/Cargo.toml                     |  1 +
 coordinator/src/p2p/libp2p/authenticate.rs |  4 +---
 coordinator/src/p2p/libp2p/gossip.rs       |  4 ++--
 coordinator/src/p2p/libp2p/mod.rs          | 14 ++++++++++----
 coordinator/src/p2p/libp2p/swarm.rs        |  6 ++++--
 6 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3902b794..49bbf65d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8339,6 +8339,7 @@ dependencies = [
  "sp-runtime",
  "tokio",
  "tributary-chain",
+ "void",
  "zalloc",
  "zeroize",
 ]
diff --git a/coordinator/Cargo.toml b/coordinator/Cargo.toml
index 2fc373aa..e0b84346 100644
--- a/coordinator/Cargo.toml
+++ b/coordinator/Cargo.toml
@@ -55,6 +55,7 @@ env_logger = { version = "0.10", default-features = false, features = ["humantim
 
 futures-util = { version = "0.3", default-features = false, features = ["std"] }
 tokio = { version = "1", default-features = false, features = ["rt-multi-thread", "sync", "time", "macros"] }
+void = { version = "1", default-features = false }
 libp2p = { version = "0.52", default-features = false, features = ["tokio", "tcp", "noise", "yamux", "ping", "request-response", "gossipsub", "macros"] }
 
 serai-cosign = { path = "./cosign" }
diff --git a/coordinator/src/p2p/libp2p/authenticate.rs b/coordinator/src/p2p/libp2p/authenticate.rs
index a8167d9e..56e5336b 100644
--- a/coordinator/src/p2p/libp2p/authenticate.rs
+++ b/coordinator/src/p2p/libp2p/authenticate.rs
@@ -1,5 +1,5 @@
 use core::{pin::Pin, future::Future};
-use std::{sync::Arc, io};
+use std::io;
 
 use zeroize::Zeroizing;
 use rand_core::{RngCore, OsRng};
@@ -9,8 +9,6 @@ use schnorrkel::{Keypair, PublicKey, Signature};
 
 use serai_client::primitives::PublicKey as Public;
 
-use tokio::sync::RwLock;
-
 use futures_util::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use libp2p::{
   core::UpgradeInfo,
diff --git a/coordinator/src/p2p/libp2p/gossip.rs b/coordinator/src/p2p/libp2p/gossip.rs
index db5af299..f64fddb5 100644
--- a/coordinator/src/p2p/libp2p/gossip.rs
+++ b/coordinator/src/p2p/libp2p/gossip.rs
@@ -5,8 +5,8 @@ use blake2::{Digest, Blake2s256};
 use borsh::{BorshSerialize, BorshDeserialize};
 
 use libp2p::gossipsub::{
-  TopicHash, IdentTopic, MessageId, MessageAuthenticity, ValidationMode, ConfigBuilder,
-  IdentityTransform, AllowAllSubscriptionFilter, Behaviour,
+  IdentTopic, MessageId, MessageAuthenticity, ValidationMode, ConfigBuilder, IdentityTransform,
+  AllowAllSubscriptionFilter, Behaviour,
 };
 pub use libp2p::gossipsub::Event;
 
diff --git a/coordinator/src/p2p/libp2p/mod.rs b/coordinator/src/p2p/libp2p/mod.rs
index fccf7ce1..bce1006d 100644
--- a/coordinator/src/p2p/libp2p/mod.rs
+++ b/coordinator/src/p2p/libp2p/mod.rs
@@ -26,6 +26,7 @@ use libp2p::{
   identity::{self, PeerId},
   tcp::Config as TcpConfig,
   yamux, allow_block_list,
+  connection_limits::{self, ConnectionLimits},
   swarm::NetworkBehaviour,
   SwarmBuilder,
 };
@@ -40,10 +41,6 @@ use validators::UpdateValidatorsTask;
 mod authenticate;
 use authenticate::OnlyValidators;
 
-/// The dial task, to find new peers to connect to
-mod dial;
-use dial::DialTask;
-
 /// The ping behavior, used to ensure connection latency is below the limit
 mod ping;
 
@@ -59,6 +56,10 @@ use gossip::Message;
 mod swarm;
 use swarm::SwarmTask;
 
+/// The dial task, to find new peers to connect to
+mod dial;
+use dial::DialTask;
+
 const PORT: u16 = 30563; // 5132 ^ (('c' << 8) | 'o')
 
 // usize::max, manually implemented, as max isn't a const fn
@@ -113,6 +114,7 @@ struct Peers {
 #[derive(NetworkBehaviour)]
 struct Behavior {
   allow_list: allow_block_list::Behaviour<allow_block_list::AllowedPeers>,
+  connection_limits: connection_limits::Behaviour,
   ping: ping::Behavior,
   reqres: reqres::Behavior,
   gossip: gossip::Behavior,
@@ -169,6 +171,10 @@ impl Libp2p {
         .unwrap()
         .with_behaviour(|_| Behavior {
           allow_list: allow_block_list::Behaviour::default(),
+          // Limit each per to a single connection
+          connection_limits: connection_limits::Behaviour::new(
+            ConnectionLimits::default().with_max_established_per_peer(Some(1)),
+          ),
           ping: ping::new_behavior(),
           reqres: reqres::new_behavior(),
           gossip: gossip::new_behavior(),
diff --git a/coordinator/src/p2p/libp2p/swarm.rs b/coordinator/src/p2p/libp2p/swarm.rs
index 10d91818..f62cb659 100644
--- a/coordinator/src/p2p/libp2p/swarm.rs
+++ b/coordinator/src/p2p/libp2p/swarm.rs
@@ -225,8 +225,10 @@ impl SwarmTask {
               }
             }
 
-            SwarmEvent::Behaviour(BehaviorEvent::AllowList(event)) => {
-              // Ensure this is an unreachable case, not an actual event
+            SwarmEvent::Behaviour(
+              BehaviorEvent::AllowList(event) | BehaviorEvent::ConnectionLimits(event)
+            ) => {
+              // Ensure these are unreachable cases, not actual events
               let _: void::Void = event;
             }
             SwarmEvent::Behaviour(