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

hns_browser_loopback_proxy/chromium/
rate_limit.rs
