nidus_http/middleware/
rate_limit.rs

1use std::{
2    collections::HashMap,
3    future::Future,
4    pin::Pin,
5    sync::{Arc, Mutex},
6    task::{Context, Poll},
7    time::{Duration, Instant},
8};
9
10use axum::{body::Body, extract::Request};
11use http::{HeaderValue, Response, StatusCode};
12use tower::{Layer, Service};
13
14use crate::context::{IdentityExtractor, RequestIdentity};
15
16type IdentityFn =
17    Arc<dyn Fn(&http::request::Parts) -> Option<RequestIdentity> + Send + Sync + 'static>;
18
19/// Error returned by rate-limit stores.
20#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
21#[error("{message}")]
22pub struct RateLimitError {
23    message: String,
24}
25
26impl RateLimitError {
27    /// Creates a rate-limit store error.
28    pub fn new(message: impl Into<String>) -> Self {
29        Self {
30            message: message.into(),
31        }
32    }
33}
34
35/// Decision returned by a rate-limit store.
36#[derive(Clone, Debug, Eq, PartialEq)]
37pub struct RateLimitDecision {
38    /// Whether the request is allowed.
39    pub allowed: bool,
40    /// Limit for the active window.
41    pub limit: u64,
42    /// Remaining requests in the active window.
43    pub remaining: u64,
44    /// Seconds until the window resets.
45    pub reset_after: Duration,
46}
47
48/// Store adapter used by rate-limit layers.
49pub trait RateLimitStore: Send + Sync + 'static {
50    /// Checks and consumes one request for an identity.
51    fn check(
52        &self,
53        identity: &RequestIdentity,
54        limit: u64,
55        window: Duration,
56    ) -> Result<RateLimitDecision, RateLimitError>;
57}
58
59/// In-memory rate-limit store intended for local development and single-process apps.
60///
61/// The store tracks counters in process memory and opportunistically removes
62/// expired identity windows during [`RateLimitStore::check`], scanning at most
63/// once per window duration so a large identity set does not add a full-map
64/// sweep to every request. Rate-limit decisions never depend on that sweep:
65/// each identity's window resets itself on its next check. The store is not
66/// shared across processes, not durable across restarts, and not a substitute
67/// for a distributed limiter at multi-instance production boundaries.
68#[derive(Clone, Default)]
69pub struct InMemoryRateLimitStore {
70    state: Arc<Mutex<StoreState>>,
71}
72
73#[derive(Default)]
74struct StoreState {
75    windows: HashMap<String, WindowState>,
76    last_prune: Option<Instant>,
77}
78
79impl InMemoryRateLimitStore {
80    /// Creates an empty in-memory rate-limit store.
81    pub fn new() -> Self {
82        Self::default()
83    }
84
85    /// Returns the number of identity windows currently retained by the store.
86    ///
87    /// Expired windows are pruned opportunistically during
88    /// [`RateLimitStore::check`] (at most one full sweep per window duration),
89    /// so this value is mainly useful for tests, diagnostics, and local tools.
90    pub fn len(&self) -> usize {
91        self.state
92            .lock()
93            .map(|state| state.windows.len())
94            .unwrap_or_default()
95    }
96
97    /// Returns whether the store currently retains no identity windows.
98    pub fn is_empty(&self) -> bool {
99        self.len() == 0
100    }
101}
102
103impl RateLimitStore for InMemoryRateLimitStore {
104    fn check(
105        &self,
106        identity: &RequestIdentity,
107        limit: u64,
108        window: Duration,
109    ) -> Result<RateLimitDecision, RateLimitError> {
110        let now = Instant::now();
111        let mut state = self
112            .state
113            .lock()
114            .map_err(|_| RateLimitError::new("rate limit store poisoned"))?;
115        // Garbage-collect expired windows at most once per window duration.
116        // Decisions never depend on this sweep: each identity's window resets
117        // itself below when its own expiry has passed.
118        if state
119            .last_prune
120            .is_none_or(|last_prune| now.duration_since(last_prune) >= window)
121        {
122            state
123                .windows
124                .retain(|_, window_state| now.duration_since(window_state.started_at) < window);
125            state.last_prune = Some(now);
126        }
127        let window_state = state
128            .windows
129            .entry(identity.as_str().to_owned())
130            .or_insert_with(|| WindowState {
131                started_at: now,
132                count: 0,
133            });
134        if now.duration_since(window_state.started_at) >= window {
135            window_state.started_at = now;
136            window_state.count = 0;
137        }
138
139        let allowed = window_state.count < limit;
140        if allowed {
141            window_state.count += 1;
142        }
143        let remaining = limit.saturating_sub(window_state.count);
144        let reset_after = window.saturating_sub(now.duration_since(window_state.started_at));
145        Ok(RateLimitDecision {
146            allowed,
147            limit,
148            remaining,
149            reset_after,
150        })
151    }
152}
153
154#[derive(Clone)]
155struct WindowState {
156    started_at: Instant,
157    count: u64,
158}
159
160/// Typed config for production-shaped rate limiting.
161#[derive(Clone)]
162pub struct RateLimitConfig {
163    limit: u64,
164    window: Duration,
165    store: Arc<dyn RateLimitStore>,
166    identity: IdentityFn,
167    fail_open: bool,
168}
169
170impl RateLimitConfig {
171    /// Creates a rate-limit config with an explicit store.
172    pub fn new(limit: u64, window: Duration, store: impl RateLimitStore) -> Self {
173        Self {
174            limit,
175            window,
176            store: Arc::new(store),
177            identity: Arc::new(|_parts| Some(RequestIdentity::new("anonymous"))),
178            fail_open: true,
179        }
180    }
181
182    /// Replaces the identity extractor.
183    pub fn identity(mut self, extractor: impl IdentityExtractor) -> Self {
184        self.identity = Arc::new(move |parts| extractor.extract(parts));
185        self
186    }
187
188    /// Allows requests when the backing store fails.
189    pub fn fail_open(mut self) -> Self {
190        self.fail_open = true;
191        self
192    }
193
194    /// Rejects requests when the backing store fails.
195    pub fn fail_closed(mut self) -> Self {
196        self.fail_open = false;
197        self
198    }
199
200    /// Creates a Tower layer from this config.
201    pub fn layer(self) -> RateLimitLayer {
202        RateLimitLayer { config: self }
203    }
204}
205
206/// Tower layer that applies configured rate limiting.
207#[derive(Clone)]
208pub struct RateLimitLayer {
209    config: RateLimitConfig,
210}
211
212impl<S> Layer<S> for RateLimitLayer {
213    type Service = RateLimitService<S>;
214
215    fn layer(&self, inner: S) -> Self::Service {
216        RateLimitService {
217            inner,
218            config: self.config.clone(),
219        }
220    }
221}
222
223/// Service produced by [`RateLimitLayer`].
224#[derive(Clone)]
225pub struct RateLimitService<S> {
226    inner: S,
227    config: RateLimitConfig,
228}
229
230impl<S> Service<Request> for RateLimitService<S>
231where
232    S: Service<Request, Response = Response<Body>> + Send + 'static,
233    S::Future: Send + 'static,
234    S::Error: Send + 'static,
235{
236    type Response = Response<Body>;
237    type Error = S::Error;
238    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
239
240    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
241        self.inner.poll_ready(cx)
242    }
243
244    fn call(&mut self, request: Request) -> Self::Future {
245        let config = self.config.clone();
246        let (parts, body) = request.into_parts();
247        let identity =
248            (config.identity)(&parts).unwrap_or_else(|| RequestIdentity::new("anonymous"));
249        let decision = config
250            .store
251            .check(&identity, config.limit, config.window)
252            .unwrap_or(RateLimitDecision {
253                allowed: config.fail_open,
254                limit: config.limit,
255                remaining: if config.fail_open { config.limit } else { 0 },
256                reset_after: config.window,
257            });
258
259        if !decision.allowed {
260            return Box::pin(async move { Ok(rate_limited_response(decision)) });
261        }
262
263        let future = self.inner.call(Request::from_parts(parts, body));
264        Box::pin(async move {
265            let mut response = future.await?;
266            insert_rate_limit_headers(response.headers_mut(), &decision);
267            Ok(response)
268        })
269    }
270}
271
272fn rate_limited_response(decision: RateLimitDecision) -> Response<Body> {
273    let mut response = Response::new(Body::from("rate limit exceeded"));
274    *response.status_mut() = StatusCode::TOO_MANY_REQUESTS;
275    insert_rate_limit_headers(response.headers_mut(), &decision);
276    response.headers_mut().insert(
277        http::header::RETRY_AFTER,
278        HeaderValue::from(decision.reset_after.as_secs().max(1)),
279    );
280    response
281}
282
283fn insert_rate_limit_headers(headers: &mut http::HeaderMap, decision: &RateLimitDecision) {
284    headers.insert("ratelimit-limit", HeaderValue::from(decision.limit));
285    headers.insert("ratelimit-remaining", HeaderValue::from(decision.remaining));
286    headers.insert(
287        "ratelimit-reset",
288        HeaderValue::from(decision.reset_after.as_secs().max(1)),
289    );
290}
nidus_http/middleware/rate_limit.rs

nidus_http/middleware/
rate_limit.rs
