nordisk/gunnar-ui

crates/gunnar-mcp/src/limit.rs

//! **Budgets**: a token bucket per connection, and a small one per client
//! address for the OAuth doors a stranger can reach. Beyond either, the answer
//! is 429 with `Retry-After`, the way the console answers a login flood.

use std::collections::HashMap;
use std::time::{Duration, Instant};

/// A token bucket: `burst` calls at once, refilled at `per_second`.
#[derive(Debug, Clone)]
pub struct Bucket {
    burst: f64,
    per_second: f64,
    tokens: f64,
    at: Instant,
}

impl Bucket {
    /// A full bucket.
    pub fn new(burst: u32, per_second: f64) -> Bucket {
        Bucket {
            burst: f64::from(burst),
            per_second,
            tokens: f64::from(burst),
            at: Instant::now(),
        }
    }

    /// Take one call at `now`, or answer how many whole seconds until one is
    /// there (at least 1).
    pub fn take_at(&mut self, now: Instant) -> Result<(), u64> {
        let elapsed = now.saturating_duration_since(self.at).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.per_second).min(self.burst);
        self.at = now;
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            Ok(())
        } else {
            let wait = (1.0 - self.tokens) / self.per_second;
            Err((wait.ceil() as u64).max(1))
        }
    }

    /// [`Bucket::take_at`] now.
    pub fn take(&mut self) -> Result<(), u64> {
        self.take_at(Instant::now())
    }
}

/// One fixed window of `max` requests per `window`, per key (a client address).
#[derive(Debug)]
pub struct Window {
    max: u32,
    window: Duration,
    seen: HashMap<String, (Instant, u32)>,
}

/// The most keys a window remembers; at the cap the oldest windows go.
const KEYS_MAX: usize = 4096;

impl Window {
    /// `max` per `window`.
    pub fn new(max: u32, window: Duration) -> Window {
        Window {
            max,
            window,
            seen: HashMap::new(),
        }
    }

    /// Count one request from `key` at `now`, or answer the seconds left in
    /// its window.
    pub fn admit_at(&mut self, key: &str, now: Instant) -> Result<(), u64> {
        if self.seen.len() >= KEYS_MAX {
            let window = self.window;
            self.seen
                .retain(|_, (start, _)| now.saturating_duration_since(*start) < window);
            if self.seen.len() >= KEYS_MAX {
                self.seen.clear();
            }
        }
        let entry = self.seen.entry(key.to_owned()).or_insert((now, 0));
        if now.saturating_duration_since(entry.0) >= self.window {
            *entry = (now, 0);
        }
        if entry.1 >= self.max {
            let left = self
                .window
                .saturating_sub(now.saturating_duration_since(entry.0));
            return Err(left.as_secs().max(1));
        }
        entry.1 += 1;
        Ok(())
    }

    /// [`Window::admit_at`] now.
    pub fn admit(&mut self, key: &str) -> Result<(), u64> {
        self.admit_at(key, Instant::now())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_bucket_gives_its_burst_then_says_when_to_come_back() {
        let t0 = Instant::now();
        let mut b = Bucket::new(3, 0.5);
        for _ in 0..3 {
            b.take_at(t0).expect("in the burst");
        }
        assert_eq!(
            b.take_at(t0),
            Err(2),
            "one token at 0.5/s is two seconds away"
        );
        b.take_at(t0 + Duration::from_secs(2)).expect("refilled");
        assert!(b.take_at(t0 + Duration::from_secs(2)).is_err());
    }

    #[test]
    fn a_window_counts_per_key_and_resets() {
        let t0 = Instant::now();
        let mut w = Window::new(2, Duration::from_secs(60));
        w.admit_at("a", t0).unwrap();
        w.admit_at("a", t0).unwrap();
        assert_eq!(w.admit_at("a", t0 + Duration::from_secs(10)), Err(50));
        w.admit_at("b", t0).expect("another key has its own window");
        w.admit_at("a", t0 + Duration::from_secs(60))
            .expect("a new window");
    }
}