Spaces:
Sleeping
Sleeping
| use axum::{ | |
| body::Body, | |
| http::{HeaderValue, Request, StatusCode}, | |
| middleware::Next, | |
| response::Response, | |
| }; | |
| use dashmap::DashMap; | |
| use once_cell::sync::Lazy; | |
| use std::sync::Arc; | |
| use std::time::{Duration, Instant}; | |
| /// High-performance rate limiting store for sensitive endpoints. | |
| pub struct SmartRateLimitStore { | |
| store: Arc<DashMap<String, Vec<Instant>>>, | |
| max_requests: usize, | |
| window_secs: u64, | |
| } | |
| impl SmartRateLimitStore { | |
| pub fn new(max_requests: usize, window_secs: u64) -> Self { | |
| let store: Arc<DashMap<String, Vec<Instant>>> = Arc::new(DashMap::new()); | |
| // Background cleanup for the smart limiter | |
| let cleanup_store = store.clone(); | |
| if let Ok(handle) = tokio::runtime::Handle::try_current() { | |
| handle.spawn(async move { | |
| let mut interval = tokio::time::interval(Duration::from_secs(300)); | |
| loop { | |
| interval.tick().await; | |
| let now = Instant::now(); | |
| let window = Duration::from_secs(window_secs); | |
| cleanup_store.retain(|_, timestamps| { | |
| timestamps.retain(|&t| now.duration_since(t) < window); | |
| !timestamps.is_empty() | |
| }); | |
| } | |
| }); | |
| } | |
| Self { | |
| store, | |
| max_requests, | |
| window_secs, | |
| } | |
| } | |
| pub fn is_allowed(&self, key: &str) -> bool { | |
| let now = Instant::now(); | |
| let window = Duration::from_secs(self.window_secs); | |
| if !self.store.contains_key(key) && self.store.len() > 50_000 { | |
| self.store.retain(|_, timestamps| { | |
| timestamps.retain(|&t| now.duration_since(t) < window); | |
| !timestamps.is_empty() | |
| }); | |
| if self.store.len() > 50_000 { | |
| return false; | |
| } | |
| } | |
| let mut entry = self.store.entry(key.to_string()).or_default(); | |
| let timestamps = entry.value_mut(); | |
| // Remove expired timestamps | |
| timestamps.retain(|&t| now.duration_since(t) < window); | |
| if timestamps.len() < self.max_requests { | |
| timestamps.push(now); | |
| true | |
| } else { | |
| false | |
| } | |
| } | |
| } | |
| /// Static store for checkout-specific rate limiting (10 requests per 60 seconds) | |
| static CHECKOUT_LIMITER: Lazy<SmartRateLimitStore> = Lazy::new(|| SmartRateLimitStore::new(10, 60)); | |
| pub async fn smart_rate_limiter(req: Request<Body>, next: Next) -> Result<Response, StatusCode> { | |
| let client_ip = crate::interfaces::http::middleware::get_client_ip(&req); | |
| if std::env::var("RENDER").is_err() { | |
| return Ok(next.run(req).await); | |
| } | |
| if !CHECKOUT_LIMITER.is_allowed(&client_ip) { | |
| tracing::warn!(ip = %client_ip, "Smart rate limit exceeded for checkout flow"); | |
| let mut response = Response::new(Body::from( | |
| "Security Alert: Rapid checkout attempts detected. Please wait 60 seconds.", | |
| )); | |
| *response.status_mut() = StatusCode::TOO_MANY_REQUESTS; | |
| response | |
| .headers_mut() | |
| .insert("Retry-After", HeaderValue::from_static("60")); | |
| return Ok(response); | |
| } | |
| Ok(next.run(req).await) | |
| } | |
| mod tests { | |
| use super::*; | |
| fn test_smart_rate_limiter_logic() { | |
| let limiter = SmartRateLimitStore::new(2, 1); // 2 reqs per 1 sec | |
| let key = "test-ip"; | |
| assert!(limiter.is_allowed(key)); | |
| assert!(limiter.is_allowed(key)); | |
| assert!(!limiter.is_allowed(key)); // Third request should be blocked | |
| // Wait for window to expire (simplified for unit test without actual sleep) | |
| // In a real test we'd use tokio::time::pause/advance | |
| } | |
| } | |