mentedb_server/
rate_limit.rs

1//! Token bucket rate limiter implemented as a tower Layer/Service.
2
3use std::collections::HashMap;
4use std::net::IpAddr;
5use std::sync::Arc;
6use std::task::{Context, Poll};
7use std::time::Instant;
8
9use axum::body::Body;
10use axum::http::{Request, StatusCode};
11use axum::response::IntoResponse;
12use serde_json::json;
13use tokio::sync::Mutex;
14use tower::{Layer, Service};
15
16/// Per-IP token bucket.
17struct Bucket {
18    tokens: u32,
19    last_refill: Instant,
20    max_tokens: u32,
21    refill_rate: u32,
22}
23
24impl Bucket {
25    fn new(max_tokens: u32, refill_rate: u32) -> Self {
26        Self {
27            tokens: max_tokens,
28            last_refill: Instant::now(),
29            max_tokens,
30            refill_rate,
31        }
32    }
33
34    /// Refill tokens based on elapsed time and try to consume one.
35    fn try_consume(&mut self) -> bool {
36        let now = Instant::now();
37        let elapsed = now.duration_since(self.last_refill);
38        let refill = (elapsed.as_secs_f64() * f64::from(self.refill_rate)) as u32;
39        if refill > 0 {
40            self.tokens = (self.tokens + refill).min(self.max_tokens);
41            self.last_refill = now;
42        }
43        if self.tokens > 0 {
44            self.tokens -= 1;
45            true
46        } else {
47            false
48        }
49    }
50}
51
52/// Shared rate limiter state.
53#[derive(Clone)]
54pub struct RateLimiter {
55    buckets: Arc<Mutex<HashMap<IpAddr, Bucket>>>,
56    max_tokens: u32,
57    refill_rate: u32,
58}
59
60impl RateLimiter {
61    /// Creates a new rate limiter with the given bucket capacity and refill rate.
62    pub fn new(max_tokens: u32, refill_rate: u32) -> Self {
63        Self {
64            buckets: Arc::new(Mutex::new(HashMap::new())),
65            max_tokens,
66            refill_rate,
67        }
68    }
69}
70
71impl Default for RateLimiter {
72    fn default() -> Self {
73        Self::new(100, 10)
74    }
75}
76
77// -- Tower Layer --
78
79impl<S> Layer<S> for RateLimiter {
80    type Service = RateLimitService<S>;
81
82    fn layer(&self, inner: S) -> Self::Service {
83        RateLimitService {
84            inner,
85            limiter: self.clone(),
86        }
87    }
88}
89
90// -- Tower Service --
91
92#[derive(Clone)]
93pub struct RateLimitService<S> {
94    inner: S,
95    limiter: RateLimiter,
96}
97
98impl<S> Service<Request<Body>> for RateLimitService<S>
99where
100    S: Service<Request<Body>, Response = axum::response::Response> + Clone + Send + 'static,
101    S::Future: Send + 'static,
102{
103    type Response = axum::response::Response;
104    type Error = S::Error;
105    type Future = std::pin::Pin<
106        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
107    >;
108
109    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
110        self.inner.poll_ready(cx)
111    }
112
113    fn call(&mut self, req: Request<Body>) -> Self::Future {
114        let limiter = self.limiter.clone();
115        let mut inner = self.inner.clone();
116
117        // Extract client IP from connection info or forwarded headers.
118        let ip = req
119            .extensions()
120            .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
121            .map(|ci| ci.0.ip())
122            .unwrap_or(IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED));
123
124        Box::pin(async move {
125            let allowed = {
126                let mut buckets = limiter.buckets.lock().await;
127                let bucket = buckets
128                    .entry(ip)
129                    .or_insert_with(|| Bucket::new(limiter.max_tokens, limiter.refill_rate));
130                bucket.try_consume()
131            };
132
133            if allowed {
134                inner.call(req).await
135            } else {
136                let body = json!({ "error": "too many requests" });
137                Ok((StatusCode::TOO_MANY_REQUESTS, axum::Json(body)).into_response())
138            }
139        })
140    }
141}
mentedb_server/rate_limit.rs

mentedb_server/
rate_limit.rs
