gitdataai/libs/git/http/rate_limit.rs

//! 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,
}

#[derive(Clone, Copy)]
enum BucketOp {
    Read,
    Write,
}

pub struct RateLimiter {
    buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
    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, BucketOp::Read, 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, BucketOp::Write, 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, BucketOp::Write, self.config.write_requests_per_window)
            .await
    }

    async fn is_allowed(&self, key: &str, op: BucketOp, 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);
        }

        let over_limit = match op {
            BucketOp::Read => bucket.read_count >= limit,
            BucketOp::Write => bucket.write_count >= limit,
        };

        if over_limit {
            return false;
        }

        match op {
            BucketOp::Read => bucket.read_count += 1,
            BucketOp::Write => bucket.write_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<Self>) -> 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);
    }
}