Custom protocol id (#163)

* Revert "Custom protocol (#158)"

This reverts commit 8f2716e.

* trait threading via functions

* clippy

* make it easier to use the default protocol id

* fix tests

Author: divagant-martian

Cargo.toml

@@ -19,7 +19,7 @@ tokio-util = { version = "0.6.9", features = ["time"] }
 libp2p-core = { version = "0.36.0", optional = true }
 zeroize = { version = "1.4.3", features = ["zeroize_derive"] }
 futures = "0.3.19"
-uint = { version = "0.9.5", default-features = false }
+uint = { version = "0.9.1", default-features = false }
 rlp = "0.5.1"
 # This version must be kept up to date do it uses the same dependencies as ENR
 hkdf = "0.12.3" 
@@ -39,7 +39,6 @@ lru = "0.7.1"
 hashlink = "0.7.0"
 delay_map = "0.1.1"
 more-asserts = "0.2.2"
-once_cell = "1.17.0"
 
 [dev-dependencies]
 rand_07 = { package = "rand", version = "0.7" }

examples/custom_executor.rs

@@ -39,7 +39,7 @@ fn main() {
     let config = Discv5ConfigBuilder::new().build();
 
     // construct the discv5 server
-    let mut discv5 = Discv5::new(enr, enr_key, config).unwrap();
+    let mut discv5: Discv5 = Discv5::new(enr, enr_key, config).unwrap();
 
     // if we know of another peer's ENR, add it known peers
     if let Some(base64_enr) = std::env::args().nth(1) {

examples/find_nodes.rs

@@ -135,7 +135,7 @@ async fn main() {
     let socket_addr = SocketAddr::new(bind_addr, port);
 
     // construct the discv5 server
-    let mut discv5 = Discv5::new(enr, enr_key, config).unwrap();
+    let mut discv5: Discv5 = Discv5::new(enr, enr_key, config).unwrap();
 
     // if we know of another peer's ENR, add it known peers
     for enr in args.remote_peer {

examples/simple_server.rs

@@ -72,7 +72,7 @@ async fn main() {
     let config = Discv5Config::default();
 
     // construct the discv5 server
-    let mut discv5 = Discv5::new(enr, enr_key, config).unwrap();
+    let mut discv5: Discv5 = Discv5::new(enr, enr_key, config).unwrap();
 
     // if we know of another peer's ENR, add it known peers
     if let Some(base64_enr) = std::env::args().nth(3) {

src/config.rs

@@ -1,17 +1,9 @@
 use crate::{
-    ipmode::IpMode,
-    kbucket::MAX_NODES_PER_BUCKET,
-    packet::{PROTOCOL_ID_LENGTH, PROTOCOL_VERSION_LENGTH},
-    Enr, Executor, PermitBanList, RateLimiter, RateLimiterBuilder,
+    ipmode::IpMode, kbucket::MAX_NODES_PER_BUCKET, Enr, Executor, PermitBanList, RateLimiter,
+    RateLimiterBuilder,
 };
 ///! A set of configuration parameters to tune the discovery protocol.
 use std::time::Duration;
-use std::{convert::TryInto, num::NonZeroU16};
-
-/// Protocol ID sent with each message.
-pub(crate) const DEFAULT_PROTOCOL_ID: [u8; PROTOCOL_ID_LENGTH] = *b"discv5";
-/// The version sent with each handshake.
-pub(crate) const DEFAULT_PROTOCOL_VERSION: [u8; PROTOCOL_VERSION_LENGTH] = 0x0001_u16.to_be_bytes();
 
 /// Configuration parameters that define the performance of the discovery network.
 #[derive(Clone)]
@@ -107,9 +99,6 @@ pub struct Discv5Config {
     /// A custom executor which can spawn the discv5 tasks. This must be a tokio runtime, with
     /// timing support. By default, the executor that created the discv5 struct will be used.
     pub executor: Option<Box<dyn Executor + Send + Sync>>,
-
-    /// The Discv5 protocol id and version, in bytes.
-    pub protocol: ([u8; PROTOCOL_ID_LENGTH], [u8; PROTOCOL_VERSION_LENGTH]),
 }
 
 impl Default for Discv5Config {
@@ -149,7 +138,6 @@ impl Default for Discv5Config {
             ban_duration: Some(Duration::from_secs(3600)), // 1 hour
             ip_mode: IpMode::default(),
             executor: None,
-            protocol: (DEFAULT_PROTOCOL_ID, DEFAULT_PROTOCOL_VERSION),
         }
     }
 }
@@ -326,30 +314,6 @@ impl Discv5ConfigBuilder {
         self
     }
 
-    /// Set the discv5 wire protocol id.
-    pub fn protocol_id(&mut self, protocol_id: &'static str) -> &mut Self {
-        let protocol_id: [u8; PROTOCOL_ID_LENGTH] =
-            protocol_id.as_bytes().try_into().unwrap_or_else(|_| {
-                panic!("The protocol id must be {} bytes long", PROTOCOL_ID_LENGTH)
-            });
-
-        self.config.protocol = (protocol_id, DEFAULT_PROTOCOL_VERSION);
-        self
-    }
-
-    /// Set the discv5 wire protocol id and the version.
-    pub fn protocol(
-        &mut self,
-        protocol_id: &'static str,
-        protocol_version: NonZeroU16,
-    ) -> &mut Self {
-        self.protocol_id(protocol_id);
-        let protocol_version = protocol_version.get().to_be_bytes();
-
-        self.config.protocol.1 = protocol_version;
-        self
-    }
-
     pub fn build(&mut self) -> Discv5Config {
         // If an executor is not provided, assume a current tokio runtime is running.
         if self.config.executor.is_none() {
@@ -383,7 +347,6 @@ impl std::fmt::Debug for Discv5Config {
             .field("incoming_bucket_limit", &self.incoming_bucket_limit)
             .field("ping_interval", &self.ping_interval)
             .field("ban_duration", &self.ban_duration)
-            .field("protocol", &self.protocol)
             .finish()
     }
 }

src/discv5.rs

@@ -19,13 +19,15 @@ use crate::{
         NodeStatus, UpdateResult,
     },
     node_info::NodeContact,
+    packet::ProtocolIdentity,
     service::{QueryKind, Service, ServiceRequest, TalkRequest},
-    Discv5Config, Enr,
+    DefaultProtocolId, Discv5Config, Enr,
 };
 use enr::{CombinedKey, EnrError, EnrKey, NodeId};
 use parking_lot::RwLock;
 use std::{
     future::Future,
+    marker::PhantomData,
     net::SocketAddr,
     sync::Arc,
     time::{Duration, Instant},
@@ -71,7 +73,10 @@ pub enum Discv5Event {
 
 /// The main Discv5 Service struct. This provides the user-level API for performing queries and
 /// interacting with the underlying service.
-pub struct Discv5 {
+pub struct Discv5<P = DefaultProtocolId>
+where
+    P: ProtocolIdentity,
+{
     config: Discv5Config,
     /// The channel to make requests from the main service.
     service_channel: Option<mpsc::Sender<ServiceRequest>>,
@@ -83,47 +88,16 @@ pub struct Discv5 {
     local_enr: Arc<RwLock<Enr>>,
     /// The key associated with the local ENR, required for updating the local ENR.
     enr_key: Arc<RwLock<CombinedKey>>,
+    /// Phantom for the protocol id.
+    _phantom: PhantomData<P>,
 }
 
-impl Discv5 {
+impl<P: ProtocolIdentity> Discv5<P> {
     pub fn new(
         local_enr: Enr,
         enr_key: CombinedKey,
         mut config: Discv5Config,
     ) -> Result<Self, &'static str> {
-        // tests use the default value, so we ignore initializing the protocol.
-        #[cfg(not(test))]
-        {
-            use crate::{
-                config::{DEFAULT_PROTOCOL_ID, DEFAULT_PROTOCOL_VERSION},
-                packet::{PROTOCOL_ID, VERSION},
-            };
-            // initialize the protocol id and version
-            let (protocol_id_bytes, protocol_version_bytes) = config.protocol;
-            PROTOCOL_ID
-                .set(protocol_id_bytes)
-                .map_err(|_old_val| "PROTOCOL_ID has already been initialized")?;
-            VERSION
-                .set(protocol_version_bytes)
-                .map_err(|_old_val| "protocol's VERSION has already been initialized")?;
-
-            if protocol_id_bytes != DEFAULT_PROTOCOL_ID
-                || protocol_version_bytes != DEFAULT_PROTOCOL_VERSION
-            {
-                let protocol_version = u16::from_be_bytes(protocol_version_bytes);
-                match std::str::from_utf8(&protocol_id_bytes) {
-                Ok(pretty_protocol_id) => tracing::info!(
-                    "Discv5 using custom protocol id and version. Id: {} Version: {}",
-                    pretty_protocol_id, protocol_version
-                ),
-                Err(_) => tracing::info!(
-                    "Discv5 using custom protocol id and version, with non utf8 protocol id. Id: {:?} Version: {}",
-                    protocol_id_bytes, protocol_version
-                ),
-            }
-            }
-        }
-
         // ensure the keypair matches the one that signed the enr.
         if local_enr.public_key() != enr_key.public() {
             return Err("Provided keypair does not match the provided ENR");
@@ -166,6 +140,7 @@ impl Discv5 {
             kbuckets,
             local_enr,
             enr_key,
+            _phantom: Default::default(),
         })
     }
 
@@ -177,7 +152,7 @@ impl Discv5 {
         }
 
         // create the main service
-        let (service_exit, service_channel) = Service::spawn(
+        let (service_exit, service_channel) = Service::spawn::<P>(
             self.local_enr.clone(),
             self.enr_key.clone(),
             self.kbuckets.clone(),
@@ -638,7 +613,7 @@ impl Discv5 {
     }
 }
 
-impl Drop for Discv5 {
+impl<P: ProtocolIdentity> Drop for Discv5<P> {
     fn drop(&mut self) {
         self.shutdown();
     }

src/discv5/test.rs

@@ -615,7 +615,7 @@ async fn test_bucket_limits() {
         .collect();
 
     let config = Discv5ConfigBuilder::new().ip_limit().build();
-    let discv5 = Discv5::new(enr, enr_key, config).unwrap();
+    let discv5: Discv5 = Discv5::new(enr, enr_key, config).unwrap();
     for enr in enrs {
         let _ = discv5.add_enr(enr.clone()); // we expect some of these to fail based on the filter.
     }

src/handler/crypto/mod.rs

@@ -223,6 +223,8 @@ pub(crate) fn encrypt_message(
 
 #[cfg(test)]
 mod tests {
+    use crate::packet::DefaultProtocolId;
+
     use super::*;
     use enr::{CombinedKey, EnrBuilder, EnrKey};
     use std::convert::TryInto;
@@ -394,7 +396,8 @@ mod tests {
         let dst_id: NodeId = node_key_2().public().into();
         let encoded_ref_packet = hex::decode("00000000000000000000000000000000088b3d4342774649325f313964a39e55ea96c005ad52be8c7560413a7008f16c9e6d2f43bbea8814a546b7409ce783d34c4f53245d08dab84102ed931f66d1492acb308fa1c6715b9d139b81acbdcc").unwrap();
         let (_packet, auth_data) =
-            crate::packet::Packet::decode(&dst_id, &encoded_ref_packet).unwrap();
+            crate::packet::Packet::decode::<DefaultProtocolId>(&dst_id, &encoded_ref_packet)
+                .unwrap();
 
         let ciphertext = hex::decode("b84102ed931f66d1492acb308fa1c6715b9d139b81acbdcc").unwrap();
         let read_key = hex::decode("00000000000000000000000000000000").unwrap();