refactor: migrate monolithic crate to Cargo Workspace with Clean Architecture
Transform the single binary crate into a 9-crate workspace monorepo: - Root Cargo.toml as [workspace] manager with resolver = "2" - zesdex-entities: Domain entity types (session, settings, store, message, etc.) - zesdex-utils: Pure utility functions (error, logger, pagination, slug, clipboard) - zesdex-dto: Data Transfer Objects for LLM provider API communication - zesdex-ipc: Unix-socket IPC layer (client/server/framing/protocol) - zesdex-iam: Identity & Access Management (Clean Architecture: domain/application/infrastructure) - zesdex-cms: Content Management (Clean Architecture: domain/application/infrastructure) - zesdex-middleware: HTTP middleware (Auth, CORS, Rate Limiting) - zesdex-libs: Composition root (AppContext, DB init, JWT, Argon2) - zesdex-backend: Main binary entry point + seed/migrate binaries - DevOps: Dockerfile, docker-compose, Nix (flake/shell/default), CI/CD updates - Remove dead root src/ and src-misc/ directories All crate re-exports maintain backward compatibility with original crate::model::*, crate::dto::*, crate::ipc::* module paths. Feature crates enforce strict layer separation: domain -> application -> infrastructure with generic trait-based dependency injection.
This commit is contained in:
asepharyana
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "zesdex-middleware"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
anyhow.workspace = true
|
||||
chrono.workspace = true
|
||||
zesdex-entities = { path = "../zesdex-entities" }
|
||||
zesdex-utils = { path = "../zesdex-utils" }
|
||||
axum.workspace = true
|
||||
tower.workspace = true
|
||||
tower-http = { workspace = true, features = ["cors", "limit"] }
|
||||
@@ -0,0 +1,294 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
//! Authentication middleware — session-lock based auth for Axum.
|
||||
//!
|
||||
//! Provides:
|
||||
//! - [`SessionAuthLayer`]: a tower [`Layer`] that injects session validation
|
||||
//! - [`SessionIdentity`]: extracted from validated requests
|
||||
//! - [`validate_session`]: low-level session existence/validity check
|
||||
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::http::header;
|
||||
use axum::http::request::Parts;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tower::{Layer, Service};
|
||||
use zesdex_entities::seaorm::common::store::Store;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SessionIdentity
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Identity extracted from a validated session token / lock.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SessionIdentity {
|
||||
/// The validated session id (from `X-Session-Id`).
|
||||
pub session_id: String,
|
||||
/// User-Agent header value, if present.
|
||||
pub user_agent: String,
|
||||
/// Unix-epoch timestamp (seconds) when the session was first seen by
|
||||
/// this middleware.
|
||||
pub connected_at: i64,
|
||||
}
|
||||
|
||||
impl SessionIdentity {
|
||||
/// Create a new identity from a validated session id.
|
||||
fn new(session_id: String, user_agent: String) -> Self {
|
||||
let connected_at = chrono::Utc::now().timestamp();
|
||||
Self {
|
||||
session_id,
|
||||
user_agent,
|
||||
connected_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extractor: pull the identity from request extensions.
|
||||
///
|
||||
/// If the identity has not been inserted by the middleware the request is
|
||||
/// rejected with 401 Unauthorized.
|
||||
impl<S: Send + Sync> FromRequestParts<S> for SessionIdentity {
|
||||
type Rejection = Response;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
|
||||
parts
|
||||
.extensions
|
||||
.get::<SessionIdentity>()
|
||||
.cloned()
|
||||
.ok_or_else(|| (StatusCode::UNAUTHORIZED, "session identity not found").into_response())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SessionAuthLayer
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Tower [`Layer`] that produces [`SessionAuthMiddleware`] services.
|
||||
///
|
||||
/// Wraps every request with session validation: if the `X-Session-Id`
|
||||
/// header points to a valid session, the request passes through and a
|
||||
/// [`SessionIdentity`] is injected into the request extensions. Otherwise
|
||||
/// a 401 response is returned immediately.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionAuthLayer {
|
||||
store: Arc<Store>,
|
||||
}
|
||||
|
||||
impl SessionAuthLayer {
|
||||
/// Create a new layer backed by the given [`Store`].
|
||||
pub fn new(store: Store) -> Self {
|
||||
Self {
|
||||
store: Arc::new(store),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience constructor using `Store::new()`.
|
||||
pub fn default() -> Self {
|
||||
Self::new(Store::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Layer<S> for SessionAuthLayer {
|
||||
type Service = SessionAuthMiddleware<S>;
|
||||
|
||||
fn layer(&self, inner: S) -> Self::Service {
|
||||
SessionAuthMiddleware {
|
||||
inner,
|
||||
store: Arc::clone(&self.store),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SessionAuthMiddleware
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Tower [`Service`] that validates `X-Session-Id` before forwarding.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SessionAuthMiddleware<S> {
|
||||
inner: S,
|
||||
store: Arc<Store>,
|
||||
}
|
||||
|
||||
impl<S, ReqBody> Service<Request<ReqBody>> for SessionAuthMiddleware<S>
|
||||
where
|
||||
S: Service<Request<ReqBody>, Response = Response> + Send + 'static,
|
||||
S::Future: Send + 'static,
|
||||
ReqBody: Send + 'static,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = S::Error;
|
||||
type Future =
|
||||
Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.inner.poll_ready(cx)
|
||||
}
|
||||
|
||||
fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
|
||||
let store = Arc::clone(&self.store);
|
||||
|
||||
// Extract session id from header.
|
||||
let session_id = req
|
||||
.headers()
|
||||
.get("X-Session-Id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let session_id = match session_id {
|
||||
Some(id) if !id.is_empty() => id,
|
||||
_ => {
|
||||
let resp = (StatusCode::UNAUTHORIZED, "missing X-Session-Id header")
|
||||
.into_response();
|
||||
return Box::pin(async move { Ok(resp) });
|
||||
}
|
||||
};
|
||||
|
||||
// Validate session.
|
||||
if let Err(e) = validate_session(&session_id, &store) {
|
||||
let resp = (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
format!("session validation failed: {e}"),
|
||||
)
|
||||
.into_response();
|
||||
return Box::pin(async move { Ok(resp) });
|
||||
}
|
||||
let user_agent = req
|
||||
.headers()
|
||||
.get(header::USER_AGENT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let identity = SessionIdentity::new(session_id, user_agent);
|
||||
req.extensions_mut().insert(identity);
|
||||
|
||||
let fut = self.inner.call(req);
|
||||
Box::pin(async move { fut.await })
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper: `require_session` (convenience middleware function)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Axum middleware function that validates `X-Session-Id` against the
|
||||
/// [`Store`] extracted from request extensions.
|
||||
///
|
||||
/// This is an alternative to [`SessionAuthLayer`] when you want to attach
|
||||
/// auth to a specific route group via `axum::middleware::from_fn_with_state`.
|
||||
pub async fn require_session(
|
||||
store: axum::extract::State<Store>,
|
||||
mut req: Request<axum::body::Body>,
|
||||
next: axum::middleware::Next,
|
||||
) -> Response {
|
||||
let session_id = req
|
||||
.headers()
|
||||
.get("X-Session-Id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let session_id = match session_id {
|
||||
Some(id) if !id.is_empty() => id,
|
||||
_ => {
|
||||
return (StatusCode::UNAUTHORIZED, "missing X-Session-Id header").into_response();
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = validate_session(&session_id, &store) {
|
||||
return (StatusCode::UNAUTHORIZED, format!("session validation failed: {e}")).into_response();
|
||||
}
|
||||
let user_agent = req
|
||||
.headers()
|
||||
.get(header::USER_AGENT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let identity = SessionIdentity::new(session_id, user_agent);
|
||||
req.extensions_mut().insert(identity);
|
||||
next.run(req).await
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Session validation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Check whether a session lock exists and is valid, returning the
|
||||
/// associated [`SessionIdentity`].
|
||||
///
|
||||
/// Validation logic:
|
||||
/// 1. Verify the session id is not a path-traversal attack.
|
||||
/// 2. Check that `<store.base_dir>/sessions/<id>/session.json` exists.
|
||||
/// 3. Deserialise the session metadata to confirm it is well-formed.
|
||||
///
|
||||
/// This is a synchronous, CPU-light check so it can be called directly
|
||||
/// inside tower service impls without spawning a blocking task.
|
||||
pub fn validate_session(session_id: &str, store: &Store) -> anyhow::Result<SessionIdentity> {
|
||||
// Directory-traversal prevention.
|
||||
if session_id.contains('/') || session_id.contains('\\') || session_id.contains("..") {
|
||||
anyhow::bail!("invalid session id: must not contain path separators");
|
||||
}
|
||||
|
||||
let session_path = store
|
||||
.base_dir
|
||||
.join("sessions")
|
||||
.join(session_id)
|
||||
.join("session.json");
|
||||
|
||||
if !session_path.exists() {
|
||||
anyhow::bail!("session not found: {session_id}");
|
||||
}
|
||||
|
||||
let _data = std::fs::read_to_string(&session_path)?;
|
||||
// We verify the JSON is well-formed by deserialising it.
|
||||
let _session: serde_json::Value = serde_json::from_str(&_data)?;
|
||||
|
||||
let user_agent = String::new();
|
||||
Ok(SessionIdentity::new(session_id.to_string(), user_agent))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_validate_session_rejects_path_traversal() {
|
||||
let store = Store::new();
|
||||
assert!(validate_session("../etc/passwd", &store).is_err());
|
||||
assert!(validate_session("foo/bar", &store).is_err());
|
||||
assert!(validate_session("foo\\bar", &store).is_err());
|
||||
assert!(validate_session("..", &store).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_session_nonexistent() {
|
||||
let store = Store::new();
|
||||
let result = validate_session("nonexistent-session-id", &store);
|
||||
assert!(result.is_err());
|
||||
assert!(result
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("session not found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_identity_creation() {
|
||||
let identity = SessionIdentity::new("sess-123".into(), "test-agent".into());
|
||||
assert_eq!(identity.session_id, "sess-123");
|
||||
assert_eq!(identity.user_agent, "test-agent");
|
||||
assert!(identity.connected_at > 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
//! CORS layer factory for the daemon HTTP (IPC) server.
|
||||
//!
|
||||
//! Since the daemon only listens on `127.0.0.1`, the CORS policy is
|
||||
//! intentionally permissive. These settings are still required because
|
||||
//! Axum rejects cross-origin requests unless a CORS layer is present.
|
||||
|
||||
use tower_http::cors::{AllowHeaders, AllowOrigin, CorsLayer};
|
||||
|
||||
/// Return a permissive [`CorsLayer`] for local daemon IPC.
|
||||
///
|
||||
/// - **Origin**: any (`*`)
|
||||
/// - **Methods**: `GET`, `POST`, `PUT`, `DELETE`, `PATCH`, `OPTIONS`
|
||||
/// - **Headers**: `Content-Type`, `Authorization`, `X-Session-Id`,
|
||||
/// `X-Request-Id`, `User-Agent`
|
||||
pub fn default_cors_layer() -> CorsLayer {
|
||||
CorsLayer::new()
|
||||
.allow_origin(AllowOrigin::any())
|
||||
.allow_methods([
|
||||
"GET".parse().unwrap(),
|
||||
"POST".parse().unwrap(),
|
||||
"PUT".parse().unwrap(),
|
||||
"DELETE".parse().unwrap(),
|
||||
"PATCH".parse().unwrap(),
|
||||
"OPTIONS".parse().unwrap(),
|
||||
])
|
||||
.allow_headers(AllowHeaders::any())
|
||||
.expose_headers([
|
||||
"Content-Type".parse().unwrap(),
|
||||
"X-Session-Id".parse().unwrap(),
|
||||
"X-Request-Id".parse().unwrap(),
|
||||
])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_cors_layer_constructs() {
|
||||
let _layer = default_cors_layer();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
|
||||
pub mod auth;
|
||||
pub mod cors;
|
||||
pub mod rate_limit;
|
||||
@@ -0,0 +1,268 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
//! Simple in-memory rate limiter for Axum.
|
||||
//!
|
||||
//! Uses a sliding-window approach: each client has a rolling list of
|
||||
//! timestamps. Requests arriving within the window that exceed the
|
||||
//! configured max are rejected.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Mutex;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::http::{Request, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use tower::{Layer, Service};
|
||||
|
||||
/// In-memory sliding-window rate limiter.
|
||||
///
|
||||
/// Thread-safe via interior mutability (`Mutex`). Each client (identified
|
||||
/// by a string key, e.g. IP address or session id) has a Vec of entry
|
||||
/// timestamps (in seconds). Old entries are cleaned on every check.
|
||||
#[derive(Debug)]
|
||||
pub struct RateLimiter {
|
||||
windows: Mutex<HashMap<String, Vec<i64>>>,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
/// Create an empty rate limiter.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
windows: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check whether a request from `client_id` should be allowed.
|
||||
///
|
||||
/// * `max_requests` — max number of requests permitted within the
|
||||
/// window.
|
||||
/// * `window_secs` — width of the sliding window in seconds.
|
||||
///
|
||||
/// Returns `Ok(true)` if the request is allowed (and records it),
|
||||
/// or `Ok(false)` if the client has exceeded the limit.
|
||||
///
|
||||
/// The window is **sliding**: only timestamps falling within
|
||||
/// `[now - window_secs, now]` are counted.
|
||||
pub fn check_rate_limit(
|
||||
&self,
|
||||
client_id: &str,
|
||||
max_requests: u32,
|
||||
window_secs: u64,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs() as i64;
|
||||
|
||||
let cutoff = now.saturating_sub(window_secs as i64);
|
||||
let mut windows = self
|
||||
.windows
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("rate limiter lock poisoned: {e}"))?;
|
||||
|
||||
let timestamps = windows.entry(client_id.to_string()).or_insert_with(Vec::new);
|
||||
|
||||
// Discard entries older than the window.
|
||||
timestamps.retain(|&ts| ts >= cutoff);
|
||||
|
||||
if timestamps.len() >= max_requests as usize {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
timestamps.push(now);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// Convenience wrapper that returns an Axum [`Response`] on rejection
|
||||
/// so it can be used directly in middleware.
|
||||
pub fn check_or_429(
|
||||
&self,
|
||||
client_id: &str,
|
||||
max_requests: u32,
|
||||
window_secs: u64,
|
||||
) -> Result<(), Response> {
|
||||
match self.check_rate_limit(client_id, max_requests, window_secs) {
|
||||
Ok(true) => Ok(()),
|
||||
Ok(false) => Err((
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
"rate limit exceeded, try again later",
|
||||
)
|
||||
.into_response()),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove all stored windows (for testing / reset).
|
||||
pub fn reset(&self) -> anyhow::Result<()> {
|
||||
let mut windows = self
|
||||
.windows
|
||||
.lock()
|
||||
.map_err(|e| anyhow::anyhow!("rate limiter lock poisoned: {e}"))?;
|
||||
windows.clear();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RateLimiter {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tower Layer / Service
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Configuration for the rate-limit middleware layer.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RateLimitLayer {
|
||||
limiter: std::sync::Arc<RateLimiter>,
|
||||
max_requests: u32,
|
||||
window_secs: u64,
|
||||
}
|
||||
|
||||
impl RateLimitLayer {
|
||||
/// Create a new layer with the given limits.
|
||||
///
|
||||
/// * `max_requests` — max requests per window per client.
|
||||
/// * `window_secs` — sliding-window width in seconds.
|
||||
pub fn new(max_requests: u32, window_secs: u64) -> Self {
|
||||
Self {
|
||||
limiter: std::sync::Arc::new(RateLimiter::new()),
|
||||
max_requests,
|
||||
window_secs,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a reference to the shared [`RateLimiter`] so callers can
|
||||
/// reset it or perform manual checks.
|
||||
pub fn limiter(&self) -> &std::sync::Arc<RateLimiter> {
|
||||
&self.limiter
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Layer<S> for RateLimitLayer {
|
||||
type Service = RateLimitMiddleware<S>;
|
||||
|
||||
fn layer(&self, inner: S) -> Self::Service {
|
||||
RateLimitMiddleware {
|
||||
inner,
|
||||
limiter: std::sync::Arc::clone(&self.limiter),
|
||||
max_requests: self.max_requests,
|
||||
window_secs: self.window_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tower [`Service`] wrapping each request with a rate-limit check.
|
||||
///
|
||||
/// Client identity is extracted from the `X-Forwarded-For` header first,
|
||||
/// falling back to the remote address, then to `"unknown"`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RateLimitMiddleware<S> {
|
||||
inner: S,
|
||||
limiter: std::sync::Arc<RateLimiter>,
|
||||
max_requests: u32,
|
||||
window_secs: u64,
|
||||
}
|
||||
|
||||
impl<S, ReqBody> Service<Request<ReqBody>> for RateLimitMiddleware<S>
|
||||
where
|
||||
S: Service<Request<ReqBody>, Response = Response> + Send + 'static,
|
||||
S::Future: Send + 'static,
|
||||
ReqBody: Send + 'static,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = S::Error;
|
||||
type Future =
|
||||
Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.inner.poll_ready(cx)
|
||||
}
|
||||
|
||||
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
|
||||
let client_id = req
|
||||
.headers()
|
||||
.get("X-Forwarded-For")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.split(',').next().unwrap_or(s).trim().to_string())
|
||||
.or_else(|| {
|
||||
req.headers()
|
||||
.get("X-Real-IP")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string())
|
||||
})
|
||||
.or_else(|| {
|
||||
req.extensions()
|
||||
.get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
|
||||
.map(|ci| ci.0.ip().to_string())
|
||||
})
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
|
||||
let limiter = std::sync::Arc::clone(&self.limiter);
|
||||
let max_requests = self.max_requests;
|
||||
let window_secs = self.window_secs;
|
||||
|
||||
match limiter.check_or_429(&client_id, max_requests, window_secs) {
|
||||
Ok(()) => {}
|
||||
Err(resp) => {
|
||||
return Box::pin(async move { Ok(resp) });
|
||||
}
|
||||
}
|
||||
|
||||
let fut = self.inner.call(req);
|
||||
Box::pin(async move { fut.await })
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_rate_limiter_allows_within_limit() {
|
||||
let limiter = RateLimiter::new();
|
||||
assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limiter_rejects_excess() {
|
||||
let limiter = RateLimiter::new();
|
||||
assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
|
||||
assert!(!limiter.check_rate_limit("client-2", 3, 60).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limiter_independent_clients() {
|
||||
let limiter = RateLimiter::new();
|
||||
assert!(limiter.check_rate_limit("alice", 2, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("alice", 2, 60).unwrap());
|
||||
assert!(!limiter.check_rate_limit("alice", 2, 60).unwrap());
|
||||
assert!(limiter.check_rate_limit("bob", 2, 60).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rate_limiter_reset() {
|
||||
let limiter = RateLimiter::new();
|
||||
assert!(limiter.check_rate_limit("client-3", 1, 60).unwrap());
|
||||
assert!(!limiter.check_rate_limit("client-3", 1, 60).unwrap());
|
||||
limiter.reset().unwrap();
|
||||
assert!(limiter.check_rate_limit("client-3", 1, 60).unwrap());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user