postrust_proxy/ratelimit/
limiter.rs

1//! Rate limiter with per-key tracking.
2
3use crate::config::RateLimitDefaults;
4use crate::ratelimit::TokenBucket;
5use dashmap::DashMap;
6use parking_lot::Mutex;
7use sha2::{Digest, Sha256};
8use std::net::{IpAddr, Ipv6Addr};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11use tokio_util::sync::CancellationToken;
12use tracing::{debug, info, warn};
13use uuid::Uuid;
14
15/// Key for rate limiting: who a bucket is counting.
16#[derive(Clone, Debug, Hash, PartialEq, Eq)]
17pub enum RateLimitKey {
18    /// Rate limit by IP address. Build it with [`RateLimitKey::client_ip`],
19    /// which groups an IPv6 client by its /64.
20    Ip(IpAddr),
21    /// Rate limit by custom header value. Build it with
22    /// [`RateLimitKey::header_value`], which stores a digest rather than the
23    /// value itself.
24    Header(String),
25    /// Rate limit by route ID
26    Route(uuid::Uuid),
27    /// Global rate limit
28    Global,
29}
30
31impl RateLimitKey {
32    /// The key for a client connecting from `ip`.
33    ///
34    /// An IPv6 client is keyed by its /64. A single host is normally handed a
35    /// whole /64, and can pick any address in it for each connection; keyed by
36    /// the full address, every request could come from a fresh bucket, which
37    /// is the same as no limit at all.
38    ///
39    /// An IPv4 client arriving on an IPv6 socket appears as `::ffff:a.b.c.d`.
40    /// That is canonicalised back to IPv4 first, or every IPv4 client would
41    /// share the one /64 that holds all the mapped addresses.
42    pub fn client_ip(ip: IpAddr) -> Self {
43        match ip.to_canonical() {
44            IpAddr::V4(v4) => RateLimitKey::Ip(IpAddr::V4(v4)),
45            IpAddr::V6(v6) => {
46                let s = v6.segments();
47                RateLimitKey::Ip(IpAddr::V6(Ipv6Addr::new(
48                    s[0], s[1], s[2], s[3], 0, 0, 0, 0,
49                )))
50            }
51        }
52    }
53
54    /// The key for a header value.
55    ///
56    /// A digest, not the value: header values are client-chosen and can be
57    /// kilobytes long, and the table holds up to
58    /// [`RateLimitDefaults::max_tracked_keys`] of them. It also keeps API keys
59    /// and the like out of the limiter's memory.
60    pub fn header_value(value: &[u8]) -> Self {
61        RateLimitKey::Header(hex::encode(Sha256::digest(value)))
62    }
63}
64
65/// Requests per window as a refill rate, in tokens per second.
66///
67/// A fraction, never rounded: see [`TokenBucket`]. A window of zero seconds is
68/// read as one, which is what the integer version effectively did.
69pub fn per_second(requests: u32, window_secs: u32) -> f64 {
70    f64::from(requests) / f64::from(window_secs.max(1))
71}
72
73/// Which limit a bucket belongs to, and whom it counts.
74///
75/// The scope is the route. Without it every route shared one bucket per
76/// client, sized by whichever route that client happened to reach first: a
77/// generous limit on one route could be spent on another, and a strict one
78/// could lock a client out of everything.
79type BucketId = (Uuid, RateLimitKey);
80
81/// Rate limiter with per-key token buckets.
82///
83/// Entries are forgotten once their bucket has refilled, which loses nothing:
84/// a full bucket answers exactly as a new one would. That happens in
85/// [`start_cleanup`](Self::start_cleanup), and on demand when the table
86/// reaches its cap. [`crate::vendored::ProxyService`] starts the cleanup task
87/// when it begins serving; anything else that holds a `RateLimiter` for long
88/// should start it too.
89pub struct RateLimiter {
90    /// Per-key token buckets
91    buckets: DashMap<BucketId, TokenBucket>,
92    /// Default configuration
93    defaults: RateLimitDefaults,
94    /// Most entries held at once.
95    max_keys: usize,
96    /// When the table was last swept because it was full, so that a flood of
97    /// new clients costs one sweep a second rather than one each.
98    last_full_sweep: Mutex<Option<Instant>>,
99}
100
101/// The scope [`RateLimiter::check`] counts in: the global defaults, which
102/// belong to no route.
103const DEFAULT_SCOPE: Uuid = Uuid::nil();
104
105impl RateLimiter {
106    /// Create a new rate limiter.
107    pub fn new(defaults: RateLimitDefaults) -> Self {
108        Self {
109            buckets: DashMap::new(),
110            max_keys: defaults.max_tracked_keys.max(1),
111            defaults,
112            last_full_sweep: Mutex::new(None),
113        }
114    }
115
116    /// Check if a request should be allowed, against the configured defaults.
117    ///
118    /// Returns `true` if the request is allowed, `false` if rate limited.
119    pub fn check(&self, key: RateLimitKey) -> bool {
120        let rate = per_second(self.defaults.requests, self.defaults.window_secs);
121        self.check_with_config(DEFAULT_SCOPE, key, rate, self.defaults.burst)
122    }
123
124    /// Check against a limit of `rate_per_sec` tokens a second and a burst of
125    /// `burst`, in the bucket for `key` within `scope`.
126    ///
127    /// `scope` keeps limits apart: pass the route's identity, so that one
128    /// route's traffic is not counted against another's.
129    pub fn check_with_config(
130        &self,
131        scope: Uuid,
132        key: RateLimitKey,
133        rate_per_sec: f64,
134        burst: u32,
135    ) -> bool {
136        let id = (scope, key);
137        if let Some(bucket) = self.buckets.get(&id) {
138            return bucket.try_acquire();
139        }
140
141        if self.buckets.len() >= self.max_keys && !self.make_room() {
142            // Refuse rather than let the request through untracked. Failing
143            // open here would mean that filling the table with junk keys turns
144            // the limit off for every client that arrives after.
145            warn!(
146                max_keys = self.max_keys,
147                "Rate limiter is tracking as many clients as it may; refusing a new one"
148            );
149            return false;
150        }
151
152        self.buckets
153            .entry(id)
154            .or_insert_with(|| TokenBucket::new(u64::from(burst), rate_per_sec))
155            .try_acquire()
156    }
157
158    /// Get remaining tokens for a key under the defaults (approximate).
159    pub fn remaining(&self, key: &RateLimitKey) -> Option<u64> {
160        self.buckets
161            .get(&(DEFAULT_SCOPE, key.clone()))
162            .map(|bucket| bucket.available())
163    }
164
165    /// How many buckets are held.
166    pub fn tracked_keys(&self) -> usize {
167        self.buckets.len()
168    }
169
170    /// Start background cleanup task.
171    pub async fn start_cleanup(self: Arc<Self>, cancel_token: CancellationToken) {
172        let cleanup_interval = Duration::from_secs(60);
173
174        info!("Rate limiter cleanup task started");
175
176        loop {
177            tokio::select! {
178                _ = cancel_token.cancelled() => {
179                    info!("Rate limiter cleanup task stopped");
180                    break;
181                }
182                _ = tokio::time::sleep(cleanup_interval) => {
183                    self.cleanup_expired();
184                }
185            }
186        }
187    }
188
189    /// Sweep a full table, at most once a second. Whether there is room now.
190    fn make_room(&self) -> bool {
191        {
192            let mut last = self.last_full_sweep.lock();
193            let now = Instant::now();
194            if last.is_some_and(|at| now.duration_since(at) < Duration::from_secs(1)) {
195                return self.buckets.len() < self.max_keys;
196            }
197            *last = Some(now);
198        }
199        self.cleanup_expired();
200        self.buckets.len() < self.max_keys
201    }
202
203    /// Forget every bucket that has refilled completely.
204    fn cleanup_expired(&self) {
205        let before = self.buckets.len();
206        self.buckets.retain(|_, bucket| !bucket.is_full());
207        let removed = before.saturating_sub(self.buckets.len());
208        if removed > 0 {
209            debug!("Rate limiter cleanup: removed {} refilled entries", removed);
210        }
211    }
212}
213
214impl Default for RateLimiter {
215    fn default() -> Self {
216        Self::new(RateLimitDefaults::default())
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use std::net::Ipv4Addr;
224
225    fn defaults(requests: u32, window_secs: u32, burst: u32) -> RateLimitDefaults {
226        RateLimitDefaults {
227            requests,
228            window_secs,
229            burst,
230            ..RateLimitDefaults::default()
231        }
232    }
233
234    #[test]
235    fn test_rate_limiter_by_ip() {
236        let limiter = RateLimiter::new(defaults(600, 60, 5));
237
238        let ip1 = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)));
239        let ip2 = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)));
240
241        // Each IP gets its own bucket
242        for _ in 0..5 {
243            assert!(limiter.check(ip1.clone()));
244            assert!(limiter.check(ip2.clone()));
245        }
246
247        // Both exhausted
248        assert!(!limiter.check(ip1.clone()));
249        assert!(!limiter.check(ip2.clone()));
250    }
251
252    #[test]
253    fn test_rate_limiter_remaining() {
254        let limiter = RateLimiter::new(defaults(600, 60, 10));
255
256        let key = RateLimitKey::Global;
257        assert!(limiter.check(key.clone()));
258
259        let remaining = limiter.remaining(&key).unwrap();
260        assert!(remaining < 10);
261    }
262
263    #[test]
264    fn the_rate_is_a_fraction() {
265        assert_eq!(per_second(3, 60), 0.05);
266        assert_eq!(per_second(600, 60), 10.0);
267        // A zero window is read as one second rather than dividing by zero.
268        assert_eq!(per_second(5, 0), 5.0);
269    }
270
271    #[test]
272    fn scopes_do_not_share_a_bucket() {
273        // The bug: one bucket per IP across every route, sized by whichever
274        // route was hit first.
275        let limiter = RateLimiter::default();
276        let ip = RateLimitKey::client_ip("10.0.0.1".parse().unwrap());
277        let strict = Uuid::new_v4();
278        let generous = Uuid::new_v4();
279
280        assert!(limiter.check_with_config(strict, ip.clone(), 0.01, 1));
281        assert!(!limiter.check_with_config(strict, ip.clone(), 0.01, 1));
282
283        // Spending the strict route's burst leaves the generous route alone,
284        // and the generous route gets its own size of bucket.
285        for _ in 0..100 {
286            assert!(limiter.check_with_config(generous, ip.clone(), 0.01, 100));
287        }
288        assert!(!limiter.check_with_config(generous, ip.clone(), 0.01, 100));
289    }
290
291    #[test]
292    fn an_ipv6_client_is_keyed_by_its_slash_64() {
293        let a = RateLimitKey::client_ip("2001:db8:1:2:aaaa::1".parse().unwrap());
294        let b = RateLimitKey::client_ip("2001:db8:1:2:ffff:1:2:3".parse().unwrap());
295        let other = RateLimitKey::client_ip("2001:db8:1:3::1".parse().unwrap());
296        assert_eq!(a, b);
297        assert_ne!(a, other);
298    }
299
300    #[test]
301    fn a_mapped_ipv4_client_is_keyed_as_ipv4() {
302        // Not by the /64 of ::ffff:0:0/96, which every IPv4 client shares.
303        let mapped = RateLimitKey::client_ip("::ffff:192.0.2.1".parse().unwrap());
304        let plain = RateLimitKey::client_ip("192.0.2.1".parse().unwrap());
305        let other = RateLimitKey::client_ip("::ffff:192.0.2.2".parse().unwrap());
306        assert_eq!(mapped, plain);
307        assert_ne!(mapped, other);
308    }
309
310    #[test]
311    fn a_header_value_is_stored_as_a_digest() {
312        let key = RateLimitKey::header_value(&[b'k'; 8192]);
313        match key {
314            RateLimitKey::Header(digest) => assert_eq!(digest.len(), 64),
315            other => panic!("expected a header key, got {other:?}"),
316        }
317        assert_eq!(
318            RateLimitKey::header_value(b"a"),
319            RateLimitKey::header_value(b"a")
320        );
321        assert_ne!(
322            RateLimitKey::header_value(b"a"),
323            RateLimitKey::header_value(b"b")
324        );
325    }
326
327    #[test]
328    fn the_table_is_capped() {
329        let limiter = RateLimiter::new(RateLimitDefaults {
330            max_tracked_keys: 3,
331            ..RateLimitDefaults::default()
332        });
333        let scope = Uuid::new_v4();
334        let key = |n: u8| RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(10, 0, 0, n)));
335
336        // Three clients, each with an empty bucket that will not refill soon.
337        for n in 0..3 {
338            assert!(limiter.check_with_config(scope, key(n), 0.0001, 1));
339        }
340        assert_eq!(limiter.tracked_keys(), 3);
341
342        // A fourth is refused rather than tracked or waved through, and the
343        // table does not grow.
344        assert!(!limiter.check_with_config(scope, key(3), 0.0001, 1));
345        assert_eq!(limiter.tracked_keys(), 3);
346
347        // Clients already tracked are still answered from their buckets.
348        assert!(!limiter.check_with_config(scope, key(0), 0.0001, 1));
349    }
350
351    #[test]
352    fn a_full_table_makes_room_by_dropping_refilled_buckets() {
353        let limiter = RateLimiter::new(RateLimitDefaults {
354            max_tracked_keys: 2,
355            ..RateLimitDefaults::default()
356        });
357        let scope = Uuid::new_v4();
358        let key = |n: u8| RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(10, 0, 0, n)));
359
360        // A burst of 2 with one spent leaves the bucket refilling; a huge
361        // rate refills it before the next check.
362        assert!(limiter.check_with_config(scope, key(0), 1e9, 2));
363        assert!(limiter.check_with_config(scope, key(1), 1e9, 2));
364        std::thread::sleep(Duration::from_millis(5));
365
366        assert!(limiter.check_with_config(scope, key(2), 1e9, 2));
367        assert!(limiter.tracked_keys() <= 2);
368    }
369
370    #[test]
371    fn cleanup_forgets_only_refilled_buckets() {
372        let limiter = RateLimiter::default();
373        let scope = Uuid::new_v4();
374        let spent = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)));
375        let refilled = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)));
376
377        assert!(limiter.check_with_config(scope, spent.clone(), 0.0001, 1));
378        assert!(limiter.check_with_config(scope, refilled, 1e9, 1));
379        std::thread::sleep(Duration::from_millis(5));
380
381        limiter.cleanup_expired();
382        assert_eq!(limiter.tracked_keys(), 1);
383        // The spent one is remembered: still refused.
384        assert!(!limiter.check_with_config(scope, spent, 0.0001, 1));
385    }
386}
postrust_proxy/ratelimit/limiter.rs

postrust_proxy/ratelimit/
limiter.rs
