diff --git a/crates/net/network/src/discovery.rs b/crates/net/network/src/discovery.rs index 38346ec1d0..07d24859ed 100644 --- a/crates/net/network/src/discovery.rs +++ b/crates/net/network/src/discovery.rs @@ -71,14 +71,17 @@ impl Discovery { /// This will spawn the [`reth_discv4::Discv4Service`] onto a new task and establish a listener /// channel to receive all discovered nodes. pub async fn new( + tcp_addr: SocketAddr, discovery_v4_addr: SocketAddr, sk: SecretKey, discv4_config: Option, discv5_config: Option, // contains discv5 listen address dns_discovery_config: Option, ) -> Result { - // setup discv4 - let local_enr = NodeRecord::from_secret_key(discovery_v4_addr, &sk); + // setup discv4 with the discovery address and tcp port + let local_enr = + NodeRecord::from_secret_key(discovery_v4_addr, &sk).with_tcp_port(tcp_addr.port()); + let discv4_future = async { let Some(disc_config) = discv4_config else { return Ok((None, None, None)) }; let (discv4, mut discv4_service) = @@ -342,6 +345,7 @@ mod tests { let (secret_key, _) = SECP256K1.generate_keypair(&mut rng); let discovery_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)); let _discovery = Discovery::new( + discovery_addr, discovery_addr, secret_key, Default::default(), @@ -370,9 +374,16 @@ mod tests { .discv5_config(discv5::ConfigBuilder::new(discv5_listen_config).build()) .build(); - Discovery::new(discv4_addr, secret_key, Some(discv4_config), Some(discv5_config), None) - .await - .expect("should build discv5 with discv4 downgrade") + Discovery::new( + discv4_addr, + discv4_addr, + secret_key, + Some(discv4_config), + Some(discv5_config), + None, + ) + .await + .expect("should build discv5 with discv4 downgrade") } #[tokio::test(flavor = "multi_thread")] diff --git a/crates/net/network/src/listener.rs b/crates/net/network/src/listener.rs index d0654042ab..9fcc15a104 100644 --- a/crates/net/network/src/listener.rs +++ b/crates/net/network/src/listener.rs @@ -1,7 +1,6 @@ //! Contains connection-oriented interfaces. use futures::{ready, Stream}; - use std::{ io, net::SocketAddr, diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs index a3fdefffab..7a4092d4a4 100644 --- a/crates/net/network/src/manager.rs +++ b/crates/net/network/src/manager.rs @@ -197,7 +197,9 @@ where let incoming = ConnectionListener::bind(listener_addr).await.map_err(|err| { NetworkError::from_io_error(err, ServiceKind::Listener(listener_addr)) })?; - let listener_address = Arc::new(Mutex::new(incoming.local_address())); + + // retrieve the tcp address of the socket + let listener_addr = incoming.local_address(); // resolve boot nodes let mut resolved_boot_nodes = vec![]; @@ -214,6 +216,7 @@ where }); let discovery = Discovery::new( + listener_addr, discovery_v4_addr, secret_key, discovery_v4_config, @@ -248,7 +251,7 @@ where let handle = NetworkHandle::new( Arc::clone(&num_active_peers), - listener_address, + Arc::new(Mutex::new(listener_addr)), to_manager_tx, secret_key, local_peer_id, @@ -314,7 +317,7 @@ where NetworkBuilder { network: self, transactions: (), request_handler: () } } - /// Returns the [`SocketAddr`] that listens for incoming connections. + /// Returns the [`SocketAddr`] that listens for incoming tcp connections. pub const fn local_addr(&self) -> SocketAddr { self.swarm.listener().local_address() } diff --git a/crates/net/network/tests/it/startup.rs b/crates/net/network/tests/it/startup.rs index 0bef94e08c..e924d322d1 100644 --- a/crates/net/network/tests/it/startup.rs +++ b/crates/net/network/tests/it/startup.rs @@ -3,7 +3,7 @@ use reth_network::{ error::{NetworkError, ServiceKind}, Discovery, NetworkConfigBuilder, NetworkManager, }; -use reth_network_api::NetworkInfo; +use reth_network_api::{NetworkInfo, PeersInfo}; use reth_provider::test_utils::NoopProvider; use secp256k1::SecretKey; use std::{ @@ -59,8 +59,46 @@ async fn test_discovery_addr_in_use() { let any_port_listener = TcpListener::bind(addr).await.unwrap(); let port = any_port_listener.local_addr().unwrap().port(); let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port)); - let _discovery = Discovery::new(addr, secret_key, Some(disc_config), None, None).await.unwrap(); + let _discovery = + Discovery::new(addr, addr, secret_key, Some(disc_config), None, None).await.unwrap(); let disc_config = Discv4Config::default(); - let result = Discovery::new(addr, secret_key, Some(disc_config), None, None).await; + let result = Discovery::new(addr, addr, secret_key, Some(disc_config), None, None).await; assert!(is_addr_in_use_kind(&result.err().unwrap(), ServiceKind::Discovery(addr))); } + +// +#[tokio::test(flavor = "multi_thread")] +async fn test_tcp_port_node_record_no_discovery() { + let secret_key = SecretKey::new(&mut rand::thread_rng()); + let config = NetworkConfigBuilder::new(secret_key) + .listener_port(0) + .disable_discovery() + .build_with_noop_provider(); + let network = NetworkManager::new(config).await.unwrap(); + + let local_addr = network.local_addr(); + // ensure we retrieved the port the OS chose + assert_ne!(local_addr.port(), 0); + + let record = network.handle().local_node_record(); + assert_eq!(record.tcp_port, local_addr.port()); +} + +// +#[tokio::test(flavor = "multi_thread")] +async fn test_tcp_port_node_record_discovery() { + let secret_key = SecretKey::new(&mut rand::thread_rng()); + let config = NetworkConfigBuilder::new(secret_key) + .listener_port(0) + .discovery_port(0) + .build_with_noop_provider(); + let network = NetworkManager::new(config).await.unwrap(); + + let local_addr = network.local_addr(); + // ensure we retrieved the port the OS chose + assert_ne!(local_addr.port(), 0); + + let record = network.handle().local_node_record(); + assert_eq!(record.tcp_port, local_addr.port()); + assert_ne!(record.udp_port, 0); +} diff --git a/crates/net/peers/src/node_record.rs b/crates/net/peers/src/node_record.rs index 99ca74b569..3b6c38170d 100644 --- a/crates/net/peers/src/node_record.rs +++ b/crates/net/peers/src/node_record.rs @@ -42,8 +42,10 @@ pub struct NodeRecord { } impl NodeRecord { + /// Derive the [`NodeRecord`] from the secret key and addr. + /// + /// Note: this will set both the TCP and UDP ports to the port of the addr. #[cfg(feature = "secp256k1")] - /// Derive the [`NodeRecord`] from the secret key and addr pub fn from_secret_key(addr: SocketAddr, sk: &secp256k1::SecretKey) -> Self { let pk = secp256k1::PublicKey::from_secret_key(secp256k1::SECP256K1, sk); let id = PeerId::from_slice(&pk.serialize_uncompressed()[1..]); @@ -73,8 +75,19 @@ impl NodeRecord { self } + /// Sets the tcp port + pub const fn with_tcp_port(mut self, port: u16) -> Self { + self.tcp_port = port; + self + } + + /// Sets the udp port + pub const fn with_udp_port(mut self, port: u16) -> Self { + self.udp_port = port; + self + } + /// Creates a new record from a socket addr and peer id. - #[allow(dead_code)] pub const fn new(addr: SocketAddr, id: PeerId) -> Self { Self { address: addr.ip(), tcp_port: addr.port(), udp_port: addr.port(), id } }