feat:disc config param (#81)

Fix #69.

Signed-off-by: Chen Kai <[email protected]>

crates/net/src/builder.rs

@@ -1,7 +1,7 @@
 //! Network Builder Module.
 
 use alloy::primitives::Address;
-use discv5::ListenConfig;
+use discv5::{Config, ListenConfig};
 use eyre::Result;
 use std::{
     net::{IpAddr, SocketAddr},
@@ -34,6 +34,9 @@ pub struct NetworkDriverBuilder {
     pub discovery_addr: Option<ListenConfig>,
     /// The [GossipConfig] constructs the config for `gossipsub`.
     pub gossip_config: Option<GossipConfig>,
+
+    /// The [Config] constructs the config for `discv5`.
+    pub discovery_config: Option<Config>,
     /// The [Keypair] for the node.
     pub keypair: Option<Keypair>,
     /// The [TcpConfig] for the swarm.
@@ -139,6 +142,39 @@ impl NetworkDriverBuilder {
         self
     }
 
+    /// Specifies the [Config] for the `discv5` configuration.
+    ///
+    /// If not set, the [NetworkDriverBuilder] will fall back to use the [ListenConfig]
+    /// to construct [Config]. These defaults can be extended by using the
+    /// [discv5::ConfigBuilder::new] method to build a custom [Config].
+    ///
+    /// ## Example
+    ///
+    /// ```rust
+    /// use discv5::{ConfigBuilder, ListenConfig};
+    /// use kona_primitives::alloy_primitives::{address, Address};
+    /// use op_net::builder::NetworkDriverBuilder;
+    /// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+    ///
+    /// let id = 10;
+    /// let signer = Address::random();
+    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
+    /// let discovery_config =
+    ///     ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
+    ///         .build();
+    /// let driver = NetworkDriverBuilder::new()
+    ///     .with_unsafe_block_signer(signer)
+    ///     .with_chain_id(id)
+    ///     .with_gossip_addr(socket)
+    ///     .with_discovery_config(discovery_config)
+    ///     .build()
+    ///     .unwrap();
+    /// ```
+    pub fn with_discovery_config(&mut self, cfg: Config) -> &mut Self {
+        self.discovery_config = Some(cfg);
+        self
+    }
+
     /// Builds the [NetworkDriver].
     ///
     /// ## Errors
@@ -231,15 +267,18 @@ impl NetworkDriverBuilder {
         let gossip = GossipDriver::new(swarm, multiaddr, handler.clone());
 
         // Build the discovery service
-        let discovery_builder =
+        let mut discovery_builder =
             DiscoveryBuilder::new().with_address(gossip_addr).with_chain_id(chain_id);
 
-        let discovery = if let Some(discovery_addr) = self.discovery_addr.take() {
-            discovery_builder.with_listen_config(discovery_addr)
-        } else {
-            discovery_builder
+        if let Some(discovery_addr) = self.discovery_addr.take() {
+            discovery_builder = discovery_builder.with_listen_config(discovery_addr);
+        }
+
+        if let Some(discovery_config) = self.discovery_config.take() {
+            discovery_builder = discovery_builder.with_discovery_config(discovery_config);
         }
-        .build()?;
+
+        let discovery = discovery_builder.build()?;
 
         Ok(NetworkDriver {
             discovery,
@@ -253,6 +292,7 @@ impl NetworkDriverBuilder {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use discv5::ConfigBuilder;
     use libp2p::gossipsub::IdentTopic;
     use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
@@ -367,4 +407,44 @@ mod tests {
 
         assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
     }
+
+    #[test]
+    fn test_build_network_driver_with_discovery_config() {
+        let id = 10;
+        let signer = Address::random();
+        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
+        let discovery_config =
+            ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
+                .build();
+        let driver = NetworkDriverBuilder::new()
+            .with_unsafe_block_signer(signer)
+            .with_chain_id(id)
+            .with_gossip_addr(socket)
+            .with_discovery_config(discovery_config)
+            .build()
+            .unwrap();
+
+        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
+    }
+
+    #[test]
+    fn test_build_network_driver_with_discovery_config_and_listen_config() {
+        let id = 10;
+        let signer = Address::random();
+        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9099);
+        let discovery_config =
+            ConfigBuilder::new(ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9098))
+                .build();
+        let discovery_addr = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 9097);
+        let driver = NetworkDriverBuilder::new()
+            .with_unsafe_block_signer(signer)
+            .with_chain_id(id)
+            .with_gossip_addr(socket)
+            .with_discovery_addr(discovery_addr)
+            .with_discovery_config(discovery_config)
+            .build()
+            .unwrap();
+
+        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9097);
+    }
 }

crates/net/src/discovery/builder.rs

@@ -3,7 +3,7 @@
 use crate::{discovery::driver::DiscoveryDriver, types::enr::OpStackEnr};
 use discv5::{
     enr::{CombinedKey, Enr},
-    ConfigBuilder, Discv5, ListenConfig,
+    Config, ConfigBuilder, Discv5, ListenConfig,
 };
 use eyre::{Report, Result};
 use std::net::SocketAddr;
@@ -19,6 +19,9 @@ pub struct DiscoveryBuilder {
     chain_id: Option<u64>,
     /// The listen config for the discovery service.
     listen_config: Option<ListenConfig>,
+
+    /// The discovery config for the discovery service.
+    discovery_config: Option<Config>,
 }
 
 impl DiscoveryBuilder {
@@ -45,6 +48,12 @@ impl DiscoveryBuilder {
         self
     }
 
+    /// Sets the discovery config for the discovery service.
+    pub fn with_discovery_config(mut self, config: Config) -> Self {
+        self.discovery_config = Some(config);
+        self
+    }
+
     /// Builds a [DiscoveryDriver].
     pub fn build(&mut self) -> Result<DiscoveryDriver> {
         let chain_id = self.chain_id.ok_or_else(|| eyre::eyre!("chain ID not set"))?;
@@ -53,17 +62,24 @@ impl DiscoveryBuilder {
         use alloy_rlp::Encodable;
         opstack.encode(&mut opstack_data);
 
+        let config = if let Some(mut discovery_config) = self.discovery_config.take() {
+            if let Some(listen_config) = self.listen_config.take() {
+                discovery_config.listen_config = listen_config;
+            }
+            Ok::<Config, Report>(discovery_config)
+        } else {
+            let listen_config = self
+                .listen_config
+                .take()
+                .or_else(|| self.address.map(ListenConfig::from))
+                .ok_or_else(|| eyre::eyre!("listen config not set"))?;
+            Ok(ConfigBuilder::new(listen_config).build())
+        }?;
+
         let key = CombinedKey::generate_secp256k1();
         let mut enr_builder = Enr::builder();
         enr_builder.add_value_rlp(OP_CL_KEY, opstack_data.into());
-        let listen_config = self.listen_config.take().map_or_else(
-            || {
-                let addr = self.address.ok_or(eyre::eyre!("address not set"))?;
-                Ok::<ListenConfig, Report>(ListenConfig::from(addr))
-            },
-            Ok,
-        )?;
-        match listen_config {
+        match config.listen_config {
             ListenConfig::Ipv4 { ip, port } => {
                 enr_builder.ip4(ip).tcp4(port);
             }
@@ -76,7 +92,6 @@ impl DiscoveryBuilder {
             }
         }
         let enr = enr_builder.build(&key)?;
-        let config = ConfigBuilder::new(listen_config).build();
 
         let disc = Discv5::new(enr, key, config)
             .map_err(|_| eyre::eyre!("could not create disc service"))?;