ckb_network/peer_store/
addr_manager.rs

1//! Address manager
2use crate::peer_store::{base_addr, types::AddrInfo};
3use ckb_logger::debug;
4use p2p::{
5    multiaddr::{Multiaddr, Protocol},
6    utils::multiaddr_to_socketaddr,
7};
8use rand::Rng;
9use std::collections::{HashMap, HashSet};
10
11/// Address manager
12#[derive(Default)]
13pub struct AddrManager {
14    next_id: u64,
15    addr_to_id: HashMap<Multiaddr, u64>,
16    id_to_info: HashMap<u64, AddrInfo>,
17    random_ids: Vec<u64>,
18}
19
20impl AddrManager {
21    /// Add an address information to address manager
22    pub fn add(&mut self, mut addr_info: AddrInfo) {
23        if let Some(&id) = self.addr_to_id.get(&addr_info.addr) {
24            let (exist_last_connected_at_ms, random_id_pos) = {
25                let info = self.id_to_info.get(&id).expect("must exists");
26                (info.last_connected_at_ms, info.random_id_pos)
27            };
28            // Get time earlier than record time, return directly
29            if addr_info.last_connected_at_ms >= exist_last_connected_at_ms {
30                addr_info.random_id_pos = random_id_pos;
31                self.id_to_info.insert(id, addr_info);
32            }
33            return;
34        }
35
36        let id = self.next_id;
37        self.addr_to_id.insert(addr_info.addr.clone(), id);
38        addr_info.random_id_pos = self.random_ids.len();
39        self.id_to_info.insert(id, addr_info);
40        self.random_ids.push(id);
41        self.next_id += 1;
42    }
43
44    /// Randomly return addrs that worth to try or connect.
45    pub fn fetch_random<F>(&mut self, count: usize, filter: F) -> Vec<AddrInfo>
46    where
47        F: Fn(&AddrInfo) -> bool,
48    {
49        let mut duplicate_ips = HashSet::new();
50        let mut addr_infos = Vec::with_capacity(count);
51        let mut rng = rand::thread_rng();
52        let now_ms = ckb_systemtime::unix_time_as_millis();
53        for i in 0..self.random_ids.len() {
54            // reuse the for loop to shuffle random ids
55            // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
56            let j = rng.gen_range(i..self.random_ids.len());
57            self.swap_random_id(j, i);
58            let addr_info: AddrInfo = self.id_to_info[&self.random_ids[i]].to_owned();
59            match multiaddr_to_socketaddr(&addr_info.addr) {
60                Some(socket_addr) => {
61                    let ip = socket_addr.ip();
62                    let is_unique_ip = !duplicate_ips.contains(&ip);
63                    // A trick to make our tests work
64                    // TODO remove this after fix the network tests.
65                    let is_test_ip = ip.is_unspecified() || ip.is_loopback();
66                    if (is_test_ip || is_unique_ip)
67                        && addr_info.is_connectable(now_ms)
68                        && filter(&addr_info)
69                    {
70                        duplicate_ips.insert(ip);
71                        addr_infos.push(addr_info);
72                    }
73                }
74                None => {
75                    if filter(&addr_info) {
76                        if addr_info.is_connectable(now_ms)
77                            || addr_info
78                                .addr
79                                .iter()
80                                .any(|p| matches!(p, Protocol::Onion3(_)))
81                        {
82                            addr_infos.push(addr_info);
83                        } else {
84                            debug!(
85                                "addr {:?} is not connectable and not an onion address",
86                                addr_info.addr
87                            );
88                        }
89                    }
90                }
91            }
92            if addr_infos.len() == count {
93                break;
94            }
95        }
96        addr_infos
97    }
98
99    /// The count of address in address manager
100    pub fn count(&self) -> usize {
101        self.addr_to_id.len()
102    }
103
104    /// Addresses iterator
105    pub fn addrs_iter(&self) -> impl Iterator<Item = &AddrInfo> {
106        self.id_to_info.values()
107    }
108
109    /// Remove an address by ip and port
110    pub fn remove(&mut self, addr: &Multiaddr) -> Option<AddrInfo> {
111        let base_addr = base_addr(addr);
112        self.addr_to_id.remove(&base_addr).and_then(|id| {
113            let random_id_pos = self.id_to_info.get(&id).expect("exists").random_id_pos;
114            // swap with last index, then remove the last index
115            self.swap_random_id(random_id_pos, self.random_ids.len() - 1);
116            self.random_ids.pop();
117            self.id_to_info.remove(&id)
118        })
119    }
120
121    /// Get an address information by ip and port
122    pub fn get(&self, addr: &Multiaddr) -> Option<&AddrInfo> {
123        let base_addr = base_addr(addr);
124        self.addr_to_id
125            .get(&base_addr)
126            .and_then(|id| self.id_to_info.get(id))
127    }
128
129    /// Get a mutable address information by ip and port
130    pub fn get_mut(&mut self, addr: &Multiaddr) -> Option<&mut AddrInfo> {
131        let base_addr = base_addr(addr);
132        if let Some(id) = self.addr_to_id.get(&base_addr) {
133            self.id_to_info.get_mut(id)
134        } else {
135            None
136        }
137    }
138
139    /// swap random_id i and j,
140    /// this function keep random_id_pos in consistency
141    fn swap_random_id(&mut self, i: usize, j: usize) {
142        if i == j {
143            return;
144        }
145        self.id_to_info
146            .get_mut(&self.random_ids[i])
147            .expect("exists")
148            .random_id_pos = j;
149        self.id_to_info
150            .get_mut(&self.random_ids[j])
151            .expect("exists")
152            .random_id_pos = i;
153        self.random_ids.swap(i, j);
154    }
155}
ckb_network/peer_store/addr_manager.rs

ckb_network/peer_store/
addr_manager.rs
