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}