Add ProtocolKind enum for overlay

ethereum · Sep 21, 2021 · b58bedc · b58bedc
1 parent 111bddf
commit b58bedc
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 11 deletions.
diff --git a/trin-core/src/portalnet/overlay.rs b/trin-core/src/portalnet/overlay.rs
@@ -2,7 +2,7 @@ use crate::utils::xor_two_values;
 
 use super::{
     discovery::Discovery,
-    types::{FindContent, FindNodes, Message, Ping, Request, SszEnr},
+    types::{FindContent, FindNodes, Message, Ping, ProtocolKind, Request, SszEnr},
     Enr, U256,
 };
 use crate::portalnet::types::HexData;
@@ -186,7 +186,7 @@ impl OverlayProtocol {
         &self,
         data_radius: U256,
         enr: Enr,
-        protocol: String,
+        protocol: ProtocolKind,
     ) -> Result<Vec<u8>, String> {
         let enr_seq = self.discovery.read().await.local_enr().seq();
         let msg = Ping {
@@ -198,7 +198,7 @@ impl OverlayProtocol {
             .await
             .send_talkreq(
                 enr,
-                protocol,
+                protocol.to_string(),
                 Message::Request(Request::Ping(msg)).to_bytes(),
             )
             .await
@@ -208,15 +208,15 @@ impl OverlayProtocol {
         &self,
         distances: Vec<u16>,
         enr: Enr,
-        protocol: String,
+        protocol: ProtocolKind,
     ) -> Result<Vec<u8>, String> {
         let msg = FindNodes { distances };
         self.discovery
             .write()
             .await
             .send_talkreq(
                 enr,
-                protocol,
+                protocol.to_string(),
                 Message::Request(Request::FindNodes(msg)).to_bytes(),
             )
             .await
@@ -226,15 +226,15 @@ impl OverlayProtocol {
         &self,
         content_key: Vec<u8>,
         enr: Enr,
-        protocol: String,
+        protocol: ProtocolKind,
     ) -> Result<Vec<u8>, String> {
         let msg = FindContent { content_key };
         self.discovery
             .write()
             .await
             .send_talkreq(
                 enr,
-                protocol,
+                protocol.to_string(),
                 Message::Request(Request::FindContent(msg)).to_bytes(),
             )
             .await

diff --git a/trin-core/src/portalnet/types.rs b/trin-core/src/portalnet/types.rs
@@ -10,6 +10,21 @@ use ssz_derive::{Decode, Encode};
 
 use super::{Enr, U256};
 
+#[derive(Debug, PartialEq, Clone)]
+pub enum ProtocolKind {
+    History,
+    State,
+}
+
+impl ToString for ProtocolKind {
+    fn to_string(&self) -> String {
+        match self {
+            ProtocolKind::History => "history".to_string(),
+            ProtocolKind::State => "state".to_string(),
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Clone)]
 pub struct ProtocolMessage {
     message_id: u8,

diff --git a/trin-history/src/network.rs b/trin-history/src/network.rs
@@ -2,10 +2,10 @@ use discv5::kbucket::KBucketsTable;
 use log::debug;
 use std::sync::Arc;
 use tokio::sync::RwLock;
-use trin_core::cli::STATE_NETWORK;
 use trin_core::portalnet::{
     discovery::Discovery,
     overlay::{OverlayConfig, OverlayProtocol, PortalnetConfig},
+    types::ProtocolKind,
     U256,
 };
 
@@ -61,7 +61,7 @@ impl HistoryNetwork {
             debug!("Pinging {} on portal history network", enr);
             let ping_result = self
                 .overlay
-                .send_ping(U256::from(u64::MAX), enr, STATE_NETWORK.to_string())
+                .send_ping(U256::from(u64::MAX), enr, ProtocolKind::History)
                 .await?;
             debug!("Portal history network Ping result: {:?}", ping_result);
         }

diff --git a/trin-state/src/network.rs b/trin-state/src/network.rs
@@ -2,10 +2,10 @@ use discv5::kbucket::KBucketsTable;
 use log::debug;
 use std::sync::Arc;
 use tokio::sync::RwLock;
-use trin_core::cli::HISTORY_NETWORK;
 use trin_core::portalnet::{
     discovery::Discovery,
     overlay::{OverlayConfig, OverlayProtocol, PortalnetConfig},
+    types::ProtocolKind,
     U256,
 };
 
@@ -61,7 +61,7 @@ impl StateNetwork {
             debug!("Pinging {} on portal state network", enr);
             let ping_result = self
                 .overlay
-                .send_ping(U256::from(u64::MAX), enr, HISTORY_NETWORK.to_string())
+                .send_ping(U256::from(u64::MAX), enr, ProtocolKind::State)
                 .await?;
             debug!("Portal state network Ping result: {:?}", ping_result);
         }