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>>, max_requests: usize, window_secs: u64, } impl SmartRateLimitStore { pub fn new(max_requests: usize, window_secs: u64) -> Self { let store: Arc>> = 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 = Lazy::new(|| SmartRateLimitStore::new(10, 60)); pub async fn smart_rate_limiter(req: Request, next: Next) -> Result { 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) } #[cfg(test)] mod tests { use super::*; #[test] 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 } }