ckb-network 1.2.0

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

/// Address manager
#[derive(Default)]
pub struct AddrManager {
    next_id: u64,
    addr_to_id: HashMap<Multiaddr, 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(&id) = self.addr_to_id.get(&addr_info.addr) {
            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(addr_info.addr.clone(), 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();
            match multiaddr_to_socketaddr(&addr_info.addr) {
                Some(socket_addr) => {
                    let ip = socket_addr.ip();
                    let is_unique_ip = !duplicate_ips.contains(&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)
                    {
                        duplicate_ips.insert(ip);
                        addr_infos.push(addr_info);
                    }
                }
                None => {
                    if filter(&addr_info) {
                        if addr_info.is_connectable(now_ms)
                            || addr_info
                                .addr
                                .iter()
                                .any(|p| matches!(p, Protocol::Onion3(_)))
                        {
                            addr_infos.push(addr_info);
                        } else {
                            debug!(
                                "addr {:?} is not connectable and not an onion address",
                                addr_info.addr
                            );
                        }
                    }
                }
            }
            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> {
        let base_addr = base_addr(addr);
        self.addr_to_id.remove(&base_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> {
        let base_addr = base_addr(addr);
        self.addr_to_id
            .get(&base_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> {
        let base_addr = base_addr(addr);
        if let Some(id) = self.addr_to_id.get(&base_addr) {
            self.id_to_info.get_mut(id)
        } 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);
    }
}