safety(web-gateway): add per-route-group rate limiting middleware

Token-bucket rate limiter (no new dependencies) with three tiers:
- Auth routes: 10 req/min, burst 5 (brute-force protection)
- Trading/read routes: 200 req/min, burst 20
- Compute routes (backtest/training/tune): 30 req/min, burst 5

Includes background cleanup task for stale entries and 7 unit tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

web-gateway/src/lib.rs

@@ -10,6 +10,7 @@ pub mod auth;
 pub mod config;
 pub mod error;
 pub mod grpc;
+pub mod rate_limit;
 pub mod routes;
 pub mod state;
 pub mod ws;

web-gateway/src/rate_limit.rs

@@ -0,0 +1,188 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use axum::extract::Request;
+use axum::http::StatusCode;
+use axum::middleware::Next;
+use axum::response::{IntoResponse, Response};
+use serde_json::json;
+use tokio::sync::Mutex;
+
+/// A simple token-bucket rate limiter keyed by client IP.
+///
+/// Each client gets `capacity` tokens, refilled at `refill_rate` per second.
+/// When tokens are exhausted, requests are rejected with 429 Too Many Requests.
+#[derive(Clone)]
+pub struct RateLimiter {
+    state: Arc<Mutex<HashMap<String, BucketState>>>,
+    capacity: u32,
+    refill_rate: f64,
+}
+
+struct BucketState {
+    tokens: f64,
+    last_refill: Instant,
+}
+
+impl RateLimiter {
+    /// Create a new rate limiter.
+    ///
+    /// - `requests_per_minute`: maximum sustained request rate per client
+    /// - `burst`: maximum burst size (bucket capacity)
+    pub fn new(requests_per_minute: u32, burst: u32) -> Self {
+        Self {
+            state: Arc::new(Mutex::new(HashMap::new())),
+            capacity: burst,
+            refill_rate: f64::from(requests_per_minute) / 60.0,
+        }
+    }
+
+    /// Try to consume a token for the given client key.
+    /// Returns `true` if allowed, `false` if rate-limited.
+    async fn try_acquire(&self, key: &str) -> bool {
+        let mut buckets = self.state.lock().await;
+        let now = Instant::now();
+
+        let bucket = buckets.entry(key.to_owned()).or_insert(BucketState {
+            tokens: f64::from(self.capacity),
+            last_refill: now,
+        });
+
+        // Refill tokens based on elapsed time
+        let elapsed = now.duration_since(bucket.last_refill).as_secs_f64();
+        bucket.tokens = (bucket.tokens + elapsed * self.refill_rate).min(f64::from(self.capacity));
+        bucket.last_refill = now;
+
+        if bucket.tokens >= 1.0 {
+            bucket.tokens -= 1.0;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+/// Extract client IP from request headers (X-Forwarded-For, X-Real-Ip)
+/// or fall back to "unknown".
+fn client_ip(request: &Request) -> String {
+    request
+        .headers()
+        .get("x-forwarded-for")
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.split(',').next())
+        .map(|s| s.trim().to_owned())
+        .or_else(|| {
+            request
+                .headers()
+                .get("x-real-ip")
+                .and_then(|v| v.to_str().ok())
+                .map(|s| s.to_owned())
+        })
+        .unwrap_or_else(|| "unknown".to_owned())
+}
+
+/// Axum middleware that enforces rate limiting using a shared [`RateLimiter`].
+///
+/// The limiter is passed via `State` (extracted from Axum's state system).
+pub async fn rate_limit_middleware(
+    axum::extract::State(limiter): axum::extract::State<RateLimiter>,
+    request: Request,
+    next: Next,
+) -> Response {
+    let ip = client_ip(&request);
+
+    if !limiter.try_acquire(&ip).await {
+        tracing::warn!(client_ip = %ip, "Rate limit exceeded");
+        return (
+            StatusCode::TOO_MANY_REQUESTS,
+            axum::Json(json!({ "error": "Too many requests" })),
+        )
+            .into_response();
+    }
+
+    next.run(request).await
+}
+
+/// Spawn a background task that periodically evicts stale entries from the
+/// rate limiter's internal state to prevent unbounded memory growth.
+pub fn spawn_cleanup_task(limiter: &RateLimiter, interval: Duration, max_idle: Duration) {
+    let state = Arc::clone(&limiter.state);
+    tokio::spawn(async move {
+        #[allow(clippy::infinite_loop)]
+        loop {
+            tokio::time::sleep(interval).await;
+            let mut buckets = state.lock().await;
+            let now = Instant::now();
+            buckets.retain(|_, bucket| now.duration_since(bucket.last_refill) < max_idle);
+        }
+    });
+}
+
+#[cfg(test)]
+#[allow(clippy::unwrap_used)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_allows_within_limit() {
+        let limiter = RateLimiter::new(60, 5); // 60/min, burst 5
+        for _ in 0..5 {
+            assert!(limiter.try_acquire("client1").await);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_rejects_over_limit() {
+        let limiter = RateLimiter::new(60, 3); // 60/min, burst 3
+        assert!(limiter.try_acquire("client1").await);
+        assert!(limiter.try_acquire("client1").await);
+        assert!(limiter.try_acquire("client1").await);
+        // 4th should be rejected (burst exhausted, no time to refill)
+        assert!(!limiter.try_acquire("client1").await);
+    }
+
+    #[tokio::test]
+    async fn test_separate_clients_independent() {
+        let limiter = RateLimiter::new(60, 1);
+        assert!(limiter.try_acquire("client1").await);
+        assert!(!limiter.try_acquire("client1").await); // exhausted
+        assert!(limiter.try_acquire("client2").await); // different client
+    }
+
+    #[tokio::test]
+    async fn test_refills_over_time() {
+        let limiter = RateLimiter::new(6000, 1); // 100/sec → refills fast
+        assert!(limiter.try_acquire("client1").await);
+        assert!(!limiter.try_acquire("client1").await);
+        tokio::time::sleep(Duration::from_millis(20)).await;
+        // After 20ms at 100/sec = ~2 tokens refilled
+        assert!(limiter.try_acquire("client1").await);
+    }
+
+    #[test]
+    fn test_client_ip_from_forwarded_for() {
+        let req = Request::builder()
+            .header("x-forwarded-for", "1.2.3.4, 5.6.7.8")
+            .body(axum::body::Body::empty())
+            .unwrap();
+        assert_eq!(client_ip(&req), "1.2.3.4");
+    }
+
+    #[test]
+    fn test_client_ip_from_real_ip() {
+        let req = Request::builder()
+            .header("x-real-ip", "10.0.0.1")
+            .body(axum::body::Body::empty())
+            .unwrap();
+        assert_eq!(client_ip(&req), "10.0.0.1");
+    }
+
+    #[test]
+    fn test_client_ip_fallback() {
+        let req = Request::builder()
+            .body(axum::body::Body::empty())
+            .unwrap();
+        assert_eq!(client_ip(&req), "unknown");
+    }
+}

web-gateway/src/routes/mod.rs

@@ -2,11 +2,13 @@ use axum::{middleware, routing, Json, Router};
 use serde_json::json;
 
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::auth::middleware::auth_middleware;
 use crate::config::GatewayConfig;
 use crate::proto::health::health_client::HealthClient;
 use crate::proto::health::HealthCheckRequest;
+use crate::rate_limit::{rate_limit_middleware, spawn_cleanup_task, RateLimiter};
 use crate::state::AppState;
 use crate::ws::handler::ws_handler;
 
@@ -21,23 +23,48 @@ pub mod training;
 pub mod tune;
 
 pub fn create_router(state: AppState) -> Router {
-    // REST endpoints require auth middleware
-    let rest_api = Router::new()
+    // Rate limiters: different tiers for different route groups
+    // Trading: 200 req/min per client, burst 20
+    let trading_limiter = RateLimiter::new(200, 20);
+    // Compute-heavy (backtest, training, tune): 30 req/min per client, burst 5
+    let compute_limiter = RateLimiter::new(30, 5);
+    // Auth (login): 10 req/min per client, burst 5
+    let auth_limiter = RateLimiter::new(10, 5);
+
+    // Spawn background cleanup tasks (evict stale entries every 5 min, idle > 10 min)
+    spawn_cleanup_task(&trading_limiter, Duration::from_secs(300), Duration::from_secs(600));
+    spawn_cleanup_task(&compute_limiter, Duration::from_secs(300), Duration::from_secs(600));
+    spawn_cleanup_task(&auth_limiter, Duration::from_secs(300), Duration::from_secs(600));
+
+    // Trading & read-only routes (moderate rate limit)
+    let trading_routes = Router::new()
         .nest("/trading", trading::router())
         .nest("/risk", risk::router())
         .nest("/ml", ml::router())
-        .nest("/training", training::router())
-        .nest("/backtest", backtesting::router())
         .nest("/performance", performance::router())
         .nest("/config", config::router())
+        .layer(middleware::from_fn_with_state(trading_limiter, rate_limit_middleware));
+
+    // Compute-heavy routes (strict rate limit)
+    let compute_routes = Router::new()
+        .nest("/training", training::router())
+        .nest("/backtest", backtesting::router())
         .nest("/tune", tune::router())
+        .layer(middleware::from_fn_with_state(compute_limiter, rate_limit_middleware));
+
+    // REST endpoints require auth middleware
+    let rest_api = Router::new()
+        .merge(trading_routes)
+        .merge(compute_routes)
         .layer(middleware::from_fn_with_state(
             Arc::<GatewayConfig>::clone(&state.config),
             auth_middleware,
         ));
 
-    // Public endpoints (no auth required)
-    let public_api = Router::new().nest("/auth", auth::router());
+    // Public endpoints (no auth required, but rate-limited)
+    let public_api = Router::new()
+        .nest("/auth", auth::router())
+        .layer(middleware::from_fn_with_state(auth_limiter, rate_limit_middleware));
 
     // WebSocket endpoint validates JWT via query param (no middleware layer)
     let ws_route = Router::new().route("/ws", routing::get(ws_handler));