ckb-network 0.105.0

ckb network implementation
Documentation
use crate::{
    multiaddr::Multiaddr,
    peer_store::{
        types::{multiaddr_to_ip_network, AddrInfo, BannedAddr},
        PeerStore,
    },
    PeerId,
};

use std::fs::File;
use std::io::Write;
use std::{collections::HashSet, fs::create_dir_all};

#[test]
fn test_peer_store_persistent() {
    let now_ms = faketime::unix_time_as_millis();
    let mut peer_store = PeerStore::default();

    // add addrs to addr manager
    let addr_manager = peer_store.mut_addr_manager();
    let addr1 = {
        let addr: Multiaddr = format!("/ip4/127.0.0.1/tcp/42/p2p/{}", PeerId::random().to_base58())
            .parse()
            .unwrap();
        AddrInfo::new(addr, 0, 60, 0)
    };
    let addr2 = {
        let addr: Multiaddr = format!("/ip4/127.0.0.5/tcp/42/p2p/{}", PeerId::random().to_base58())
            .parse()
            .unwrap();
        let mut addr_info = AddrInfo::new(addr, 100, 30, 0);
        addr_info.mark_tried(now_ms);
        addr_info
    };
    addr_manager.add(addr1.clone());
    addr_manager.add(addr2.clone());

    // add addrs to ban list
    let ban_list = peer_store.mut_ban_list();
    let addr3 = multiaddr_to_ip_network(&"/ip4/127.0.0.1/tcp/42".parse().unwrap()).unwrap();
    let addr4 = multiaddr_to_ip_network(&"/ip4/127.0.0.2/tcp/42".parse().unwrap()).unwrap();
    let addr5 = multiaddr_to_ip_network(&"/ip4/255.0.0.1/tcp/42".parse().unwrap()).unwrap();
    let ban1 = BannedAddr {
        address: addr3,
        ban_until: now_ms + 10_000,
        ban_reason: "test1".into(),
        created_at: now_ms,
    };
    let ban2 = BannedAddr {
        address: addr4,
        ban_until: now_ms + 20_000,
        ban_reason: "test2".into(),
        created_at: now_ms + 1,
    };
    let ban3 = BannedAddr {
        address: addr5,
        ban_until: now_ms + 30_000,
        ban_reason: "test3".into(),
        created_at: now_ms + 2,
    };
    ban_list.ban(ban1.clone());
    ban_list.ban(ban2.clone());
    ban_list.ban(ban3.clone());

    // dump and load
    let dir = tempfile::tempdir().unwrap();
    peer_store.dump_to_dir(&dir.path()).unwrap();
    let peer_store2 = PeerStore::load_from_dir_or_default(&dir.path());

    // check addr manager
    let addr_manager2 = peer_store2.addr_manager();
    // set random_id_pos to default, this field is internal used only
    let addrs = addr_manager2.addrs_iter().cloned().map(|mut paddr| {
        paddr.random_id_pos = 0;
        paddr
    });
    assert_eq!(
        addrs.collect::<HashSet<_>>(),
        vec![addr1, addr2].into_iter().collect::<HashSet<_>>()
    );

    // check ban list
    let ban_list2 = peer_store2.ban_list();
    assert_eq!(
        ban_list2
            .get_banned_addrs()
            .into_iter()
            .collect::<HashSet<_>>(),
        vec![ban1, ban2, ban3].into_iter().collect::<HashSet<_>>()
    );
}

#[test]
fn test_peer_store_load_from_dir_should_not_panic() {
    // should return an empty store when dir does not exist
    {
        let peer_store = PeerStore::load_from_dir_or_default("/tmp/a_directory_does_not_exist");
        assert_eq!(0, peer_store.addr_manager().count());
        assert_eq!(0, peer_store.ban_list().get_banned_addrs().len());
    }

    // should return an empty store when AddrManager db is empty or broken
    {
        let dir = tempfile::tempdir().unwrap();
        let file_path = dir.path().join("addr_manager.db");
        let mut file = File::create(file_path).unwrap();
        writeln!(file).unwrap();
        let peer_store = PeerStore::load_from_dir_or_default(dir);
        assert_eq!(0, peer_store.addr_manager().count());
    }

    {
        let dir = tempfile::tempdir().unwrap();
        let file_path = dir.path().join("addr_manager.db");
        let mut file = File::create(file_path).unwrap();
        writeln!(file, "broken").unwrap();
        let peer_store = PeerStore::load_from_dir_or_default(dir);
        assert_eq!(0, peer_store.addr_manager().count());
    }

    // should return an empty store when BanList db is empty or broken
    {
        let dir = tempfile::tempdir().unwrap();
        let file_path = dir.path().join("ban_list.db");
        let mut file = File::create(file_path).unwrap();
        writeln!(file).unwrap();
        let peer_store = PeerStore::load_from_dir_or_default(dir);
        assert_eq!(0, peer_store.ban_list().get_banned_addrs().len());
    }
    {
        let dir = tempfile::tempdir().unwrap();
        let file_path = dir.path().join("ban_list.db");
        let mut file = File::create(file_path).unwrap();
        writeln!(file, "broken").unwrap();
        let peer_store = PeerStore::load_from_dir_or_default(dir);
        assert_eq!(0, peer_store.ban_list().get_banned_addrs().len());
    }
}

#[test]
fn test_peer_store_dump_with_broken_tmp_file_should_be_ok() {
    let dir = tempfile::tempdir().unwrap();
    create_dir_all(dir.path().join("tmp")).unwrap();
    // write a truncated json with 8 peers to tmp file
    {
        let tmp_addr_file_path = dir.path().join("tmp/addr_manager.db");
        let mut file = File::create(tmp_addr_file_path).unwrap();
        let truncated_json = r#"[{"peer_id":"QmZDfQQPzQmPXW8DjoCw9QVLJU85rnSxgH3j3u4j19hq4o","ip_port":{"ip":"127.0.0.1","port":438},"addr":"/ip4/127.0.0.1/tcp/438","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":438},{"peer_id":"Qmb34eHAHK4BgSCNQ5KV3Jc6iqh3FBRWEzkT6M4Yztoac6","ip_port":{"ip":"127.0.0.1","port":46},"addr":"/ip4/127.0.0.1/tcp/46","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":46},{"peer_id":"QmeKJZ2AgE3tbP5hjeKeoLGArApQDcgXenKJ3w5eG48esg","ip_port":{"ip":"127.0.0.1","port":538},"addr":"/ip4/127.0.0.1/tcp/538","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":538},{"peer_id":"QmVshsmSHSjk84tcac1wdncMN4hcSXpM5LSQvYkjig3YYS","ip_port":{"ip":"127.0.0.1","port":773},"addr":"/ip4/127.0.0.1/tcp/773","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":773},{"peer_id":"QmbPjmp1rQ1M4G533YPs6CvB5aP6suH92qnhJ6eA1QJJC7","ip_port":{"ip":"127.0.0.1","port":156},"addr":"/ip4/127.0.0.1/tcp/156","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":156},{"peer_id":"Qmf3xG6wJuhXP1QQQMtwAzvz5oCsMtLkyZAJEVe9oW6hae","ip_port":{"ip":"127.0.0.1","port":217},"addr":"/ip4/127.0.0.1/tcp/217","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":217},{"peer_id":"QmdAkE4i1V7ikhnKenYMQ21955q58i9b7XnUcKVRs9TEtP","ip_port":{"ip":"127.0.0.1","port":594},"addr":"/ip4/127.0.0.1/tcp/594","score":60,"last_connected_at_ms":0,"last_tried_at_ms":0,"attempts_count":0,"random_id_pos":594},{"peer_id":"QmawsN1dHNMMx1sDWXdp2r8H8rXanu"#;
        writeln!(file, "{}", truncated_json).unwrap();
        file.sync_all().unwrap();
    }
    // write a truncated json with 3 ban records to tmp file
    {
        let tmp_ban_file_path = dir.path().join("tmp/ban_list.db");
        let mut file = File::create(tmp_ban_file_path).unwrap();
        let truncated_json = r#"[{"address":"192.168.0.2/32","ban_until":31061427677740,"ban_reason":"test","created_at":1612678877739},{"address":"192.168.0.3/32","ban_until":472792659688893,"ban_reason":"test","created_at":1612678888636},{"address":"192.168.0.4/32","ban_until":472792659688893,"ban_reason"#;
        writeln!(file, "{}", truncated_json).unwrap();
        file.sync_all().unwrap();
    }

    // dump with 3 peers and 1 ban list
    let mut peer_store = PeerStore::default();
    let addr_manager = peer_store.mut_addr_manager();
    for i in 0..3 {
        let addr: Multiaddr = format!(
            "/ip4/127.0.0.1/tcp/{}/p2p/{}",
            i,
            PeerId::random().to_base58()
        )
        .parse()
        .unwrap();
        addr_manager.add(AddrInfo::new(addr, 0, 60, 0));
    }
    let ban_list = peer_store.mut_ban_list();
    let now_ms = faketime::unix_time_as_millis();
    ban_list.ban(BannedAddr {
        address: multiaddr_to_ip_network(&"/ip4/127.0.0.1/tcp/42".parse().unwrap()).unwrap(),
        ban_until: now_ms + 10_000,
        ban_reason: "test".into(),
        created_at: now_ms,
    });
    peer_store.dump_to_dir(dir.as_ref()).unwrap();

    // reload from dumped data should be OK
    let peer_store = PeerStore::load_from_dir_or_default(dir.as_ref());
    assert_eq!(1, peer_store.ban_list().count());
    assert_eq!(3, peer_store.addr_manager().count());
}