d5939db
use std::sync::Arc;
use axum::extract::{Request, State};
use axum::http::header::{HeaderValue, SET_COOKIE};
use axum::http::{HeaderMap, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::Json;
use axum_extra::extract::cookie::{Cookie, SameSite};
use cache::RedisCache;
use serde_json::json;
use time::OffsetDateTime;
use uuid::Uuid;
use crate::config::RateLimitConfig;
pub const X_RATE_LIMIT_LIMIT: &str = "x-ratelimit-limit";
pub const X_RATE_LIMIT_REMAINING: &str = "x-ratelimit-remaining";
pub const X_RATE_LIMIT_RESET: &str = "x-ratelimit-reset";
pub const RETRY_AFTER: &str = "retry-after";
/// Supplies the Redis handle and settings used by the rate-limit middleware.
///
/// Implemented by the domain services (e.g. `AuthService`) so the middleware
/// can run with `from_fn_with_state` against their router state while keeping
/// an explicit, per-API partition key derived from the request path.
pub trait RateLimitProvider: Send + Sync + Clone + 'static {
fn rate_limit_kv(&self) -> Arc<RedisCache>;
fn rate_limit_settings(&self) -> Arc<RateLimitConfig>;
}
/// The resolved rate-limit client id (the `rl_id` cookie value) for the
/// current request. Injected into request extensions by [`ratelimit`] so
/// handlers can bind issued PoW challenges to the same client.
#[derive(Debug, Clone)]
pub struct ClientId(pub String);
fn read_cookie(headers: &HeaderMap, name: &str) -> Option<String> {
let header = headers.get("cookie")?;
let value = header.to_str().ok()?;
for c in Cookie::split_parse(value).flatten() {
if c.name() == name {
return Some(c.value().to_string());
}
}
None
}
fn write_cookie(headers: &mut HeaderMap, name: &str, value: &str, secure: bool) {
let cookie = Cookie::build((name.to_string(), value.to_string()))
.http_only(true)
.same_site(SameSite::Strict)
.path("/")
.secure(secure)
.build();
if let Ok(hv) = HeaderValue::from_str(&cookie.to_string()) {
headers.insert(SET_COOKIE, hv);
}
}
/// Sliding-window rate-limit middleware.
///
/// Redis ZSET layout: `rate:{api}:{rl_id}` where `api` is the request path
/// (per-API isolation, no global limit) and `rl_id` is the dedicated
/// rate-limit cookie value. ZSET member = `{request_id}:{url}`, score =
/// `unix_seconds(req_time)`. Entries older than `window_secs` are pruned
/// before counting, so only requests within the last window are considered.
pub async fn ratelimit<S: RateLimitProvider>(
State(svc): State<S>,
mut request: Request,
next: Next,
) -> Response {
let kv = svc.rate_limit_kv();
let settings = svc.rate_limit_settings();
let now = OffsetDateTime::now_utc().unix_timestamp();
let window_start = now - settings.window_secs;
let existing = read_cookie(request.headers(), &settings.cookie_name);
let rl_id = existing.clone().unwrap_or_else(|| Uuid::new_v4().to_string());
let api = request.uri().path();
let key = format!("rate:{}:{}", api, rl_id);
let url = request
.uri()
.path_and_query()
.map(|p| p.as_str().to_string())
.unwrap_or_default();
let member = format!("{}:{}", Uuid::new_v4(), url);
let ttl = (settings.window_secs + 60) as u64;
let _ = kv.zadd(&key, &member, now as f64, ttl).await;
let _ = kv.zremrange_by_score_lt(&key, window_start as f64).await;
let count = kv.zcard(&key).await.unwrap_or(0);
let reset = now + settings.window_secs;
let remaining = count.saturating_sub(settings.default_limit).saturating_sub(1);
let mut response: Response;
if count > settings.default_limit {
response = (StatusCode::TOO_MANY_REQUESTS, Json(json!({ "error": "rate limit exceeded" })))
.into_response();
} else {
request.extensions_mut().insert(ClientId(rl_id.clone()));
response = next.run(request).await;
}
let headers = response.headers_mut();
if let Ok(v) = HeaderValue::from_str(&settings.default_limit.to_string()) {
headers.insert(X_RATE_LIMIT_LIMIT, v);
}
if let Ok(v) = HeaderValue::from_str(&remaining.to_string()) {
headers.insert(X_RATE_LIMIT_REMAINING, v);
}
if let Ok(v) = HeaderValue::from_str(&reset.to_string()) {
headers.insert(X_RATE_LIMIT_RESET, v);
}
if count > settings.default_limit
&& let Ok(v) = HeaderValue::from_str(&settings.window_secs.to_string())
{
headers.insert(RETRY_AFTER, v);
}
if existing.is_none() {
write_cookie(response.headers_mut(), &settings.cookie_name, &rl_id, settings.cookie_secure);
}
response
}