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) {
        #[cfg(unix)]
        {
            use std::io::Write;
            use std::os::unix::fs::OpenOptionsExt;
            if let Ok(mut file) = std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .mode(0o600)
                .open(&path)
            {
                let _ = file.write_all(json.as_bytes());
            }
        }
        #[cfg(not(unix))]
        {
            let _ = std::fs::write(path, json);
        }
    }
}

/// Run `f` while holding an exclusive file lock on the rate-limit state
/// lock file. All state-mutating paths go through this helper so that
/// every reader/writer observes the same serialization guarantee — and
/// so lock-acquisition failures cannot silently degrade to an unlocked
/// read-modify-write (R6-H26).
///
/// Previously `check_and_record_atomic` propagated lock errors while
/// `end_session`/`record_session` used a helper that returned
/// `Option<File>` and dropped failures on the floor. Under disk
/// pressure or EMFILE conditions, `end_session` could proceed without a
/// lock and race a concurrent `check_and_record_atomic` holding a
/// legitimate exclusive lock, corrupting the rate-limit count.
fn with_lock<F, R>(f: F) -> Result<R>
where
    F: FnOnce() -> Result<R>,
{
    use fs4::fs_std::FileExt;

    let path = state_path();
    if let Some(parent) = path.parent() {
        let _ = std::fs::create_dir_all(parent);
    }

    let lock_path = path.with_extension("lock");
    let lock_file = std::fs::OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(false)
        .open(&lock_path)
        .map_err(|e| AvError::InvalidPolicy(format!("Rate limit lock open failed: {}", e)))?;

    lock_file
        .lock_exclusive()
        .map_err(|e| AvError::InvalidPolicy(format!("Rate limit lock failed: {}", e)))?;

    let result = f();
    let _ = FileExt::unlock(&lock_file);
    result
}

/// Perform a rate limit check-and-record atomically with file locking.
/// Prevents concurrent CLI invocations from racing past the limit.
///
/// The state file is loaded, checked, and (on success) updated within a
/// single exclusive lock acquisition so no concurrent process can slip
/// through between the check and the record (TOCTOU-safe).
pub fn check_and_record_atomic(
    config: &RateLimitConfig,
    identity: &str,
    session_id: &str,
) -> Result<()> {
    with_lock(|| {
        // Under a single lock: load once, check, then record — no second load.
        let mut state = load_state();
        let now = Utc::now();

        // --- check phase (mirrors check()) ---
        if let Some(records) = state.sessions.get(identity) {
            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
                    )));
                }
            }
            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
                    )));
                }
            }
            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
                    )));
                }
            }
        }

        // --- record phase (mirrors record_session()) ---
        let records = state.sessions.entry(identity.to_string()).or_default();
        records.push(SessionRecord {
            timestamp: now,
            session_id: session_id.to_string(),
            active: true,
        });
        save_state(&state);
        Ok(())
    })
}

/// 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) -> Result<()> {
    with_lock(|| {
        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);
        Ok(())
    })
}

/// Mark a session as ended (no longer active).
pub fn end_session(identity: &str, session_id: &str) -> Result<()> {
    with_lock(|| {
        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);
        Ok(())
    })
}

#[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"));
    }
}