pacenotes_server/reports/
rate_limiter.rs

1/// Token buckets: each client may burst `client_burst` reports and then earns
2/// `client_per_hour`, the server as a whole likewise, and `max_clients` bounds
3/// what a flood of addresses can take.
4#[derive(Clone, Debug)]
5pub struct Limits {
6    pub client_burst: f64,
7    pub client_per_hour: f64,
8    pub global_burst: f64,
9    pub global_per_hour: f64,
10    pub max_clients: usize,
11}
12
13impl Default for Limits {
14    /// A healthy installer reports at most once per run, and 600 an hour is far
15    /// beyond any real user base while still bounding a broken client.
16    fn default() -> Self {
17        Limits {
18            client_burst: 10.0,
19            client_per_hour: 10.0,
20            global_burst: 600.0,
21            global_per_hour: 600.0,
22            max_clients: 10_000,
23        }
24    }
25}
26
27struct Bucket {
28    tokens: f64,
29    updated: std::time::Instant,
30}
31
32impl Bucket {
33    fn full(capacity: f64, now: std::time::Instant) -> Bucket {
34        Bucket {
35            tokens: capacity,
36            updated: now,
37        }
38    }
39
40    fn refill(&mut self, now: std::time::Instant, capacity: f64, per_hour: f64) {
41        let earned = now.duration_since(self.updated).as_secs_f64() * per_hour / 3_600.0;
42        self.tokens = (self.tokens + earned).min(capacity);
43        self.updated = now;
44    }
45
46    /// Takes a token, or says how many seconds until one is earned.
47    fn take(&mut self, now: std::time::Instant, capacity: f64, per_hour: f64) -> Result<(), u64> {
48        self.refill(now, capacity, per_hour);
49        if self.tokens >= 1.0 {
50            self.tokens -= 1.0;
51            Ok(())
52        } else {
53            Err(((1.0 - self.tokens) * 3_600.0 / per_hour).ceil() as u64)
54        }
55    }
56}
57
58pub struct RateLimiter {
59    limits: Limits,
60    clients: std::sync::Mutex<std::collections::HashMap<std::net::IpAddr, Bucket>>,
61    global: std::sync::Mutex<Bucket>,
62}
63
64impl RateLimiter {
65    pub fn new(limits: Limits) -> RateLimiter {
66        let now = std::time::Instant::now();
67        RateLimiter {
68            global: std::sync::Mutex::new(Bucket::full(limits.global_burst, now)),
69            clients: std::sync::Mutex::new(std::collections::HashMap::new()),
70            limits,
71        }
72    }
73
74    /// `Err` carries a Retry-After in seconds. The client's own bucket goes
75    /// first, so a client over its limit cannot drain the global one.
76    pub fn check(&self, client: std::net::IpAddr, now: std::time::Instant) -> Result<(), u64> {
77        let limits = &self.limits;
78        {
79            let mut clients = self.clients.lock().unwrap_or_else(|e| e.into_inner());
80            if !clients.contains_key(&client) && clients.len() >= limits.max_clients {
81                // A refilled bucket is indistinguishable from a new client.
82                clients.retain(|_, bucket| {
83                    bucket.refill(now, limits.client_burst, limits.client_per_hour);
84                    bucket.tokens < limits.client_burst
85                });
86                if clients.len() >= limits.max_clients {
87                    return Err(60);
88                }
89            }
90            clients
91                .entry(client)
92                .or_insert_with(|| Bucket::full(limits.client_burst, now))
93                .take(now, limits.client_burst, limits.client_per_hour)?;
94        }
95        self.global.lock().unwrap_or_else(|e| e.into_inner()).take(
96            now,
97            limits.global_burst,
98            limits.global_per_hour,
99        )
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    fn tight() -> super::Limits {
106        super::Limits {
107            client_burst: 2.0,
108            client_per_hour: 2.0,
109            global_burst: 3.0,
110            global_per_hour: 3.0,
111            max_clients: 2,
112        }
113    }
114
115    fn ip(text: &str) -> std::net::IpAddr {
116        text.parse().unwrap()
117    }
118
119    #[test]
120    fn the_limiter_holds_per_client_and_globally() {
121        let limiter = super::RateLimiter::new(tight());
122        let now = std::time::Instant::now();
123        let a = ip("198.51.100.1");
124        let b = ip("198.51.100.2");
125
126        assert!(limiter.check(a, now).is_ok());
127        assert!(limiter.check(a, now).is_ok());
128        let retry = limiter.check(a, now).unwrap_err();
129        assert!((1_700..=1_800).contains(&retry), "{retry}");
130        assert!(limiter.check(b, now).is_ok());
131        assert!(limiter.check(b, now).is_err(), "the global bucket ran out");
132
133        let an_hour_later = now + std::time::Duration::from_secs(3_600);
134        assert!(limiter.check(a, an_hour_later).is_ok());
135    }
136
137    #[test]
138    fn a_flood_of_addresses_cannot_grow_memory_without_bound() {
139        let limiter = super::RateLimiter::new(tight());
140        let now = std::time::Instant::now();
141        assert!(limiter.check(ip("198.51.100.1"), now).is_ok());
142        assert!(limiter.check(ip("198.51.100.2"), now).is_ok());
143        assert!(limiter.check(ip("198.51.100.3"), now).is_err());
144
145        let once_refilled = now + std::time::Duration::from_secs(7_200);
146        assert!(limiter.check(ip("198.51.100.3"), once_refilled).is_ok());
147    }
148}
pacenotes_server/reports/rate_limiter.rs

pacenotes_server/reports/
rate_limiter.rs
