Add NetworkID to kadcast configuration

Resolves #123
dusk-network · Oct 17, 2023 · 9e8861f · 9e8861f
1 parent 1500cd3
commit 9e8861f
Show file tree

Hide file tree

Showing 13 changed files with 124 additions and 58 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -37,6 +37,9 @@ pub const DEFAULT_BLOCKLIST_REFRESH_SECS: u64 = 10;
 
 #[derive(Clone, Serialize, Deserialize)]
 pub struct Config {
+    /// KadcastID
+    pub kadcast_id: Option<u8>,
+
     /// Public `SocketAddress` of the [Peer]. No domain name allowed
     ///
     /// This is the address where other peers can contact you.
@@ -86,6 +89,7 @@ impl Default for Config {
     fn default() -> Self {
         Self {
             public_address: "127.0.0.1:9000".to_string(),
+            kadcast_id: None,
             listen_address: None,
             bootstrapping_nodes: vec![],
             auto_propagate: ENABLE_BROADCAST_PROPAGATION,

diff --git a/src/encoding.rs b/src/encoding.rs
@@ -35,32 +35,34 @@ mod tests {
 
     #[test]
     fn test_encode_ping() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let a = Message::Ping(peer.to_header());
         test_kadkast_marshal(a)
     }
     #[test]
     fn test_encode_pong() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let a = Message::Pong(peer.to_header());
         test_kadkast_marshal(a)
     }
 
     #[test]
     fn test_encode_find_nodes() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
-        let target = *PeerNode::generate("192.168.1.1:666")?.id().as_binary();
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
+        let target =
+            *PeerNode::generate("192.168.1.1:666", 0)?.id().as_binary();
         let a = Message::FindNodes(peer.to_header(), target);
         test_kadkast_marshal(a)
     }
 
     #[test]
     fn test_encode_nodes() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let nodes = vec![
-            PeerNode::generate("192.168.1.1:666")?,
+            PeerNode::generate("192.168.1.1:666", 0)?,
             PeerNode::generate(
                 "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:666",
+                0,
             )?,
         ]
         .iter()
@@ -72,13 +74,13 @@ mod tests {
 
     #[test]
     fn test_encode_empty_nodes() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let a = Message::Nodes(peer.to_header(), NodePayload { peers: vec![] });
         test_kadkast_marshal(a)
     }
     #[test]
     fn test_encode_broadcast() -> Result<()> {
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let a = Message::Broadcast(
             peer.to_header(),
             BroadcastPayload {

diff --git a/src/encoding/header.rs b/src/encoding/header.rs
@@ -13,6 +13,7 @@ use crate::{kbucket::BinaryID, K_ID_LEN_BYTES, K_NONCE_LEN};
 pub struct Header {
     pub(crate) binary_id: BinaryID,
     pub(crate) sender_port: u16,
+    pub(crate) network_id: u8,
     pub(crate) reserved: [u8; 2],
 }
 
@@ -30,6 +31,7 @@ impl Marshallable for Header {
         writer.write_all(self.binary_id.as_binary())?;
         writer.write_all(self.binary_id.nonce())?;
         writer.write_all(&self.sender_port.to_le_bytes())?;
+        writer.write_all(&[self.network_id])?;
         writer.write_all(&self.reserved)?;
         Ok(())
     }
@@ -50,12 +52,19 @@ impl Marshallable for Header {
         let mut port_buffer = [0; 2];
         reader.read_exact(&mut port_buffer)?;
         let port = u16::from_le_bytes(port_buffer);
+
+        let mut network_id = [0; 1];
+        reader.read_exact(&mut network_id)?;
+        let network_id = network_id[0];
+
         let mut reserved = [0; 2];
         reader.read_exact(&mut reserved)?;
+
         Ok(Header {
             binary_id,
             sender_port: port,
             reserved,
+            network_id,
         })
     }
 }
diff --git a/src/handling.rs b/src/handling.rs
@@ -97,6 +97,7 @@ impl MessageHandler {
                 let remote_peer = PeerNode::from_socket(
                     remote_peer_addr,
                     *message.header().binary_id(),
+                    message.header().network_id,
                 );
 
                 match handler.handle_peer(remote_peer).await {
@@ -114,6 +115,14 @@ impl MessageHandler {
                         );
                         continue;
                     }
+                    Err(NodeInsertError::MismatchNetwork(n)) => {
+                        error!(
+                            "Unable to insert node - NETWORK MISMATCH {} - {}",
+                            n.value().address(),
+                            n.network_id,
+                        );
+                        continue;
+                    }
                 };
 
                 handler.handle_message(message, remote_peer_addr).await;

diff --git a/src/kbucket.rs b/src/kbucket.rs
@@ -36,6 +36,9 @@ impl<V> Tree<V> {
         &mut self,
         node: Node<V>,
     ) -> Result<InsertOk<V>, InsertError<V>> {
+        if self.root().network_id != node.network_id {
+            return Err(NodeInsertError::MismatchNetwork(node));
+        }
         match self.root.calculate_distance(&node) {
             None => Err(NodeInsertError::Invalid(node)),
             Some(height) => self.get_or_create_bucket(height).insert(node),
@@ -172,29 +175,24 @@ mod tests {
     use crate::tests::Result;
     #[test]
     fn test_buckets() -> Result<()> {
-        let root = PeerNode::generate("192.168.0.1:666")?;
+        let root = PeerNode::generate("192.168.0.1:666", 0)?;
         let mut config = BucketConfig::default();
         config.node_evict_after = Duration::from_millis(5000);
         config.node_ttl = Duration::from_secs(60);
 
         let mut route_table = Tree::new(root, config);
         for i in 2..255 {
-            let res = route_table.insert(PeerNode::generate(
-                &format!("192.168.0.{}:666", i)[..],
-            )?);
+            let node = PeerNode::generate(format!("192.168.0.{}:666", i), 0)?;
+            let res = route_table.insert(node);
             match res {
-                Ok(_) => {}
-                Err(error) => match error {
-                    NodeInsertError::Invalid(_) => {
-                        assert!(false)
-                    }
-                    _ => {}
-                },
+                Ok(_) | Err(NodeInsertError::Full(_)) => {}
+                _ => panic!("Node must be valid"),
             }
         }
         let res = route_table
-            .insert(PeerNode::generate("192.168.0.1:666")?)
+            .insert(PeerNode::generate("192.168.0.1:666", 0)?)
             .expect_err("this should be an error");
+
         assert!(if let NodeInsertError::Invalid(_) = res {
             true
         } else {

diff --git a/src/kbucket/bucket.rs b/src/kbucket/bucket.rs
@@ -38,6 +38,7 @@ pub enum NodeInsertOk<'a, TNode> {
 pub enum NodeInsertError<TNode> {
     Invalid(TNode),
     Full(TNode),
+    MismatchNetwork(TNode),
 }
 
 impl<'a, TNode> NodeInsertOk<'a, TNode> {
@@ -256,17 +257,17 @@ mod tests {
 
     #[test]
     fn test_lru_base_5secs() -> Result<()> {
-        let root = PeerNode::generate("127.0.0.1:666")?;
+        let root = PeerNode::generate("127.0.0.1:666", 0)?;
         let mut config = BucketConfig::default();
         config.node_evict_after = Duration::from_millis(1000);
         config.node_ttl = Duration::from_secs(5);
 
         let mut route_table = Tree::new(root, config);
 
         let bucket = route_table.bucket_for_test();
-        let node1 = PeerNode::generate("192.168.1.1:8080")?;
+        let node1 = PeerNode::generate("192.168.1.1:8080", 0)?;
         let id_node1 = node1.id().as_binary().clone();
-        let node1_copy = PeerNode::generate("192.168.1.1:8080")?;
+        let node1_copy = PeerNode::generate("192.168.1.1:8080", 0)?;
         match bucket.insert(node1).expect("This should return an ok()") {
             NodeInsertOk::Inserted { .. } => {}
             _ => assert!(false),
@@ -282,7 +283,7 @@ mod tests {
             _ => assert!(false),
         }
         assert_eq!(Some(&id_node1), bucket.last_id());
-        let node2 = PeerNode::generate("192.168.1.2:8080")?;
+        let node2 = PeerNode::generate("192.168.1.2:8080", 0)?;
         let id_node2 = node2.id().as_binary().clone();
 
         match bucket.insert(node2).expect("This should return an ok()") {
@@ -295,7 +296,7 @@ mod tests {
         assert_eq!(Some(&id_node1), bucket.least_used_id());
 
         match bucket
-            .insert(PeerNode::generate("192.168.1.1:8080")?)
+            .insert(PeerNode::generate("192.168.1.1:8080", 0)?)
             .expect("This should return an ok()")
         {
             NodeInsertOk::Updated { .. } => {}
@@ -313,6 +314,7 @@ mod tests {
             match bucket
                 .insert(PeerNode::generate(
                     &format!("192.168.1.{}:8080", i)[..],
+                    0,
                 )?)
                 .expect("This should return an ok()")
             {
@@ -323,7 +325,7 @@ mod tests {
             }
         }
         assert_eq!(bucket.pick::<K_BETA>().count(), K_BETA);
-        let pending = PeerNode::generate("192.168.1.21:8080")?;
+        let pending = PeerNode::generate("192.168.1.21:8080", 0)?;
         let pending_id = pending.id().as_binary().clone();
         match bucket.insert(pending).expect_err("this should be error") {
             NodeInsertError::Full(pending) => {
@@ -339,7 +341,8 @@ mod tests {
                             &pending_id
                         );
                         thread::sleep(Duration::from_secs(1));
-                        let pending = PeerNode::generate("192.168.1.21:8080")?;
+                        let pending =
+                            PeerNode::generate("192.168.1.21:8080", 0)?;
                         match bucket.insert(pending).expect("this should be ok")
                         {
                             NodeInsertOk::Inserted { inserted: _ } => {}

diff --git a/src/kbucket/key.rs b/src/kbucket/key.rs
@@ -169,12 +169,12 @@ mod tests {
 
     #[test]
     fn test_distance() -> Result<()> {
-        let n1 = PeerNode::generate("192.168.0.1:666")?;
-        let n2 = PeerNode::generate("192.168.0.1:666")?;
+        let n1 = PeerNode::generate("192.168.0.1:666", 0)?;
+        let n2 = PeerNode::generate("192.168.0.1:666", 0)?;
         assert_eq!(n1.calculate_distance(&n2), None);
         assert_eq!(n1.id().calculate_distance_native(n2.id()), None);
         for i in 2..255 {
-            let n_in = PeerNode::generate(&format!("192.168.0.{}:666", i)[..])?;
+            let n_in = PeerNode::generate(format!("192.168.0.{}:666", i), 0)?;
             assert_eq!(
                 n1.calculate_distance(&n_in),
                 n1.id().calculate_distance_native(n_in.id())
@@ -185,7 +185,7 @@ mod tests {
 
     #[test]
     fn test_id_nonce() -> Result<()> {
-        let root = PeerNode::generate("192.168.0.1:666")?;
+        let root = PeerNode::generate("192.168.0.1:666", 0)?;
         println!("Nonce is {:?}", root.id().nonce());
         assert!(root.id().verify_nonce());
         Ok(())

diff --git a/src/kbucket/node.rs b/src/kbucket/node.rs
@@ -15,6 +15,7 @@ pub struct Node<TValue> {
     value: TValue,
     pub(super) eviction_status: NodeEvictionStatus,
     pub(super) seen_at: Instant,
+    pub(crate) network_id: u8,
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
@@ -24,12 +25,13 @@ pub enum NodeEvictionStatus {
 }
 
 impl<TValue> Node<TValue> {
-    pub fn new(id: BinaryID, value: TValue) -> Self {
+    pub fn new(id: BinaryID, value: TValue, network_id: u8) -> Self {
         Node {
             id,
             value,
             seen_at: Instant::now(),
             eviction_status: NodeEvictionStatus::None,
+            network_id,
         }
     }
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -81,8 +81,9 @@ impl Peer {
         config: Config,
         listener: L,
     ) -> Result<Self, AddrParseError> {
+        let network_id = config.kadcast_id.unwrap_or_default();
         let tree = Tree::new(
-            PeerNode::generate(&config.public_address[..])?,
+            PeerNode::generate(&config.public_address[..], network_id)?,
             config.bucket,
         );
 

diff --git a/src/peer.rs b/src/peer.rs
@@ -27,18 +27,25 @@ impl PeerInfo {
 }
 
 impl PeerNode {
-    pub fn generate(address: &str) -> Result<Self, AddrParseError> {
-        let address: SocketAddr = address.parse()?;
+    pub fn generate(
+        address: impl AsRef<str>,
+        network_id: u8,
+    ) -> Result<Self, AddrParseError> {
+        let address: SocketAddr = address.as_ref().parse()?;
         let info = PeerInfo { address };
         let binary =
             PeerNode::compute_id(&info.address.ip(), info.address.port());
         let id = BinaryID::generate(binary);
-        Ok(Node::new(id, info))
+        Ok(Node::new(id, info, network_id))
     }
 
-    pub fn from_socket(address: SocketAddr, id: BinaryID) -> Self {
+    pub fn from_socket(
+        address: SocketAddr,
+        id: BinaryID,
+        network_id: u8,
+    ) -> Self {
         let info = PeerInfo { address };
-        Node::new(id, info)
+        Node::new(id, info, network_id)
     }
 
     pub(crate) fn verify_header(header: &Header, ip: &IpAddr) -> bool {
@@ -64,6 +71,7 @@ impl PeerNode {
             binary_id: *self.id(),
             sender_port: self.value().address.port(),
             reserved: [0; 2],
+            network_id: self.network_id,
         }
     }
 
@@ -86,13 +94,14 @@ mod tests {
     use crate::tests::Result;
     #[test]
     fn test_verify_header() -> Result<()> {
-        let wrong_header = PeerNode::generate("10.0.0.1:333")?.to_header();
+        let wrong_header = PeerNode::generate("10.0.0.1:333", 0)?.to_header();
         let wrong_header_sameport =
-            PeerNode::generate("10.0.0.1:666")?.to_header();
+            PeerNode::generate("10.0.0.1:666", 0)?.to_header();
         vec![
-            PeerNode::generate("192.168.1.1:666")?,
+            PeerNode::generate("192.168.1.1:666", 0)?,
             PeerNode::generate(
                 "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:666",
+                0,
             )?,
         ]
         .iter()

diff --git a/src/transport/encoding/raptorq.rs b/src/transport/encoding/raptorq.rs
@@ -120,7 +120,7 @@ mod tests {
         for i in 0..data.len() {
             data[i] = rand::Rng::gen(&mut rand::thread_rng());
         }
-        let peer = PeerNode::generate("192.168.0.1:666")?;
+        let peer = PeerNode::generate("192.168.0.1:666", 0)?;
         let header = peer.to_header();
         let payload = BroadcastPayload {
             height: 255,

diff --git a/src/transport/encoding/raptorq/decoder.rs b/src/transport/encoding/raptorq/decoder.rs
@@ -176,7 +176,7 @@ mod tests {
 
     #[test]
     fn test_expiring_cache() -> Result<()> {
-        let root = PeerNode::generate("192.168.0.1:666")?;
+        let root = PeerNode::generate("192.168.0.1:666", 0)?;
         let enc =
             RaptorQEncoder::configure(&RaptorQEncoder::default_configuration());
         let mut conf = RaptorQDecoder::default_configuration();