From 37ca7e0b33662190cf437a66050e6a7634548c04 Mon Sep 17 00:00:00 2001
From: Ikechukwu Ahiara Marvellous <ahiaraikechukwu@gmail.com>
Date: Wed, 18 Jan 2023 02:53:15 +0100
Subject: [PATCH] feat: add --nat flag to reth node (#888)

Co-authored-by: Georgios Konstantopoulos <me@gakonst.com>
---
 Cargo.lock               | 2 ++
 bin/reth/Cargo.toml      | 2 ++
 bin/reth/src/config.rs   | 3 +++
 bin/reth/src/node/mod.rs | 5 +++++
 4 files changed, 12 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 390af41439..e2cf3c1d1f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3928,9 +3928,11 @@ dependencies = [
  "reth-cli-utils",
  "reth-consensus",
  "reth-db",
+ "reth-discv4",
  "reth-downloaders",
  "reth-executor",
  "reth-interfaces",
+ "reth-net-nat",
  "reth-network",
  "reth-network-api",
  "reth-primitives",
diff --git a/bin/reth/Cargo.toml b/bin/reth/Cargo.toml
index 0aff9f0506..7772bb52ad 100644
--- a/bin/reth/Cargo.toml
+++ b/bin/reth/Cargo.toml
@@ -24,6 +24,8 @@ reth-network-api = {path = "../../crates/net/network-api" }
 reth-downloaders = {path = "../../crates/net/downloaders" }
 reth-cli-utils = { path = "../../crates/cli/utils" }
 reth-tracing = { path = "../../crates/tracing" }
+reth-net-nat = { path = "../../crates/net/nat" }
+reth-discv4 = { path = "../../crates/net/discv4" }
 
 # tracing
 tracing = "0.1"
diff --git a/bin/reth/src/config.rs b/bin/reth/src/config.rs
index 1e7be4cc6b..e4d4e641e4 100644
--- a/bin/reth/src/config.rs
+++ b/bin/reth/src/config.rs
@@ -2,6 +2,7 @@
 use std::sync::Arc;
 
 use reth_db::database::Database;
+use reth_discv4::Discv4Config;
 use reth_network::{
     config::{mainnet_nodes, rng_secret_key},
     NetworkConfig, PeersConfig,
@@ -29,6 +30,7 @@ impl Config {
         chain_spec: ChainSpec,
         disable_discovery: bool,
         bootnodes: Option<Vec<NodeRecord>>,
+        nat_resolution_method: reth_net_nat::NatResolver,
     ) -> NetworkConfig<ProviderImpl<DB>> {
         let peer_config = reth_network::PeersConfig::default()
             .with_trusted_nodes(self.peers.trusted_nodes.clone())
@@ -36,6 +38,7 @@ impl Config {
         NetworkConfig::builder(Arc::new(ProviderImpl::new(db)), rng_secret_key())
             .boot_nodes(bootnodes.unwrap_or_else(mainnet_nodes))
             .peer_config(peer_config)
+            .discovery(Discv4Config::builder().external_ip_resolver(Some(nat_resolution_method)))
             .chain_spec(chain_spec)
             .set_discovery(disable_discovery)
             .build()
diff --git a/bin/reth/src/node/mod.rs b/bin/reth/src/node/mod.rs
index 0bf549c305..7169993283 100644
--- a/bin/reth/src/node/mod.rs
+++ b/bin/reth/src/node/mod.rs
@@ -16,6 +16,7 @@ use reth_cli_utils::init::init_genesis;
 use reth_consensus::BeaconConsensus;
 use reth_downloaders::{bodies, headers};
 use reth_interfaces::consensus::ForkchoiceState;
+use reth_net_nat::NatResolver;
 use reth_network::NetworkEvent;
 use reth_network_api::NetworkInfo;
 use reth_primitives::{BlockNumber, ChainSpec, NodeRecord, H256};
@@ -83,6 +84,9 @@ pub struct Command {
 
     #[arg(long, value_delimiter = ',')]
     bootnodes: Option<Vec<NodeRecord>>,
+
+    #[arg(long, default_value = "any")]
+    nat: NatResolver,
 }
 
 impl Command {
@@ -126,6 +130,7 @@ impl Command {
                 self.chain.clone(),
                 self.network.disable_discovery,
                 self.bootnodes.clone(),
+                self.nat,
             )
             .start_network()
             .await?;