From 238eea37cf23bd212823d76d2d07065a59ce7e2a Mon Sep 17 00:00:00 2001
From: Bjerg <onbjerg@users.noreply.github.com>
Date: Sat, 20 May 2023 09:24:31 +0200
Subject: [PATCH] feat: add node identity to networking stack (#2758)

---
 bin/reth/src/args/network_args.rs | 16 +++++++++++++---
 crates/net/eth-wire/src/hello.rs  |  2 +-
 crates/net/network/src/lib.rs     |  2 +-
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/bin/reth/src/args/network_args.rs b/bin/reth/src/args/network_args.rs
index 23a995f0f72d..94ed5d90e93f 100644
--- a/bin/reth/src/args/network_args.rs
+++ b/bin/reth/src/args/network_args.rs
@@ -1,8 +1,9 @@
 //! clap [Args](clap::Args) for network related arguments.
 
+use crate::version::P2P_VERSION;
 use clap::Args;
 use reth_net_nat::NatResolver;
-use reth_network::NetworkConfigBuilder;
+use reth_network::{HelloMessage, NetworkConfigBuilder};
 use reth_primitives::{mainnet_nodes, ChainSpec, NodeRecord};
 use reth_staged_sync::Config;
 use secp256k1::SecretKey;
@@ -36,6 +37,10 @@ pub struct NetworkArgs {
     #[arg(long, value_name = "FILE", verbatim_doc_comment, conflicts_with = "no_persist_peers")]
     pub peers_file: Option<PathBuf>,
 
+    /// Custom node identity
+    #[arg(long, value_name = "IDENTITY", default_value = P2P_VERSION)]
+    pub identity: String,
+
     /// Secret key to use for this node.
     ///
     /// This will also deterministically set the peer ID. If not specified, it will be set in the
@@ -71,14 +76,19 @@ impl NetworkArgs {
         default_peers_file: PathBuf,
     ) -> NetworkConfigBuilder {
         let chain_bootnodes = chain_spec.chain.bootnodes().unwrap_or_else(mainnet_nodes);
-
         let peers_file = self.peers_file.clone().unwrap_or(default_peers_file);
 
-        let network_config_builder = config
+        // Configure basic network stack.
+        let mut network_config_builder = config
             .network_config(self.nat, self.persistent_peers_file(peers_file), secret_key)
             .boot_nodes(self.bootnodes.clone().unwrap_or(chain_bootnodes))
             .chain_spec(chain_spec);
 
+        // Configure node identity
+        let peer_id = network_config_builder.get_peer_id();
+        network_config_builder = network_config_builder
+            .hello_message(HelloMessage::builder(peer_id).client_version(&self.identity).build());
+
         self.discovery.apply_to_builder(network_config_builder)
     }
 }
diff --git a/crates/net/eth-wire/src/hello.rs b/crates/net/eth-wire/src/hello.rs
index 4ab36b05f5d4..c14759fd40ba 100644
--- a/crates/net/eth-wire/src/hello.rs
+++ b/crates/net/eth-wire/src/hello.rs
@@ -84,7 +84,7 @@ impl HelloMessageBuilder {
         self
     }
 
-    /// Sets client version.
+    /// Sets protocol version.
     pub fn protocol_version(mut self, protocol_version: ProtocolVersion) -> Self {
         self.protocol_version = Some(protocol_version);
         self
diff --git a/crates/net/network/src/lib.rs b/crates/net/network/src/lib.rs
index 79971e769936..f6e424eb9b09 100644
--- a/crates/net/network/src/lib.rs
+++ b/crates/net/network/src/lib.rs
@@ -147,4 +147,4 @@ pub use network::NetworkHandle;
 pub use peers::PeersConfig;
 pub use session::{PeerInfo, SessionsConfig};
 
-pub use reth_eth_wire::DisconnectReason;
+pub use reth_eth_wire::{DisconnectReason, HelloBuilder, HelloMessage};