Spaces:
Sleeping
Sleeping
File size: 3,789 Bytes
7856e60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | 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)
}
#[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
}
}
|