Custom protocol (#158)

* add OnceCell protocol bytes

* store explicit arrays

* include pr changes

* adjust config to use bytes

* initialize oncecells

* adjust for tests

* comment for clarity

* fmt

* remove everything goes allow in kbucket key

* clippy

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.1", default-features = false }
+uint = { version = "0.9.5", 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,6 +39,7 @@ 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" }

src/config.rs

@@ -1,9 +1,17 @@
 use crate::{
-    ipmode::IpMode, kbucket::MAX_NODES_PER_BUCKET, Enr, Executor, PermitBanList, RateLimiter,
-    RateLimiterBuilder,
+    ipmode::IpMode,
+    kbucket::MAX_NODES_PER_BUCKET,
+    packet::{PROTOCOL_ID_LENGTH, PROTOCOL_VERSION_LENGTH},
+    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)]
@@ -99,6 +107,9 @@ 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 {
@@ -138,6 +149,7 @@ 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),
         }
     }
 }
@@ -314,6 +326,30 @@ 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() {
@@ -347,6 +383,7 @@ 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

@@ -91,6 +91,39 @@ impl Discv5 {
         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");

src/kbucket/key.rs

@@ -21,8 +21,7 @@
 // This basis of this file has been taken from the rust-libp2p codebase:
 // https://github.com/libp2p/rust-libp2p
 
-#![allow(clippy::all)]
-
+#![allow(clippy::assign_op_pattern)]
 use enr::{
     k256::sha2::digest::generic_array::{typenum::U32, GenericArray},
     NodeId,

src/packet/mod.rs

@@ -15,12 +15,17 @@ use aes::{
     Aes128Ctr,
 };
 use enr::NodeId;
+use once_cell::sync::OnceCell;
 use rand::Rng;
 use std::convert::TryInto;
 use zeroize::Zeroize;
 
 /// The packet IV length (u128).
 pub const IV_LENGTH: usize = 16;
+/// Length of the protocol id.
+pub const PROTOCOL_ID_LENGTH: usize = 6;
+/// Length of the protocol version.
+pub const PROTOCOL_VERSION_LENGTH: usize = 2;
 /// The length of the static header. (6 byte protocol id, 2 bytes version, 1 byte kind, 12 byte
 /// message nonce and a 2 byte authdata-size).
 pub const STATIC_HEADER_LENGTH: usize = 23;
@@ -29,10 +34,18 @@ pub const MESSAGE_NONCE_LENGTH: usize = 12;
 /// The Id nonce length (in bytes).
 pub const ID_NONCE_LENGTH: usize = 16;
 
-/// Protocol ID sent with each message.
-const PROTOCOL_ID: &str = "discv5";
-/// The version sent with each handshake.
-const VERSION: u16 = 0x0001;
+/// Protocol ID bytes sent with each message.
+#[cfg(not(test))]
+pub(crate) static PROTOCOL_ID: OnceCell<[u8; PROTOCOL_ID_LENGTH]> = OnceCell::new();
+#[cfg(test)]
+pub(crate) static PROTOCOL_ID: OnceCell<[u8; PROTOCOL_ID_LENGTH]> =
+    OnceCell::with_value(crate::config::DEFAULT_PROTOCOL_ID);
+/// The version bytes sent with each handshake.
+#[cfg(not(test))]
+pub(crate) static VERSION: OnceCell<[u8; PROTOCOL_VERSION_LENGTH]> = OnceCell::new();
+#[cfg(test)]
+pub(crate) static VERSION: OnceCell<[u8; PROTOCOL_VERSION_LENGTH]> =
+    OnceCell::with_value(crate::config::DEFAULT_PROTOCOL_VERSION);
 
 pub(crate) const MAX_PACKET_SIZE: usize = 1280;
 // The smallest packet must be at least this large
@@ -95,8 +108,8 @@ impl PacketHeader {
     pub fn encode(&self) -> Vec<u8> {
         let auth_data = self.kind.encode();
         let mut buf = Vec::with_capacity(auth_data.len() + STATIC_HEADER_LENGTH);
-        buf.extend_from_slice(PROTOCOL_ID.as_bytes());
-        buf.extend_from_slice(&VERSION.to_be_bytes());
+        buf.extend_from_slice(PROTOCOL_ID.wait());
+        buf.extend_from_slice(VERSION.wait());
         let kind: u8 = (&self.kind).into();
         buf.extend_from_slice(&kind.to_be_bytes());
         buf.extend_from_slice(&self.message_nonce);
@@ -440,17 +453,16 @@ impl Packet {
         }
 
         // Check the protocol id
-        if &static_header[..6] != PROTOCOL_ID.as_bytes() {
+        if &static_header[..PROTOCOL_ID_LENGTH] != PROTOCOL_ID.wait() {
             return Err(PacketError::HeaderDecryptionFailed);
         }
 
         // Check the version matches
-        let version = u16::from_be_bytes(
-            static_header[6..8]
-                .try_into()
-                .expect("Must be correct size"),
-        );
-        if version != VERSION {
+        let version_bytes =
+            &static_header[PROTOCOL_ID_LENGTH..PROTOCOL_ID_LENGTH + PROTOCOL_VERSION_LENGTH];
+        if version_bytes != VERSION.wait() {
+            let version =
+                u16::from_be_bytes(version_bytes.try_into().expect("Must be correct size"));
             return Err(PacketError::InvalidVersion(version));
         }