From d2ec304bd38fd5f1d1e0e2a1f74ef37a2166ddf3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien?= <aurelien.catinon@gmail.com>
Date: Thu, 16 Feb 2023 19:30:47 +0100
Subject: [PATCH] feat: allows to disable dns and discv4 discovery separately
 (#1406)

---
 bin/reth/src/args/mod.rs          |  2 +-
 bin/reth/src/args/network_args.rs | 44 +++++++++++++++++++++++++++----
 bin/reth/src/p2p/mod.rs           | 19 ++++++++-----
 crates/net/network/src/config.rs  | 15 +++++------
 4 files changed, 58 insertions(+), 22 deletions(-)

diff --git a/bin/reth/src/args/mod.rs b/bin/reth/src/args/mod.rs
index c8f539f0c1..3b7e530c17 100644
--- a/bin/reth/src/args/mod.rs
+++ b/bin/reth/src/args/mod.rs
@@ -2,7 +2,7 @@
 
 /// NetworkArg struct
 mod network_args;
-pub use network_args::NetworkArgs;
+pub use network_args::{DiscoveryArgs, NetworkArgs};
 
 /// RpcServerArg struct
 mod rpc_server_args;
diff --git a/bin/reth/src/args/network_args.rs b/bin/reth/src/args/network_args.rs
index ee63b7329f..6bc3212df0 100644
--- a/bin/reth/src/args/network_args.rs
+++ b/bin/reth/src/args/network_args.rs
@@ -14,8 +14,8 @@ use std::path::PathBuf;
 #[command(next_help_heading = "Networking")]
 pub struct NetworkArgs {
     /// Disable the discovery service.
-    #[arg(short, long)]
-    pub disable_discovery: bool,
+    #[command(flatten)]
+    pub discovery: DiscoveryArgs,
 
     /// Target trusted peer enodes
     /// --trusted-peers enode://abcd@192.168.0.1:30303
@@ -52,11 +52,12 @@ impl NetworkArgs {
     /// values in this option struct.
     pub fn network_config(&self, config: &Config, chain_spec: ChainSpec) -> NetworkConfigBuilder {
         let peers_file = (!self.no_persist_peers).then_some(&self.peers_file);
-        config
+        let network_config_builder = config
             .network_config(self.nat, peers_file.map(|f| f.as_ref().to_path_buf()))
             .boot_nodes(self.bootnodes.clone().unwrap_or_else(mainnet_nodes))
-            .chain_spec(chain_spec)
-            .set_discovery(self.disable_discovery)
+            .chain_spec(chain_spec);
+
+        self.discovery.apply_to_builder(network_config_builder)
     }
 }
 
@@ -71,3 +72,36 @@ impl NetworkArgs {
         Some(self.peers_file.clone().into())
     }
 }
+
+/// Arguments to setup discovery
+#[derive(Debug, Args)]
+pub struct DiscoveryArgs {
+    /// Disable the discovery service.
+    #[arg(short, long)]
+    disable_discovery: bool,
+
+    /// Disable the DNS discovery.
+    #[arg(long, conflicts_with = "disable_discovery")]
+    disable_dns_discovery: bool,
+
+    /// Disable Discv4 discovery.
+    #[arg(long, conflicts_with = "disable_discovery")]
+    disable_discv4_discovery: bool,
+}
+
+impl DiscoveryArgs {
+    /// Apply the discovery settings to the given [NetworkConfigBuilder]
+    pub fn apply_to_builder(
+        &self,
+        mut network_config_builder: NetworkConfigBuilder,
+    ) -> NetworkConfigBuilder {
+        if self.disable_discovery || self.disable_dns_discovery {
+            network_config_builder = network_config_builder.disable_dns_discovery();
+        }
+
+        if self.disable_discovery || self.disable_discv4_discovery {
+            network_config_builder = network_config_builder.disable_discv4_discovery();
+        }
+        network_config_builder
+    }
+}
diff --git a/bin/reth/src/p2p/mod.rs b/bin/reth/src/p2p/mod.rs
index 84fa90436d..1bc0007912 100644
--- a/bin/reth/src/p2p/mod.rs
+++ b/bin/reth/src/p2p/mod.rs
@@ -1,5 +1,8 @@
 //! P2P Debugging tool
-use crate::dirs::{ConfigPath, PlatformPath};
+use crate::{
+    args::DiscoveryArgs,
+    dirs::{ConfigPath, PlatformPath},
+};
 use backon::{ConstantBackoff, Retryable};
 use clap::{Parser, Subcommand};
 use reth_db::mdbx::{Env, EnvKind, WriteMap};
@@ -42,8 +45,8 @@ pub struct Command {
     chain: ChainSpec,
 
     /// Disable the discovery service.
-    #[arg(short, long)]
-    disable_discovery: bool,
+    #[command(flatten)]
+    pub discovery: DiscoveryArgs,
 
     /// Target trusted peer
     #[arg(long)]
@@ -98,10 +101,12 @@ impl Command {
 
         config.peers.connect_trusted_nodes_only = self.trusted_only;
 
-        let network = config
-            .network_config(self.nat, None)
-            .set_discovery(self.disable_discovery)
-            .chain_spec(self.chain.clone())
+        let mut network_config_builder =
+            config.network_config(self.nat, None).chain_spec(self.chain.clone());
+
+        network_config_builder = self.discovery.apply_to_builder(network_config_builder);
+
+        let network = network_config_builder
             .build(Arc::new(ShareableDatabase::new(noop_db)))
             .start_network()
             .await?;
diff --git a/crates/net/network/src/config.rs b/crates/net/network/src/config.rs
index 7142d41f26..b31e53f2ea 100644
--- a/crates/net/network/src/config.rs
+++ b/crates/net/network/src/config.rs
@@ -274,19 +274,16 @@ impl NetworkConfigBuilder {
         self
     }
 
-    /// Sets the discovery service off on true.
-    // TODO(onbjerg): This name does not imply `true` = disable
-    pub fn set_discovery(mut self, disable_discovery: bool) -> Self {
-        if disable_discovery {
-            self.disable_discovery();
-        }
+    /// Disable the DNS discovery.
+    pub fn disable_dns_discovery(mut self) -> Self {
+        self.dns_discovery_config = None;
         self
     }
 
-    /// Disables all discovery services.
-    pub fn disable_discovery(&mut self) {
+    /// Disable the Discv4 discovery.
+    pub fn disable_discv4_discovery(mut self) -> Self {
         self.discovery_v4_builder = None;
-        self.dns_discovery_config = None;
+        self
     }
 
     /// Consumes the type and creates the actual [`NetworkConfig`]