ckb_network/peer_store/
peer_store_db.rs

1use crate::{
2    errors::{Error, PeerStoreError},
3    peer_store::{
4        PeerStore,
5        addr_manager::AddrManager,
6        ban_list::BanList,
7        types::{AddrInfo, BannedAddr},
8    },
9};
10use ckb_logger::{debug, error};
11use std::path::Path;
12use std::{
13    fs::{File, OpenOptions, copy, create_dir_all, remove_file, rename},
14    io::{Read, Write},
15};
16
17const DEFAULT_ADDR_MANAGER_DB: &str = "addr_manager.db";
18const DEFAULT_BAN_LIST_DB: &str = "ban_list.db";
19
20impl AddrManager {
21    /// Load address list from disk
22    pub fn load<R: Read>(r: R) -> Result<Self, Error> {
23        let addrs: Vec<AddrInfo> = serde_json::from_reader(r).map_err(PeerStoreError::Serde)?;
24        let mut addr_manager = AddrManager::default();
25        addrs.into_iter().for_each(|addr| addr_manager.add(addr));
26        Ok(addr_manager)
27    }
28
29    /// Dump address list to disk
30    pub fn dump(&self, mut file: File) -> Result<(), Error> {
31        let addrs: Vec<_> = self.addrs_iter().collect();
32        debug!("Dump {} addrs", addrs.len());
33        // empty file and dump the json string to it
34        file.set_len(0)
35            .and_then(|_| serde_json::to_string(&addrs).map_err(Into::into))
36            .and_then(|json_string| file.write_all(json_string.as_bytes()))
37            .and_then(|_| file.sync_all())
38            .map_err(Into::into)
39    }
40
41    #[cfg(target_family = "wasm")]
42    pub fn dump_data(&self) -> Vec<u8> {
43        let addrs: Vec<_> = self.addrs_iter().collect();
44        serde_json::to_string(&addrs).unwrap().into_bytes()
45    }
46}
47
48impl BanList {
49    /// Load ban list from disk
50    pub fn load<R: Read>(r: R) -> Result<Self, Error> {
51        let banned_addrs: Vec<BannedAddr> =
52            serde_json::from_reader(r).map_err(PeerStoreError::Serde)?;
53        let mut ban_list = BanList::default();
54        banned_addrs
55            .into_iter()
56            .for_each(|banned_addr| ban_list.ban(banned_addr));
57        Ok(ban_list)
58    }
59
60    /// Dump ban list to disk
61    pub fn dump(&self, mut file: File) -> Result<(), Error> {
62        let banned_addrs = self.get_banned_addrs();
63        debug!("Dump {} banned addrs", banned_addrs.len());
64        // empty file and dump the json string to it
65        file.set_len(0)
66            .and_then(|_| serde_json::to_string(&banned_addrs).map_err(Into::into))
67            .and_then(|json_string| file.write_all(json_string.as_bytes()))
68            .and_then(|_| file.sync_all())
69            .map_err(Into::into)
70    }
71
72    #[cfg(target_family = "wasm")]
73    pub fn dump_data(&self) -> Vec<u8> {
74        let banned_addrs = self.get_banned_addrs();
75        serde_json::to_string(&banned_addrs).unwrap().into_bytes()
76    }
77}
78
79impl PeerStore {
80    /// Init peer store from disk
81    pub fn load_from_dir_or_default<P: AsRef<Path>>(path: P) -> Self {
82        let addr_manager_path = path.as_ref().join(DEFAULT_ADDR_MANAGER_DB);
83        let ban_list_path = path.as_ref().join(DEFAULT_BAN_LIST_DB);
84
85        let addr_manager = File::open(&addr_manager_path)
86            .map_err(|err| {
87                debug!(
88                    "Failed to open AddrManager db, file: {:?}, error: {:?}",
89                    addr_manager_path, err
90                )
91            })
92            .and_then(|file| {
93                AddrManager::load(std::io::BufReader::new(file)).map_err(|err| {
94                    error!(
95                        "Failed to load AddrManager db, file: {:?}, error: {:?}",
96                        addr_manager_path, err
97                    )
98                })
99            })
100            .unwrap_or_default();
101
102        let ban_list = File::open(&ban_list_path)
103            .map_err(|err| {
104                debug!(
105                    "Failed to open BanList db, file: {:?}, error: {:?}",
106                    ban_list_path, err
107                )
108            })
109            .and_then(|file| {
110                BanList::load(std::io::BufReader::new(file)).map_err(|err| {
111                    error!(
112                        "Failed to load BanList db, file: {:?}, error: {:?}",
113                        ban_list_path, err
114                    )
115                })
116            })
117            .unwrap_or_default();
118
119        PeerStore::new(addr_manager, ban_list)
120    }
121
122    #[cfg(target_family = "wasm")]
123    pub async fn load_from_idb<P: AsRef<Path>>(path: P) -> Self {
124        use crate::peer_store::browser::get_db;
125
126        let addr_manager_path = path
127            .as_ref()
128            .join(DEFAULT_ADDR_MANAGER_DB)
129            .to_str()
130            .unwrap()
131            .to_owned()
132            .into_bytes();
133        let ban_list_path = path
134            .as_ref()
135            .join(DEFAULT_BAN_LIST_DB)
136            .to_str()
137            .unwrap()
138            .to_owned()
139            .into_bytes();
140
141        let db = get_db(path).await;
142
143        let addr_manager = db
144            .get(&addr_manager_path)
145            .await
146            .map_err(|err| debug!("Failed to get indexdb value, error: {:?}", err))
147            .and_then(|data| {
148                AddrManager::load(std::io::Cursor::new(data.unwrap_or_default()))
149                    .map_err(|err| debug!("Failed to load peer store value, error: {:?}", err))
150            })
151            .unwrap_or_default();
152
153        let ban_list = db
154            .get(&ban_list_path)
155            .await
156            .map_err(|err| debug!("Failed to get indexdb value, error: {:?}", err))
157            .and_then(|data| {
158                BanList::load(std::io::Cursor::new(data.unwrap_or_default()))
159                    .map_err(|err| debug!("Failed to load BanList value, error: {:?}", err))
160            })
161            .unwrap_or_default();
162        PeerStore::new(addr_manager, ban_list)
163    }
164
165    /// Dump all info to disk
166    pub fn dump_to_dir<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
167        // create dir
168        create_dir_all(&path)?;
169        // dump file to a temporary sub-directory
170        let tmp_dir = path.as_ref().join("tmp");
171        create_dir_all(&tmp_dir)?;
172        let tmp_addr_manager = tmp_dir.join(DEFAULT_ADDR_MANAGER_DB);
173        let tmp_ban_list = tmp_dir.join(DEFAULT_BAN_LIST_DB);
174        self.addr_manager().dump(
175            OpenOptions::new()
176                .write(true)
177                .create(true)
178                .truncate(true)
179                .append(false)
180                .open(&tmp_addr_manager)?,
181        )?;
182        move_file(
183            tmp_addr_manager,
184            path.as_ref().join(DEFAULT_ADDR_MANAGER_DB),
185        )?;
186        self.ban_list().dump(
187            OpenOptions::new()
188                .write(true)
189                .create(true)
190                .truncate(true)
191                .append(false)
192                .open(&tmp_ban_list)?,
193        )?;
194        move_file(tmp_ban_list, path.as_ref().join(DEFAULT_BAN_LIST_DB))?;
195        Ok(())
196    }
197
198    #[cfg(target_family = "wasm")]
199    pub fn dump_to_idb<P: AsRef<Path>>(
200        &self,
201        path: P,
202    ) -> impl std::future::Future<Output = ()> + use<P> {
203        use crate::peer_store::browser::get_db;
204        let ban_list = self.ban_list().dump_data();
205        let addr_manager = self.addr_manager().dump_data();
206        let addr_manager_path = path
207            .as_ref()
208            .join(DEFAULT_ADDR_MANAGER_DB)
209            .to_str()
210            .unwrap()
211            .to_owned();
212        let ban_list_path = path
213            .as_ref()
214            .join(DEFAULT_BAN_LIST_DB)
215            .to_str()
216            .unwrap()
217            .to_owned();
218        async {
219            let db = get_db(path).await;
220
221            let _ignore = db.put(addr_manager_path.into_bytes(), addr_manager).await;
222            let _ignore = db.put(ban_list_path.into_bytes(), ban_list).await;
223        }
224    }
225}
226
227/// This function use `copy` then `remove_file` as a fallback when `rename` failed,
228/// this maybe happen when src and dst on different file systems.
229fn move_file<P: AsRef<Path>>(src: P, dst: P) -> Result<(), Error> {
230    if rename(&src, &dst).is_err() {
231        copy(&src, &dst)?;
232        remove_file(&src)?;
233    }
234    Ok(())
235}