claude-utils 0.2.0

use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct RateLimiter {
    requests: Arc<RwLock<HashMap<IpAddr, Vec<Instant>>>>,
    max_requests: usize,
    window: Duration,
}

impl RateLimiter {
    pub fn new(max_requests: usize, window_seconds: u64) -> Self {
        Self {
            requests: Arc::new(RwLock::new(HashMap::new())),
            max_requests,
            window: Duration::from_secs(window_seconds),
        }
    }

    pub async fn check_rate_limit(&self, ip: IpAddr) -> bool {
        let now = Instant::now();
        let mut requests = self.requests.write().await;
        
        let timestamps = requests.entry(ip).or_insert_with(Vec::new);
        
        // Remove old requests outside the window
        timestamps.retain(|&timestamp| now.duration_since(timestamp) < self.window);
        
        // Check if under limit
        if timestamps.len() < self.max_requests {
            timestamps.push(now);
            true
        } else {
            false
        }
    }

    pub async fn cleanup_old_entries(&self) {
        let now = Instant::now();
        let mut requests = self.requests.write().await;
        
        requests.retain(|_, timestamps| {
            timestamps.retain(|&timestamp| now.duration_since(timestamp) < self.window);
            !timestamps.is_empty()
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[tokio::test]
    async fn test_rate_limiter() {
        let limiter = RateLimiter::new(3, 1);
        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));

        // First 3 requests should pass
        assert!(limiter.check_rate_limit(ip).await);
        assert!(limiter.check_rate_limit(ip).await);
        assert!(limiter.check_rate_limit(ip).await);

        // 4th request should fail
        assert!(!limiter.check_rate_limit(ip).await);

        // Wait for window to expire
        tokio::time::sleep(Duration::from_secs(1)).await;

        // Should pass again
        assert!(limiter.check_rate_limit(ip).await);
    }
}