GitRust
rate_limit.rs
321 lignes · 11370 octets
//! Rate Limiting Middleware and Helpers //! //! Feature 009: Admin & Settings — per-route tower layers for blunt-instrument //! rate limiting on admin endpoints. //! //! Feature 041: Files Sync & Sharing (T074) — fine-grained per-IP limits for //! public share endpoints and authentication endpoints, routed through the //! generic [`RateLimitService`](crate::services::rate_limit_service::RateLimitService). //! //! # Thresholds //! //! | Scope | Default | //! |-------|---------| //! | Admin (tower `RateLimitLayer`) | 100 req/min | //! | User management (tower) | 20 req/min | //! | Reset password (tower) | 10 req/min | //! | Public share per IP | 30 req/min | //! | Auth per IP | env `RATE_LIMIT_AUTH_PER_MINUTE`, default 10 req/min | //! //! Public share and auth limits are backed by the in-process //! [`RateLimitService`](crate::services::rate_limit_service::RateLimitService) //! so they share bounding, cleanup and the canonical //! `TooManyRequests` error (HTTP 429 + `Retry-After`). use axum::{ extract::{Request, State}, middleware::Next, response::{IntoResponse, Response}, }; use sea_orm::DatabaseConnection; use std::{sync::OnceLock, time::Duration}; use tower::limit::RateLimitLayer; use crate::error::AppError; use crate::services::rate_limit_service::RateLimitService; /// Rate limiting configuration for admin routes pub struct AdminRateLimitConfig; impl AdminRateLimitConfig { /// Create rate limiter for admin endpoints (100 requests per minute) pub fn admin_limit() -> RateLimitLayer { RateLimitLayer::new(100, Duration::from_secs(60)) } /// Create rate limiter for user management endpoints (20 requests per minute) pub fn user_management_limit() -> RateLimitLayer { RateLimitLayer::new(20, Duration::from_secs(60)) } /// Create rate limiter for reset password endpoint (10 requests per minute) pub fn reset_password_limit() -> RateLimitLayer { RateLimitLayer::new(10, Duration::from_secs(60)) } } /// Rate limiting configuration for public share endpoints /// Feature 041: Files Sync & Sharing pub struct PublicShareRateLimitConfig; impl PublicShareRateLimitConfig { /// Create rate limiter for public share endpoints (30 requests per minute per IP) /// /// This limit is more restrictive than authenticated endpoints to prevent abuse /// of public links (brute force password attempts, token enumeration, etc.) pub fn public_share_limit() -> RateLimitLayer { RateLimitLayer::new(30, Duration::from_secs(60)) } } /// Placeholder middleware — real admin rate limiting is applied via /// [`AdminRateLimitConfig::admin_limit`] in the router layer. /// /// Kept for composability if a future version routes admin traffic through /// [`RateLimitService`] (for per-user quotas). pub async fn admin_rate_limit( State(_db): State<DatabaseConnection>, request: Request, next: Next, ) -> Result<Response, AppError> { Ok(next.run(request).await) } const PUBLIC_SHARE_MAX_PER_MINUTE: u32 = 30; const DEFAULT_AUTH_MAX_PER_MINUTE: u32 = 10; static PUBLIC_SHARE_LIMITER: OnceLock<RateLimitService> = OnceLock::new(); static AUTH_LIMITER: OnceLock<RateLimitService> = OnceLock::new(); static AUTH_MAX_PER_MINUTE: OnceLock<u32> = OnceLock::new(); fn public_share_limiter() -> &'static RateLimitService { PUBLIC_SHARE_LIMITER.get_or_init(RateLimitService::new) } fn auth_limiter() -> &'static RateLimitService { AUTH_LIMITER.get_or_init(RateLimitService::new) } fn auth_max_per_minute() -> u32 { *AUTH_MAX_PER_MINUTE.get_or_init(|| { std::env::var("RATE_LIMIT_AUTH_PER_MINUTE") .ok() .and_then(|v| v.parse().ok()) .unwrap_or(DEFAULT_AUTH_MAX_PER_MINUTE) }) } /// SEC-H3: Number of trusted reverse proxy hops (0 = direct connection, ignore XFF). /// Read once from TRUSTED_PROXY_HOPS environment variable. static TRUSTED_PROXY_HOPS: OnceLock<usize> = OnceLock::new(); fn trusted_proxy_hops() -> usize { *TRUSTED_PROXY_HOPS.get_or_init(|| { std::env::var("TRUSTED_PROXY_HOPS") .ok() .and_then(|v| v.parse().ok()) .unwrap_or(0) }) } /// Pick the client IP from `X-Forwarded-For` and/or the direct connection. /// /// `hops` is the number of trusted reverse-proxy hops (0 = ignore XFF). /// When `hops > 0`, takes the Nth IP from the **right** of the XFF list. fn extract_ip(xff: Option<&str>, connect: Option<String>, hops: usize) -> String { if hops > 0 { if let Some(xff_value) = xff { let ips: Vec<&str> = xff_value.split(',').map(|s| s.trim()).collect(); if ips.len() >= hops { return ips[ips.len() - hops].to_string(); } } } connect.unwrap_or_else(|| "unknown".to_string()) } /// SEC-H3: Extract the real client IP, respecting TRUSTED_PROXY_HOPS. /// When TRUSTED_PROXY_HOPS > 0, takes the Nth IP from the right of X-Forwarded-For. /// Otherwise falls back to the direct connection IP. fn extract_ip_from_xff_or_connect(xff: Option<&str>, connect_ip: Option<String>) -> String { extract_ip(xff, connect_ip, trusted_proxy_hops()) } /// Extract IP address from Request (uses ConnectInfo when available) pub fn extract_client_ip(request: &Request) -> String { let xff = request .headers() .get("x-forwarded-for") .and_then(|h| h.to_str().ok()); let connect_ip = request .extensions() .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>() .map(|ci| ci.0.ip().to_string()) .or_else(|| { // `MockConnectInfo` stores itself, not `ConnectInfo` (axum 0.8). request .extensions() .get::<axum::extract::connect_info::MockConnectInfo<std::net::SocketAddr>>() .map(|mock| mock.0.ip().to_string()) }); extract_ip_from_xff_or_connect(xff, connect_ip) } /// Extract IP address from HeaderMap (used in handlers where Request is consumed) pub fn extract_client_ip_from_headers(headers: &axum::http::HeaderMap) -> String { let xff = headers.get("x-forwarded-for").and_then(|h| h.to_str().ok()); extract_ip_from_xff_or_connect(xff, None) } /// Check rate limit for a public share endpoint request. /// /// Feature 041: Files Sync & Sharing (T074). Returns an [`AppError`] /// response (HTTP 429 with `Retry-After`) if the limit is exceeded. #[allow(clippy::result_large_err)] pub fn check_public_share_rate_limit(request: &Request) -> Result<(), Response> { let ip = extract_client_ip(request); match public_share_limiter().check(&format!("share:{ip}"), PUBLIC_SHARE_MAX_PER_MINUTE) { Ok(()) => Ok(()), Err(err) => { tracing::warn!(ip = %ip, "Public share rate limit exceeded"); Err(err.into_response()) } } } /// Check rate limit for a public share endpoint using `HeaderMap`. /// /// Feature 041: Files Sync & Sharing (T074). Prefer this variant when the /// `Request` has already been destructured by an extractor. #[allow(clippy::result_large_err)] pub fn check_public_share_rate_limit_from_headers( headers: &axum::http::HeaderMap, ) -> Result<(), Response> { let ip = extract_client_ip_from_headers(headers); match public_share_limiter().check(&format!("share:{ip}"), PUBLIC_SHARE_MAX_PER_MINUTE) { Ok(()) => Ok(()), Err(err) => { tracing::warn!(ip = %ip, "Public share rate limit exceeded"); Err(err.into_response()) } } } /// SEC-C2: Check auth rate limit for an incoming request. /// /// Protects login, register, refresh, and forgot-password from brute-force /// attacks. Returns 429 `Too Many Requests` with a `Retry-After` header when /// the per-IP budget is exceeded. #[allow(clippy::result_large_err)] pub fn check_auth_rate_limit(headers: &axum::http::HeaderMap) -> Result<(), Response> { let ip = extract_client_ip_from_headers(headers); match auth_limiter().check(&format!("auth:{ip}"), auth_max_per_minute()) { Ok(()) => Ok(()), Err(err) => Err(err.into_response()), } } /// Per-IP rate-limit layer for authentication endpoints. /// /// Uses [`extract_client_ip`] so the key comes from `ConnectInfo` when /// `TRUSTED_PROXY_HOPS` is 0 (the default), rather than collapsing every /// request into a global `auth:unknown` bucket. pub(crate) async fn auth_ip_rate_limit(request: Request, next: Next) -> Response { let ip = extract_client_ip(&request); match auth_limiter().check(&format!("auth:{ip}"), auth_max_per_minute()) { Ok(()) => next.run(request).await, Err(err) => err.into_response(), } } #[cfg(test)] mod tests { use super::*; #[test] fn test_rate_limit_config_creation() { let _admin = AdminRateLimitConfig::admin_limit(); let _user_mgmt = AdminRateLimitConfig::user_management_limit(); let _reset = AdminRateLimitConfig::reset_password_limit(); } #[test] fn auth_limiter_isolates_budget_per_ip() { let rl = RateLimitService::new(); for _ in 0..3 { rl.check("auth:1.2.3.4", 3).unwrap(); } assert!(rl.check("auth:1.2.3.4", 3).is_err()); // Different IP has its own budget. rl.check("auth:5.6.7.8", 3) .expect("distinct IP must not share the other IP's budget"); } #[test] fn share_limiter_isolates_budget_per_ip() { let rl = RateLimitService::new(); for _ in 0..3 { rl.check("share:1.2.3.4", 3).unwrap(); } assert!(rl.check("share:1.2.3.4", 3).is_err()); rl.check("share:9.9.9.9", 3) .expect("distinct IP must not share the other IP's budget"); } #[test] fn extract_client_ip_uses_connect_info_without_xff() { let mut request = Request::builder() .uri("/") .body(axum::body::Body::empty()) .expect("request"); request .extensions_mut() .insert(axum::extract::ConnectInfo(std::net::SocketAddr::from(( [203, 0, 113, 1], 12345, )))); assert_eq!(extract_client_ip(&request), "203.0.113.1"); } #[test] fn extract_client_ip_uses_mock_connect_info_without_xff() { let mut request = Request::builder() .uri("/") .body(axum::body::Body::empty()) .expect("request"); request .extensions_mut() .insert(axum::extract::connect_info::MockConnectInfo( std::net::SocketAddr::from(([203, 0, 113, 1], 12345)), )); assert_eq!(extract_client_ip(&request), "203.0.113.1"); } #[test] fn hops_zero_uses_connect_not_xff() { assert_eq!( extract_ip(Some("1.2.3.4"), Some("10.0.0.9".into()), 0), "10.0.0.9" ); } #[test] fn hops_one_takes_rightmost_xff() { assert_eq!( extract_ip(Some("1.2.3.4, 5.6.7.8"), Some("10.0.0.9".into()), 1), "5.6.7.8" ); } #[test] fn hops_two_takes_second_from_right() { assert_eq!( extract_ip( Some("1.2.3.4, 5.6.7.8, 9.9.9.9"), Some("10.0.0.9".into()), 2 ), "5.6.7.8" ); } }