hypercall_api/caches/

rate_limit.rs

//! Rate limiting cache using fixed window algorithm.
//!
//! This module provides per-wallet rate limiting for:
//! - Order placements
//! - Order cancellations
//! - API requests
//!
//! Supports both Redis (for distributed rate limiting across multiple instances)
//! and in-memory backends (for local development without Redis).

use crate::boundary::read_models::TierCacheApi;
use crate::models::TradingLimits;
use hypercall_types::WalletAddress;
use redis::aio::ConnectionManager;
use redis::Script;
use std::collections::HashMap;
use std::sync::{Arc, RwLock as StdRwLock};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

/// Lua script for atomic check-and-increment in Redis.
/// Returns {current_count, allowed (1=yes, 0=no)}
const CHECK_AND_INCREMENT_LUA: &str = r#"
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
local limit = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])

if current < limit then
    current = redis.call('INCR', KEYS[1])
    redis.call('EXPIRE', KEYS[1], ttl)
    return {current, 1}
end
return {current, 0}
"#;

/// Extra TTL buffer to prevent race conditions at window boundaries (in seconds).
const REDIS_TTL_BUFFER_SECS: u64 = 10;

/// Lua script to get current count without incrementing.
const GET_COUNT_LUA: &str = r#"
local current = redis.call('GET', KEYS[1])
if current == false then
    return 0
end
return tonumber(current)
"#;

/// Fixed window counter for rate limiting (in-memory backend).
#[derive(Debug, Clone)]
struct FixedWindowCounter {
    count: u32,
    window_start: Instant,
}

impl FixedWindowCounter {
    fn new() -> Self {
        Self {
            count: 0,
            window_start: Instant::now(),
        }
    }

    /// Check if the window has expired and reset if necessary.
    /// Returns true if the window was reset.
    fn maybe_reset(&mut self, window_duration: Duration) -> bool {
        if self.window_start.elapsed() >= window_duration {
            self.count = 0;
            self.window_start = Instant::now();
            true
        } else {
            false
        }
    }

    /// Increment counter if under limit, return true if allowed.
    fn try_increment(&mut self, limit: u32, window_duration: Duration) -> bool {
        self.maybe_reset(window_duration);
        if self.count < limit {
            self.count += 1;
            true
        } else {
            false
        }
    }

    /// Get time remaining until window reset.
    fn time_until_reset(&self, window_duration: Duration) -> Duration {
        let elapsed = self.window_start.elapsed();
        if elapsed >= window_duration {
            Duration::ZERO
        } else {
            window_duration - elapsed
        }
    }

    /// Get current count.
    fn current_count(&self) -> u32 {
        self.count
    }
}

/// Type of rate-limited action.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RateLimitAction {
    /// Order placement
    OrderPlacement,
    /// Order cancellation
    OrderCancellation,
    /// General API request
    ApiRequest,
    /// RFQ submission — per-wallet throttle to prevent a single taker
    /// from fan-out-DOS-ing every connected QP via rapid-fire
    /// /rfq/request calls. Default: 10/minute (see get_limit_for_action).
    RfqSubmit,
}

impl RateLimitAction {
    /// Get the string key for this action (used in Redis keys).
    fn as_key_str(&self) -> &'static str {
        match self {
            RateLimitAction::OrderPlacement => "order",
            RateLimitAction::OrderCancellation => "cancel",
            RateLimitAction::ApiRequest => "api",
            RateLimitAction::RfqSubmit => "rfq_submit",
        }
    }
}

/// Rate limit check result with metadata.
#[derive(Debug, Clone)]
pub struct RateLimitInfo {
    /// The limit for this action
    pub limit: u32,
    /// Remaining requests in current window
    pub remaining: u32,
    /// Unix timestamp when the window resets
    pub reset_at: u64,
}

/// Error returned when rate limit is exceeded or service is unavailable.
#[derive(Debug, Clone)]
pub enum RateLimitError {
    /// Rate limit exceeded - client should retry after the specified time
    Exceeded {
        /// Seconds until the limit resets
        retry_after_secs: u32,
        /// The limit that was exceeded
        limit: u32,
        /// Type of action that was rate limited
        action: RateLimitAction,
    },
    /// Redis service unavailable - fail closed
    ServiceUnavailable {
        /// Error message for logging
        message: String,
    },
}

impl std::fmt::Display for RateLimitError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RateLimitError::Exceeded {
                retry_after_secs,
                limit,
                action,
            } => {
                write!(
                    f,
                    "Rate limit exceeded for {:?}: {} per minute, retry after {} seconds",
                    action, limit, retry_after_secs
                )
            }
            RateLimitError::ServiceUnavailable { message } => {
                write!(f, "Rate limit service unavailable: {}", message)
            }
        }
    }
}

impl std::error::Error for RateLimitError {}

/// Backend type for rate limiting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RateLimitBackend {
    /// Redis backend for distributed rate limiting
    Redis,
    /// In-memory backend for single-instance rate limiting
    InMemory,
}

/// Rate limit cache using fixed window algorithm.
///
/// Supports both Redis (distributed) and in-memory (single-instance) backends.
/// When Redis is configured, rate limits are shared across all API server instances.
/// When Redis is not configured, falls back to in-memory rate limiting.
pub struct RateLimitCache {
    /// Redis connection manager (None if using in-memory backend)
    redis: Option<ConnectionManager>,
    /// In-memory counters for order placements per wallet
    order_counts: StdRwLock<HashMap<WalletAddress, FixedWindowCounter>>,
    /// In-memory counters for cancellations per wallet
    cancel_counts: StdRwLock<HashMap<WalletAddress, FixedWindowCounter>>,
    /// In-memory counters for API requests per wallet
    api_counts: StdRwLock<HashMap<WalletAddress, FixedWindowCounter>>,
    /// In-memory counters for RFQ submissions per wallet
    rfq_submit_counts: StdRwLock<HashMap<WalletAddress, FixedWindowCounter>>,
    /// Reference to tier cache for getting limits (TierCache handles its own locking)
    tier_cache: Arc<dyn TierCacheApi>,
    /// Window duration (default: 60 seconds)
    window_duration: Duration,
    /// Lua script for atomic check-and-increment
    check_and_increment_script: Script,
    /// Lua script for getting current count
    get_count_script: Script,
}

impl RateLimitCache {
    /// Create a new rate limit cache.
    ///
    /// If `redis_url` is provided and connection succeeds, uses Redis backend.
    /// Otherwise falls back to in-memory backend.
    pub async fn new(tier_cache: Arc<dyn TierCacheApi>, redis_url: Option<&str>) -> Self {
        let redis = if let Some(url) = redis_url {
            match Self::connect_redis(url).await {
                Ok(conn) => {
                    tracing::info!("RateLimitCache using Redis backend at {}", url);
                    Some(conn)
                }
                Err(e) => {
                    tracing::warn!(
                        "Failed to connect to Redis at {}, using in-memory backend: {}",
                        url,
                        e
                    );
                    None
                }
            }
        } else {
            tracing::info!("REDIS_URL not set, using in-memory rate limiting");
            None
        };

        Self {
            redis,
            order_counts: StdRwLock::new(HashMap::new()),
            cancel_counts: StdRwLock::new(HashMap::new()),
            api_counts: StdRwLock::new(HashMap::new()),
            rfq_submit_counts: StdRwLock::new(HashMap::new()),
            tier_cache,
            window_duration: Duration::from_secs(60),
            check_and_increment_script: Script::new(CHECK_AND_INCREMENT_LUA),
            get_count_script: Script::new(GET_COUNT_LUA),
        }
    }

    /// Connect to Redis and return a connection manager.
    async fn connect_redis(url: &str) -> Result<ConnectionManager, redis::RedisError> {
        let client = redis::Client::open(url)?;
        ConnectionManager::new(client).await
    }

    /// Get the backend type being used.
    pub fn backend(&self) -> RateLimitBackend {
        if self.redis.is_some() {
            RateLimitBackend::Redis
        } else {
            RateLimitBackend::InMemory
        }
    }

    /// Generate Redis key for a wallet/action/window.
    fn redis_key(wallet: &WalletAddress, action: RateLimitAction, window_start: u64) -> String {
        format!(
            "ratelimit:{}:{}:{}",
            wallet,
            action.as_key_str(),
            window_start
        )
    }

    /// Get the window start timestamp (truncated to minute boundary).
    fn get_window_start(now_secs: u64, window_secs: u64) -> u64 {
        (now_secs / window_secs) * window_secs
    }

    /// Get the counters map for the given action type (in-memory backend).
    fn get_counters(
        &self,
        action: RateLimitAction,
    ) -> &StdRwLock<HashMap<WalletAddress, FixedWindowCounter>> {
        match action {
            RateLimitAction::OrderPlacement => &self.order_counts,
            RateLimitAction::OrderCancellation => &self.cancel_counts,
            RateLimitAction::ApiRequest => &self.api_counts,
            RateLimitAction::RfqSubmit => &self.rfq_submit_counts,
        }
    }

    /// Get the limit for the given action from trading limits.
    fn get_limit_for_action(limits: &TradingLimits, action: RateLimitAction) -> u32 {
        match action {
            RateLimitAction::OrderPlacement => limits.orders_per_minute as u32,
            RateLimitAction::OrderCancellation => limits.cancels_per_minute as u32,
            RateLimitAction::ApiRequest => limits.api_requests_per_minute as u32,
            // RFQ submit: 10/minute per wallet. Each submit fans out to
            // every connected QP, so a higher rate would amplify into
            // proportionally more WS traffic + quote generation load.
            // Not yet in TradingLimits — hardcoded until we have a
            // tier-aware RFQ limit field.
            RateLimitAction::RfqSubmit => 10,
        }
    }

    /// Check and increment the rate limit counter for a wallet/action.
    ///
    /// Returns Ok(RateLimitInfo) if allowed, Err(RateLimitError) if exceeded or unavailable.
    pub async fn check_and_increment(
        &self,
        wallet: &WalletAddress,
        action: RateLimitAction,
    ) -> Result<RateLimitInfo, RateLimitError> {
        // Wait for the tier-cache lock instead of using the sync try_read path,
        // where transient lock contention falls through to missing-tier defaults.
        let limits = self.tier_cache.get_trading_limits_async(wallet).await;
        let limit = Self::get_limit_for_action(&limits, action);

        if let Some(ref redis) = self.redis {
            self.check_and_increment_redis(redis.clone(), wallet, action, limit)
                .await
        } else {
            self.check_and_increment_memory(wallet, action, limit)
        }
    }

    /// Check and increment using Redis backend.
    async fn check_and_increment_redis(
        &self,
        mut redis: ConnectionManager,
        wallet: &WalletAddress,
        action: RateLimitAction,
        limit: u32,
    ) -> Result<RateLimitInfo, RateLimitError> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let window_secs = self.window_duration.as_secs();
        let window_start = Self::get_window_start(now, window_secs);
        let window_end = window_start + window_secs;
        let key = Self::redis_key(wallet, action, window_start);

        // TTL is window duration plus buffer to prevent race conditions
        let ttl = window_secs + REDIS_TTL_BUFFER_SECS;

        // Execute Lua script atomically
        let result: Result<(i64, i64), redis::RedisError> = self
            .check_and_increment_script
            .key(&key)
            .arg(limit)
            .arg(ttl)
            .invoke_async(&mut redis)
            .await;

        match result {
            Ok((count, allowed)) => {
                if allowed == 1 {
                    Ok(RateLimitInfo {
                        limit,
                        remaining: limit.saturating_sub(count as u32),
                        reset_at: window_end,
                    })
                } else {
                    let retry_after = (window_end.saturating_sub(now)).max(1) as u32;
                    Err(RateLimitError::Exceeded {
                        retry_after_secs: retry_after,
                        limit,
                        action,
                    })
                }
            }
            Err(e) => {
                tracing::error!("Redis rate limit check failed: {}", e);
                // Fail closed: reject requests when rate limiting is unavailable.
                // This is a security decision to prevent abuse during Redis outages.
                Err(RateLimitError::ServiceUnavailable {
                    message: format!("Redis error: {}", e),
                })
            }
        }
    }

    /// Check and increment using in-memory backend.
    fn check_and_increment_memory(
        &self,
        wallet: &WalletAddress,
        action: RateLimitAction,
        limit: u32,
    ) -> Result<RateLimitInfo, RateLimitError> {
        let counters = self.get_counters(action);

        // Use write lock to update counters
        let mut map = counters.write().unwrap();
        let counter = map.entry(*wallet).or_insert_with(FixedWindowCounter::new);

        if counter.try_increment(limit, self.window_duration) {
            let remaining = limit.saturating_sub(counter.current_count());
            let reset_at = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs()
                + counter.time_until_reset(self.window_duration).as_secs();

            Ok(RateLimitInfo {
                limit,
                remaining,
                reset_at,
            })
        } else {
            let retry_after = counter.time_until_reset(self.window_duration).as_secs() as u32;
            Err(RateLimitError::Exceeded {
                retry_after_secs: retry_after.max(1),
                limit,
                action,
            })
        }
    }

    /// Get current rate limit info for a wallet/action without incrementing.
    pub async fn get_info(
        &self,
        wallet: &WalletAddress,
        action: RateLimitAction,
    ) -> Result<RateLimitInfo, RateLimitError> {
        let limits = self.tier_cache.get_trading_limits_async(wallet).await;
        let limit = Self::get_limit_for_action(&limits, action);

        if let Some(ref redis) = self.redis {
            self.get_info_redis(redis.clone(), wallet, action, limit)
                .await
        } else {
            Ok(self.get_info_memory(wallet, action, limit))
        }
    }

    /// Get info using Redis backend.
    async fn get_info_redis(
        &self,
        mut redis: ConnectionManager,
        wallet: &WalletAddress,
        action: RateLimitAction,
        limit: u32,
    ) -> Result<RateLimitInfo, RateLimitError> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let window_secs = self.window_duration.as_secs();
        let window_start = Self::get_window_start(now, window_secs);
        let window_end = window_start + window_secs;
        let key = Self::redis_key(wallet, action, window_start);

        let result: Result<i64, redis::RedisError> = self
            .get_count_script
            .key(&key)
            .invoke_async(&mut redis)
            .await;

        match result {
            Ok(count) => {
                let remaining = limit.saturating_sub(count as u32);
                Ok(RateLimitInfo {
                    limit,
                    remaining,
                    reset_at: window_end,
                })
            }
            Err(e) => {
                tracing::error!("Redis rate limit get_info failed: {}", e);
                Err(RateLimitError::ServiceUnavailable {
                    message: format!("Redis error: {}", e),
                })
            }
        }
    }

    /// Get info using in-memory backend.
    fn get_info_memory(
        &self,
        wallet: &WalletAddress,
        action: RateLimitAction,
        limit: u32,
    ) -> RateLimitInfo {
        let counters = self.get_counters(action);

        let (current_count, time_until_reset) = {
            let mut map = counters.write().unwrap();
            if let Some(counter) = map.get_mut(wallet) {
                counter.maybe_reset(self.window_duration);
                (
                    counter.current_count(),
                    counter.time_until_reset(self.window_duration),
                )
            } else {
                (0, self.window_duration)
            }
        };

        let remaining = limit.saturating_sub(current_count);
        let reset_at = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
            + time_until_reset.as_secs();

        RateLimitInfo {
            limit,
            remaining,
            reset_at,
        }
    }

    /// Clean up expired entries to prevent memory growth (in-memory backend only).
    /// Call this periodically (e.g., every 5 minutes).
    pub fn cleanup_expired(&self) {
        let threshold = self.window_duration * 2;

        let cleanup_map = |map: &StdRwLock<HashMap<WalletAddress, FixedWindowCounter>>| {
            let mut guard = map.write().unwrap();
            guard.retain(|_, counter| counter.window_start.elapsed() < threshold);
        };

        cleanup_map(&self.order_counts);
        cleanup_map(&self.cancel_counts);
        cleanup_map(&self.api_counts);
    }

    /// Get the number of tracked wallets (for metrics, in-memory backend only).
    pub fn tracked_wallet_count(&self) -> usize {
        let order_len = self.order_counts.read().unwrap().len();
        let cancel_len = self.cancel_counts.read().unwrap().len();
        let api_len = self.api_counts.read().unwrap().len();
        order_len.max(cancel_len).max(api_len)
    }
}

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

    #[test]
    fn test_fixed_window_counter_increments() {
        let mut counter = FixedWindowCounter::new();
        let window = Duration::from_secs(60);

        assert!(counter.try_increment(10, window));
        assert_eq!(counter.current_count(), 1);

        assert!(counter.try_increment(10, window));
        assert_eq!(counter.current_count(), 2);
    }

    #[test]
    fn test_fixed_window_counter_blocks_at_limit() {
        let mut counter = FixedWindowCounter::new();
        let window = Duration::from_secs(60);

        // Fill up to limit
        for _ in 0..10 {
            assert!(counter.try_increment(10, window));
        }

        // 11th request should fail
        assert!(!counter.try_increment(10, window));
        assert_eq!(counter.current_count(), 10);
    }

    #[test]
    fn test_rate_limit_error_display() {
        let err = RateLimitError::Exceeded {
            retry_after_secs: 30,
            limit: 60,
            action: RateLimitAction::OrderPlacement,
        };

        let display = format!("{}", err);
        assert!(display.contains("60"));
        assert!(display.contains("30"));

        let err = RateLimitError::ServiceUnavailable {
            message: "Redis error".to_string(),
        };
        let display = format!("{}", err);
        assert!(display.contains("unavailable"));
    }

    #[test]
    fn test_get_limit_for_action() {
        let limits = TradingLimits {
            max_open_orders: 100,
            max_open_positions: 50,
            orders_per_minute: 60,
            cancels_per_minute: 120,
            api_requests_per_minute: 600,
        };

        assert_eq!(
            RateLimitCache::get_limit_for_action(&limits, RateLimitAction::OrderPlacement),
            60
        );
        assert_eq!(
            RateLimitCache::get_limit_for_action(&limits, RateLimitAction::OrderCancellation),
            120
        );
        assert_eq!(
            RateLimitCache::get_limit_for_action(&limits, RateLimitAction::ApiRequest),
            600
        );
    }

    #[test]
    fn test_redis_key_generation() {
        use std::str::FromStr;
        let wallet = WalletAddress::from_str("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266").unwrap();
        let key = RateLimitCache::redis_key(&wallet, RateLimitAction::OrderPlacement, 1704067200);
        assert!(key.starts_with("ratelimit:"));
        assert!(key.contains(":order:"));
        assert!(key.contains("1704067200"));
    }

    #[test]
    fn test_get_window_start() {
        // 1704067234 should round down to 1704067200 (60-second windows)
        assert_eq!(RateLimitCache::get_window_start(1704067234, 60), 1704067200);
        assert_eq!(RateLimitCache::get_window_start(1704067200, 60), 1704067200);
        assert_eq!(RateLimitCache::get_window_start(1704067259, 60), 1704067200);
        assert_eq!(RateLimitCache::get_window_start(1704067260, 60), 1704067260);
    }

    #[test]
    fn test_action_key_str() {
        assert_eq!(RateLimitAction::OrderPlacement.as_key_str(), "order");
        assert_eq!(RateLimitAction::OrderCancellation.as_key_str(), "cancel");
        assert_eq!(RateLimitAction::ApiRequest.as_key_str(), "api");
    }
}