From 65f77db97a6fd91c51869aeea27a43f95a60434a Mon Sep 17 00:00:00 2001
From: Milos Stankovic <82043364+morph-dev@users.noreply.github.com>
Date: Wed, 25 Sep 2024 16:23:27 +0300
Subject: [PATCH] refactor: use NodeId instead of String for discv5 json types

---
 ethportal-api/src/types/discv5.rs         | 12 ++++--------
 ethportal-peertest/src/scenarios/basic.rs |  6 ++++--
 portalnet/src/discovery.rs                |  6 +++---
 portalnet/src/overlay/protocol.rs         |  2 +-
 4 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/ethportal-api/src/types/discv5.rs b/ethportal-api/src/types/discv5.rs
index 28138df7f..98bb903ad 100644
--- a/ethportal-api/src/types/discv5.rs
+++ b/ethportal-api/src/types/discv5.rs
@@ -1,5 +1,4 @@
 use super::enr::Enr;
-use crate::utils::bytes::hex_encode;
 use serde::{Deserialize, Serialize};
 
 use discv5::enr::NodeId;
@@ -9,7 +8,7 @@ use discv5::enr::NodeId;
 #[serde(transparent)]
 pub struct Bucket {
     #[serde(skip_serializing_if = "Vec::is_empty")]
-    pub node_ids: Vec<String>,
+    pub node_ids: Vec<NodeId>,
 }
 
 /// Represents a discv5 kbuckets table
@@ -24,7 +23,7 @@ pub struct KBucketsTable {
 #[serde(rename_all = "camelCase")]
 pub struct NodeInfo {
     pub enr: Enr,
-    pub node_id: String,
+    pub node_id: NodeId,
     pub ip: Option<String>,
 }
 
@@ -32,7 +31,7 @@ pub struct NodeInfo {
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct RoutingTableInfo {
-    pub local_node_id: String,
+    pub local_node_id: NodeId,
     pub buckets: KBucketsTable,
 }
 
@@ -41,10 +40,7 @@ impl<TVal: Eq> From<discv5::kbucket::KBucketsTable<NodeId, TVal>> for KBucketsTa
         let buckets = table
             .buckets_iter()
             .map(|bucket| Bucket {
-                node_ids: bucket
-                    .iter()
-                    .map(|node| hex_encode(*node.key.preimage()))
-                    .collect(),
+                node_ids: bucket.iter().map(|node| *node.key.preimage()).collect(),
             })
             .collect();
 
diff --git a/ethportal-peertest/src/scenarios/basic.rs b/ethportal-peertest/src/scenarios/basic.rs
index 96fc14d90..5277d6ad9 100644
--- a/ethportal-peertest/src/scenarios/basic.rs
+++ b/ethportal-peertest/src/scenarios/basic.rs
@@ -25,12 +25,14 @@ pub async fn test_discv5_node_info(peertest: &Peertest) {
 
 pub async fn test_discv5_routing_table_info(target: &Client) {
     info!("Testing discv5_routingTableInfo");
+    let node_info = target.node_info().await.unwrap();
     let result = Discv5ApiClient::routing_table_info(target).await.unwrap();
-    assert!(result.local_node_id.starts_with("0x"));
+    assert_eq!(result.local_node_id, node_info.node_id);
 }
 
 pub async fn test_routing_table_info(subnetwork: Subnetwork, target: &Client) {
     info!("Testing routing_table_info for {subnetwork}");
+    let node_info = target.node_info().await.unwrap();
     let result = match subnetwork {
         Subnetwork::Beacon => BeaconNetworkApiClient::routing_table_info(target),
         Subnetwork::History => HistoryNetworkApiClient::routing_table_info(target),
@@ -39,7 +41,7 @@ pub async fn test_routing_table_info(subnetwork: Subnetwork, target: &Client) {
     }
     .await
     .unwrap();
-    assert!(result.local_node_id.starts_with("0x"));
+    assert_eq!(result.local_node_id, node_info.node_id);
 }
 
 pub async fn test_radius(subnetwork: Subnetwork, target: &Client) {
diff --git a/portalnet/src/discovery.rs b/portalnet/src/discovery.rs
index 10b8d49fa..25cb7681e 100644
--- a/portalnet/src/discovery.rs
+++ b/portalnet/src/discovery.rs
@@ -26,7 +26,7 @@ use super::config::PortalnetConfig;
 use crate::socket;
 use ethportal_api::{
     types::{discv5::RoutingTableInfo, enr::Enr, network::Subnetwork, portal_wire::NetworkSpec},
-    utils::bytes::{hex_decode, hex_encode},
+    utils::bytes::hex_decode,
     NodeInfo,
 };
 use trin_utils::version::get_trin_version;
@@ -279,7 +279,7 @@ impl Discovery {
         Ok(NodeInfo {
             enr: Enr::from_str(&self.discv5.local_enr().to_base64())
                 .map_err(|err| anyhow!("{err}"))?,
-            node_id: hex_encode(self.discv5.local_enr().node_id().raw()),
+            node_id: self.discv5.local_enr().node_id(),
             ip: self
                 .discv5
                 .local_enr()
@@ -292,7 +292,7 @@ impl Discovery {
     /// k-buckets.
     pub fn routing_table_info(&self) -> RoutingTableInfo {
         RoutingTableInfo {
-            local_node_id: hex_encode(self.discv5.local_enr().node_id().raw()),
+            local_node_id: self.discv5.local_enr().node_id(),
             buckets: self.discv5.kbuckets().into(),
         }
     }
diff --git a/portalnet/src/overlay/protocol.rs b/portalnet/src/overlay/protocol.rs
index c63339794..7d898b0fc 100644
--- a/portalnet/src/overlay/protocol.rs
+++ b/portalnet/src/overlay/protocol.rs
@@ -254,7 +254,7 @@ where
     /// Returns the node-id and a nested array of node-ids to represent this node's k-buckets table.
     pub fn routing_table_info(&self) -> RoutingTableInfo {
         RoutingTableInfo {
-            local_node_id: hex_encode(self.local_enr().node_id().raw()),
+            local_node_id: self.local_enr().node_id(),
             buckets: self.kbuckets.read().clone().into(),
         }
     }