diff --git a/crates/net/README.md b/crates/net/README.md index e484424..206c232 100644 --- a/crates/net/README.md +++ b/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"); diff --git a/crates/net/src/builder.rs b/crates/net/src/builder.rs index ff90911..d049d81 100644 --- a/crates/net/src/builder.rs +++ b/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, /// The unsafe block signer. pub unsafe_block_signer: Option
, - /// The socket address that the service is listening on. - pub socket: Option, + /// The socket address that the gossip service is listening on. + pub gossip_addr: Option, + /// The listen config that the discovery service is listening on. + pub discovery_addr: Option, /// The [GossipConfig] constructs the config for `gossipsub`. pub gossip_config: Option, /// 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); + } } diff --git a/crates/net/src/discovery/builder.rs b/crates/net/src/discovery/builder.rs index 1a092b4..b86c262 100644 --- a/crates/net/src/discovery/builder.rs +++ b/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, /// The chain ID of the network. chain_id: Option, + /// The listen config for the discovery service. + listen_config: Option, } impl DiscoveryBuilder { @@ -37,9 +39,14 @@ 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 { - 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(); @@ -47,8 +54,28 @@ impl DiscoveryBuilder { 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::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)