From b46101afb5e549d40b7b2537fff9b67e05ad4448 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Altu=C4=9F=20Bakan?= <mail@alt.ug>
Date: Tue, 1 Aug 2023 20:44:39 +0200
Subject: [PATCH] feat(cli): add max peer args (#4024)

---
 bin/reth/src/args/network_args.rs       | 41 ++++++++++++++++++++++---
 crates/net/network/src/peers/manager.rs | 18 +++++++++++
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/bin/reth/src/args/network_args.rs b/bin/reth/src/args/network_args.rs
index c945fdaf3..086373687 100644
--- a/bin/reth/src/args/network_args.rs
+++ b/bin/reth/src/args/network_args.rs
@@ -60,6 +60,14 @@ pub struct NetworkArgs {
     /// Network listening port. default: 30303
     #[arg(long = "port", value_name = "PORT")]
     pub port: Option<u16>,
+
+    /// Maximum number of outbound requests. default: 100
+    #[arg(long)]
+    pub max_outbound_peers: Option<usize>,
+
+    /// Maximum number of inbound requests. default: 30
+    #[arg(long)]
+    pub max_inbound_peers: Option<usize>,
 }
 
 impl NetworkArgs {
@@ -78,9 +86,17 @@ impl NetworkArgs {
         let chain_bootnodes = chain_spec.chain.bootnodes().unwrap_or_else(mainnet_nodes);
         let peers_file = self.peers_file.clone().unwrap_or(default_peers_file);
 
-        // Configure basic network stack.
+        // Configure peer connections
+        let peer_config = config
+            .peers
+            .clone()
+            .with_max_inbound_opt(self.max_inbound_peers)
+            .with_max_outbound_opt(self.max_outbound_peers);
+
+        // Configure basic network stack
         let mut network_config_builder = config
             .network_config(self.nat, self.persistent_peers_file(peers_file), secret_key)
+            .peer_config(peer_config)
             .boot_nodes(self.bootnodes.clone().unwrap_or(chain_bootnodes))
             .chain_spec(chain_spec);
 
@@ -91,11 +107,7 @@ impl NetworkArgs {
 
         self.discovery.apply_to_builder(network_config_builder)
     }
-}
 
-// === impl NetworkArgs ===
-
-impl NetworkArgs {
     /// If `no_persist_peers` is true then this returns the path to the persistent peers file path.
     pub fn persistent_peers_file(&self, peers_file: PathBuf) -> Option<PathBuf> {
         if self.no_persist_peers {
@@ -163,4 +175,23 @@ mod tests {
             CommandParser::<NetworkArgs>::parse_from(["reth", "--nat", "extip:0.0.0.0"]).args;
         assert_eq!(args.nat, NatResolver::ExternalIp("0.0.0.0".parse().unwrap()));
     }
+
+    #[test]
+    fn parse_peer_args() {
+        let args =
+            CommandParser::<NetworkArgs>::parse_from(["reth", "--max-outbound-peers", "50"]).args;
+        assert_eq!(args.max_outbound_peers, Some(50));
+        assert_eq!(args.max_inbound_peers, None);
+
+        let args = CommandParser::<NetworkArgs>::parse_from([
+            "reth",
+            "--max-outbound-peers",
+            "75",
+            "--max-inbound-peers",
+            "15",
+        ])
+        .args;
+        assert_eq!(args.max_outbound_peers, Some(75));
+        assert_eq!(args.max_inbound_peers, Some(15));
+    }
 }
diff --git a/crates/net/network/src/peers/manager.rs b/crates/net/network/src/peers/manager.rs
index 8b192f1d7..1955603c7 100644
--- a/crates/net/network/src/peers/manager.rs
+++ b/crates/net/network/src/peers/manager.rs
@@ -1134,11 +1134,29 @@ impl PeersConfig {
         self.connection_info.num_inbound = num_inbound;
         self
     }
+
     /// Maximum allowed outbound connections.
     pub fn with_max_outbound(mut self, max_outbound: usize) -> Self {
         self.connection_info.max_outbound = max_outbound;
         self
     }
+
+    /// Maximum allowed inbound connections with optional update.
+    pub fn with_max_inbound_opt(mut self, max_inbound: Option<usize>) -> Self {
+        if let Some(max_inbound) = max_inbound {
+            self.connection_info.max_inbound = max_inbound;
+        }
+        self
+    }
+
+    /// Maximum allowed outbound connections with optional update.
+    pub fn with_max_outbound_opt(mut self, max_outbound: Option<usize>) -> Self {
+        if let Some(max_outbound) = max_outbound {
+            self.connection_info.max_outbound = max_outbound;
+        }
+        self
+    }
+
     /// Maximum allowed inbound connections.
     pub fn with_max_inbound(mut self, max_inbound: usize) -> Self {
         self.connection_info.max_inbound = max_inbound;