fraiseql-auth 2.0.0-rc.13

// Rate limiting for brute-force protection
// Uses an in-memory approach with Arc and Mutex for simplicity
//
// # Threading Model
//
// All rate limiting operations are **atomic** with respect to concurrent access:
// - Each call to check() holds a lock for its entire duration
// - Check-and-update operations cannot be interleaved with other threads
// - This prevents race conditions where multiple threads simultaneously exceed limits
// - The lock is held while reading current time, reading record, and updating counter
// - This ensures that the decision to allow/deny a request is consistent

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
    time::{SystemTime, UNIX_EPOCH},
};

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

/// Rate limit configuration for an endpoint
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Whether rate limiting is enabled for this endpoint
    pub enabled:      bool,
    /// Maximum number of requests allowed in the window
    pub max_requests: u32,
    /// Window duration in seconds
    pub window_secs:  u64,
}

impl RateLimitConfig {
    /// IP-based rate limiting for public endpoints
    /// 100 requests per 60 seconds (typical for auth/start, auth/callback)
    pub fn per_ip_standard() -> Self {
        Self {
            enabled:      true,
            max_requests: 100,
            window_secs:  60,
        }
    }

    /// Stricter IP-based rate limiting for sensitive endpoints
    /// 50 requests per 60 seconds
    pub fn per_ip_strict() -> Self {
        Self {
            enabled:      true,
            max_requests: 50,
            window_secs:  60,
        }
    }

    /// User-based rate limiting for authenticated endpoints
    /// 10 requests per 60 seconds
    pub fn per_user_standard() -> Self {
        Self {
            enabled:      true,
            max_requests: 10,
            window_secs:  60,
        }
    }

    /// Failed login attempt limiting
    /// 5 failed attempts per 3600 seconds (1 hour)
    pub fn failed_login_attempts() -> Self {
        Self {
            enabled:      true,
            max_requests: 5,
            window_secs:  3600,
        }
    }
}

/// Request record for tracking
#[derive(Debug, Clone)]
struct RequestRecord {
    /// Number of requests in current window
    count:        u32,
    /// Unix timestamp of window start
    window_start: u64,
}

/// Per-key rate limiter using in-memory tracking
/// Maintains separate rate limits for each key (IP, user ID, etc.)
pub struct KeyedRateLimiter {
    records: Arc<Mutex<HashMap<String, RequestRecord>>>,
    config:  RateLimitConfig,
}

impl KeyedRateLimiter {
    /// Create a new keyed rate limiter
    pub fn new(config: RateLimitConfig) -> Self {
        Self {
            records: Arc::new(Mutex::new(HashMap::new())),
            config,
        }
    }

    /// Get current Unix timestamp in seconds
    ///
    /// # Error Handling (Fail-Safe)
    ///
    /// If system time cannot be determined (e.g., clock moved backward), returns u64::MAX.
    ///
    /// This value guarantees that:
    /// - The check `now >= record.window_start + window_secs` will always be true
    /// - This resets the rate limit window for this request (allowing it through)
    /// - Subsequent requests will establish a new window starting at u64::MAX
    /// - Until the system time issue is fixed, rate limiting effectively allows all requests
    ///
    /// Rationale: It's safer to fail open (allow requests) during system time errors than to
    /// permanently block all requests (fail closed). System time errors are usually temporary
    /// (clock adjustments, NTP sync issues) and this allows the service to remain available.
    /// Once the system time is fixed, normal rate limiting resumes.
    ///
    /// # Security Implications
    ///
    /// This means a compromised or misconfigured system (with broken clock) can temporarily
    /// bypass rate limiting. This is an acceptable trade-off for availability. The system
    /// should be monitored for persistent time errors which could indicate tampering.
    fn current_timestamp() -> u64 {
        match SystemTime::now().duration_since(UNIX_EPOCH) {
            Ok(duration) => duration.as_secs(),
            Err(e) => {
                // CRITICAL: System time error - fail-safe to allow requests during time issues
                // This resets rate limit windows and allows subsequent requests through
                eprintln!(
                    "CRITICAL: System time error in rate limiter: {}. \
                     Rate limiting is temporarily disabled. \
                     System clock may have moved backward or time source is unavailable.",
                    e
                );
                u64::MAX
            },
        }
    }

    /// Check if a request should be allowed for the given key
    ///
    /// # Atomicity
    ///
    /// This operation is **atomic** - the entire check-and-update sequence happens atomically:
    /// 1. Acquires exclusive lock on rate limit records
    /// 2. Gets current timestamp
    /// 3. Loads or creates request record for this key
    /// 4. Decides: allow, reset window, or deny
    /// 5. Updates counter/window only if request is allowed
    /// 6. Releases lock
    ///
    /// No concurrent thread can observe a partial state. This prevents classic
    /// time-of-check-time-of-use (TOCTOU) race conditions where multiple threads
    /// simultaneously exceed the rate limit.
    ///
    /// # Returns
    ///
    /// - Ok(()) if the request is allowed and counter has been incremented
    /// - Err(AuthError::RateLimited) if the key has exceeded the rate limit
    ///
    /// # Panics
    ///
    /// Panics if the Mutex is poisoned (another thread panicked while holding the lock).
    /// This is acceptable because a poisoned lock indicates a thread panic, suggesting
    /// the system is already in an inconsistent state and should be restarted.
    pub fn check(&self, key: &str) -> Result<()> {
        // If rate limiting is disabled, always allow the request
        // Note: This check is outside the lock for efficiency, but there's a benign race:
        // if another thread changes config.enabled between this check and acquiring the lock,
        // we still proceed to update the counter. This is safe because we only update counters
        // and don't depend on the enabled flag for correctness (counter updates are idempotent).
        if !self.config.enabled {
            return Ok(());
        }

        // CRITICAL: Acquire lock - this ensures all operations below are atomic
        let mut records = self
            .records
            .lock()
            .expect("rate limiter mutex poisoned - system in critical state");
        let now = Self::current_timestamp();

        // Get or create record for this key (first request from this key)
        let record = records.entry(key.to_string()).or_insert_with(|| RequestRecord {
            count:        0,
            window_start: now,
        });

        // Thread-safe decision: all branches update state atomically while holding the lock
        if now >= record.window_start + self.config.window_secs {
            // CASE 1: Window has expired - start a new window
            // This request is the first in the new window, so it's allowed
            record.count = 1;
            record.window_start = now;
            Ok(())
        } else if record.count < self.config.max_requests {
            // CASE 2: Window is active and we haven't exceeded the limit
            // This request is allowed - increment the counter atomically
            record.count += 1;
            Ok(())
        } else {
            // CASE 3: Window is active and we've reached the limit
            // This request is NOT allowed - counter is not incremented
            // Subsequent requests will also fail until the window expires
            Err(AuthError::RateLimited {
                retry_after_secs: self.config.window_secs,
            })
        }
    }

    /// Get the number of active rate limiters (for monitoring)
    pub fn active_limiters(&self) -> usize {
        let records = self
            .records
            .lock()
            .expect("rate limiter mutex poisoned - system in critical state");
        records.len()
    }

    /// Clear all rate limiters (for testing or reset)
    pub fn clear(&self) {
        let mut records = self
            .records
            .lock()
            .expect("rate limiter mutex poisoned - system in critical state");
        records.clear();
    }

    /// Create a copy for independent testing
    pub fn clone_config(&self) -> RateLimitConfig {
        self.config.clone()
    }
}

/// Global rate limiters for different endpoints
pub struct RateLimiters {
    /// auth/start: per-IP, 100 req/min
    pub auth_start:    KeyedRateLimiter,
    /// auth/callback: per-IP, 50 req/min
    pub auth_callback: KeyedRateLimiter,
    /// auth/refresh: per-user, 10 req/min
    pub auth_refresh:  KeyedRateLimiter,
    /// auth/logout: per-user, 20 req/min
    pub auth_logout:   KeyedRateLimiter,
    /// Failed login tracking: per-user, 5 attempts/hour
    pub failed_logins: KeyedRateLimiter,
}

impl RateLimiters {
    /// Create default rate limiters for all endpoints
    pub fn new() -> Self {
        Self {
            auth_start:    KeyedRateLimiter::new(RateLimitConfig::per_ip_standard()),
            auth_callback: KeyedRateLimiter::new(RateLimitConfig::per_ip_strict()),
            auth_refresh:  KeyedRateLimiter::new(RateLimitConfig::per_user_standard()),
            auth_logout:   KeyedRateLimiter::new(RateLimitConfig::per_user_standard()),
            failed_logins: KeyedRateLimiter::new(RateLimitConfig::failed_login_attempts()),
        }
    }

    /// Create with custom configurations
    pub fn with_configs(
        start_cfg: RateLimitConfig,
        callback_cfg: RateLimitConfig,
        refresh_cfg: RateLimitConfig,
        logout_cfg: RateLimitConfig,
        failed_cfg: RateLimitConfig,
    ) -> Self {
        Self {
            auth_start:    KeyedRateLimiter::new(start_cfg),
            auth_callback: KeyedRateLimiter::new(callback_cfg),
            auth_refresh:  KeyedRateLimiter::new(refresh_cfg),
            auth_logout:   KeyedRateLimiter::new(logout_cfg),
            failed_logins: KeyedRateLimiter::new(failed_cfg),
        }
    }
}

impl Default for RateLimiters {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_rate_limiter_allows_within_limit() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 3,
            window_secs:  60,
        });

        // Should allow up to max_requests
        for i in 0..3 {
            let result = limiter.check("key");
            assert!(result.is_ok(), "Request {} should be allowed", i);
        }
    }

    #[test]
    fn test_rate_limiter_rejects_over_limit() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 2,
            window_secs:  60,
        });

        limiter.check("key").ok();
        limiter.check("key").ok();

        // Third should fail
        let result = limiter.check("key");
        assert!(result.is_err(), "Request over limit should fail");
    }

    #[test]
    fn test_rate_limiter_per_key() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 2,
            window_secs:  60,
        });

        // Key 1: use allowance
        limiter.check("key1").ok();
        limiter.check("key1").ok();

        // Key 2: should have fresh allowance
        let result = limiter.check("key2");
        assert!(result.is_ok(), "Different key should have independent limit");
    }

    #[test]
    fn test_rate_limiter_error_contains_retry_after() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 1,
            window_secs:  60,
        });

        limiter.check("key").ok();
        let result = limiter.check("key");

        match result {
            Err(AuthError::RateLimited { retry_after_secs }) => {
                assert_eq!(retry_after_secs, 60);
            },
            _ => panic!("Expected RateLimited error"),
        }
    }

    #[test]
    fn test_rate_limiter_active_limiters_count() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 100,
            window_secs:  60,
        });

        assert_eq!(limiter.active_limiters(), 0);

        limiter.check("key1").ok();
        assert_eq!(limiter.active_limiters(), 1);

        limiter.check("key2").ok();
        assert_eq!(limiter.active_limiters(), 2);
    }

    #[test]
    fn test_rate_limiters_default() {
        let limiters = RateLimiters::new();

        // auth/start should allow requests
        let result = limiters.auth_start.check("ip_1");
        assert!(result.is_ok());

        // auth/refresh should track per-user
        let result = limiters.auth_refresh.check("user_1");
        assert!(result.is_ok());
    }

    #[test]
    fn test_rate_limit_config_presets() {
        let standard_ip = RateLimitConfig::per_ip_standard();
        assert_eq!(standard_ip.max_requests, 100);
        assert_eq!(standard_ip.window_secs, 60);

        let strict_ip = RateLimitConfig::per_ip_strict();
        assert_eq!(strict_ip.max_requests, 50);

        let user_limit = RateLimitConfig::per_user_standard();
        assert_eq!(user_limit.max_requests, 10);

        let failed = RateLimitConfig::failed_login_attempts();
        assert_eq!(failed.max_requests, 5);
        assert_eq!(failed.window_secs, 3600);
    }

    #[test]
    fn test_ip_based_rate_limiting() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig::per_ip_standard());

        let ip = "203.0.113.1";

        // Should allow up to 100 requests
        for _ in 0..100 {
            let result = limiter.check(ip);
            assert!(result.is_ok());
        }

        // 101st should fail
        let result = limiter.check(ip);
        assert!(result.is_err());
    }

    #[test]
    fn test_rejected_login_tracking() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig::failed_login_attempts());

        let user = "alice@example.com";

        // Should allow 5 failed attempts
        for _ in 0..5 {
            let result = limiter.check(user);
            assert!(result.is_ok());
        }

        // 6th should fail
        let result = limiter.check(user);
        assert!(result.is_err());
    }

    #[test]
    fn test_multiple_users_independent() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig::failed_login_attempts());

        // User 1 uses attempts
        for _ in 0..5 {
            limiter.check("user1").ok();
        }

        // User 1 blocked
        let result = limiter.check("user1");
        assert!(result.is_err());

        // User 2 should have fresh attempts
        let result = limiter.check("user2");
        assert!(result.is_ok());
    }

    #[test]
    fn test_clear_limiters() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 1,
            window_secs:  60,
        });

        limiter.check("key").ok();
        let result = limiter.check("key");
        assert!(result.is_err());

        limiter.clear();

        // After clear, should allow again
        let result = limiter.check("key");
        assert!(result.is_ok());
    }

    #[test]
    fn test_thread_safe_rate_limiting() {
        use std::sync::Arc as StdArc;

        let limiter = StdArc::new(KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 100,
            window_secs:  60,
        }));

        let mut handles = vec![];

        for _ in 0..10 {
            let limiter_clone = StdArc::clone(&limiter);
            let handle = std::thread::spawn(move || {
                for _ in 0..10 {
                    let _ = limiter_clone.check("concurrent");
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().ok();
        }

        // After 100 concurrent requests, next should fail
        let result = limiter.check("concurrent");
        assert!(result.is_err());
    }

    #[test]
    fn test_rate_limiting_many_keys() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 10,
            window_secs:  60,
        });

        // Simulate 1000 different IPs, each with requests
        for i in 0..1000 {
            let key = format!("192.168.{}.{}", i / 256, i % 256);
            let result = limiter.check(&key);
            assert!(result.is_ok());
        }

        assert_eq!(limiter.active_limiters(), 1000);
    }

    #[test]
    fn test_endpoint_combinations() {
        let limiters = RateLimiters::new();

        let ip = "203.0.113.1";
        let user = "bob@example.com";

        // Complete flow
        let result = limiters.auth_start.check(ip);
        assert!(result.is_ok());

        let result = limiters.auth_callback.check(ip);
        assert!(result.is_ok());

        let result = limiters.auth_refresh.check(user);
        assert!(result.is_ok());

        let result = limiters.auth_logout.check(user);
        assert!(result.is_ok());

        let result = limiters.failed_logins.check(user);
        assert!(result.is_ok());
    }

    #[test]
    fn test_attack_prevention_scenario() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 10,
            window_secs:  60,
        });

        let target = "admin@example.com";

        // Attacker tries 10 failed attempts
        for _ in 0..10 {
            let _ = limiter.check(target);
        }

        // 11th blocked
        let result = limiter.check(target);
        assert!(result.is_err());
    }

    #[test]
    fn test_rate_limiter_disabled() {
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      false,
            max_requests: 1,
            window_secs:  60,
        });

        // Even with max_requests = 1, should allow many requests when disabled
        for i in 0..100 {
            let result = limiter.check("key");
            assert!(result.is_ok(), "Request {} should be allowed when rate limiting disabled", i);
        }
    }

    // CONCURRENCY AND ATOMICITY TESTS
    // These tests verify that the rate limiter is thread-safe and atomic

    #[test]
    fn test_concurrent_requests_from_same_key_respects_limit() {
        // RACE CONDITION CHECK: Multiple threads simultaneously checking the same key
        // This verifies that atomic operations prevent exceeding the limit
        use std::{sync::Arc, thread};

        let limiter = Arc::new(KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 50,
            window_secs:  60,
        }));

        let key = "shared_key";
        let allowed_count = Arc::new(std::sync::atomic::AtomicU32::new(0));
        let rejected_count = Arc::new(std::sync::atomic::AtomicU32::new(0));

        let mut handles = vec![];

        // Spawn 100 concurrent threads, all checking the same key
        for _ in 0..100 {
            let limiter = Arc::clone(&limiter);
            let allowed = Arc::clone(&allowed_count);
            let rejected = Arc::clone(&rejected_count);

            let handle = thread::spawn(move || {
                match limiter.check(key) {
                    Ok(_) => allowed.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
                    Err(_) => rejected.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
                };
            });
            handles.push(handle);
        }

        // Wait for all threads to complete
        for handle in handles {
            handle.join().unwrap();
        }

        let allowed = allowed_count.load(std::sync::atomic::Ordering::SeqCst);
        let rejected = rejected_count.load(std::sync::atomic::Ordering::SeqCst);

        // CRITICAL: Due to atomicity, at most 50 requests should be allowed
        assert_eq!(allowed, 50, "Atomic operations should limit to max_requests");
        assert_eq!(rejected, 50, "Remaining requests should be rejected");
        assert_eq!(allowed + rejected, 100, "All requests should be accounted for");
    }

    #[test]
    fn test_concurrent_requests_different_keys_independent() {
        // RACE CONDITION CHECK: Multiple threads checking different keys
        // This verifies that per-key isolation works under concurrent access
        use std::{sync::Arc, thread};

        let limiter = Arc::new(KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 10,
            window_secs:  60,
        }));

        let mut handles = vec![];

        // Spawn 10 threads, each using a different key and making 15 requests
        for thread_id in 0..10 {
            let limiter = Arc::clone(&limiter);
            let handle = thread::spawn(move || {
                let key = format!("key_{}", thread_id);
                let mut allowed = 0;
                let mut rejected = 0;

                for _ in 0..15 {
                    match limiter.check(&key) {
                        Ok(_) => allowed += 1,
                        Err(_) => rejected += 1,
                    }
                }

                (allowed, rejected)
            });
            handles.push(handle);
        }

        // Collect results from all threads
        let mut total_allowed = 0;
        let mut total_rejected = 0;

        for handle in handles {
            let (allowed, rejected) = handle.join().unwrap();
            total_allowed += allowed;
            total_rejected += rejected;
        }

        // CRITICAL: Each key gets independent limit of 10 requests
        assert_eq!(total_allowed, 100, "Each of 10 keys should allow 10 requests");
        assert_eq!(total_rejected, 50, "Each of 10 keys should reject 5 requests");
    }

    #[test]
    fn test_atomic_check_and_update_not_interleaved() {
        // This test verifies that the check-and-update sequence is atomic
        // by ensuring the counter never gets into an inconsistent state
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 3,
            window_secs:  60,
        });

        let key = "test_key";

        // Make 3 allowed requests
        assert!(limiter.check(key).is_ok());
        assert!(limiter.check(key).is_ok());
        assert!(limiter.check(key).is_ok());

        // Verify counter is at 3 (not less, not more)
        assert_eq!(limiter.active_limiters(), 1);

        // 4th request should be rejected
        assert!(limiter.check(key).is_err());

        // 5th request should also be rejected (counter didn't change)
        assert!(limiter.check(key).is_err());

        // Counter should still be at 3 (not decremented on rejection)
        // This verifies that rejected requests didn't partially update state
    }

    #[test]
    fn test_concurrent_window_reset_safety() {
        // Verify that window reset (when window expires) is atomic
        // even under concurrent access
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 2,
            window_secs:  3600, // 1 hour - won't expire in test
        });

        let key = "reset_key";

        // Fill the window
        limiter.check(key).ok();
        limiter.check(key).ok();

        // Further requests should fail
        assert!(limiter.check(key).is_err());
        assert!(limiter.check(key).is_err());

        // Verify state is consistent by clearing and re-checking
        limiter.clear();
        assert_eq!(limiter.active_limiters(), 0);

        // After clear, new requests should be allowed
        assert!(limiter.check(key).is_ok());
    }

    #[test]
    fn test_no_toctou_race_condition() {
        // Time-of-Check-Time-of-Use (TOCTOU) race condition test
        // Verifies that checking the limit and updating the counter happen atomically
        let limiter = KeyedRateLimiter::new(RateLimitConfig {
            enabled:      true,
            max_requests: 1, // Very strict: only 1 request allowed
            window_secs:  60,
        });

        let key = "single_key";

        // First request is allowed
        assert!(limiter.check(key).is_ok());

        // Due to atomic check-and-update, the second request must fail
        // There's no window where both can check and both succeed
        let result = limiter.check(key);
        assert!(
            result.is_err(),
            "Second request must fail - check-and-update is atomic so no TOCTOU race"
        );
    }
}