//! HTTP rate limiting for git operations. //! //! Uses a token-bucket approach with per-IP and per-repo-write limits. //! Cleanup runs every 5 minutes to prevent unbounded memory growth. use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::RwLock; use tokio::time::interval; #[derive(Debug, Clone)] pub struct RateLimitConfig { /// Requests allowed per window for read operations. pub read_requests_per_window: u32, /// Requests allowed per window for write operations. pub write_requests_per_window: u32, /// Window duration in seconds. pub window_secs: u64, } impl Default for RateLimitConfig { fn default() -> Self { Self { read_requests_per_window: 120, write_requests_per_window: 30, window_secs: 60, } } } #[derive(Debug)] struct RateLimitBucket { read_count: u32, write_count: u32, reset_time: Instant, } pub struct RateLimiter { buckets: Arc>>, config: RateLimitConfig, } impl RateLimiter { pub fn new(config: RateLimitConfig) -> Self { Self { buckets: Arc::new(RwLock::new(HashMap::new())), config, } } pub async fn is_ip_read_allowed(&self, ip: &str) -> bool { let key = format!("ip:read:{}", ip); self.is_allowed(&key, self.config.read_requests_per_window) .await } pub async fn is_ip_write_allowed(&self, ip: &str) -> bool { let key = format!("ip:write:{}", ip); self.is_allowed(&key, self.config.write_requests_per_window) .await } pub async fn is_repo_write_allowed(&self, ip: &str, repo_path: &str) -> bool { let key = format!("repo:write:{}:{}", ip, repo_path); self.is_allowed(&key, self.config.write_requests_per_window) .await } async fn is_allowed(&self, key: &str, limit: u32) -> bool { let now = Instant::now(); let mut buckets = self.buckets.write().await; let bucket = buckets .entry(key.to_string()) .or_insert_with(|| RateLimitBucket { read_count: 0, write_count: 0, reset_time: now + Duration::from_secs(self.config.window_secs), }); if now >= bucket.reset_time { bucket.read_count = 0; bucket.write_count = 0; bucket.reset_time = now + Duration::from_secs(self.config.window_secs); } // Use read_count for both read/write since we don't distinguish in bucket if bucket.read_count >= limit { return false; } bucket.read_count += 1; true } pub async fn retry_after(&self, ip: &str) -> u64 { let key_read = format!("ip:read:{}", ip); let now = Instant::now(); let buckets = self.buckets.read().await; if let Some(bucket) = buckets.get(&key_read) { if now < bucket.reset_time { return bucket.reset_time.saturating_duration_since(now).as_secs() as u64; } } 0 } /// Start a background cleanup task that removes expired entries. /// Should be spawned once at startup. pub fn start_cleanup(self: Arc) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let mut ticker = interval(Duration::from_secs(300)); // every 5 minutes loop { ticker.tick().await; let now = Instant::now(); let mut buckets = self.buckets.write().await; buckets.retain(|_, bucket| now < bucket.reset_time); } }) } } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_rate_limit_allows_requests_up_to_limit() { let limiter = Arc::new(RateLimiter::new(RateLimitConfig { read_requests_per_window: 3, write_requests_per_window: 1, window_secs: 60, })); for _ in 0..3 { assert!(limiter.is_ip_read_allowed("1.2.3.4").await); } assert!(!limiter.is_ip_read_allowed("1.2.3.4").await); } }