cp_tor/

rate_limit.rs

//! Token bucket rate limiter for Canon search protocol.
//!
//! Per CP-013 §14: 10 requests per minute per requester public key.
//! Uses a token bucket algorithm with max 10 tokens refilling at 10/60 per second.

use std::collections::HashMap;

/// Token bucket rate limiter keyed by requester public key.
pub struct RateLimiter {
    buckets: HashMap<[u8; 32], TokenBucket>,
    max_tokens: f64,
    refill_rate: f64,
}

struct TokenBucket {
    tokens: f64,
    last_refill: i64,
}

impl RateLimiter {
    /// Create a new rate limiter with the default CP-013 parameters:
    /// 10 tokens max, refilling at 10/60 tokens per second.
    pub fn new() -> Self {
        Self {
            buckets: HashMap::new(),
            max_tokens: 10.0,
            refill_rate: 10.0 / 60.0,
        }
    }

    /// Create a rate limiter with custom parameters (for testing).
    pub fn with_params(max_tokens: f64, refill_rate: f64) -> Self {
        Self {
            buckets: HashMap::new(),
            max_tokens,
            refill_rate,
        }
    }

    /// Check if a request from the given public key is allowed.
    /// Consumes one token if allowed.
    pub fn check(&mut self, public_key: &[u8; 32], now_ms: i64) -> bool {
        let bucket = self.buckets.entry(*public_key).or_insert(TokenBucket {
            tokens: self.max_tokens,
            last_refill: now_ms,
        });

        // Refill tokens based on elapsed time
        let elapsed_secs = (now_ms - bucket.last_refill) as f64 / 1000.0;
        if elapsed_secs > 0.0 {
            bucket.tokens = (bucket.tokens + elapsed_secs * self.refill_rate).min(self.max_tokens);
            bucket.last_refill = now_ms;
        }

        // Try to consume a token
        if bucket.tokens >= 1.0 {
            bucket.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Remove stale entries older than the given threshold (milliseconds).
    /// Call periodically to prevent unbounded memory growth.
    pub fn cleanup(&mut self, now_ms: i64, stale_threshold_ms: i64) {
        self.buckets
            .retain(|_, bucket| (now_ms - bucket.last_refill) < stale_threshold_ms);
    }
}

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

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

    #[test]
    fn test_rate_limiter_allows_up_to_max() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let now = 1000000i64;

        // First 10 should be allowed
        for _ in 0..10 {
            assert!(limiter.check(&key, now));
        }

        // 11th should be rejected
        assert!(!limiter.check(&key, now));
    }

    #[test]
    fn test_rate_limiter_refills_over_time() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let start = 1000000i64;

        // Exhaust all tokens
        for _ in 0..10 {
            assert!(limiter.check(&key, start));
        }
        assert!(!limiter.check(&key, start));

        // After 6 seconds, one token should have refilled (10/60 * 6 = 1.0)
        let after_6s = start + 6000;
        assert!(limiter.check(&key, after_6s));
        assert!(!limiter.check(&key, after_6s));
    }

    #[test]
    fn test_rate_limiter_full_refill() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let start = 1000000i64;

        // Exhaust all tokens
        for _ in 0..10 {
            limiter.check(&key, start);
        }

        // After 60 seconds, all 10 tokens should be refilled
        let after_60s = start + 60_000;
        for _ in 0..10 {
            assert!(limiter.check(&key, after_60s));
        }
        assert!(!limiter.check(&key, after_60s));
    }

    #[test]
    fn test_rate_limiter_independent_keys() {
        let mut limiter = RateLimiter::new();
        let key_a = [1u8; 32];
        let key_b = [2u8; 32];
        let now = 1000000i64;

        // Exhaust key_a
        for _ in 0..10 {
            limiter.check(&key_a, now);
        }
        assert!(!limiter.check(&key_a, now));

        // key_b should still have tokens
        assert!(limiter.check(&key_b, now));
    }

    #[test]
    fn test_rate_limiter_cleanup() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let start = 1_000_000i64;

        limiter.check(&key, start);
        assert_eq!(limiter.buckets.len(), 1);

        // Cleanup with threshold: entries older than 1 hour
        let stale_threshold = 3_600_000;
        limiter.cleanup(start + stale_threshold + 1, stale_threshold);
        assert_eq!(limiter.buckets.len(), 0);
    }

    #[test]
    fn test_rate_limiter_no_negative_tokens() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let now = 1000000i64;

        // Exhaust tokens then try more
        for _ in 0..20 {
            limiter.check(&key, now);
        }

        // After 6 seconds should have exactly 1 token, not negative + refill
        let after_6s = now + 6000;
        assert!(limiter.check(&key, after_6s));
        assert!(!limiter.check(&key, after_6s));
    }

    #[test]
    fn test_rate_limiter_tokens_cap_at_max() {
        let mut limiter = RateLimiter::new();
        let key = [1u8; 32];
        let start = 1000000i64;

        // Use one token
        limiter.check(&key, start);

        // Wait a very long time (should cap at 10, not exceed)
        let much_later = start + 1_000_000;
        for _ in 0..10 {
            assert!(limiter.check(&key, much_later));
        }
        assert!(!limiter.check(&key, much_later));
    }

    #[test]
    fn test_rate_limiter_custom_params() {
        let mut limiter = RateLimiter::with_params(3.0, 1.0); // 3 max, 1/sec refill
        let key = [1u8; 32];
        let now = 1000000i64;

        assert!(limiter.check(&key, now));
        assert!(limiter.check(&key, now));
        assert!(limiter.check(&key, now));
        assert!(!limiter.check(&key, now)); // 4th rejected

        // After 1 second, 1 token refilled
        assert!(limiter.check(&key, now + 1000));
        assert!(!limiter.check(&key, now + 1000));
    }
}