dht-crawler 0.2.1

高性能的 Rust DHT (Distributed Hash Table) 爬虫库 | A high-performance Rust DHT crawler library for fetching torrent information from the BitTorrent DHT network
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use crate::types::NodeTuple;
use rand::seq::SliceRandom;

#[derive(Default)]
pub(crate) struct RoutingSnapshot {
    v4: Vec<NodeTuple>,
    v6: Vec<NodeTuple>,
}

impl RoutingSnapshot {
    pub(crate) fn from_nodes(nodes: Vec<NodeTuple>, limit: usize) -> Self {
        let mut v4 = Vec::new();
        let mut v6 = Vec::new();
        for node in nodes.into_iter().take(limit) {
            if node.addr.is_ipv6() {
                v6.push(node);
            } else {
                v4.push(node);
            }
        }
        Self { v4, v6 }
    }

    pub(crate) fn random_nodes(&self, count: usize, filter_ipv6: Option<bool>) -> Vec<NodeTuple> {
        let mut rng = rand::thread_rng();
        match filter_ipv6 {
            Some(true) => self.v6.choose_multiple(&mut rng, count).cloned().collect(),
            Some(false) => self.v4.choose_multiple(&mut rng, count).cloned().collect(),
            None => {
                let mut all = Vec::with_capacity(self.v4.len() + self.v6.len());
                all.extend_from_slice(&self.v4);
                all.extend_from_slice(&self.v6);
                all.choose_multiple(&mut rng, count).cloned().collect()
            }
        }
    }

    pub(crate) fn closest_nodes(
        &self,
        target: &[u8; 20],
        count: usize,
        filter_ipv6: Option<bool>,
    ) -> Vec<NodeTuple> {
        if count == 0 {
            return Vec::new();
        }
        let mut nodes = match filter_ipv6 {
            Some(true) => self.v6.clone(),
            Some(false) => self.v4.clone(),
            None => {
                let mut all = Vec::with_capacity(self.v4.len() + self.v6.len());
                all.extend_from_slice(&self.v4);
                all.extend_from_slice(&self.v6);
                all
            }
        };
        let compare =
            |left: &NodeTuple, right: &NodeTuple| xor_distance_cmp(&left.id, &right.id, target);
        if nodes.len() > count {
            nodes.select_nth_unstable_by(count, compare);
            nodes.truncate(count);
        }
        nodes.sort_unstable_by(compare);
        nodes
    }
}

pub(crate) fn xor_distance_cmp(
    left: &[u8; 20],
    right: &[u8; 20],
    target: &[u8; 20],
) -> std::cmp::Ordering {
    for index in 0..20 {
        let ordering = (left[index] ^ target[index]).cmp(&(right[index] ^ target[index]));
        if ordering != std::cmp::Ordering::Equal {
            return ordering;
        }
    }
    std::cmp::Ordering::Equal
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::SocketAddr;

    fn node(id: u8, port: u16) -> NodeTuple {
        NodeTuple {
            id: [id; 20],
            addr: SocketAddr::from(([8, 8, 8, 8], port)),
        }
    }

    #[test]
    fn closest_nodes_orders_by_xor_distance_and_limits_results() {
        let snapshot =
            RoutingSnapshot::from_nodes(vec![node(0xf0, 1), node(0x01, 2), node(0x10, 3)], 3);
        let closest = snapshot.closest_nodes(&[0; 20], 2, Some(false));
        assert_eq!(
            closest
                .iter()
                .map(|node| node.addr.port())
                .collect::<Vec<_>>(),
            vec![2, 3]
        );
    }
}