rustclamp 1.0.0

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

use super::{Next, Request, Response, error};

/// Most clients tracked at once by default; see [`Throttle::capacity`].
const CAPACITY: usize = 10_000;

/// A fixed-window rate limiter: at most `max` hits per key in each window.
///
/// Usable on its own (for example from a console command) or as middleware
/// through [`Throttle::middleware`] and [`throttle`].
#[derive(Debug)]
pub struct Throttle {
    max: u32,
    window: Duration,
    trust_forwarded: bool,
    capacity: usize,
    hits: Mutex<HashMap<String, (Instant, u32)>>,
    #[cfg(feature = "redis")]
    shared: Option<(crate::redis::Redis, String)>,
}

impl Throttle {
    /// At most `max` hits per key in every `window`.
    pub fn new(max: u32, window: Duration) -> Self {
        Self {
            max,
            window,
            trust_forwarded: false,
            capacity: CAPACITY,
            hits: Mutex::default(),
            #[cfg(feature = "redis")]
            shared: None,
        }
    }

    /// At most `max` hits per key per minute.
    pub fn per_minute(max: u32) -> Self {
        Self::new(max, Duration::from_secs(60))
    }

    /// `max` requests per client per minute, read from config `key` (default
    /// `max`), trusting `X-Forwarded-For` when `TRUST_PROXY=true`, as
    /// [`trust_forwarded`](Self::trust_forwarded) describes.
    ///
    /// ```
    /// use rustclamp::config::Config;
    /// use rustclamp::web::Throttle;
    ///
    /// let contact = Throttle::from_config(&Config::parse("CONTACT_PER_MINUTE=2"), "CONTACT_PER_MINUTE", 5);
    /// ```
    pub fn from_config(config: &crate::config::Config, key: &str, max: u32) -> Self {
        Self::per_minute(config.get_or(key, max))
            .trust_forwarded(config.get_or("TRUST_PROXY", false))
    }

    /// When `trust` is set, identifies clients by the last `X-Forwarded-For`
    /// address, the one the proxy appended, instead of the connecting peer.
    /// Enable this only behind exactly one proxy that sets the header, such as
    /// nginx; otherwise clients can pick their own key.
    #[must_use]
    pub fn trust_forwarded(mut self, trust: bool) -> Self {
        self.trust_forwarded = trust;
        self
    }

    /// Tracks at most `clients` keys at once (10,000 by default), so memory
    /// stays bounded when addresses rotate. When a new key arrives at the cap,
    /// expired windows are dropped first; only when every window is still live
    /// is the oldest one dropped, and that client starts a fresh window early.
    /// Set it above your peak number of clients per window to rule that out.
    #[must_use]
    pub fn capacity(mut self, clients: usize) -> Self {
        self.capacity = clients.max(1);
        self
    }

    /// Counts hits in `redis` under keys starting with `prefix` (such as
    /// `throttle:login:`), so every app instance sharing that Redis shares
    /// the limit. While Redis cannot be reached, each instance counts on its
    /// own, as without this.
    #[cfg(feature = "redis")]
    #[must_use]
    pub fn shared(mut self, redis: crate::redis::Redis, prefix: &str) -> Self {
        self.shared = Some((redis, prefix.to_owned()));
        self
    }

    /// Counts one hit for `key`. Over the limit, returns how long until the
    /// window resets.
    pub fn hit(&self, key: &str) -> Result<(), Duration> {
        #[cfg(feature = "redis")]
        if let Some((redis, prefix)) = &self.shared
            && let Ok((count, left)) = redis.hit(&format!("{prefix}{key}"), self.window)
        {
            return if count > i64::from(self.max) {
                Err(left)
            } else {
                Ok(())
            };
        }
        let now = Instant::now();
        let mut hits = self
            .hits
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        // ponytail: O(n) scans, only at capacity; a timing wheel if that ever shows in profiles
        if hits.len() >= self.capacity && !hits.contains_key(key) {
            hits.retain(|_, (start, _)| now.duration_since(*start) < self.window);
            if hits.len() >= self.capacity
                && let Some(oldest) = hits
                    .iter()
                    .min_by_key(|(_, (start, _))| *start)
                    .map(|(k, _)| k.clone())
            {
                hits.remove(&oldest);
            }
        }
        let entry = hits.entry(key.to_owned()).or_insert((now, 0));
        if now.duration_since(entry.0) >= self.window {
            *entry = (now, 0);
        }
        entry.1 += 1;
        if entry.1 > self.max {
            Err(self.window - now.duration_since(entry.0))
        } else {
            Ok(())
        }
    }

    /// The client key for `request`: its forwarded or peer address.
    pub(super) fn client(&self, request: &Request) -> String {
        let forwarded = self
            .trust_forwarded
            .then(|| request.header("x-forwarded-for"))
            .flatten()
            .and_then(|list| list.rsplit(',').next())
            .map(str::trim);
        match (forwarded, request.peer) {
            (Some(address), _) => address.to_owned(),
            (None, Some(peer)) => peer.to_string(),
            (None, None) => "unknown".to_owned(),
        }
    }

    /// Middleware answering `429` with `Retry-After` once a client is over the limit.
    pub fn middleware(self) -> impl Fn(&Request, Next) -> Response + Send + Sync + 'static {
        move |request, next| match self.hit(&self.client(request)) {
            Ok(()) => next(request),
            Err(wait) => {
                let seconds = wait.as_secs() + u64::from(wait.subsec_nanos() > 0);
                error(429).with_header("Retry-After", &seconds.max(1).to_string())
            }
        }
    }
}

/// Middleware allowing each client `max` requests per `window`, keyed by peer address.
pub fn throttle(
    max: u32,
    window: Duration,
) -> impl Fn(&Request, Next) -> Response + Send + Sync + 'static {
    Throttle::new(max, window).middleware()
}

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

    #[cfg(feature = "redis")]
    #[test]
    fn shared_counts_in_redis_and_falls_back_to_its_own_count() {
        use std::io::{BufRead, BufReader, Write};
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        // Answers each EVAL with the next count and 30 s left in the window.
        let server = std::thread::spawn(move || {
            let (stream, _) = listener.accept().unwrap();
            let mut reader = BufReader::new(stream.try_clone().unwrap());
            let mut writer = stream;
            for count in 1..=3 {
                // *5 header, then five $len + value line pairs.
                for _ in 0..11 {
                    reader.read_line(&mut String::new()).unwrap();
                }
                write!(writer, "*2\r\n:{count}\r\n:30000\r\n").unwrap();
            }
        });
        let redis = crate::redis::Redis::connect(&format!("redis://{address}")).unwrap();
        let limiter = Throttle::new(2, Duration::from_secs(60)).shared(redis, "throttle:");
        assert!(limiter.hit("a").is_ok());
        assert!(limiter.hit("a").is_ok());
        assert_eq!(limiter.hit("a"), Err(Duration::from_secs(30)));
        server.join().unwrap();
        // Redis gone: this instance counts on its own, from zero.
        assert!(limiter.hit("a").is_ok());
        assert!(limiter.hit("a").is_ok());
        assert!(limiter.hit("a").is_err());
    }

    #[test]
    fn limits_each_client_per_window() {
        let limiter = Throttle::new(2, Duration::from_secs(60));
        assert!(limiter.hit("a").is_ok());
        assert!(limiter.hit("a").is_ok());
        let wait = limiter.hit("a").unwrap_err();
        assert!(wait > Duration::from_secs(58));
        assert!(limiter.hit("b").is_ok());

        let short = Throttle::new(1, Duration::from_millis(20));
        assert!(short.hit("a").is_ok());
        assert!(short.hit("a").is_err());
        std::thread::sleep(Duration::from_millis(25));
        assert!(short.hit("a").is_ok());
    }

    #[test]
    fn memory_stays_bounded() {
        let limiter = Throttle::new(1, Duration::from_millis(50)).capacity(3);
        for key in ["a", "b", "c"] {
            limiter.hit(key).unwrap();
        }
        limiter.hit("d").unwrap();
        let keys = |limiter: &Throttle| limiter.hits.lock().unwrap().len();
        assert_eq!(keys(&limiter), 3, "the oldest window makes room");
        std::thread::sleep(Duration::from_millis(60));
        limiter.hit("e").unwrap();
        assert_eq!(keys(&limiter), 1, "expired windows are pruned first");
    }

    #[test]
    fn middleware_answers_429_with_retry_after() {
        let routes = Router::new()
            .middleware(throttle(1, Duration::from_secs(30)))
            .get("/", |_| Response::text(200, "ok"));
        let peer = "10.0.0.1".parse().unwrap();
        assert_eq!(
            routes.handle(&Request::get("/").with_peer(peer)).status,
            200
        );
        let limited = routes.handle(&Request::get("/").with_peer(peer));
        assert_eq!(limited.status, 429);
        assert_eq!(limited.header("retry-after"), Some("30"));
        let other = "10.0.0.2".parse().unwrap();
        assert_eq!(
            routes.handle(&Request::get("/").with_peer(other)).status,
            200
        );
    }

    #[test]
    fn forwarded_address_is_opt_in() {
        let proxy = "127.0.0.1".parse().unwrap();
        let request = |client: &str| {
            Request::get("/")
                .with_peer(proxy)
                .with_header("X-Forwarded-For", client)
        };
        let direct = Throttle::new(1, Duration::from_secs(60));
        assert_eq!(direct.client(&request("1.1.1.1")), "127.0.0.1");
        let proxied = Throttle::new(1, Duration::from_secs(60)).trust_forwarded(true);
        // nginx appends the real address; anything before it came from the client.
        assert_eq!(proxied.client(&request("6.6.6.6, 1.1.1.1")), "1.1.1.1");
        assert_eq!(proxied.client(&request("1.1.1.1")), "1.1.1.1");
        let routes = crate::web::Router::new()
            .middleware(proxied.middleware())
            .get("/", |_| Response::text(200, "ok"));
        assert_eq!(routes.handle(&request("fake-1, 1.1.1.1")).status, 200);
        assert_eq!(
            routes.handle(&request("fake-2, 1.1.1.1")).status,
            429,
            "spoofed prefixes share one bucket"
        );
    }
}