chore: integrate discv5 config builder in networkconfig builder (#7856)

Co-authored-by: Emilia Hane <elsaemiliaevahane@gmail.com>

Cargo.lock

@@ -6376,6 +6376,7 @@ dependencies = [
  "reth-consensus-common",
  "reth-db",
  "reth-discv4",
+ "reth-discv5",
  "reth-downloaders",
  "reth-ethereum-payload-builder",
  "reth-evm",

bin/reth/Cargo.toml

@@ -44,6 +44,7 @@ reth-payload-builder.workspace = true
 reth-payload-validator.workspace = true
 reth-basic-payload-builder.workspace = true
 reth-discv4.workspace = true
+reth-discv5.workspace = true
 reth-static-file = { workspace = true }
 reth-trie = { workspace = true, features = ["metrics"] }
 reth-nippy-jar.workspace = true

bin/reth/src/commands/p2p/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     args::{
         get_secret_key,
         utils::{chain_help, chain_spec_value_parser, hash_or_num_value_parser, SUPPORTED_CHAINS},
-        DatabaseArgs, DiscoveryArgs,
+        DatabaseArgs, DiscoveryArgs, NetworkArgs,
     },
     dirs::{DataDirPath, MaybePlatformPath},
     utils::get_single_header,
@@ -14,12 +14,11 @@ use clap::{Parser, Subcommand};
 use discv5::ListenConfig;
 use reth_config::Config;
 use reth_db::create_db;
-use reth_discv4::NatResolver;
 use reth_interfaces::p2p::bodies::client::BodiesClient;
-use reth_primitives::{BlockHashOrNumber, ChainSpec, NodeRecord};
+use reth_primitives::{BlockHashOrNumber, ChainSpec};
 use reth_provider::ProviderFactory;
 use std::{
-    net::{SocketAddrV4, SocketAddrV6},
+    net::{IpAddr, SocketAddrV4, SocketAddrV6},
     path::PathBuf,
     sync::Arc,
 };
@@ -53,31 +52,14 @@ pub struct Command {
     #[arg(long, value_name = "DATA_DIR", verbatim_doc_comment, default_value_t)]
     datadir: MaybePlatformPath<DataDirPath>,
 
-    /// Secret key to use for this node.
-    ///
-    /// This also will deterministically set the peer ID.
-    #[arg(long, value_name = "PATH")]
-    p2p_secret_key: Option<PathBuf>,
-
     /// Disable the discovery service.
     #[command(flatten)]
-    pub discovery: DiscoveryArgs,
+    pub network: NetworkArgs,
-
-    /// Target trusted peer
-    #[arg(long)]
-    trusted_peer: Option<NodeRecord>,
-
-    /// Connect only to trusted peers
-    #[arg(long)]
-    trusted_only: bool,
 
     /// The number of retries per request
     #[arg(long, default_value = "5")]
     retries: usize,
 
-    #[arg(long, default_value = "any")]
-    nat: NatResolver,
-
     #[command(flatten)]
     db: DatabaseArgs,
 
@@ -113,65 +95,75 @@ impl Command {
 
         let mut config: Config = confy::load_path(&config_path).unwrap_or_default();
 
-        if let Some(peer) = self.trusted_peer {
+        for &peer in &self.network.trusted_peers {
             config.peers.trusted_nodes.insert(peer);
         }
 
-        if config.peers.trusted_nodes.is_empty() && self.trusted_only {
+        if config.peers.trusted_nodes.is_empty() && self.network.trusted_only {
             eyre::bail!("No trusted nodes. Set trusted peer with `--trusted-peer <enode record>` or set `--trusted-only` to `false`")
         }
 
-        config.peers.trusted_nodes_only = self.trusted_only;
+        config.peers.trusted_nodes_only = self.network.trusted_only;
 
         let default_secret_key_path = data_dir.p2p_secret();
-        let secret_key_path = self.p2p_secret_key.clone().unwrap_or(default_secret_key_path);
+        let secret_key_path =
+            self.network.p2p_secret_key.clone().unwrap_or(default_secret_key_path);
         let p2p_secret_key = get_secret_key(&secret_key_path)?;
+        let rlpx_socket = (self.network.addr, self.network.port).into();
+        let boot_nodes = self.chain.bootnodes().unwrap_or_default();
 
         let mut network_config_builder = config
-            .network_config(self.nat, None, p2p_secret_key)
+            .network_config(self.network.nat, None, p2p_secret_key)
             .chain_spec(self.chain.clone())
             .disable_discv4_discovery_if(self.chain.chain.is_optimism())
-            .boot_nodes(self.chain.bootnodes().unwrap_or_default());
+            .boot_nodes(boot_nodes.clone());
 
-        network_config_builder = self.discovery.apply_to_builder(network_config_builder);
+        network_config_builder = self
-
+            .network
-        let mut network_config = network_config_builder.build(Arc::new(ProviderFactory::new(
+            .discovery
-            noop_db,
+            .apply_to_builder(network_config_builder, rlpx_socket)
-            self.chain.clone(),
+            .map_discv5_config_builder(|builder| {
-            data_dir.static_files(),
-        )?));
-
-        if !self.discovery.disable_discovery &&
-            (self.discovery.enable_discv5_discovery ||
-                network_config.chain_spec.chain.is_optimism())
-        {
-            network_config = network_config.discovery_v5_with_config_builder(|builder| {
                 let DiscoveryArgs {
-                    discv5_addr: discv5_addr_ipv4,
+                    discv5_addr,
                     discv5_addr_ipv6,
-                    discv5_port: discv5_port_ipv4,
+                    discv5_port,
                     discv5_port_ipv6,
                     discv5_lookup_interval,
                     discv5_bootstrap_lookup_interval,
                     discv5_bootstrap_lookup_countdown,
                     ..
-                } = self.discovery;
+                } = self.network.discovery;
+
+                // Use rlpx address if none given
+                let discv5_addr_ipv4 = discv5_addr.or(match self.network.addr {
+                    IpAddr::V4(ip) => Some(ip),
+                    IpAddr::V6(_) => None,
+                });
+                let discv5_addr_ipv6 = discv5_addr_ipv6.or(match self.network.addr {
+                    IpAddr::V4(_) => None,
+                    IpAddr::V6(ip) => Some(ip),
+                });
 
                 builder
                     .discv5_config(
                         discv5::ConfigBuilder::new(ListenConfig::from_two_sockets(
-                            discv5_addr_ipv4.map(|addr| SocketAddrV4::new(addr, discv5_port_ipv4)),
+                            discv5_addr_ipv4.map(|addr| SocketAddrV4::new(addr, discv5_port)),
                             discv5_addr_ipv6
                                 .map(|addr| SocketAddrV6::new(addr, discv5_port_ipv6, 0, 0)),
                         ))
                         .build(),
                     )
+                    .add_unsigned_boot_nodes(boot_nodes.into_iter())
                     .lookup_interval(discv5_lookup_interval)
                     .bootstrap_lookup_interval(discv5_bootstrap_lookup_interval)
                     .bootstrap_lookup_countdown(discv5_bootstrap_lookup_countdown)
-                    .build()
             });
-        }
+
+        let network_config = network_config_builder.build(Arc::new(ProviderFactory::new(
+            noop_db,
+            self.chain.clone(),
+            data_dir.static_files(),
+        )?));
 
         let network = network_config.start_network().await?;
         let fetch_client = network.fetch_client().await?;

crates/net/discv5/src/config.rs

@@ -250,7 +250,7 @@ impl ConfigBuilder {
 }
 
 /// Config used to bootstrap [`discv5::Discv5`].
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct Config {
     /// Config used by [`discv5::Discv5`]. Contains the [`ListenConfig`], with the discovery listen
     /// socket.
@@ -296,9 +296,7 @@ impl Config {
             discovered_peer_filter: None,
         }
     }
-}
 
-impl Config {
     /// Returns the discovery (UDP) socket contained in the [`discv5::Config`]. Returns the IPv6
     /// socket, if both IPv4 and v6 are configured. This socket will be advertised to peers in the
     /// local [`Enr`](discv5::enr::Enr).
@@ -416,7 +414,7 @@ pub fn discv5_sockets_wrt_rlpx_addr(
 
 /// A boot node can be added either as a string in either 'enode' URL scheme or serialized from
 /// [`Enr`](discv5::Enr) type.
-#[derive(Debug, PartialEq, Eq, Hash, Display)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Display)]
 pub enum BootNode {
     /// An unsigned node record.
     #[display(fmt = "{_0}")]

crates/net/network/src/config.rs

@@ -104,43 +104,6 @@ impl<C> NetworkConfig<C> {
         self
     }
 
-    /// Sets the config to use for the discovery v5 protocol, with help of the
-    /// [`reth_discv5::ConfigBuilder`].
-    /// ```
-    /// use reth_network::NetworkConfigBuilder;
-    /// use secp256k1::{rand::thread_rng, SecretKey};
-    ///
-    /// let sk = SecretKey::new(&mut thread_rng());
-    /// let network_config = NetworkConfigBuilder::new(sk).build(());
-    /// let fork_id = network_config.status.forkid;
-    /// let network_config = network_config
-    ///     .discovery_v5_with_config_builder(|builder| builder.fork(b"eth", fork_id).build());
-    /// ```
-
-    pub fn discovery_v5_with_config_builder(
-        self,
-        f: impl FnOnce(reth_discv5::ConfigBuilder) -> reth_discv5::Config,
-    ) -> Self {
-        let network_stack_id = NetworkStackId::id(&self.chain_spec);
-        let fork_id = self.chain_spec.latest_fork_id();
-        let boot_nodes = self.boot_nodes.clone();
-
-        let mut builder = reth_discv5::Config::builder(self.listener_addr)
-            .add_unsigned_boot_nodes(boot_nodes.into_iter());
-
-        if let Some(id) = network_stack_id {
-            builder = builder.fork(id, fork_id)
-        }
-
-        self.set_discovery_v5(f(builder))
-    }
-
-    /// Sets the config to use for the discovery v5 protocol.
-    pub fn set_discovery_v5(mut self, discv5_config: reth_discv5::Config) -> Self {
-        self.discovery_v5_config = Some(discv5_config);
-        self
-    }
-
     /// Sets the address for the incoming RLPx connection listener.
     pub fn set_listener_addr(mut self, listener_addr: SocketAddr) -> Self {
         self.listener_addr = listener_addr;
@@ -180,6 +143,9 @@ pub struct NetworkConfigBuilder {
     dns_discovery_config: Option<DnsDiscoveryConfig>,
     /// How to set up discovery version 4.
     discovery_v4_builder: Option<Discv4ConfigBuilder>,
+    /// How to set up discovery version 5.
+    #[serde(skip)]
+    discovery_v5_builder: Option<reth_discv5::ConfigBuilder>,
     /// All boot nodes to start network discovery with.
     boot_nodes: HashSet<NodeRecord>,
     /// Address to use for discovery
@@ -222,6 +188,7 @@ impl NetworkConfigBuilder {
             secret_key,
             dns_discovery_config: Some(Default::default()),
             discovery_v4_builder: Some(Default::default()),
+            discovery_v5_builder: None,
             boot_nodes: Default::default(),
             discovery_addr: None,
             listener_addr: None,
@@ -348,12 +315,17 @@ impl NetworkConfigBuilder {
     }
 
     /// Sets the discv4 config to use.
-    //
     pub fn discovery(mut self, builder: Discv4ConfigBuilder) -> Self {
         self.discovery_v4_builder = Some(builder);
         self
     }
 
+    /// Sets the discv5 config to use.
+    pub fn discovery_v5(mut self, builder: reth_discv5::ConfigBuilder) -> Self {
+        self.discovery_v5_builder = Some(builder);
+        self
+    }
+
     /// Sets the dns discovery config to use.
     pub fn dns_discovery(mut self, config: DnsDiscoveryConfig) -> Self {
         self.dns_discovery_config = Some(config);
@@ -420,6 +392,36 @@ impl NetworkConfigBuilder {
         }
     }
 
+    /// Calls a closure on [`reth_discv5::ConfigBuilder`], if discv5 discovery is enabled and the
+    /// builder has been set.
+    /// ```
+    /// use reth_network::NetworkConfigBuilder;
+    /// use reth_primitives::MAINNET;
+    /// use reth_provider::test_utils::NoopProvider;
+    /// use secp256k1::{rand::thread_rng, SecretKey};
+    ///
+    /// let sk = SecretKey::new(&mut thread_rng());
+    /// let fork_id = MAINNET.latest_fork_id();
+    /// let network_config = NetworkConfigBuilder::new(sk)
+    ///     .map_discv5_config_builder(|builder| builder.fork(b"eth", fork_id))
+    ///     .build(NoopProvider::default());
+    /// ```
+    pub fn map_discv5_config_builder(
+        mut self,
+        f: impl FnOnce(reth_discv5::ConfigBuilder) -> reth_discv5::ConfigBuilder,
+    ) -> Self {
+        if let Some(mut builder) = self.discovery_v5_builder {
+            if let Some(network_stack_id) = NetworkStackId::id(&self.chain_spec) {
+                let fork_id = self.chain_spec.latest_fork_id();
+                builder = builder.fork(network_stack_id, fork_id);
+            }
+
+            self.discovery_v5_builder = Some(f(builder));
+        }
+
+        self
+    }
+
     /// Adds a new additional protocol to the RLPx sub-protocol list.
     pub fn add_rlpx_sub_protocol(mut self, protocol: impl IntoRlpxSubProtocol) -> Self {
         self.extra_protocols.push(protocol);
@@ -458,6 +460,7 @@ impl NetworkConfigBuilder {
             secret_key,
             mut dns_discovery_config,
             discovery_v4_builder,
+            discovery_v5_builder,
             boot_nodes,
             discovery_addr,
             listener_addr,
@@ -511,7 +514,7 @@ impl NetworkConfigBuilder {
             boot_nodes,
             dns_discovery_config,
             discovery_v4_config: discovery_v4_builder.map(|builder| builder.build()),
-            discovery_v5_config: None,
+            discovery_v5_config: discovery_v5_builder.map(|builder| builder.build()),
             discovery_v4_addr: discovery_addr.unwrap_or(DEFAULT_DISCOVERY_ADDRESS),
             listener_addr,
             peers_config: peers_config.unwrap_or_default(),

crates/node-core/src/args/network.rs

@@ -20,7 +20,7 @@ use reth_network::{
 use reth_primitives::{mainnet_nodes, ChainSpec, NodeRecord};
 use secp256k1::SecretKey;
 use std::{
-    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
     ops::Not,
     path::PathBuf,
     sync::Arc,
@@ -119,7 +119,10 @@ impl NetworkArgs {
         secret_key: SecretKey,
         default_peers_file: PathBuf,
     ) -> NetworkConfigBuilder {
-        let chain_bootnodes = chain_spec.bootnodes().unwrap_or_else(mainnet_nodes);
+        let boot_nodes = self
+            .bootnodes
+            .clone()
+            .unwrap_or_else(|| chain_spec.bootnodes().unwrap_or_else(mainnet_nodes));
         let peers_file = self.peers_file.clone().unwrap_or(default_peers_file);
 
         // Configure peer connections
@@ -136,7 +139,6 @@ impl NetworkArgs {
                 self.soft_limit_byte_size_pooled_transactions_response_on_pack_request,
             ),
         };
-
         // Configure basic network stack
         let mut network_config_builder = config
             .network_config(self.nat, self.persistent_peers_file(peers_file), secret_key)
@@ -144,8 +146,8 @@ impl NetworkArgs {
                 SessionsConfig::default().with_upscaled_event_buffer(peers_config.max_peers()),
             )
             .peer_config(peers_config)
-            .boot_nodes(self.bootnodes.clone().unwrap_or(chain_bootnodes))
+            .boot_nodes(boot_nodes.clone())
-            .chain_spec(chain_spec)
+            .chain_spec(chain_spec.clone())
             .transactions_manager_config(transactions_manager_config);
 
         // Configure node identity
@@ -154,7 +156,29 @@ impl NetworkArgs {
             HelloMessageWithProtocols::builder(peer_id).client_version(&self.identity).build(),
         );
 
-        self.discovery.apply_to_builder(network_config_builder)
+        let rlpx_socket = (self.addr, self.port).into();
+        network_config_builder =
+            self.discovery.apply_to_builder(network_config_builder, rlpx_socket);
+
+        if chain_spec.is_optimism() && !self.discovery.disable_discovery {
+            network_config_builder =
+                network_config_builder.discovery_v5(reth_discv5::Config::builder(rlpx_socket));
+        }
+
+        network_config_builder.map_discv5_config_builder(|builder| {
+            let DiscoveryArgs {
+                discv5_lookup_interval,
+                discv5_bootstrap_lookup_interval,
+                discv5_bootstrap_lookup_countdown,
+                ..
+            } = self.discovery;
+
+            builder
+                .add_unsigned_boot_nodes(boot_nodes.into_iter())
+                .lookup_interval(discv5_lookup_interval)
+                .bootstrap_lookup_interval(discv5_bootstrap_lookup_interval)
+                .bootstrap_lookup_countdown(discv5_bootstrap_lookup_countdown)
+        })
     }
 
     /// If `no_persist_peers` is false then this returns the path to the persistent peers file path.
@@ -228,11 +252,13 @@ pub struct DiscoveryArgs {
     #[arg(id = "discovery.port", long = "discovery.port", value_name = "DISCOVERY_PORT", default_value_t = DEFAULT_DISCOVERY_PORT)]
     pub port: u16,
 
-    /// The UDP IPv4 address to use for devp2p peer discovery version 5.
+    /// The UDP IPv4 address to use for devp2p peer discovery version 5. Overwritten by RLPx
+    /// address, if it's also IPv4.
     #[arg(id = "discovery.v5.addr", long = "discovery.v5.addr", value_name = "DISCOVERY_V5_ADDR", default_value = None)]
     pub discv5_addr: Option<Ipv4Addr>,
 
-    /// The UDP IPv6 address to use for devp2p peer discovery version 5.
+    /// The UDP IPv6 address to use for devp2p peer discovery version 5. Overwritten by RLPx
+    /// address, if it's also IPv6.
     #[arg(id = "discovery.v5.addr.ipv6", long = "discovery.v5.addr.ipv6", value_name = "DISCOVERY_V5_ADDR_IPV6", default_value = None)]
     pub discv5_addr_ipv6: Option<Ipv6Addr>,
 
@@ -270,6 +296,7 @@ impl DiscoveryArgs {
     pub fn apply_to_builder(
         &self,
         mut network_config_builder: NetworkConfigBuilder,
+        rlpx_tcp_socket: SocketAddr,
     ) -> NetworkConfigBuilder {
         if self.disable_discovery || self.disable_dns_discovery {
             network_config_builder = network_config_builder.disable_dns_discovery();
@@ -279,6 +306,11 @@ impl DiscoveryArgs {
             network_config_builder = network_config_builder.disable_discv4_discovery();
         }
 
+        if !self.disable_discovery && self.enable_discv5_discovery {
+            network_config_builder =
+                network_config_builder.discovery_v5(reth_discv5::Config::builder(rlpx_tcp_socket));
+        }
+
         network_config_builder
     }
 
@@ -295,8 +327,8 @@ impl Default for DiscoveryArgs {
         Self {
             disable_discovery: false,
             disable_dns_discovery: false,
-            disable_discv4_discovery: cfg!(feature = "optimism"),
+            disable_discv4_discovery: false,
-            enable_discv5_discovery: cfg!(feature = "optimism"),
+            enable_discv5_discovery: false,
             addr: DEFAULT_DISCOVERY_ADDR,
             port: DEFAULT_DISCOVERY_PORT,
             discv5_addr: None,

crates/node-core/src/node_config.rs

@@ -27,7 +27,7 @@ use reth_provider::{
 use reth_tasks::TaskExecutor;
 use secp256k1::SecretKey;
 use std::{
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
+    net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6},
     path::PathBuf,
     sync::Arc,
 };
@@ -456,8 +456,7 @@ impl NodeConfig {
         secret_key: SecretKey,
         default_peers_path: PathBuf,
     ) -> NetworkConfig<C> {
-        let cfg_builder = self
+        self.network
-            .network
             .network_config(config, self.chain.clone(), secret_key, default_peers_path)
             .with_task_executor(Box::new(executor))
             .set_head(head)
@@ -471,39 +470,30 @@ impl NodeConfig {
                 self.network.discovery.addr,
                 // set discovery port based on instance number
                 self.network.discovery.port + self.instance - 1,
-            ));
+            ))
+            .map_discv5_config_builder(|builder| {
+                let DiscoveryArgs {
+                    discv5_addr,
+                    discv5_addr_ipv6,
+                    discv5_port,
+                    discv5_port_ipv6,
+                    ..
+                } = self.network.discovery;
 
-        let config = cfg_builder.build(client);
+                // Use rlpx address if none given
+                let discv5_addr_ipv4 = discv5_addr.or(match self.network.addr {
+                    IpAddr::V4(ip) => Some(ip),
+                    IpAddr::V6(_) => None,
+                });
+                let discv5_addr_ipv6 = discv5_addr_ipv6.or(match self.network.addr {
+                    IpAddr::V4(_) => None,
+                    IpAddr::V6(ip) => Some(ip),
+                });
 
-        if self.network.discovery.disable_discovery ||
+                let discv5_port_ipv4 = discv5_port + self.instance - 1;
-            !self.network.discovery.enable_discv5_discovery &&
+                let discv5_port_ipv6 = discv5_port_ipv6 + self.instance - 1;
-                !config.chain_spec.chain.is_optimism()
-        {
-            return config
-        }
 
-        let rlpx_addr = config.listener_addr().ip();
+                builder.discv5_config(
-        // work around since discv5 config builder can't be integrated into network config builder
-        // due to unsatisfied trait bounds
-        config.discovery_v5_with_config_builder(|builder| {
-            let DiscoveryArgs {
-                discv5_addr,
-                discv5_addr_ipv6,
-                discv5_port,
-                discv5_port_ipv6,
-                discv5_lookup_interval,
-                discv5_bootstrap_lookup_interval,
-                discv5_bootstrap_lookup_countdown,
-                ..
-            } = self.network.discovery;
-
-            let discv5_addr_ipv4 = discv5_addr.or_else(|| ipv4(rlpx_addr));
-            let discv5_addr_ipv6 = discv5_addr_ipv6.or_else(|| ipv6(rlpx_addr));
-            let discv5_port_ipv4 = discv5_port + self.instance - 1;
-            let discv5_port_ipv6 = discv5_port_ipv6 + self.instance - 1;
-
-            builder
-                .discv5_config(
                     discv5::ConfigBuilder::new(ListenConfig::from_two_sockets(
                         discv5_addr_ipv4.map(|addr| SocketAddrV4::new(addr, discv5_port_ipv4)),
                         discv5_addr_ipv6
@@ -511,11 +501,8 @@ impl NodeConfig {
                     ))
                     .build(),
                 )
-                .lookup_interval(discv5_lookup_interval)
+            })
-                .bootstrap_lookup_interval(discv5_bootstrap_lookup_interval)
+            .build(client)
-                .bootstrap_lookup_countdown(discv5_bootstrap_lookup_countdown)
-                .build()
-        })
     }
 
     /// Change rpc port numbers based on the instance number, using the inner
@@ -551,19 +538,3 @@ impl Default for NodeConfig {
         }
     }
 }
-
-/// Returns the address if this is an [`Ipv4Addr`].
-pub fn ipv4(ip: IpAddr) -> Option<Ipv4Addr> {
-    match ip {
-        IpAddr::V4(ip) => Some(ip),
-        IpAddr::V6(_) => None,
-    }
-}
-
-/// Returns the address if this is an [`Ipv6Addr`].
-pub fn ipv6(ip: IpAddr) -> Option<Ipv6Addr> {
-    match ip {
-        IpAddr::V4(_) => None,
-        IpAddr::V6(ip) => Some(ip),
-    }
-}