lxc-streamutils/apps/server-rust/src/rate_limit.rs

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

use tokio::sync::Mutex;

#[derive(Clone)]
pub struct AuthRateLimiter {
    attempts: Arc<Mutex<HashMap<String, AttemptBucket>>>,
    window: Duration,
    block_for: Duration,
    max_failures: usize,
}

#[derive(Default)]
struct AttemptBucket {
    failures: VecDeque<Instant>,
    blocked_until: Option<Instant>,
}

#[derive(Clone, Copy, Debug)]
pub struct RateLimited {
    pub retry_after: Duration,
}

impl Default for AuthRateLimiter {
    fn default() -> Self {
        Self::new(5, Duration::from_secs(5 * 60), Duration::from_secs(10 * 60))
    }
}

impl AuthRateLimiter {
    pub fn new(max_failures: usize, window: Duration, block_for: Duration) -> Self {
        Self {
            attempts: Arc::new(Mutex::new(HashMap::new())),
            window,
            block_for,
            max_failures: max_failures.max(1),
        }
    }

    /// Check both account and network dimensions. Callers intentionally receive
    /// one generic result so this cannot be used to enumerate usernames.
    pub async fn check(&self, username: &str, ip: &str) -> Result<(), RateLimited> {
        let now = Instant::now();
        let mut attempts = self.attempts.lock().await;
        for key in keys(username, ip) {
            let bucket = attempts.entry(key).or_default();
            prune(bucket, now, self.window);
            if let Some(until) = bucket.blocked_until.filter(|until| *until > now) {
                return Err(RateLimited {
                    retry_after: until.duration_since(now),
                });
            }
        }
        Ok(())
    }

    pub async fn failure(&self, username: &str, ip: &str) {
        let now = Instant::now();
        let mut attempts = self.attempts.lock().await;
        for key in keys(username, ip) {
            let bucket = attempts.entry(key).or_default();
            prune(bucket, now, self.window);
            bucket.failures.push_back(now);
            if bucket.failures.len() >= self.max_failures {
                bucket.blocked_until = Some(now + self.block_for);
                bucket.failures.clear();
            }
        }
        // Opportunistic pruning bounds memory for a public login endpoint.
        if attempts.len() > 8_192 {
            attempts.retain(|_, bucket| {
                prune(bucket, now, self.window);
                !bucket.failures.is_empty() || bucket.blocked_until.is_some_and(|until| until > now)
            });
        }
    }

    pub async fn success(&self, username: &str, ip: &str) {
        let mut attempts = self.attempts.lock().await;
        // A successful account verification clears the account bucket. Keep
        // the IP bucket so one valid account cannot reset an attack on others.
        attempts.remove(&format!("account:{}", normalize_username(username)));
        let _ = ip;
    }

    /// Consume one request from an IP-scoped budget. This is used for costly
    /// anonymous enrollment work even when a request would otherwise succeed.
    pub async fn consume_ip(&self, namespace: &str, ip: &str) -> Result<(), RateLimited> {
        let identity = format!("{namespace}:{}", normalize_ip(ip));
        self.check(&identity, ip).await?;
        self.failure(&identity, ip).await;
        Ok(())
    }
}

fn keys(username: &str, ip: &str) -> [String; 2] {
    [
        format!("account:{}", normalize_username(username)),
        format!("network:{}", normalize_ip(ip)),
    ]
}

fn normalize_username(value: &str) -> String {
    value.trim().to_lowercase()
}

fn normalize_ip(value: &str) -> String {
    let value = value.trim();
    if value.is_empty() {
        "unknown".into()
    } else {
        value.chars().take(96).collect()
    }
}

fn prune(bucket: &mut AttemptBucket, now: Instant, window: Duration) {
    while bucket
        .failures
        .front()
        .is_some_and(|timestamp| now.duration_since(*timestamp) >= window)
    {
        bucket.failures.pop_front();
    }
    if bucket.blocked_until.is_some_and(|until| until <= now) {
        bucket.blocked_until = None;
    }
}

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

    #[tokio::test]
    async fn blocks_account_and_network_after_threshold() {
        let limiter = AuthRateLimiter::new(2, Duration::from_secs(60), Duration::from_secs(60));
        assert!(limiter.check("Streamer", "127.0.0.1").await.is_ok());
        limiter.failure("Streamer", "127.0.0.1").await;
        limiter.failure("Streamer", "127.0.0.1").await;
        assert!(limiter.check("streamer", "127.0.0.1").await.is_err());
        assert!(limiter.check("another", "127.0.0.1").await.is_err());
    }

    #[tokio::test]
    async fn request_budget_counts_successful_anonymous_work() {
        let limiter = AuthRateLimiter::new(2, Duration::from_secs(60), Duration::from_secs(60));
        assert!(limiter.consume_ip("enroll", "127.0.0.1").await.is_ok());
        assert!(limiter.consume_ip("enroll", "127.0.0.1").await.is_ok());
        assert!(limiter.consume_ip("enroll", "127.0.0.1").await.is_err());
    }
}