proper productionize project

apps/server-rust/src/rate_limit.rs

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