ListenConfig parametrization

crates/net/README.md

@@ -15,7 +15,7 @@ let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9099);
 let mut driver = NetworkDriver::builder()
     .with_chain_id(10) // op mainnet chain id
     .with_unsafe_block_signer(signer)
-    .with_socket(socket)
+    .with_gossip_addr(socket)
     .build()
     .expect("Failed to builder network driver");
 

crates/net/src/builder.rs

@@ -1,6 +1,7 @@
 //! Network Builder Module.
 
 use alloy::primitives::Address;
+use discv5::ListenConfig;
 use eyre::Result;
 use std::net::{IpAddr, SocketAddr};
 use tokio::sync::watch::channel;
@@ -24,8 +25,10 @@ pub struct NetworkDriverBuilder {
     pub chain_id: Option<u64>,
     /// The unsafe block signer.
     pub unsafe_block_signer: Option<Address>,
-    /// The socket address that the service is listening on.
-    pub socket: Option<SocketAddr>,
+    /// The socket address that the gossip service is listening on.
+    pub gossip_addr: Option<SocketAddr>,
+    /// The listen config that the discovery service is listening on.
+    pub discovery_addr: Option<ListenConfig>,
     /// The [GossipConfig] constructs the config for `gossipsub`.
     pub gossip_config: Option<GossipConfig>,
     /// The [Keypair] for the node.
@@ -56,9 +59,15 @@ impl NetworkDriverBuilder {
         self
     }
 
-    /// Specifies the socket address that the service is listening on.
-    pub fn with_socket(&mut self, socket: SocketAddr) -> &mut Self {
-        self.socket = Some(socket);
+    /// Specifies the socket address that the gossip service is listening on.
+    pub fn with_gossip_addr(&mut self, socket: SocketAddr) -> &mut Self {
+        self.gossip_addr = Some(socket);
+        self
+    }
+
+    /// Specifies the listen config that the discovery service is listening on.
+    pub fn with_discovery_addr(&mut self, listen_config: ListenConfig) -> &mut Self {
+        self.discovery_addr = Some(listen_config);
         self
     }
 
@@ -109,7 +118,7 @@ impl NetworkDriverBuilder {
     /// let mut builder = NetworkDriverBuilder::new()
     ///    .with_unsafe_block_signer(signer)
     ///    .with_chain_id(chain_id)
-    ///    .with_socket(socket)
+    ///    .with_gossip_addr(socket)
     ///    .with_gossip_config(cfg);
     ///    .build()
     ///    .unwrap();
@@ -126,7 +135,11 @@ impl NetworkDriverBuilder {
     /// Returns an error if any of the following required fields are not set:
     /// - [NetworkDriverBuilder::unsafe_block_signer]
     /// - [NetworkDriverBuilder::chain_id]
-    /// - [NetworkDriverBuilder::socket]
+    /// - [NetworkDriverBuilder::gossip_addr]
+    ///
+    /// If explicitly set, the following fields are used for discovery address, otherwise the gossip
+    /// address is used:
+    /// - [NetworkDriverBuilder::discovery_addr]
     ///
     /// Set these fields using the respective methods on the [NetworkDriverBuilder]
     /// before calling this method.
@@ -143,7 +156,21 @@ impl NetworkDriverBuilder {
     /// let driver = NetworkDriverBuilder::new()
     ///    .with_unsafe_block_signer(signer)
     ///    .with_chain_id(chain_id)
-    ///    .with_socket(socket)
+    ///    .with_gossip_addr(socket)
+    ///    .build()
+    ///    .unwrap();
+    ///
+    /// Or if you want to use a different discovery address:
+    ///
+    /// let chain_id = 10;
+    /// let signer = Address::random();
+    /// let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9099);
+    /// let listen_config = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9999);
+    /// let driver = NetworkDriverBuilder::new()
+    ///    .with_unsafe_block_signer(signer)
+    ///    .with_chain_id(chain_id)
+    ///    .with_gossip_addr(socket)
+    ///    .with_discovery_addr(listen_config)
     ///    .build()
     ///    .unwrap();
     /// ```
@@ -179,19 +206,27 @@ impl NetworkDriverBuilder {
             )?
             .with_behaviour(|_| behaviour)?
             .build();
-        let socket = self.socket.take().ok_or(eyre::eyre!("socket address not set"))?;
+
+        let gossip_addr =
+            self.gossip_addr.take().ok_or(eyre::eyre!("gossip_addr address not set"))?;
         let mut multiaddr = Multiaddr::empty();
-        match socket.ip() {
+        match gossip_addr.ip() {
             IpAddr::V4(ip) => multiaddr.push(Protocol::Ip4(ip)),
             IpAddr::V6(ip) => multiaddr.push(Protocol::Ip6(ip)),
         }
-        multiaddr.push(Protocol::Tcp(socket.port()));
-        let gossip = GossipDriver::new(swarm, multiaddr, handler);
+        multiaddr.push(Protocol::Tcp(gossip_addr.port()));
+        let gossip = GossipDriver::new(swarm, multiaddr, handler.clone());
 
         // Build the discovery service
-        let discovery =
-            DiscoveryBuilder::new().with_address(socket).with_chain_id(chain_id).build()?;
+        let 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
+        }
+        .build()?;
         Ok(NetworkDriver { unsafe_block_recv, unsafe_block_signer_sender, gossip, discovery })
     }
 }
@@ -227,7 +262,7 @@ mod tests {
         else {
             panic!("expected error when building NetworkDriver without socket");
         };
-        assert_eq!(err.to_string(), "socket address not set");
+        assert_eq!(err.to_string(), "gossip_addr address not set");
     }
 
     #[test]
@@ -239,7 +274,7 @@ mod tests {
         let driver = NetworkDriverBuilder::new()
             .with_unsafe_block_signer(signer)
             .with_chain_id(id)
-            .with_socket(socket)
+            .with_gossip_addr(socket)
             .with_gossip_config(cfg)
             .build()
             .unwrap();
@@ -272,7 +307,7 @@ mod tests {
         let driver = NetworkDriverBuilder::new()
             .with_unsafe_block_signer(signer)
             .with_chain_id(id)
-            .with_socket(socket)
+            .with_gossip_addr(socket)
             .build()
             .unwrap();
         let mut multiaddr = Multiaddr::empty();
@@ -285,6 +320,7 @@ mod tests {
         // Driver Assertions
         assert_eq!(driver.gossip.addr, multiaddr);
         assert_eq!(driver.discovery.chain_id, id);
+        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9099);
 
         // Block Handler Assertions
         assert_eq!(driver.gossip.handler.chain_id, id);
@@ -295,4 +331,21 @@ mod tests {
         let v3 = IdentTopic::new(format!("/optimism/{}/2/blocks", id));
         assert_eq!(driver.gossip.handler.blocks_v3_topic.hash(), v3.hash());
     }
+
+    #[test]
+    fn test_build_network_driver_with_discovery_addr() {
+        let id = 10;
+        let signer = Address::random();
+        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9099);
+        let discovery_addr = ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9098);
+        let driver = NetworkDriverBuilder::new()
+            .with_unsafe_block_signer(signer)
+            .with_chain_id(id)
+            .with_gossip_addr(socket)
+            .with_discovery_addr(discovery_addr)
+            .build()
+            .unwrap();
+
+        assert_eq!(driver.discovery.disc.local_enr().tcp4().unwrap(), 9098);
+    }
 }

crates/net/src/discovery/builder.rs

@@ -5,7 +5,7 @@ use discv5::{
     enr::{CombinedKey, Enr},
     ConfigBuilder, Discv5, ListenConfig,
 };
-use eyre::Result;
+use eyre::{Report, Result};
 use std::net::SocketAddr;
 
 use crate::types::enr::OP_CL_KEY;
@@ -17,6 +17,8 @@ pub struct DiscoveryBuilder {
     address: Option<SocketAddr>,
     /// The chain ID of the network.
     chain_id: Option<u64>,
+    /// The listen config for the discovery service.
+    listen_config: Option<ListenConfig>,
 }
 
 impl DiscoveryBuilder {
@@ -37,18 +39,43 @@ impl DiscoveryBuilder {
         self
     }
 
+    /// Sets the listen config for the discovery service.
+    pub fn with_listen_config(mut self, listen_config: ListenConfig) -> Self {
+        self.listen_config = Some(listen_config);
+        self
+    }
+
     /// Builds a [DiscoveryDriver].
     pub fn build(&mut self) -> Result<DiscoveryDriver> {
-        let addr = self.address.ok_or_else(|| eyre::eyre!("address not set"))?;
         let chain_id = self.chain_id.ok_or_else(|| eyre::eyre!("chain ID not set"))?;
         let opstack = OpStackEnr::new(chain_id, 0);
         let mut opstack_data = Vec::new();
         use alloy_rlp::Encodable;
         opstack.encode(&mut opstack_data);
 
         let key = CombinedKey::generate_secp256k1();
-        let enr = Enr::builder().add_value_rlp(OP_CL_KEY, opstack_data.into()).build(&key)?;
-        let listen_config = ListenConfig::from_ip(addr.ip(), addr.port());
+        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 {
+            ListenConfig::Ipv4 { ip, port } => {
+                enr_builder.ip4(ip).tcp4(port);
+            }
+            ListenConfig::Ipv6 { ip, port } => {
+                enr_builder.ip6(ip).tcp6(port);
+            }
+            ListenConfig::DualStack { ipv4, ipv4_port, ipv6, ipv6_port } => {
+                enr_builder.ip4(ipv4).tcp4(ipv4_port);
+                enr_builder.ip6(ipv6).tcp6(ipv6_port);
+            }
+        }
+        let enr = enr_builder.build(&key)?;
         let config = ConfigBuilder::new(listen_config).build();
 
         let disc = Discv5::new(enr, key, config)