tryaudex_core/

ratelimit.rs

use std::collections::HashMap;
use std::path::PathBuf;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::error::{AvError, Result};

/// Rate limiting configuration in `[ratelimit]` config section.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct RateLimitConfig {
    /// Max sessions per hour per identity (default: unlimited)
    pub max_per_hour: Option<u32>,
    /// Max sessions per day per identity (default: unlimited)
    pub max_per_day: Option<u32>,
    /// Max concurrent active sessions per identity (default: unlimited)
    pub max_concurrent: Option<u32>,
}

/// Rate limit state stored on disk.
#[derive(Debug, Serialize, Deserialize, Default)]
struct RateLimitState {
    /// Map of identity -> list of session timestamps
    sessions: HashMap<String, Vec<SessionRecord>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct SessionRecord {
    timestamp: DateTime<Utc>,
    session_id: String,
    active: bool,
}

fn state_path() -> PathBuf {
    dirs::data_local_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("audex")
        .join("ratelimit.json")
}

fn load_state() -> RateLimitState {
    let path = state_path();
    std::fs::read_to_string(&path)
        .ok()
        .and_then(|s| serde_json::from_str(&s).ok())
        .unwrap_or_default()
}

fn save_state(state: &RateLimitState) {
    let path = state_path();
    if let Some(parent) = path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }
    if let Ok(json) = serde_json::to_string(state) {
        let _ = std::fs::write(path, json);
    }
}

/// Check if a new session is allowed under rate limits.
/// Returns Ok(()) if allowed, Err with details if rate limited.
pub fn check(config: &RateLimitConfig, identity: &str) -> Result<()> {
    let state = load_state();
    let now = Utc::now();

    let records = match state.sessions.get(identity) {
        Some(r) => r,
        None => return Ok(()), // No history = no limits hit
    };

    // Check hourly limit
    if let Some(max) = config.max_per_hour {
        let one_hour_ago = now - chrono::Duration::hours(1);
        let count = records
            .iter()
            .filter(|r| r.timestamp >= one_hour_ago)
            .count() as u32;
        if count >= max {
            return Err(AvError::InvalidPolicy(format!(
                "Rate limit exceeded: {} sessions in the last hour (max: {}). Try again later.",
                count, max
            )));
        }
    }

    // Check daily limit
    if let Some(max) = config.max_per_day {
        let one_day_ago = now - chrono::Duration::days(1);
        let count = records
            .iter()
            .filter(|r| r.timestamp >= one_day_ago)
            .count() as u32;
        if count >= max {
            return Err(AvError::InvalidPolicy(format!(
                "Rate limit exceeded: {} sessions in the last 24 hours (max: {}). Try again tomorrow.",
                count, max
            )));
        }
    }

    // Check concurrent limit
    if let Some(max) = config.max_concurrent {
        let active_count = records.iter().filter(|r| r.active).count() as u32;
        if active_count >= max {
            return Err(AvError::InvalidPolicy(format!(
                "Rate limit exceeded: {} active sessions (max: {}). End an existing session first.",
                active_count, max
            )));
        }
    }

    Ok(())
}

/// Record a new session for rate limiting.
pub fn record_session(identity: &str, session_id: &str) {
    let mut state = load_state();
    let records = state.sessions.entry(identity.to_string()).or_default();
    records.push(SessionRecord {
        timestamp: Utc::now(),
        session_id: session_id.to_string(),
        active: true,
    });
    save_state(&state);
}

/// Mark a session as ended (no longer active).
pub fn end_session(identity: &str, session_id: &str) {
    let mut state = load_state();
    if let Some(records) = state.sessions.get_mut(identity) {
        for record in records.iter_mut() {
            if record.session_id == session_id {
                record.active = false;
            }
        }
        // Prune old records (older than 48 hours) to prevent unbounded growth
        let cutoff = Utc::now() - chrono::Duration::hours(48);
        records.retain(|r| r.timestamp >= cutoff || r.active);
    }
    save_state(&state);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_default() {
        let config = RateLimitConfig::default();
        assert!(config.max_per_hour.is_none());
        assert!(config.max_per_day.is_none());
        assert!(config.max_concurrent.is_none());
    }

    #[test]
    fn test_config_deserialize() {
        let toml_str = r#"
max_per_hour = 10
max_per_day = 50
max_concurrent = 3
"#;
        let config: RateLimitConfig = toml::from_str(toml_str).unwrap();
        assert_eq!(config.max_per_hour, Some(10));
        assert_eq!(config.max_per_day, Some(50));
        assert_eq!(config.max_concurrent, Some(3));
    }

    #[test]
    fn test_check_no_limits() {
        let config = RateLimitConfig::default();
        assert!(check(&config, "anyone").is_ok());
    }

    #[test]
    fn test_check_no_history() {
        let config = RateLimitConfig {
            max_per_hour: Some(5),
            max_per_day: Some(20),
            max_concurrent: Some(2),
        };
        // No history for this identity = allowed
        assert!(check(&config, "new-user-test-ratelimit").is_ok());
    }

    #[test]
    fn test_state_serialization() {
        let mut state = RateLimitState::default();
        state.sessions.insert(
            "test@example.com".to_string(),
            vec![SessionRecord {
                timestamp: Utc::now(),
                session_id: "sess-001".to_string(),
                active: true,
            }],
        );
        let json = serde_json::to_string(&state).unwrap();
        let parsed: RateLimitState = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.sessions.len(), 1);
        assert!(parsed.sessions.contains_key("test@example.com"));
    }
}