Turn ProtocolIdentity into a struct (#285)

Co-authored-by: Age Manning <[email protected]>

Author: dknopik

src/config.rs

@@ -1,7 +1,7 @@
 //! A set of configuration parameters to tune the discovery protocol.
 use crate::{
-    kbucket::MAX_NODES_PER_BUCKET, socket::ListenConfig, Enr, Executor, PermitBanList, RateLimiter,
-    RateLimiterBuilder,
+    kbucket::MAX_NODES_PER_BUCKET, socket::ListenConfig, Enr, Executor, PermitBanList,
+    ProtocolIdentity, RateLimiter, RateLimiterBuilder,
 };
 use std::time::Duration;
 
@@ -110,6 +110,9 @@ pub struct Config {
 
     /// Configuration for the sockets to listen on.
     pub listen_config: ListenConfig,
+
+    /// The protocol identity to use in network messages.
+    pub protocol_identity: ProtocolIdentity,
 }
 
 #[derive(Debug)]
@@ -156,6 +159,7 @@ impl ConfigBuilder {
             auto_nat_listen_duration: Some(Duration::from_secs(300)), // 5 minutes
             executor: None,
             listen_config,
+            protocol_identity: ProtocolIdentity::default(),
         };
 
         ConfigBuilder { config }
@@ -333,6 +337,13 @@ impl ConfigBuilder {
         self
     }
 
+    /// Configures use of a custom protocol identifier to use in network messages. Any packets not
+    /// matching the configured protocol identity will be ignored.
+    pub fn protocol_identity(&mut self, protocol_identity: ProtocolIdentity) -> &mut Self {
+        self.config.protocol_identity = protocol_identity;
+        self
+    }
+
     pub fn build(&mut self) -> Config {
         // If an executor is not provided, assume a current tokio runtime is running.
         if self.config.executor.is_none() {

src/discv5.rs

@@ -19,15 +19,13 @@ use crate::{
         NodeStatus, UpdateResult,
     },
     node_info::NodeContact,
-    packet::ProtocolIdentity,
     service::{QueryKind, Service, ServiceRequest, TalkRequest},
-    Config, DefaultProtocolId, Enr, IpMode,
+    Config, Enr, IpMode,
 };
 use enr::{CombinedKey, EnrKey, Error as EnrError, NodeId};
 use parking_lot::RwLock;
 use std::{
     future::Future,
-    marker::PhantomData,
     net::SocketAddr,
     sync::Arc,
     time::{Duration, Instant},
@@ -82,10 +80,7 @@ pub enum Event {
 
 /// The main Discv5 Service struct. This provides the user-level API for performing queries and
 /// interacting with the underlying service.
-pub struct Discv5<P = DefaultProtocolId>
-where
-    P: ProtocolIdentity,
-{
+pub struct Discv5 {
     config: Config,
     /// The channel to make requests from the main service.
     service_channel: Option<mpsc::Sender<ServiceRequest>>,
@@ -99,11 +94,9 @@ where
     enr_key: Arc<RwLock<CombinedKey>>,
     // Type of socket we are using
     ip_mode: IpMode,
-    /// Phantom for the protocol id.
-    _phantom: PhantomData<P>,
 }
 
-impl<P: ProtocolIdentity> Discv5<P> {
+impl Discv5 {
     pub fn new(
         local_enr: Enr,
         enr_key: CombinedKey,
@@ -154,7 +147,6 @@ impl<P: ProtocolIdentity> Discv5<P> {
             local_enr,
             enr_key,
             ip_mode,
-            _phantom: Default::default(),
         })
     }
 
@@ -166,7 +158,7 @@ impl<P: ProtocolIdentity> Discv5<P> {
         }
 
         // create the main service
-        let (service_exit, service_channel) = Service::spawn::<P>(
+        let (service_exit, service_channel) = Service::spawn(
             self.local_enr.clone(),
             self.enr_key.clone(),
             self.kbuckets.clone(),
@@ -739,7 +731,7 @@ impl<P: ProtocolIdentity> Discv5<P> {
     }
 }
 
-impl<P: ProtocolIdentity> Drop for Discv5<P> {
+impl Drop for Discv5 {
     fn drop(&mut self) {
         self.shutdown();
     }

src/handler/crypto/mod.rs

@@ -220,9 +220,8 @@ pub(crate) fn encrypt_message(
 
 #[cfg(test)]
 mod tests {
-    use crate::packet::DefaultProtocolId;
-
     use super::*;
+    use crate::ProtocolIdentity;
     use enr::{CombinedKey, Enr, EnrKey};
     use std::convert::TryInto;
 
@@ -392,9 +391,12 @@ mod tests {
     fn decrypt_ref_test_ping() {
         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::<DefaultProtocolId>(&dst_id, &encoded_ref_packet)
-                .unwrap();
+        let (_packet, auth_data) = crate::packet::Packet::decode(
+            &dst_id,
+            ProtocolIdentity::default(),
+            &encoded_ref_packet,
+        )
+        .unwrap();
 
         let ciphertext = hex::decode("b84102ed931f66d1492acb308fa1c6715b9d139b81acbdcc").unwrap();
         let read_key = hex::decode("00000000000000000000000000000000").unwrap();