Merge pull request #5 from Nuhvi/bep_0042

feat: add bep_0042 support

README.md

@@ -16,7 +16,7 @@ Check the [Examples](./examples).
 
 Supported BEPs:
 - [x] [BEP0005 DHT Protocol](https://www.bittorrent.org/beps/bep_0005.html)
-- [ ] [BEP0042 DHT Security extension](https://www.bittorrent.org/beps/bep_0042.html)
+- [X] [BEP0042 DHT Security extension](https://www.bittorrent.org/beps/bep_0042.html)
 - [x] [BEP0043 Read-only DHT Nodes](https://www.bittorrent.org/beps/bep_0043.html)
 - [x] [BEP0044 Storing arbitrary data in the DHT](https://www.bittorrent.org/beps/bep_0044.html)
 
@@ -27,3 +27,8 @@ Supported BEPs:
 - [ ] [BEP0042 DHT Security extension](https://www.bittorrent.org/beps/bep_0042.html)
 - [x] [BEP0043 Read-only DHT Nodes](https://www.bittorrent.org/beps/bep_0043.html)
 - [ ] [BEP0044 Storing arbitrary data in the DHT](https://www.bittorrent.org/beps/bep_0044.html)
+
+
+## Acknowledgment
+
+This implementation was possible thanks to [Webtorrent's Bittorrent-dht](https://github.com/webtorrent/bittorrent-dht) as a reference, and [Rustydht-lib](https://github.com/raptorswing/rustydht-lib) that saved me a lot of time, especially at the serialization and deserialization of Bencode messages.

src/common/id.rs

@@ -1,7 +1,10 @@
 //! Kademlia node Id or a lookup target
+use crc::{Crc, CRC_32_ISCSI};
 use rand::Rng;
+use std::convert::TryInto;
 use std::{
     fmt::{self, Debug, Display, Formatter},
+    net::{IpAddr, Ipv4Addr, SocketAddr},
     str::FromStr,
 };
 
@@ -11,6 +14,9 @@ use crate::{Error, Result};
 pub const ID_SIZE: usize = 20;
 pub const MAX_DISTANCE: u8 = ID_SIZE as u8 * 8;
 
+const IPV4_MASK: u32 = 0x030f3fff;
+const CASTAGNOLI: Crc<u32> = Crc::<u32>::new(&CRC_32_ISCSI);
+
 #[derive(Clone, Copy, PartialEq, Ord, PartialOrd, Eq, Hash)]
 /// Kademlia node Id or a lookup target
 pub struct Id {
@@ -61,6 +67,80 @@ impl Id {
     pub fn to_vec(self) -> Vec<u8> {
         self.bytes.to_vec()
     }
+
+    /// Create a new Id according to [BEP0042](http://bittorrent.org/beps/bep_0042.html).
+    pub fn from_addr(addr: &SocketAddr) -> Id {
+        let ip = addr.ip();
+
+        let mut rng = rand::thread_rng();
+        let r: u8 = rng.gen();
+
+        let mut bytes: [u8; 20] = rng.gen();
+
+        match ip {
+            IpAddr::V4(addr) => from_ipv4_and_r(bytes, addr, r),
+            IpAddr::V6(_addr) => unimplemented!(),
+        }
+    }
+
+    /// Validate that this Id is valid with respect to [BEP0042](http://bittorrent.org/beps/bep_0042.html).
+    pub fn is_valid_for_ip(&self, ip: &IpAddr) -> bool {
+        match ip {
+            IpAddr::V4(ipv4) => {
+                if ipv4.is_private() {
+                    return true;
+                }
+
+                let expected = first_21_bits(&id_prefix_ipv4(ipv4, self.bytes[ID_SIZE - 1]));
+                let actual = first_21_bits(&self.bytes);
+
+                expected == actual
+            }
+            IpAddr::V6(ipv6) => {
+                unimplemented!()
+
+                // // For IPv6, checking the ULA range fc00::/7
+                // if (ipv6.segments()[0] & 0xFE00 == 0xFC00) {
+                //     return true;
+                // }
+            }
+        }
+    }
+}
+
+fn first_21_bits(bytes: &[u8]) -> [u8; 3] {
+    [bytes[0], bytes[1], bytes[2] & 0xf8]
+}
+
+fn from_ipv4_and_r(bytes: [u8; 20], ip: Ipv4Addr, r: u8) -> Id {
+    let mut bytes = bytes;
+    let prefix = id_prefix_ipv4(&ip, r);
+
+    // Set first 21 bits to the prefix
+    bytes[0] = prefix[0];
+    bytes[1] = prefix[1];
+    // set the first 5 bits of the 3 byte to the remaining 5 bits of the prefix
+    bytes[2] = (prefix[2] & 0xf8) | (bytes[2] & 0x7);
+
+    // Set the last byte to the random r
+    bytes[ID_SIZE - 1] = r;
+
+    Id { bytes }
+}
+
+fn id_prefix_ipv4(ip: &Ipv4Addr, r: u8) -> [u8; 3] {
+    let r32: u32 = r.into();
+    let ip_int: u32 = u32::from_be_bytes(ip.octets());
+    let nonsense: u32 = (ip_int & IPV4_MASK) | (r32 << 29);
+
+    let mut digest = CASTAGNOLI.digest();
+    digest.update(&nonsense.to_be_bytes());
+
+    let crc = digest.finalize();
+
+    crc.to_be_bytes()[..3]
+        .try_into()
+        .expect("Failed to convert bytes 0-2 of the crc into a 3-byte array")
 }
 
 impl Display for Id {
@@ -168,4 +248,64 @@ mod test {
 
         assert_eq!(id.bytes, bytes);
     }
+
+    #[test]
+    fn from_ipv4() {
+        let vectors = vec![
+            (Ipv4Addr::new(124, 31, 75, 21), 1, [0x5f, 0xbf, 0xbf]),
+            (Ipv4Addr::new(21, 75, 31, 124), 86, [0x5a, 0x3c, 0xe9]),
+            (Ipv4Addr::new(65, 23, 51, 170), 22, [0xa5, 0xd4, 0x32]),
+            (Ipv4Addr::new(84, 124, 73, 14), 65, [0x1b, 0x03, 0x21]),
+            (Ipv4Addr::new(43, 213, 53, 83), 90, [0xe5, 0x6f, 0x6c]),
+        ];
+
+        for vector in vectors {
+            test(vector.0, vector.1, vector.2);
+        }
+
+        fn test(ip: Ipv4Addr, r: u8, expected_prefix: [u8; 3]) {
+            let id = Id::random();
+            let result = from_ipv4_and_r(id.bytes, ip, r);
+            let prefix = first_21_bits(&result.bytes);
+
+            assert_eq!(prefix, first_21_bits(&expected_prefix));
+            assert_eq!(result.bytes[ID_SIZE - 1], r);
+        }
+    }
+
+    #[test]
+    fn is_valid_for_ipv4() {
+        let valid_vectors = vec![
+            (
+                Ipv4Addr::new(124, 31, 75, 21),
+                "5fbfbff10c5d6a4ec8a88e4c6ab4c28b95eee401",
+            ),
+            (
+                Ipv4Addr::new(21, 75, 31, 124),
+                "5a3ce9c14e7a08645677bbd1cfe7d8f956d53256",
+            ),
+            (
+                Ipv4Addr::new(65, 23, 51, 170),
+                "a5d43220bc8f112a3d426c84764f8c2a1150e616",
+            ),
+            (
+                Ipv4Addr::new(84, 124, 73, 14),
+                "1b0321dd1bb1fe518101ceef99462b947a01ff41",
+            ),
+            (
+                Ipv4Addr::new(43, 213, 53, 83),
+                "e56f6cbf5b7c4be0237986d5243b87aa6d51305a",
+            ),
+        ];
+
+        for vector in valid_vectors {
+            test(vector.0, vector.1);
+        }
+
+        fn test(ip: Ipv4Addr, hex: &str) {
+            let id = Id::from_str(hex).unwrap();
+
+            assert!(id.is_valid_for_ip(&IpAddr::V4(ip)));
+        }
+    }
 }

src/common/node.rs

@@ -73,12 +73,17 @@ impl Node {
         self.last_seen.elapsed() > STALE_TIME
     }
 
-    pub(crate) fn _should_ping(&self) -> bool {
+    pub(crate) fn should_ping(&self) -> bool {
         self.last_seen.elapsed() > MIN_PING_BACKOFF_INTERVAL
     }
 
     /// Returns true if both nodes have the same id and address
     pub fn same_adress(&self, other: &Self) -> bool {
         self.address == other.address
     }
+
+    /// Check [BEP0042](https://www.bittorrent.org/beps/bep_0042.html).
+    pub(crate) fn is_secure(&self) -> bool {
+        self.id.is_valid_for_ip(&self.address.ip())
+    }
 }

src/routing_table.rs

@@ -57,8 +57,12 @@ impl RoutingTable {
         }
     }
 
+    /// Return the closest nodes to the target while prioritizing secure nodes,
+    /// as defined in [BEP_0042](https://www.bittorrent.org/beps/bep_0042.html)
     pub fn closest(&self, target: &Id) -> Vec<Node> {
         let mut result = Vec::with_capacity(MAX_BUCKET_SIZE_K);
+        let mut unsecure = Vec::with_capacity(MAX_BUCKET_SIZE_K);
+
         let distance = self.id.distance(target);
 
         for i in
@@ -71,7 +75,11 @@ impl RoutingTable {
                 Some(bucket) => {
                     for node in bucket.iter() {
                         if result.len() < 20 {
-                            result.push(node.clone());
+                            if node.is_secure() {
+                                result.push(node.clone());
+                            } else {
+                                unsecure.push(node.clone())
+                            }
                         } else {
                             return result;
                         }
@@ -81,6 +89,16 @@ impl RoutingTable {
             }
         }
 
+        if result.len() < 20 {
+            for node in unsecure {
+                if result.len() < 20 {
+                    result.push(node);
+                } else {
+                    break;
+                }
+            }
+        }
+
         result
     }
 

src/rpc.rs

@@ -536,7 +536,7 @@ impl Rpc {
             for node in self.routing_table.to_vec() {
                 if node.is_stale() {
                     self.routing_table.remove(&node.id);
-                } else if node._should_ping() {
+                } else if node.should_ping() {
                     self.ping(node.address);
                 }
             }