Serve the web client from yaak-web (#582)

crates-server/yaak-web/src/limits.rs

@@ -0,0 +1,90 @@
+//! 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());
+        }
+    }
+}