Broadcast external IP from NAT in enode record (#10274)

Co-authored-by: Matthias Seitz <[email protected]>

book/cli/reth/debug/execution.md

@@ -82,6 +82,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/debug/in-memory-merkle.md

@@ -82,6 +82,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/debug/merkle.md

@@ -82,6 +82,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/debug/replay-engine.md

@@ -82,6 +82,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/node.md

@@ -74,6 +74,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/p2p.md

@@ -59,6 +59,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/stage/run.md

@@ -125,6 +125,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

book/cli/reth/stage/unwind.md

@@ -87,6 +87,9 @@ Networking:
       --enable-discv5-discovery
           Enable Discv5 discovery
 
+      --disable-nat
+          Disable Nat discovery
+
       --discovery.addr <DISCOVERY_ADDR>
           The UDP address to use for devp2p peer discovery version 4
 

crates/net/nat/src/lib.rs

@@ -61,6 +61,14 @@ impl NatResolver {
     pub async fn external_addr(self) -> Option<IpAddr> {
         external_addr_with(self).await
     }
+
+    /// Returns the external ip, if it is [`NatResolver::ExternalIp`]
+    pub const fn as_external_ip(self) -> Option<IpAddr> {
+        match self {
+            Self::ExternalIp(ip) => Some(ip),
+            _ => None,
+        }
+    }
 }
 
 impl fmt::Display for NatResolver {

crates/net/network/src/config.rs

@@ -81,6 +81,8 @@ pub struct NetworkConfig<C> {
     pub tx_gossip_disabled: bool,
     /// How to instantiate transactions manager.
     pub transactions_manager_config: TransactionsManagerConfig,
+    /// The NAT resolver for external IP
+    pub nat: Option<NatResolver>,
 }
 
 // === impl NetworkConfig ===
@@ -197,6 +199,8 @@ pub struct NetworkConfigBuilder {
     block_import: Option<Box<dyn BlockImport>>,
     /// How to instantiate transactions manager.
     transactions_manager_config: TransactionsManagerConfig,
+    /// The NAT resolver for external IP
+    nat: Option<NatResolver>,
 }
 
 // === impl NetworkConfigBuilder ===
@@ -228,6 +232,7 @@ impl NetworkConfigBuilder {
             tx_gossip_disabled: false,
             block_import: None,
             transactions_manager_config: Default::default(),
+            nat: None,
         }
     }
 
@@ -369,6 +374,7 @@ impl NetworkConfigBuilder {
         self.discovery_v4_builder
             .get_or_insert_with(Discv4Config::builder)
             .external_ip_resolver(Some(resolver));
+        self.nat = Some(resolver);
         self
     }
 
@@ -417,9 +423,15 @@ impl NetworkConfigBuilder {
         self
     }
 
+    // Disable nat
+    pub const fn disable_nat(mut self) -> Self {
+        self.nat = None;
+        self
+    }
+
     /// Disables all discovery.
     pub fn disable_discovery(self) -> Self {
-        self.disable_discv4_discovery().disable_dns_discovery()
+        self.disable_discv4_discovery().disable_dns_discovery().disable_nat()
     }
 
     /// Disables all discovery if the given condition is true.
@@ -485,6 +497,12 @@ impl NetworkConfigBuilder {
         self.build(NoopBlockReader::new(chain_spec))
     }
 
+    /// Sets the NAT resolver for external IP.
+    pub const fn add_nat(mut self, nat: Option<NatResolver>) -> Self {
+        self.nat = nat;
+        self
+    }
+
     /// Consumes the type and creates the actual [`NetworkConfig`]
     /// for the given client type that can interact with the chain.
     ///
@@ -515,6 +533,7 @@ impl NetworkConfigBuilder {
             tx_gossip_disabled,
             block_import,
             transactions_manager_config,
+            nat,
         } = self;
 
         discovery_v5_builder = discovery_v5_builder.map(|mut builder| {
@@ -581,6 +600,7 @@ impl NetworkConfigBuilder {
             fork_filter,
             tx_gossip_disabled,
             transactions_manager_config,
+            nat,
         }
     }
 }

crates/net/network/src/manager.rs

@@ -188,6 +188,7 @@ impl NetworkManager {
             extra_protocols,
             tx_gossip_disabled,
             transactions_manager_config: _,
+            nat,
         } = config;
 
         let peers_manager = PeersManager::new(peers_config);
@@ -267,6 +268,7 @@ impl NetworkManager {
             discv4,
             discv5,
             event_sender.clone(),
+            nat,
         );
 
         Ok(Self {

crates/net/network/src/network.rs

@@ -9,7 +9,7 @@ use std::{
 use alloy_primitives::B256;
 use enr::Enr;
 use parking_lot::Mutex;
-use reth_discv4::Discv4;
+use reth_discv4::{Discv4, NatResolver};
 use reth_discv5::Discv5;
 use reth_eth_wire::{DisconnectReason, NewBlock, NewPooledTransactionHashes, SharedTransactions};
 use reth_network_api::{
@@ -65,6 +65,7 @@ impl NetworkHandle {
         discv4: Option<Discv4>,
         discv5: Option<Discv5>,
         event_sender: EventSender<NetworkEvent>,
+        nat: Option<NatResolver>,
     ) -> Self {
         let inner = NetworkInner {
             num_active_peers,
@@ -81,6 +82,7 @@ impl NetworkHandle {
             discv4,
             discv5,
             event_sender,
+            nat,
         };
         Self { inner: Arc::new(inner) }
     }
@@ -216,10 +218,13 @@ impl PeersInfo for NetworkHandle {
         } else if let Some(record) = self.inner.discv5.as_ref().and_then(|d| d.node_record()) {
             record
         } else {
-            let id = *self.peer_id();
-            let mut socket_addr = *self.inner.listener_address.lock();
+            let external_ip = self.inner.nat.and_then(|nat| nat.as_external_ip());
 
-            if socket_addr.ip().is_unspecified() {
+            let mut socket_addr = *self.inner.listener_address.lock();
+            if let Some(ip) = external_ip {
+                // if able to resolve external ip, use it instead and also set the local address
+                socket_addr.set_ip(ip)
+            } else if socket_addr.ip().is_unspecified() {
                 // zero address is invalid
                 if socket_addr.ip().is_ipv4() {
                     socket_addr.set_ip(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST));
@@ -228,7 +233,7 @@ impl PeersInfo for NetworkHandle {
                 }
             }
 
-            NodeRecord::new(socket_addr, id)
+            NodeRecord::new(socket_addr, *self.peer_id())
         }
     }
 
@@ -432,6 +437,8 @@ struct NetworkInner {
     discv5: Option<Discv5>,
     /// Sender for high level network events.
     event_sender: EventSender<NetworkEvent>,
+    /// The NAT resolver
+    nat: Option<NatResolver>,
 }
 
 /// Provides access to modify the network's additional protocol handlers.

crates/net/network/tests/it/startup.rs

@@ -1,10 +1,10 @@
 use std::{
     io,
-    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
 };
 
 use reth_chainspec::MAINNET;
-use reth_discv4::Discv4Config;
+use reth_discv4::{Discv4Config, NatResolver};
 use reth_network::{
     error::{NetworkError, ServiceKind},
     Discovery, NetworkConfigBuilder, NetworkManager,
@@ -105,3 +105,32 @@ async fn test_tcp_port_node_record_discovery() {
     assert_eq!(record.tcp_port, local_addr.port());
     assert_ne!(record.udp_port, 0);
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_node_record_address_with_nat() {
+    let secret_key = SecretKey::new(&mut rand::thread_rng());
+    let config = NetworkConfigBuilder::new(secret_key)
+        .add_nat(Some(NatResolver::ExternalIp("10.1.1.1".parse().unwrap())))
+        .disable_discv4_discovery()
+        .disable_dns_discovery()
+        .build_with_noop_provider(MAINNET.clone());
+
+    let network = NetworkManager::new(config).await.unwrap();
+    let record = network.handle().local_node_record();
+
+    assert_eq!(record.address, IpAddr::V4(Ipv4Addr::new(10, 1, 1, 1)));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_node_record_address_with_nat_disable_discovery() {
+    let secret_key = SecretKey::new(&mut rand::thread_rng());
+    let config = NetworkConfigBuilder::new(secret_key)
+        .add_nat(Some(NatResolver::ExternalIp("10.1.1.1".parse().unwrap())))
+        .disable_discovery()
+        .build_with_noop_provider(MAINNET.clone());
+
+    let network = NetworkManager::new(config).await.unwrap();
+    let record = network.handle().local_node_record();
+
+    assert_eq!(record.address, IpAddr::V4(std::net::Ipv4Addr::LOCALHOST));
+}

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

@@ -172,7 +172,7 @@ impl NetworkArgs {
 
                     DEFAULT_DISCOVERY_ADDR
                 }
-            }
+            };
         }
 
         self.addr
@@ -355,6 +355,10 @@ pub struct DiscoveryArgs {
     #[arg(long, conflicts_with = "disable_discovery")]
     pub enable_discv5_discovery: bool,
 
+    /// Disable Nat discovery.
+    #[arg(long, conflicts_with = "disable_discovery")]
+    pub disable_nat: bool,
+
     /// The UDP address to use for devp2p peer discovery version 4.
     #[arg(id = "discovery.addr", long = "discovery.addr", value_name = "DISCOVERY_ADDR", default_value_t = DEFAULT_DISCOVERY_ADDR)]
     pub addr: IpAddr,
@@ -418,6 +422,10 @@ impl DiscoveryArgs {
             network_config_builder = network_config_builder.disable_discv4_discovery();
         }
 
+        if self.disable_discovery || self.disable_nat {
+            network_config_builder = network_config_builder.disable_nat();
+        }
+
         if !self.disable_discovery && self.enable_discv5_discovery {
             network_config_builder = network_config_builder
                 .discovery_v5(self.discovery_v5_builder(rlpx_tcp_socket, boot_nodes));
@@ -494,6 +502,7 @@ impl Default for DiscoveryArgs {
             disable_dns_discovery: false,
             disable_discv4_discovery: false,
             enable_discv5_discovery: false,
+            disable_nat: false,
             addr: DEFAULT_DISCOVERY_ADDR,
             port: DEFAULT_DISCOVERY_PORT,
             discv5_addr: None,