diff --git a/network/src/network.rs b/network/src/network.rs index 9fb5a14cc8..5824c0e141 100644 --- a/network/src/network.rs +++ b/network/src/network.rs @@ -105,7 +105,9 @@ impl NetworkState { .chain(config.public_addresses.iter()) .map(|addr| (addr.to_owned(), std::u8::MAX)) .collect(); - let peer_store = Mutex::new(PeerStore::load_from_dir(config.peer_store_path())?); + let peer_store = Mutex::new(PeerStore::load_from_dir_or_default( + config.peer_store_path(), + )); let bootnodes = config.bootnodes()?; let whitelist_peers = config diff --git a/network/src/peer_store/peer_store_db.rs b/network/src/peer_store/peer_store_db.rs index 54ca837d18..bfa19627e2 100644 --- a/network/src/peer_store/peer_store_db.rs +++ b/network/src/peer_store/peer_store_db.rs @@ -7,7 +7,7 @@ use crate::{ PeerStore, }, }; -use ckb_logger::debug; +use ckb_logger::{debug, error}; use std::fs::{copy, create_dir_all, remove_file, rename, File, OpenOptions}; use std::io::{Read, Write}; use std::path::Path; @@ -49,25 +49,45 @@ impl BanList { } impl PeerStore { - pub fn load_from_dir>(path: P) -> Result { + pub fn load_from_dir_or_default>(path: P) -> Self { let addr_manager_path = path.as_ref().join(DEFAULT_ADDR_MANAGER_DB); let ban_list_path = path.as_ref().join(DEFAULT_BAN_LIST_DB); - let addr_manager = if addr_manager_path.exists() { - AddrManager::load(File::open(addr_manager_path)?)? - } else { - debug!("Failed to load addr manager from {:?}", addr_manager_path); - AddrManager::default() - }; + let addr_manager = File::open(&addr_manager_path) + .map_err(|err| { + debug!( + "Failed to open AddrManager db, file: {:?}, error: {:?}", + addr_manager_path, err + ) + }) + .and_then(|file| { + AddrManager::load(file).map_err(|err| { + error!( + "Failed to load AddrManager db, file: {:?}, error: {:?}", + addr_manager_path, err + ) + }) + }) + .unwrap_or_default(); - let ban_list = if ban_list_path.exists() { - BanList::load(File::open(ban_list_path)?)? - } else { - debug!("Failed to load ban list from {:?}", ban_list_path); - BanList::default() - }; + let ban_list = File::open(&ban_list_path) + .map_err(|err| { + debug!( + "Failed to open BanList db, file: {:?}, error: {:?}", + ban_list_path, err + ) + }) + .and_then(|file| { + BanList::load(file).map_err(|err| { + error!( + "Failed to load BanList db, file: {:?}, error: {:?}", + ban_list_path, err + ) + }) + }) + .unwrap_or_default(); - Ok(PeerStore::new(addr_manager, ban_list)) + PeerStore::new(addr_manager, ban_list) } pub fn dump_to_dir>(&self, path: P) -> Result<(), Error> { diff --git a/network/src/tests/peer_store_db.rs b/network/src/tests/peer_store_db.rs index 918e9d1baf..1b4b9a04b4 100644 --- a/network/src/tests/peer_store_db.rs +++ b/network/src/tests/peer_store_db.rs @@ -8,6 +8,8 @@ use crate::{ }; use std::collections::HashSet; +use std::fs::File; +use std::io::Write; #[test] fn test_peer_store_persistent() { @@ -71,7 +73,7 @@ fn test_peer_store_persistent() { // dump and load let dir = tempfile::tempdir().unwrap(); peer_store.dump_to_dir(&dir.path()).unwrap(); - let peer_store2 = PeerStore::load_from_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(); @@ -95,3 +97,50 @@ fn test_peer_store_persistent() { vec![ban1, ban2, ban3].into_iter().collect::>() ); } + +#[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()); + } +}