cuprate_address_book/
peer_list.rs

Help
1use std::collections::{BTreeMap, HashMap, HashSet};
2
3use indexmap::IndexMap;
4use rand::prelude::*;
5
6use cuprate_constants::block::MAX_BLOCK_HEIGHT_USIZE;
7use cuprate_p2p_core::{services::ZoneSpecificPeerListEntryBase, NetZoneAddress, NetworkZone};
8use cuprate_pruning::PruningSeed;
9
10#[cfg(test)]
11pub(crate) mod tests;
12
13/// A Peer list in the address book.
14///
15/// This could either be the white list or gray list.
16#[derive(Debug)]
17pub(crate) struct PeerList<Z: NetworkZone> {
18    /// The peers with their peer data.
19    pub peers: IndexMap<Z::Addr, ZoneSpecificPeerListEntryBase<Z::Addr>>,
20    /// An index of Pruning seed to address, so can quickly grab peers with the blocks
21    /// we want.
22    ///
23    /// Pruning seeds are sorted by first their `log_stripes` and then their stripe.
24    /// This means the first peers in this list will store more blocks than peers
25    /// later on. So when we need a peer with a certain block we look at the peers
26    /// storing more blocks first then work our way to the peers storing less.
27    ///  
28    pruning_seeds: BTreeMap<PruningSeed, Vec<Z::Addr>>,
29    /// A hashmap linking `ban_ids` to addresses.
30    ban_ids: HashMap<<Z::Addr as NetZoneAddress>::BanID, Vec<Z::Addr>>,
31}
32
33impl<Z: NetworkZone> PeerList<Z> {
34    /// Creates a new peer list.
35    pub(crate) fn new(list: Vec<ZoneSpecificPeerListEntryBase<Z::Addr>>) -> Self {
36        let mut peers = IndexMap::with_capacity(list.len());
37        let mut pruning_seeds = BTreeMap::new();
38        let mut ban_ids = HashMap::with_capacity(list.len());
39
40        for peer in list {
41            pruning_seeds
42                .entry(peer.pruning_seed)
43                .or_insert_with(Vec::new)
44                .push(peer.adr);
45
46            ban_ids
47                .entry(peer.adr.ban_id())
48                .or_insert_with(Vec::new)
49                .push(peer.adr);
50
51            peers.insert(peer.adr, peer);
52        }
53        Self {
54            peers,
55            pruning_seeds,
56            ban_ids,
57        }
58    }
59
60    /// Gets the length of the peer list
61    pub(crate) fn len(&self) -> usize {
62        self.peers.len()
63    }
64
65    /// Adds a new peer to the peer list
66    pub(crate) fn add_new_peer(&mut self, peer: ZoneSpecificPeerListEntryBase<Z::Addr>) {
67        if self.peers.insert(peer.adr, peer).is_none() {
68            #[expect(clippy::unwrap_or_default, reason = "It's more clear with this")]
69            self.pruning_seeds
70                .entry(peer.pruning_seed)
71                .or_insert_with(Vec::new)
72                .push(peer.adr);
73
74            #[expect(clippy::unwrap_or_default)]
75            self.ban_ids
76                .entry(peer.adr.ban_id())
77                .or_insert_with(Vec::new)
78                .push(peer.adr);
79        }
80    }
81
82    /// Returns a random peer.
83    /// If the pruning seed is specified then we will get a random peer with
84    /// that pruning seed otherwise we will just get a random peer in the whole
85    /// list.
86    ///
87    /// The given peer will be removed from the peer list.
88    pub(crate) fn take_random_peer<R: Rng>(
89        &mut self,
90        r: &mut R,
91        block_needed: Option<usize>,
92        must_keep_peers: &HashSet<Z::Addr>,
93    ) -> Option<ZoneSpecificPeerListEntryBase<Z::Addr>> {
94        // Take a random peer and see if it's in the list of must_keep_peers, if it is try again.
95        // TODO: improve this
96
97        for _ in 0..3 {
98            if let Some(needed_height) = block_needed {
99                let (_, addresses_with_block) = self.pruning_seeds.iter().find(|(seed, _)| {
100                    // TODO: factor in peer blockchain height?
101                    seed.get_next_unpruned_block(needed_height, MAX_BLOCK_HEIGHT_USIZE)
102                        .expect("Block needed is higher than max block allowed.")
103                        == needed_height
104                })?;
105                let n = r.gen_range(0..addresses_with_block.len());
106                let peer = addresses_with_block[n];
107                if must_keep_peers.contains(&peer) {
108                    continue;
109                }
110
111                return self.remove_peer(&peer);
112            }
113            let len = self.len();
114
115            if len == 0 {
116                return None;
117            }
118
119            let n = r.gen_range(0..len);
120
121            let (&key, _) = self.peers.get_index(n).unwrap();
122            if !must_keep_peers.contains(&key) {
123                return self.remove_peer(&key);
124            }
125        }
126
127        None
128    }
129
130    pub(crate) fn get_random_peers<R: Rng>(
131        &self,
132        r: &mut R,
133        len: usize,
134    ) -> Vec<ZoneSpecificPeerListEntryBase<Z::Addr>> {
135        let mut peers = self.peers.values().copied().choose_multiple(r, len);
136        // Order of the returned peers is not random, I am unsure of the impact of this, potentially allowing someone to make guesses about which peers
137        // were connected first.
138        // So to mitigate this shuffle the result.
139        peers.shuffle(r);
140        peers.drain(len.min(peers.len())..peers.len());
141        peers
142    }
143
144    /// Returns a mutable reference to a peer.
145    pub(crate) fn get_peer_mut(
146        &mut self,
147        peer: &Z::Addr,
148    ) -> Option<&mut ZoneSpecificPeerListEntryBase<Z::Addr>> {
149        self.peers.get_mut(peer)
150    }
151
152    /// Returns true if the list contains this peer.
153    pub(crate) fn contains_peer(&self, peer: &Z::Addr) -> bool {
154        self.peers.contains_key(peer)
155    }
156
157    /// Removes a peer from the pruning idx
158    ///
159    /// MUST NOT BE USED ALONE
160    fn remove_peer_pruning_idx(&mut self, peer: &ZoneSpecificPeerListEntryBase<Z::Addr>) {
161        remove_peer_idx::<Z>(self.pruning_seeds.get_mut(&peer.pruning_seed), &peer.adr);
162        if self
163            .pruning_seeds
164            .get(&peer.pruning_seed)
165            .expect("There must be a peer with this id")
166            .is_empty()
167        {
168            self.pruning_seeds.remove(&peer.pruning_seed);
169        }
170    }
171
172    /// Removes a peer from the ban idx
173    ///
174    /// MUST NOT BE USED ALONE
175    fn remove_peer_ban_idx(&mut self, peer: &ZoneSpecificPeerListEntryBase<Z::Addr>) {
176        remove_peer_idx::<Z>(self.ban_ids.get_mut(&peer.adr.ban_id()), &peer.adr);
177        if self
178            .ban_ids
179            .get(&peer.adr.ban_id())
180            .expect("There must be a peer with this id")
181            .is_empty()
182        {
183            self.ban_ids.remove(&peer.adr.ban_id());
184        }
185    }
186
187    /// Removes a peer from all the indexes
188    ///
189    /// MUST NOT BE USED ALONE
190    fn remove_peer_from_all_idxs(&mut self, peer: &ZoneSpecificPeerListEntryBase<Z::Addr>) {
191        self.remove_peer_pruning_idx(peer);
192        self.remove_peer_ban_idx(peer);
193    }
194
195    /// Removes a peer from the peer list
196    pub(crate) fn remove_peer(
197        &mut self,
198        peer: &Z::Addr,
199    ) -> Option<ZoneSpecificPeerListEntryBase<Z::Addr>> {
200        let peer_eb = self.peers.swap_remove(peer)?;
201        self.remove_peer_from_all_idxs(&peer_eb);
202        Some(peer_eb)
203    }
204
205    /// Removes all peers with a specific ban id.
206    pub(crate) fn remove_peers_with_ban_id(&mut self, ban_id: &<Z::Addr as NetZoneAddress>::BanID) {
207        let Some(addresses) = self.ban_ids.get(ban_id) else {
208            // No peers to ban
209            return;
210        };
211
212        for addr in addresses.clone() {
213            self.remove_peer(&addr);
214        }
215    }
216
217    /// Tries to reduce the peer list to `new_len`.
218    ///
219    /// This function could keep the list bigger than `new_len` if `must_keep_peers`s length
220    /// is larger than `new_len`, in that case we will remove as much as we can.
221    pub(crate) fn reduce_list(&mut self, must_keep_peers: &HashSet<Z::Addr>, new_len: usize) {
222        if new_len >= self.len() {
223            return;
224        }
225
226        let target_removed = self.len() - new_len;
227        let mut removed_count = 0;
228        let mut peers_to_remove: Vec<Z::Addr> = Vec::with_capacity(target_removed);
229
230        for peer_adr in self.peers.keys() {
231            if removed_count >= target_removed {
232                break;
233            }
234            if !must_keep_peers.contains(peer_adr) {
235                peers_to_remove.push(*peer_adr);
236                removed_count += 1;
237            }
238        }
239
240        for peer_adr in peers_to_remove {
241            let _ = self.remove_peer(&peer_adr);
242        }
243    }
244}
245
246/// Remove a peer from an index.
247fn remove_peer_idx<Z: NetworkZone>(peer_list: Option<&mut Vec<Z::Addr>>, addr: &Z::Addr) {
248    if let Some(peer_list) = peer_list {
249        if let Some(idx) = peer_list.iter().position(|peer_adr| peer_adr == addr) {
250            peer_list.swap_remove(idx);
251        } else {
252            unreachable!("This function will only be called when the peer exists.");
253        }
254    } else {
255        unreachable!("Index must exist if a peer has that index");
256    }
257}
cuprate_address_book/peer_list.rs

cuprate_address_book/
peer_list.rs
