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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! Address manager
use crate::peer_store::types::AddrInfo;
use p2p::{multiaddr::Multiaddr, utils::multiaddr_to_socketaddr};
use rand::Rng;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;

/// Address manager
#[derive(Default)]
pub struct AddrManager {
    next_id: u64,
    addr_to_id: HashMap<SocketAddr, u64>,
    id_to_info: HashMap<u64, AddrInfo>,
    random_ids: Vec<u64>,
}

impl AddrManager {
    /// Add an address information to address manager
    pub fn add(&mut self, mut addr_info: AddrInfo) {
        if let Some(key) = multiaddr_to_socketaddr(&addr_info.addr) {
            if let Some(&id) = self.addr_to_id.get(&key) {
                let (exist_last_connected_at_ms, random_id_pos) = {
                    let info = self.id_to_info.get(&id).expect("must exists");
                    (info.last_connected_at_ms, info.random_id_pos)
                };
                // Get time earlier than record time, return directly
                if addr_info.last_connected_at_ms >= exist_last_connected_at_ms {
                    addr_info.random_id_pos = random_id_pos;
                    self.id_to_info.insert(id, addr_info);
                }
                return;
            }

            let id = self.next_id;
            self.addr_to_id.insert(key, id);
            addr_info.random_id_pos = self.random_ids.len();
            self.id_to_info.insert(id, addr_info);
            self.random_ids.push(id);
            self.next_id += 1;
        }
    }

    /// Randomly return addrs that worth to try or connect.
    pub fn fetch_random<F>(&mut self, count: usize, filter: F) -> Vec<AddrInfo>
    where
        F: Fn(&AddrInfo) -> bool,
    {
        let mut duplicate_ips = HashSet::new();
        let mut addr_infos = Vec::with_capacity(count);
        let mut rng = rand::thread_rng();
        let now_ms = ckb_systemtime::unix_time_as_millis();
        for i in 0..self.random_ids.len() {
            // reuse the for loop to shuffle random ids
            // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
            let j = rng.gen_range(i, self.random_ids.len());
            self.swap_random_id(j, i);
            let addr_info: AddrInfo = self.id_to_info[&self.random_ids[i]].to_owned();
            if let Some(socket_addr) = multiaddr_to_socketaddr(&addr_info.addr) {
                let ip = socket_addr.ip();
                let is_unique_ip = duplicate_ips.insert(ip);
                // A trick to make our tests work
                // TODO remove this after fix the network tests.
                let is_test_ip = ip.is_unspecified() || ip.is_loopback();
                if (is_test_ip || is_unique_ip)
                    && addr_info.is_connectable(now_ms)
                    && filter(&addr_info)
                {
                    addr_infos.push(addr_info);
                }
                if addr_infos.len() == count {
                    break;
                }
            }
        }
        addr_infos
    }

    /// The count of address in address manager
    pub fn count(&self) -> usize {
        self.addr_to_id.len()
    }

    /// Addresses iterator
    pub fn addrs_iter(&self) -> impl Iterator<Item = &AddrInfo> {
        self.id_to_info.values()
    }

    /// Remove an address by ip and port
    pub fn remove(&mut self, addr: &Multiaddr) -> Option<AddrInfo> {
        multiaddr_to_socketaddr(addr).and_then(|addr| {
            self.addr_to_id.remove(&addr).and_then(|id| {
                let random_id_pos = self.id_to_info.get(&id).expect("exists").random_id_pos;
                // swap with last index, then remove the last index
                self.swap_random_id(random_id_pos, self.random_ids.len() - 1);
                self.random_ids.pop();
                self.id_to_info.remove(&id)
            })
        })
    }

    /// Get an address information by ip and port
    pub fn get(&self, addr: &Multiaddr) -> Option<&AddrInfo> {
        multiaddr_to_socketaddr(addr).and_then(|addr| {
            self.addr_to_id
                .get(&addr)
                .and_then(|id| self.id_to_info.get(id))
        })
    }

    /// Get a mutable address information by ip and port
    pub fn get_mut(&mut self, addr: &Multiaddr) -> Option<&mut AddrInfo> {
        if let Some(addr) = multiaddr_to_socketaddr(addr) {
            if let Some(id) = self.addr_to_id.get(&addr) {
                self.id_to_info.get_mut(id)
            } else {
                None
            }
        } else {
            None
        }
    }

    /// swap random_id i and j,
    /// this function keep random_id_pos in consistency
    fn swap_random_id(&mut self, i: usize, j: usize) {
        if i == j {
            return;
        }
        self.id_to_info
            .get_mut(&self.random_ids[i])
            .expect("exists")
            .random_id_pos = j;
        self.id_to_info
            .get_mut(&self.random_ids[j])
            .expect("exists")
            .random_id_pos = i;
        self.random_ids.swap(i, j);
    }
}