leonwww/protocol/cli/src/security.rs

use crate::config::GurtConfig;
use gurtlib::{prelude::*, GurtMethod, GurtStatusCode};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tracing::{warn, debug};

#[derive(Debug)]
pub struct RateLimitData {
    requests: Vec<Instant>,
    connections: u32,
}

impl RateLimitData {
    fn new() -> Self {
        Self {
            requests: Vec::new(),
            connections: 0,
        }
    }

    fn cleanup_old_requests(&mut self, window: Duration) {
        let cutoff = Instant::now() - window;
        self.requests.retain(|&request_time| request_time > cutoff);
    }

    fn add_request(&mut self) {
        self.requests.push(Instant::now());
    }

    fn request_count(&self) -> usize {
        self.requests.len()
    }

    fn increment_connections(&mut self) {
        self.connections += 1;
    }

    fn decrement_connections(&mut self) {
        if self.connections > 0 {
            self.connections -= 1;
        }
    }

    fn connection_count(&self) -> u32 {
        self.connections
    }
}

pub struct SecurityMiddleware {
    config: Arc<GurtConfig>,
    rate_limit_data: Arc<Mutex<HashMap<IpAddr, RateLimitData>>>,
}

impl SecurityMiddleware {
    pub fn new(config: Arc<GurtConfig>) -> Self {
        Self {
            config,
            rate_limit_data: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn is_method_allowed(&self, method: &GurtMethod) -> bool {
        if let Some(security) = &self.config.security {
            let method_str = method.to_string();
            security.allowed_methods.contains(&method_str)
        } else {
            true
        }
    }

    pub fn check_rate_limit(&self, client_ip: IpAddr) -> bool {
        if let Some(security) = &self.config.security {
            let mut data = self.rate_limit_data.lock().unwrap();
            let rate_data = data.entry(client_ip).or_insert_with(RateLimitData::new);
            
            rate_data.cleanup_old_requests(Duration::from_secs(60));
            
            if rate_data.request_count() >= security.rate_limit_requests as usize {
                warn!("Rate limit exceeded for IP {}: {} requests in the last minute", 
                      client_ip, rate_data.request_count());
                return false;
            }
            
            rate_data.add_request();
            debug!("Request from {}: {}/{} requests in the last minute", 
                   client_ip, rate_data.request_count(), security.rate_limit_requests);
        }
        
        true
    }

    pub fn check_connection_limit(&self, client_ip: IpAddr) -> bool {
        if let Some(security) = &self.config.security {
            let mut data = self.rate_limit_data.lock().unwrap();
            let rate_data = data.entry(client_ip).or_insert_with(RateLimitData::new);
            
            if rate_data.connection_count() >= security.rate_limit_connections {
                warn!("Connection limit exceeded for IP {}: {} concurrent connections", 
                      client_ip, rate_data.connection_count());
                return false;
            }
        }
        
        true
    }

    pub fn register_connection(&self, client_ip: IpAddr) {
        if self.config.security.is_some() {
            let mut data = self.rate_limit_data.lock().unwrap();
            let rate_data = data.entry(client_ip).or_insert_with(RateLimitData::new);
            rate_data.increment_connections();
            debug!("Connection registered for {}: {} concurrent connections", 
                   client_ip, rate_data.connection_count());
        }
    }

    pub fn unregister_connection(&self, client_ip: IpAddr) {
        if self.config.security.is_some() {
            let mut data = self.rate_limit_data.lock().unwrap();
            if let Some(rate_data) = data.get_mut(&client_ip) {
                rate_data.decrement_connections();
                debug!("Connection unregistered for {}: {} concurrent connections remaining", 
                       client_ip, rate_data.connection_count());
            }
        }
    }

    pub fn create_method_not_allowed_response(&self) -> std::result::Result<GurtResponse, GurtError> {
        let response = GurtResponse::new(GurtStatusCode::MethodNotAllowed)
            .with_header("Content-Type", "text/html");
        Ok(response)
    }

    pub fn create_rate_limit_response(&self) -> std::result::Result<GurtResponse, GurtError> {
        let response = GurtResponse::new(GurtStatusCode::TooManyRequests)
            .with_header("Content-Type", "text/html")
            .with_header("Retry-After", "60");
        Ok(response)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::IpAddr;
    use std::sync::Arc;
    use std::time::Duration;

    fn create_test_config() -> Arc<GurtConfig> {
        let mut config = crate::config::GurtConfig::default();
        config.security = Some(crate::config::SecurityConfig {
            deny_files: vec!["*.secret".to_string(), "private/*".to_string()],
            allowed_methods: vec!["GET".to_string(), "POST".to_string()],
            rate_limit_requests: 5,
            rate_limit_connections: 2,
        });
        Arc::new(config)
    }

    #[test]
    fn test_rate_limit_data_initialization() {
        let data = RateLimitData::new();
        
        assert_eq!(data.request_count(), 0);
        assert_eq!(data.connection_count(), 0);
    }

    #[test]
    fn test_rate_limit_data_request_tracking() {
        let mut data = RateLimitData::new();
        
        data.add_request();
        data.add_request();
        assert_eq!(data.request_count(), 2);
        
        data.cleanup_old_requests(Duration::from_secs(0));
        assert_eq!(data.request_count(), 0);
    }

    #[test]
    fn test_rate_limit_data_connection_tracking() {
        let mut data = RateLimitData::new();
        
        data.increment_connections();
        data.increment_connections();
        assert_eq!(data.connection_count(), 2);
        
        data.decrement_connections();
        assert_eq!(data.connection_count(), 1);
        
        data.decrement_connections();
        data.decrement_connections();
        assert_eq!(data.connection_count(), 0);
    }

    #[test]
    fn test_security_middleware_initialization() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        
        assert!(middleware.rate_limit_data.lock().unwrap().is_empty());
    }

    #[test]
    fn test_connection_tracking() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        let ip: IpAddr = "127.0.0.1".parse().unwrap();
        
        middleware.register_connection(ip);
        {
            let data = middleware.rate_limit_data.lock().unwrap();
            assert_eq!(data.get(&ip).unwrap().connection_count(), 1);
        }
        
        middleware.unregister_connection(ip);
        {
            let data = middleware.rate_limit_data.lock().unwrap();
            assert_eq!(data.get(&ip).unwrap().connection_count(), 0);
        }
    }

    #[test]
    fn test_rate_limiting_requests() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        let ip: IpAddr = "127.0.0.1".parse().unwrap();
        
        for _ in 0..5 {
            assert!(middleware.check_rate_limit(ip));
        }
        
        assert!(!middleware.check_rate_limit(ip));
    }

    #[test]
    fn test_connection_limiting() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        let ip: IpAddr = "127.0.0.1".parse().unwrap();
        
        middleware.register_connection(ip);
        middleware.register_connection(ip);
        
        assert!(!middleware.check_connection_limit(ip));
    }

    #[test]
    fn test_method_validation() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        
        assert!(middleware.is_method_allowed(&GurtMethod::GET));
        assert!(middleware.is_method_allowed(&GurtMethod::POST));
        
        assert!(!middleware.is_method_allowed(&GurtMethod::PUT));
        assert!(!middleware.is_method_allowed(&GurtMethod::DELETE));
    }

    #[test]
    fn test_multiple_ips_isolation() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        let ip1: IpAddr = "127.0.0.1".parse().unwrap();
        let ip2: IpAddr = "127.0.0.2".parse().unwrap();
        
        for _ in 0..6 {
            middleware.check_rate_limit(ip1);
        }
        
        assert!(middleware.check_rate_limit(ip2));
        assert!(!middleware.check_rate_limit(ip1));
    }

    #[test]
    fn test_response_creation() {
        let config = create_test_config();
        let middleware = SecurityMiddleware::new(config.clone());
        
        let response = middleware.create_method_not_allowed_response().unwrap();
        assert_eq!(response.status_code, 405);
        
        let response = middleware.create_rate_limit_response().unwrap();
        assert_eq!(response.status_code, 429);
    }
}