mirror of
https://github.com/mountain-loop/yaak.git
synced 2026-08-20 18:34:04 +02:00
91 lines
2.9 KiB
Rust
91 lines
2.9 KiB
Rust
//! Per-client rate limiting, kept deliberately small.
|
|
//!
|
|
//! One token bucket per client IP, refilled continuously, in a mutex-guarded
|
|
//! map that is swept of idle entries as it goes. Good enough to keep one
|
|
//! caller from monopolising a hosted instance; not a substitute for whatever
|
|
//! sits in front of it in production.
|
|
|
|
use std::collections::HashMap;
|
|
use std::net::IpAddr;
|
|
use std::sync::Mutex;
|
|
use std::time::{Duration, Instant};
|
|
|
|
pub struct RateLimiter {
|
|
per_minute: u32,
|
|
buckets: Mutex<HashMap<IpAddr, Bucket>>,
|
|
}
|
|
|
|
struct Bucket {
|
|
tokens: f64,
|
|
last: Instant,
|
|
}
|
|
|
|
impl RateLimiter {
|
|
/// `per_minute == 0` disables limiting.
|
|
pub fn new(per_minute: u32) -> Self {
|
|
Self { per_minute, buckets: Mutex::new(HashMap::new()) }
|
|
}
|
|
|
|
/// Take one token for `client`, or say how long until one is available.
|
|
pub fn check(&self, client: IpAddr) -> Result<(), Duration> {
|
|
if self.per_minute == 0 {
|
|
return Ok(());
|
|
}
|
|
let capacity = self.per_minute as f64;
|
|
let per_second = capacity / 60.0;
|
|
let now = Instant::now();
|
|
|
|
let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
|
|
|
|
// Sweep buckets that have been idle long enough to be full again; there is nothing
|
|
// to remember about them.
|
|
if buckets.len() > 1024 {
|
|
buckets.retain(|_, b| now.duration_since(b.last).as_secs_f64() * per_second < capacity);
|
|
}
|
|
|
|
let bucket = buckets.entry(client).or_insert(Bucket { tokens: capacity, last: now });
|
|
let elapsed = now.duration_since(bucket.last).as_secs_f64();
|
|
bucket.tokens = (bucket.tokens + elapsed * per_second).min(capacity);
|
|
bucket.last = now;
|
|
|
|
if bucket.tokens >= 1.0 {
|
|
bucket.tokens -= 1.0;
|
|
Ok(())
|
|
} else {
|
|
let wait = (1.0 - bucket.tokens) / per_second;
|
|
Err(Duration::from_secs_f64(wait.max(0.001)))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn a_full_bucket_then_a_wait() {
|
|
let limiter = RateLimiter::new(3);
|
|
let ip: IpAddr = "203.0.113.5".parse().unwrap();
|
|
assert!(limiter.check(ip).is_ok());
|
|
assert!(limiter.check(ip).is_ok());
|
|
assert!(limiter.check(ip).is_ok());
|
|
let wait = limiter.check(ip).expect_err("fourth call in a burst should wait");
|
|
assert!(wait > Duration::ZERO && wait <= Duration::from_secs(20));
|
|
}
|
|
|
|
#[test]
|
|
fn clients_are_independent_and_zero_disables() {
|
|
let limiter = RateLimiter::new(1);
|
|
let a: IpAddr = "203.0.113.5".parse().unwrap();
|
|
let b: IpAddr = "203.0.113.6".parse().unwrap();
|
|
assert!(limiter.check(a).is_ok());
|
|
assert!(limiter.check(a).is_err());
|
|
assert!(limiter.check(b).is_ok());
|
|
|
|
let unlimited = RateLimiter::new(0);
|
|
for _ in 0..1000 {
|
|
assert!(unlimited.check(a).is_ok());
|
|
}
|
|
}
|
|
}
|