hns_browser_loopback_proxy/chromium/
rate_limit.rs

1//! Bounded concurrency and request-rate admission for the loopback listener.
2
3use std::collections::{HashMap, VecDeque, hash_map::RandomState};
4use std::hash::{BuildHasher, Hasher};
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::{Arc, Mutex, MutexGuard};
7use std::time::{Duration, Instant};
8use thiserror::Error;
9
10pub const DEFAULT_MAX_ACTIVE_CLIENTS: usize = 64;
11pub const DEFAULT_GLOBAL_REQUESTS_PER_WINDOW: usize = 240;
12pub const DEFAULT_HOST_REQUESTS_PER_WINDOW: usize = DEFAULT_GLOBAL_REQUESTS_PER_WINDOW;
13pub const DEFAULT_MAX_TRACKED_HOSTS: usize = 256;
14pub const DEFAULT_RATE_WINDOW: Duration = Duration::from_secs(10);
15
16/// Independently configurable request limits for a proxy generation.
17///
18/// Keeping this type in the admission module lets platform configuration map
19/// into it without making the limiter depend on the server or configuration
20/// modules.
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub struct RateLimitConfig {
23    pub global_requests: usize,
24    pub per_host_requests: usize,
25    pub window: Duration,
26    pub max_tracked_hosts: usize,
27}
28
29impl Default for RateLimitConfig {
30    fn default() -> Self {
31        Self {
32            global_requests: DEFAULT_GLOBAL_REQUESTS_PER_WINDOW,
33            per_host_requests: DEFAULT_HOST_REQUESTS_PER_WINDOW,
34            window: DEFAULT_RATE_WINDOW,
35            max_tracked_hosts: DEFAULT_MAX_TRACKED_HOSTS,
36        }
37    }
38}
39
40#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
41pub enum RateLimitConfigError {
42    #[error("the global request limit must be non-zero")]
43    ZeroGlobalLimit,
44    #[error("the per-host request limit must be non-zero")]
45    ZeroHostLimit,
46    #[error("the rate-limit window must be non-zero")]
47    ZeroWindow,
48    #[error("the tracked-host limit must be non-zero")]
49    ZeroTrackedHosts,
50    #[error("the per-host request limit cannot exceed the global request limit")]
51    HostLimitExceedsGlobal,
52    #[error("a rate-limit value exceeds the hard proxy maximum")]
53    ExceedsMaximum,
54}
55
56impl RateLimitConfig {
57    pub fn validate(self) -> Result<Self, RateLimitConfigError> {
58        if self.global_requests == 0 {
59            return Err(RateLimitConfigError::ZeroGlobalLimit);
60        }
61        if self.per_host_requests == 0 {
62            return Err(RateLimitConfigError::ZeroHostLimit);
63        }
64        if self.window.is_zero() {
65            return Err(RateLimitConfigError::ZeroWindow);
66        }
67        if self.max_tracked_hosts == 0 {
68            return Err(RateLimitConfigError::ZeroTrackedHosts);
69        }
70        if self.per_host_requests > self.global_requests {
71            return Err(RateLimitConfigError::HostLimitExceedsGlobal);
72        }
73        if self.global_requests > DEFAULT_GLOBAL_REQUESTS_PER_WINDOW
74            || self.per_host_requests > DEFAULT_HOST_REQUESTS_PER_WINDOW
75            || self.window > DEFAULT_RATE_WINDOW
76            || self.max_tracked_hosts > DEFAULT_MAX_TRACKED_HOSTS
77        {
78            return Err(RateLimitConfigError::ExceedsMaximum);
79        }
80        Ok(self)
81    }
82}
83
84/// Shared active-client admission counter.
85#[derive(Clone, Debug)]
86pub struct ActiveClientLimiter {
87    inner: Arc<ActiveClientLimiterInner>,
88}
89
90#[derive(Debug)]
91struct ActiveClientLimiterInner {
92    max_active: usize,
93    active: AtomicUsize,
94}
95
96impl ActiveClientLimiter {
97    pub fn new(max_active: usize) -> Self {
98        Self {
99            inner: Arc::new(ActiveClientLimiterInner {
100                max_active,
101                active: AtomicUsize::new(0),
102            }),
103        }
104    }
105
106    /// Acquires one client slot. Dropping the returned permit releases it.
107    pub fn try_acquire(&self) -> Option<ActiveClientPermit> {
108        let mut active = self.inner.active.load(Ordering::Acquire);
109        loop {
110            if active >= self.inner.max_active {
111                return None;
112            }
113            match self.inner.active.compare_exchange_weak(
114                active,
115                active + 1,
116                Ordering::AcqRel,
117                Ordering::Acquire,
118            ) {
119                Ok(_) => {
120                    return Some(ActiveClientPermit {
121                        inner: Arc::clone(&self.inner),
122                    });
123                }
124                Err(observed) => active = observed,
125            }
126        }
127    }
128
129    pub fn active(&self) -> usize {
130        self.inner.active.load(Ordering::Acquire)
131    }
132
133    pub fn max_active(&self) -> usize {
134        self.inner.max_active
135    }
136}
137
138/// A single non-cloneable active-client slot.
139#[derive(Debug)]
140pub struct ActiveClientPermit {
141    inner: Arc<ActiveClientLimiterInner>,
142}
143
144impl Drop for ActiveClientPermit {
145    fn drop(&mut self) {
146        self.inner.active.fetch_sub(1, Ordering::AcqRel);
147    }
148}
149
150#[derive(Clone, Copy, Debug, Eq, PartialEq)]
151pub enum RateLimitScope {
152    Global,
153    Host,
154}
155
156#[derive(Clone, Copy, Debug, Eq, PartialEq)]
157pub enum RateLimitDecision {
158    Allowed,
159    Limited {
160        scope: RateLimitScope,
161        retry_after: Duration,
162    },
163}
164
165/// An exact sliding-window limiter with bounded per-host state.
166///
167/// Host names are hashed with a per-instance random key and never retained.
168/// Hash collisions can only make admission stricter. Callers should pass the
169/// canonical host so aliases do not receive separate windows.
170#[derive(Clone, Debug)]
171pub struct RequestRateLimiter {
172    inner: Arc<RequestRateLimiterInner>,
173}
174
175#[derive(Debug)]
176struct RequestRateLimiterInner {
177    config: RateLimitConfig,
178    host_hasher: RandomState,
179    state: Mutex<RateLimitState>,
180}
181
182#[derive(Debug, Default)]
183struct RateLimitState {
184    global: VecDeque<Instant>,
185    hosts: HashMap<u64, HostWindow>,
186}
187
188#[derive(Debug)]
189struct HostWindow {
190    requests: VecDeque<Instant>,
191}
192
193impl RequestRateLimiter {
194    pub fn new(config: RateLimitConfig) -> Result<Self, RateLimitConfigError> {
195        let config = config.validate()?;
196        Ok(Self {
197            inner: Arc::new(RequestRateLimiterInner {
198                config,
199                host_hasher: RandomState::new(),
200                state: Mutex::new(RateLimitState::default()),
201            }),
202        })
203    }
204
205    pub fn config(&self) -> RateLimitConfig {
206        self.inner.config
207    }
208
209    /// Atomically admits and records one request, or reports the limiting
210    /// window and the earliest useful retry time.
211    pub fn check(&self, canonical_host: &str, now: Instant) -> RateLimitDecision {
212        let config = self.inner.config;
213        let host_key = self.host_key(canonical_host);
214        let mut state = lock_recover(&self.inner.state);
215
216        expire(&mut state.global, now, config.window);
217        for window in state.hosts.values_mut() {
218            expire(&mut window.requests, now, config.window);
219        }
220        state.hosts.retain(|_, window| !window.requests.is_empty());
221
222        if state.global.len() >= config.global_requests {
223            return RateLimitDecision::Limited {
224                scope: RateLimitScope::Global,
225                retry_after: retry_after(&state.global, now, config.window),
226            };
227        }
228
229        if let Some(host) = state.hosts.get_mut(&host_key) {
230            if host.requests.len() >= config.per_host_requests {
231                return RateLimitDecision::Limited {
232                    scope: RateLimitScope::Host,
233                    retry_after: retry_after(&host.requests, now, config.window),
234                };
235            }
236        } else {
237            if state.hosts.len() >= config.max_tracked_hosts {
238                let retry_after = state
239                    .hosts
240                    .values()
241                    .filter_map(|window| window.requests.front())
242                    .filter_map(|request| request.checked_add(config.window))
243                    .map(|expires| expires.saturating_duration_since(now))
244                    .min()
245                    .unwrap_or(config.window);
246                return RateLimitDecision::Limited {
247                    scope: RateLimitScope::Host,
248                    retry_after,
249                };
250            }
251            state.hosts.insert(
252                host_key,
253                HostWindow {
254                    requests: VecDeque::new(),
255                },
256            );
257        }
258
259        state.global.push_back(now);
260        if let Some(host) = state.hosts.get_mut(&host_key) {
261            host.requests.push_back(now);
262        }
263        RateLimitDecision::Allowed
264    }
265
266    pub fn tracked_hosts(&self) -> usize {
267        lock_recover(&self.inner.state).hosts.len()
268    }
269
270    fn host_key(&self, canonical_host: &str) -> u64 {
271        let mut hasher = self.inner.host_hasher.build_hasher();
272        for byte in canonical_host.bytes() {
273            hasher.write_u8(byte.to_ascii_lowercase());
274        }
275        hasher.finish()
276    }
277}
278
279fn lock_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
280    mutex
281        .lock()
282        .unwrap_or_else(|poisoned| poisoned.into_inner())
283}
284
285fn expire(requests: &mut VecDeque<Instant>, now: Instant, window: Duration) {
286    while requests.front().is_some_and(|request| {
287        request
288            .checked_add(window)
289            .is_some_and(|expires| expires <= now)
290    }) {
291        requests.pop_front();
292    }
293}
294
295fn retry_after(requests: &VecDeque<Instant>, now: Instant, window: Duration) -> Duration {
296    requests
297        .front()
298        .and_then(|request| request.checked_add(window))
299        .map_or(window, |expires| expires.saturating_duration_since(now))
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use std::sync::Arc;
306    use std::sync::mpsc;
307    use std::thread;
308
309    fn test_config() -> RateLimitConfig {
310        RateLimitConfig {
311            global_requests: 4,
312            per_host_requests: 2,
313            window: Duration::from_secs(10),
314            max_tracked_hosts: 2,
315        }
316    }
317
318    #[test]
319    fn defaults_match_the_proxy_budget() {
320        assert_eq!(
321            RateLimitConfig::default(),
322            RateLimitConfig {
323                global_requests: 240,
324                per_host_requests: 240,
325                window: Duration::from_secs(10),
326                max_tracked_hosts: 256,
327            }
328        );
329        assert_eq!(DEFAULT_MAX_ACTIVE_CLIENTS, 64);
330    }
331
332    #[test]
333    fn permit_releases_exactly_one_slot_on_drop() {
334        let limiter = ActiveClientLimiter::new(1);
335        let permit = limiter.try_acquire().expect("first slot");
336        assert_eq!(limiter.active(), 1);
337        assert!(limiter.try_acquire().is_none());
338        drop(permit);
339        assert_eq!(limiter.active(), 0);
340        assert!(limiter.try_acquire().is_some());
341    }
342
343    #[test]
344    fn concurrent_permits_never_exceed_the_limit() {
345        let limiter = Arc::new(ActiveClientLimiter::new(8));
346        let (sender, receiver) = mpsc::channel();
347        let mut threads = Vec::new();
348        for _ in 0..32 {
349            let limiter = Arc::clone(&limiter);
350            let sender = sender.clone();
351            threads.push(thread::spawn(move || {
352                sender
353                    .send(limiter.try_acquire())
354                    .expect("receiver remains alive");
355            }));
356        }
357        drop(sender);
358        for handle in threads {
359            handle.join().expect("thread joins");
360        }
361        let permits = receiver.into_iter().flatten().collect::<Vec<_>>();
362        assert_eq!(permits.len(), 8);
363        assert_eq!(limiter.active(), 8);
364        drop(permits);
365        assert_eq!(limiter.active(), 0);
366    }
367
368    #[test]
369    fn host_window_is_exact_and_case_insensitive() {
370        let limiter = RequestRateLimiter::new(test_config()).expect("valid config");
371        let now = Instant::now();
372        assert_eq!(limiter.check("Example", now), RateLimitDecision::Allowed);
373        assert_eq!(
374            limiter.check("example", now + Duration::from_secs(1)),
375            RateLimitDecision::Allowed
376        );
377        assert_eq!(
378            limiter.check("EXAMPLE", now + Duration::from_secs(2)),
379            RateLimitDecision::Limited {
380                scope: RateLimitScope::Host,
381                retry_after: Duration::from_secs(8),
382            }
383        );
384        assert_eq!(
385            limiter.check("example", now + Duration::from_secs(10)),
386            RateLimitDecision::Allowed
387        );
388    }
389
390    #[test]
391    fn global_window_applies_across_hosts() {
392        let limiter = RequestRateLimiter::new(RateLimitConfig {
393            max_tracked_hosts: 8,
394            ..test_config()
395        })
396        .expect("valid config");
397        let now = Instant::now();
398        for host in ["one", "two", "three", "four"] {
399            assert_eq!(limiter.check(host, now), RateLimitDecision::Allowed);
400        }
401        assert_eq!(
402            limiter.check("five", now),
403            RateLimitDecision::Limited {
404                scope: RateLimitScope::Global,
405                retry_after: Duration::from_secs(10),
406            }
407        );
408    }
409
410    #[test]
411    fn default_window_allows_same_host_asset_bursts_up_to_global_limit() {
412        let limiter = RequestRateLimiter::new(RateLimitConfig::default()).expect("valid config");
413        let now = Instant::now();
414
415        for _ in 0..DEFAULT_GLOBAL_REQUESTS_PER_WINDOW {
416            assert_eq!(limiter.check("app.example", now), RateLimitDecision::Allowed);
417        }
418        assert_eq!(
419            limiter.check("app.example", now),
420            RateLimitDecision::Limited {
421                scope: RateLimitScope::Global,
422                retry_after: DEFAULT_RATE_WINDOW,
423            }
424        );
425    }
426
427    #[test]
428    fn host_state_stays_capped_under_rotation() {
429        let limiter = RequestRateLimiter::new(test_config()).expect("valid config");
430        let now = Instant::now();
431        assert_eq!(limiter.check("one", now), RateLimitDecision::Allowed);
432        assert_eq!(limiter.check("two", now), RateLimitDecision::Allowed);
433        assert_eq!(
434            limiter.check("three", now),
435            RateLimitDecision::Limited {
436                scope: RateLimitScope::Host,
437                retry_after: Duration::from_secs(10),
438            }
439        );
440        assert_eq!(limiter.tracked_hosts(), 2);
441    }
442
443    #[test]
444    fn invalid_config_is_rejected() {
445        let mut config = test_config();
446        config.max_tracked_hosts = 0;
447        assert!(matches!(
448            RequestRateLimiter::new(config),
449            Err(RateLimitConfigError::ZeroTrackedHosts)
450        ));
451    }
452}
hns_browser_loopback_proxy/chromium/rate_limit.rs

hns_browser_loopback_proxy/chromium/
rate_limit.rs
