zebra_network/peer_set/initialize/recent_by_ip.rs
1//! A set of IPs from recent connection attempts.
2
3use std::{
4 collections::{HashMap, VecDeque},
5 net::IpAddr,
6 time::{Duration, Instant},
7};
8
9use crate::constants;
10
11#[cfg(test)]
12mod tests;
13
14#[derive(Debug)]
15/// Stores IPs of recently attempted inbound connections.
16pub struct RecentByIp {
17 /// The list of IPs in decreasing connection age order.
18 pub by_time: VecDeque<(IpAddr, Instant)>,
19
20 /// Stores IPs for recently attempted inbound connections.
21 pub by_ip: HashMap<IpAddr, usize>,
22
23 /// The maximum number of peer connections Zebra will keep for a given IP address
24 /// before it drops any additional peer connections with that IP.
25 pub max_connections_per_ip: usize,
26
27 /// The duration to wait after an entry is added before removing it.
28 pub time_limit: Duration,
29}
30
31impl Default for RecentByIp {
32 fn default() -> Self {
33 Self::new(None, None)
34 }
35}
36
37impl RecentByIp {
38 /// Creates a new [`RecentByIp`]
39 pub fn new(time_limit: Option<Duration>, max_connections_per_ip: Option<usize>) -> Self {
40 let (by_time, by_ip) = Default::default();
41 Self {
42 by_time,
43 by_ip,
44 time_limit: time_limit.unwrap_or(constants::MIN_PEER_RECONNECTION_DELAY),
45 max_connections_per_ip: max_connections_per_ip
46 .unwrap_or(constants::DEFAULT_MAX_CONNS_PER_IP),
47 }
48 }
49
50 /// Prunes outdated entries, checks if there's a recently attempted inbound connection with
51 /// this IP, and adds the entry to `by_time`, and `by_ip` if needed.
52 ///
53 /// Returns true if the recently attempted inbound connection count is past the configured limit.
54 pub fn is_past_limit_or_add(&mut self, ip: IpAddr) -> bool {
55 let now = Instant::now();
56 self.prune_by_time(now);
57
58 let count = self.by_ip.entry(ip).or_default();
59 if *count >= self.max_connections_per_ip {
60 true
61 } else {
62 *count += 1;
63 self.by_time.push_back((ip, now));
64 false
65 }
66 }
67
68 /// Prunes entries older than `time_limit`, decrementing or removing their counts in `by_ip`.
69 fn prune_by_time(&mut self, now: Instant) {
70 // Currently saturates to zero:
71 // <https://doc.rust-lang.org/std/time/struct.Instant.html#monotonicity>
72 //
73 // This discards the whole structure if the time limit is very large,
74 // which is unexpected, but stops this list growing without limit.
75 // After the handshake, the peer set will remove any duplicate connections over the limit.
76 let age_limit = now - self.time_limit;
77
78 // `by_time` must be sorted for this to work.
79 let split_off_idx = self.by_time.partition_point(|&(_, time)| time <= age_limit);
80
81 let updated_by_time = self.by_time.split_off(split_off_idx);
82
83 for (ip, _) in &self.by_time {
84 if let Some(count) = self.by_ip.get_mut(ip) {
85 *count -= 1;
86 if *count == 0 {
87 self.by_ip.remove(ip);
88 }
89 }
90 }
91
92 self.by_time = updated_by_time;
93 }
94}