From 3053748321c0d6967818316723100331e8feb5ea Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Sun, 10 Sep 2023 10:59:42 +0200
Subject: [PATCH] feat: add discovery port function

---
 crates/net/discv4/src/lib.rs     |  8 +++++++-
 crates/net/network/src/config.rs | 35 ++++++++++++++++----------------
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs
index fa7077bc3122..36dde7c6e96a 100644
--- a/crates/net/discv4/src/lib.rs
+++ b/crates/net/discv4/src/lib.rs
@@ -53,7 +53,7 @@ use std::{
     cell::RefCell,
     collections::{btree_map, hash_map::Entry, BTreeMap, HashMap, VecDeque},
     io,
-    net::{IpAddr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
     pin::Pin,
     rc::Rc,
     sync::Arc,
@@ -96,6 +96,12 @@ pub use reth_net_nat::{external_ip, NatResolver};
 /// Note: the default TCP port is the same.
 pub const DEFAULT_DISCOVERY_PORT: u16 = 30303;
 
+/// The default address for discv4 via UDP: "0.0.0.0:30303"
+///
+/// Note: The default TCP address is the same.
+pub const DEFAULT_DISCOVERY_ADDRESS: SocketAddr =
+    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT));
+
 /// The maximum size of any packet is 1280 bytes.
 const MAX_PACKET_SIZE: usize = 1280;
 
diff --git a/crates/net/network/src/config.rs b/crates/net/network/src/config.rs
index 0b093e8f1250..c1492f8ae6b0 100644
--- a/crates/net/network/src/config.rs
+++ b/crates/net/network/src/config.rs
@@ -7,7 +7,7 @@ use crate::{
     session::SessionsConfig,
     NetworkHandle, NetworkManager,
 };
-use reth_discv4::{Discv4Config, Discv4ConfigBuilder, DEFAULT_DISCOVERY_PORT};
+use reth_discv4::{Discv4Config, Discv4ConfigBuilder, DEFAULT_DISCOVERY_ADDRESS};
 use reth_dns_discovery::DnsDiscoveryConfig;
 use reth_ecies::util::pk2id;
 use reth_eth_wire::{HelloMessage, Status};
@@ -17,11 +17,7 @@ use reth_primitives::{
 use reth_provider::{BlockReader, HeaderProvider};
 use reth_tasks::{TaskSpawner, TokioTaskExecutor};
 use secp256k1::SECP256K1;
-use std::{
-    collections::HashSet,
-    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
-    sync::Arc,
-};
+use std::{collections::HashSet, net::SocketAddr, sync::Arc};
 // re-export for convenience
 pub use secp256k1::SecretKey;
 
@@ -244,14 +240,15 @@ impl NetworkConfigBuilder {
     /// This is a convenience function for both [NetworkConfigBuilder::listener_addr] and
     /// [NetworkConfigBuilder::discovery_addr].
     ///
-    /// By default, both are on the same port: [DEFAULT_DISCOVERY_PORT]
+    /// By default, both are on the same port:
+    /// [DEFAULT_DISCOVERY_PORT](reth_discv4::DEFAULT_DISCOVERY_PORT)
     pub fn set_addrs(self, addr: SocketAddr) -> Self {
         self.listener_addr(addr).discovery_addr(addr)
     }
 
     /// Sets the socket address the network will listen on.
     ///
-    /// By default, this is [Ipv4Addr::UNSPECIFIED] on [DEFAULT_DISCOVERY_PORT]
+    /// By default, this is [DEFAULT_DISCOVERY_ADDRESS]
     pub fn listener_addr(mut self, listener_addr: SocketAddr) -> Self {
         self.listener_addr = Some(listener_addr);
         self
@@ -259,11 +256,9 @@ impl NetworkConfigBuilder {
 
     /// Sets the port of the address the network will listen on.
     ///
-    /// By default, this is [DEFAULT_DISCOVERY_PORT]
+    /// By default, this is [DEFAULT_DISCOVERY_PORT](reth_discv4::DEFAULT_DISCOVERY_PORT)
     pub fn listener_port(mut self, port: u16) -> Self {
-        self.listener_addr
-            .get_or_insert(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT).into())
-            .set_port(port);
+        self.listener_addr.get_or_insert(DEFAULT_DISCOVERY_ADDRESS).set_port(port);
         self
     }
 
@@ -273,6 +268,14 @@ impl NetworkConfigBuilder {
         self
     }
 
+    /// Sets the port of the address the discovery network will listen on.
+    ///
+    /// By default, this is [DEFAULT_DISCOVERY_PORT](reth_discv4::DEFAULT_DISCOVERY_PORT)
+    pub fn discovery_port(mut self, port: u16) -> Self {
+        self.discovery_addr.get_or_insert(DEFAULT_DISCOVERY_ADDRESS).set_port(port);
+        self
+    }
+
     /// Sets the discv4 config to use.
     pub fn discovery(mut self, builder: Discv4ConfigBuilder) -> Self {
         self.discovery_v4_builder = Some(builder);
@@ -369,9 +372,7 @@ impl NetworkConfigBuilder {
             head,
         } = self;
 
-        let listener_addr = listener_addr.unwrap_or_else(|| {
-            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT))
-        });
+        let listener_addr = listener_addr.unwrap_or(DEFAULT_DISCOVERY_ADDRESS);
 
         let mut hello_message =
             hello_message.unwrap_or_else(|| HelloMessage::builder(peer_id).build());
@@ -408,9 +409,7 @@ impl NetworkConfigBuilder {
             boot_nodes,
             dns_discovery_config,
             discovery_v4_config: discovery_v4_builder.map(|builder| builder.build()),
-            discovery_addr: discovery_addr.unwrap_or_else(|| {
-                SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, DEFAULT_DISCOVERY_PORT))
-            }),
+            discovery_addr: discovery_addr.unwrap_or(DEFAULT_DISCOVERY_ADDRESS),
             listener_addr,
             peers_config: peers_config.unwrap_or_default(),
             sessions_config: sessions_config.unwrap_or_default(),