rustydht_lib/common/
id.rs

1use anyhow::anyhow;
2
3extern crate crc;
4use crc::crc32;
5
6use rand::prelude::*;
7
8use std::convert::TryInto;
9use std::net::IpAddr;
10
11use crate::errors::RustyDHTError;
12
13/// The length (in bytes) of BitTorrent info hashes and DHT node ids.
14pub const ID_SIZE: usize = 20;
15
16/// Represents the id of a [Node](crate::common::Node) or a BitTorrent info-hash. Basically, it's a
17/// 20-byte identifier.
18#[derive(Eq, PartialEq, Hash, Copy, Clone)]
19pub struct Id {
20    bytes: [u8; ID_SIZE],
21}
22
23impl Id {
24    /// Create a new Id from some bytes. Returns Err if `bytes` is not of length
25    /// [ID_SIZE](crate::common::ID_SIZE).
26    pub fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Id, RustyDHTError> {
27        let bytes = bytes.as_ref();
28        if bytes.len() != ID_SIZE {
29            return Err(RustyDHTError::PacketParseError(anyhow!(
30                "Wrong number of bytes"
31            )));
32        }
33
34        let mut tmp: [u8; ID_SIZE] = [0; ID_SIZE];
35        tmp[..ID_SIZE].clone_from_slice(&bytes[..ID_SIZE]);
36
37        Ok(Id { bytes: tmp })
38    }
39
40    /// Generates a random Id for a mainline DHT node with the provided IP address.
41    /// The generated Id will be valid with respect to [BEP0042](http://bittorrent.org/beps/bep_0042.html).
42    pub fn from_ip(ip: &IpAddr) -> Id {
43        let mut rng = thread_rng();
44        let r: u8 = rng.gen();
45
46        let magic_prefix = IdPrefixMagic::from_ip(ip, r);
47
48        let mut bytes = [0; ID_SIZE];
49
50        bytes[0] = magic_prefix.prefix[0];
51        bytes[1] = magic_prefix.prefix[1];
52        bytes[2] = (magic_prefix.prefix[2] & 0xf8) | (rng.gen::<u8>() & 0x7);
53        for item in bytes.iter_mut().take(ID_SIZE - 1).skip(3) {
54            *item = rng.gen();
55        }
56        bytes[ID_SIZE - 1] = r;
57
58        Id { bytes }
59    }
60
61    /// Generates a completely random Id. The returned Id is *not* guaranteed to be
62    /// valid with respect to [BEP0042](http://bittorrent.org/beps/bep_0042.html).
63    pub fn from_random<T: rand::RngCore>(rng: &mut T) -> Id {
64        let mut bytes: [u8; ID_SIZE] = [0; ID_SIZE];
65        rng.fill_bytes(&mut bytes);
66
67        Id { bytes }
68    }
69
70    /// Copies the byes that make up the Id and returns them in a Vec
71    pub fn to_vec(&self) -> Vec<u8> {
72        self.bytes.to_vec()
73    }
74
75    /// Evaluates the Id and decides if it's a valid Id for a DHT node with the
76    /// provided IP address (based on [BEP0042](http://bittorrent.org/beps/bep_0042.html)).
77    /// Note: the current implementation does not handle non-globally-routable address space
78    /// properly. It will likely return false for non-routable IPv4 address space (against the spec).
79    pub fn is_valid_for_ip(&self, ip: &IpAddr) -> bool {
80        // TODO return true if ip is not globally routable
81        if ip.is_loopback() {
82            return true;
83        }
84        let expected = IdPrefixMagic::from_ip(ip, self.bytes[ID_SIZE - 1]);
85        let actual = IdPrefixMagic::from_id(self);
86
87        expected == actual
88    }
89
90    /// Returns the number of bits of prefix that the two ids have in common.
91    ///
92    /// Consider two Ids with binary values `10100000` and `10100100`. This function
93    /// would return `5` because the Ids share the common 5-bit prefix `10100`.
94    pub fn matching_prefix_bits(&self, other: &Self) -> usize {
95        let xored = self.xor(other);
96        let mut to_ret: usize = 0;
97
98        for i in 0..ID_SIZE {
99            let leading_zeros: usize = xored.bytes[i]
100                .leading_zeros()
101                .try_into()
102                .expect("this should never fail");
103            to_ret += leading_zeros;
104
105            if leading_zeros < 8 {
106                break;
107            }
108        }
109
110        to_ret
111    }
112
113    /// Creates an Id from a hex string.
114    ///
115    /// For example: `let id = Id::from_hex("88ffb73943354a00dc2dadd14c54d28020a513c8").unwrap();`
116    pub fn from_hex(h: &str) -> Result<Id, RustyDHTError> {
117        let bytes =
118            hex::decode(h).map_err(|hex_err| RustyDHTError::PacketParseError(hex_err.into()))?;
119
120        Id::from_bytes(&bytes)
121    }
122
123    /// Computes the exclusive or (XOR) of this Id with another. The BitTorrent DHT
124    /// uses XOR as its distance metric.
125    ///
126    /// Example: `let distance_between_nodes = id.xor(other_id);`
127    pub fn xor(&self, other: &Id) -> Id {
128        let mut bytes: [u8; ID_SIZE] = [0; ID_SIZE];
129        for (i, item) in bytes.iter_mut().enumerate().take(ID_SIZE) {
130            *item = self.bytes[i] ^ other.bytes[i];
131        }
132
133        Id::from_bytes(&bytes).expect("Wrong number of bytes for id")
134    }
135
136    /// Makes a new id that's similar to this one.
137    /// `identical_bytes` specifies how many bytes of the resulting Id should be the same as `this`.
138    /// `identical_bytes` must be in the range (0, [ID_SIZE](crate::common::ID_SIZE)) otherwise Err
139    /// is returned.
140    pub fn make_mutant(&self, identical_bytes: usize) -> Result<Id, RustyDHTError> {
141        if identical_bytes == 0 || identical_bytes >= ID_SIZE {
142            return Err(RustyDHTError::GeneralError(anyhow!(
143                "identical_bytes must be in range (0, ID_SIZE)"
144            )));
145        }
146        let mut mutant = Id::from_random(&mut thread_rng());
147        for i in 0..identical_bytes {
148            mutant.bytes[i] = self.bytes[i];
149        }
150        Ok(mutant)
151    }
152
153    /// An Id with all of its bits set to 0
154    pub const ZERO: Self = Id {
155        bytes: [0; ID_SIZE],
156    };
157}
158
159impl std::fmt::Display for Id {
160    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        f.write_str(&hex::encode(&self.bytes))
162    }
163}
164
165impl std::fmt::Debug for Id {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        f.write_str(&hex::encode(&self.bytes))
168    }
169}
170
171impl PartialOrd for Id {
172    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
173        for i in 0..self.bytes.len() {
174            match self.bytes[i].cmp(&other.bytes[i]) {
175                std::cmp::Ordering::Less => return Some(std::cmp::Ordering::Less),
176                std::cmp::Ordering::Greater => return Some(std::cmp::Ordering::Greater),
177                _ => continue,
178            };
179        }
180
181        Some(std::cmp::Ordering::Equal)
182    }
183}
184
185#[derive(Debug)]
186struct IdPrefixMagic {
187    prefix: [u8; 3],
188    suffix: u8,
189}
190
191impl IdPrefixMagic {
192    // Populates an IdPrefixMagic from the bytes of an id.
193    // This isn't a way to generate a valid IdPrefixMagic from a random id.
194    // For that, use Id::from_ip
195    fn from_id(id: &Id) -> IdPrefixMagic {
196        IdPrefixMagic {
197            prefix: id.bytes[..3]
198                .try_into()
199                .expect("Failed to grab first three bytes of id"),
200            suffix: id.bytes[ID_SIZE - 1],
201        }
202    }
203
204    fn from_ip(ip: &IpAddr, seed_r: u8) -> IdPrefixMagic {
205        match ip {
206            IpAddr::V4(ipv4) => {
207                let r32: u32 = seed_r.into();
208                let magic: u32 = 0x030f3fff;
209                let ip_int: u32 = u32::from_be_bytes(ipv4.octets());
210                let nonsense: u32 = (ip_int & magic) | (r32 << 29);
211                let crc: u32 = crc32::checksum_castagnoli(&nonsense.to_be_bytes());
212                IdPrefixMagic {
213                    prefix: crc.to_be_bytes()[..3]
214                        .try_into()
215                        .expect("Failed to convert bytes 0-2 of the crc into a 3-byte array"),
216                    suffix: seed_r,
217                }
218            }
219            IpAddr::V6(ipv6) => {
220                let r64: u64 = seed_r.into();
221                let magic: u64 = 0x0103070f1f3f7fff;
222                let ip_int: u64 = u64::from_be_bytes(
223                    ipv6.octets()[8..]
224                        .try_into()
225                        .expect("Failed to get IPv6 bytes"),
226                );
227                let nonsense: u64 = ip_int & magic | (r64 << 61);
228                let crc: u32 = crc32::checksum_castagnoli(&nonsense.to_be_bytes());
229                IdPrefixMagic {
230                    prefix: crc.to_be_bytes()[..2].try_into().expect("Failed to poop"),
231                    suffix: seed_r,
232                }
233            }
234        }
235    }
236}
237
238impl PartialEq for IdPrefixMagic {
239    fn eq(&self, other: &Self) -> bool {
240        self.prefix[0] == other.prefix[0]
241            && self.prefix[1] == other.prefix[1]
242            && self.prefix[2] & 0xf8 == other.prefix[2] & 0xf8
243            && self.suffix == other.suffix
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use std::net::Ipv4Addr;
251
252    #[test]
253    fn test_from_ip_v4() {
254        assert_eq!(
255            IdPrefixMagic::from_ip(&IpAddr::V4(Ipv4Addr::new(124, 31, 75, 21)), 1),
256            IdPrefixMagic {
257                prefix: [0x5f, 0xbf, 0xbf],
258                suffix: 1
259            }
260        );
261        assert_eq!(
262            IdPrefixMagic::from_ip(&IpAddr::V4(Ipv4Addr::new(21, 75, 31, 124)), 86),
263            IdPrefixMagic {
264                prefix: [0x5a, 0x3c, 0xe9],
265                suffix: 0x56
266            }
267        );
268
269        assert_eq!(
270            IdPrefixMagic::from_ip(&IpAddr::V4(Ipv4Addr::new(65, 23, 51, 170)), 22),
271            IdPrefixMagic {
272                prefix: [0xa5, 0xd4, 0x32],
273                suffix: 0x16
274            }
275        );
276
277        assert_eq!(
278            IdPrefixMagic::from_ip(&IpAddr::V4(Ipv4Addr::new(84, 124, 73, 14)), 65),
279            IdPrefixMagic {
280                prefix: [0x1b, 0x03, 0x21],
281                suffix: 0x41
282            }
283        );
284
285        assert_eq!(
286            IdPrefixMagic::from_ip(&IpAddr::V4(Ipv4Addr::new(43, 213, 53, 83)), 90),
287            IdPrefixMagic {
288                prefix: [0xe5, 0x6f, 0x6c],
289                suffix: 0x5a
290            }
291        );
292    }
293
294    #[test]
295    fn test_generate_valid_id() {
296        let ip = IpAddr::V4(Ipv4Addr::new(124, 31, 75, 21));
297        let id = Id::from_ip(&ip);
298        assert!(id.is_valid_for_ip(&ip));
299    }
300
301    #[test]
302    fn test_id_xor() {
303        let h1 = Id::from_hex("0000000000000000000000000000000000000001").unwrap();
304        let h2 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
305        let h3 = h1.xor(&h2);
306        assert_eq!(h3, h1);
307
308        let h1 = Id::from_hex("0000000000000000000000000000000000000001").unwrap();
309        let h2 = Id::from_hex("0000000000000000000000000000000000000001").unwrap();
310        let h3 = h1.xor(&h2);
311        assert_eq!(
312            h3,
313            Id::from_hex("0000000000000000000000000000000000000000").unwrap()
314        );
315
316        let h1 = Id::from_hex("1010101010101010101010101010101010101010").unwrap();
317        let h2 = Id::from_hex("0101010101010101010101010101010101010101").unwrap();
318        let h3 = h1.xor(&h2);
319        assert_eq!(
320            h3,
321            Id::from_hex("1111111111111111111111111111111111111111").unwrap()
322        );
323
324        let h1 = Id::from_hex("fefefefefefefefefefefefefefefefefefefefe").unwrap();
325        let h2 = Id::from_hex("0505050505050505050505050505050505050505").unwrap();
326        let h3 = h1.xor(&h2);
327        assert_eq!(
328            h3,
329            Id::from_hex("fbfbfbfbfbfbfbfbfbfbfbfbfbfbfbfbfbfbfbfb").unwrap()
330        );
331    }
332
333    #[test]
334    fn test_id_ordering() {
335        let h1 = Id::from_hex("0000000000000000000000000000000000000001").unwrap();
336        let h2 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
337        assert!(h1 > h2);
338        assert!(h2 < h1);
339        assert_ne!(h1, h2);
340
341        let h1 = Id::from_hex("00000000000000000000f0000000000000000000").unwrap();
342        let h2 = Id::from_hex("000000000000000000000fffffffffffffffffff").unwrap();
343        assert!(h1 > h2);
344        assert!(h2 < h1);
345        assert_ne!(h1, h2);
346
347        let h1 = Id::from_hex("1000000000000000000000000000000000000000").unwrap();
348        let h2 = Id::from_hex("0fffffffffffffffffffffffffffffffffffffff").unwrap();
349        assert!(h1 > h2);
350        assert!(h2 < h1);
351        assert_ne!(h1, h2);
352
353        let h1 = Id::from_hex("1010101010101010101010101010101010101010").unwrap();
354        let h2 = Id::from_hex("1010101010101010101010101010101010101010").unwrap();
355        assert!(h1 <= h2);
356        assert!(h2 <= h1);
357        assert_eq!(h1, h2);
358
359        let h1 = Id::from_hex("0000000000000000000000000000000000000010").unwrap();
360        let h2 = Id::from_hex("0000000000000000000000000000000000000001").unwrap();
361        assert!(h1 > h2);
362        assert!(h2 < h1);
363        assert_ne!(h1, h2);
364    }
365
366    #[test]
367    fn test_matching_prefix_bits() {
368        let h1 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
369        let h2 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
370        assert_eq!(h1.matching_prefix_bits(&h2), 160);
371
372        let h1 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
373        let h2 = Id::from_hex("f000000000000000000000000000000000000000").unwrap();
374        assert_eq!(h1.matching_prefix_bits(&h2), 0);
375
376        let h1 = Id::from_hex("0000000000000000000000000000000000000000").unwrap();
377        let h2 = Id::from_hex("1000000000000000000000000000000000000000").unwrap();
378        assert_eq!(h1.matching_prefix_bits(&h2), 3);
379    }
380}
rustydht_lib/common/id.rs

rustydht_lib/common/
id.rs
