diff --git a/web-gateway/src/lib.rs b/web-gateway/src/lib.rs index 155b8a939..f0c857662 100644 --- a/web-gateway/src/lib.rs +++ b/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; diff --git a/web-gateway/src/rate_limit.rs b/web-gateway/src/rate_limit.rs new file mode 100644 index 000000000..be55d854c --- /dev/null +++ b/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>>, + 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, + 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"); + } +} diff --git a/web-gateway/src/routes/mod.rs b/web-gateway/src/routes/mod.rs index 7b7785cf6..dc4bbfcea 100644 --- a/web-gateway/src/routes/mod.rs +++ b/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::::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));