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        if let Some(window_state) = state.windows.get_mut(identity.as_str()) {
128            return Ok(window_state.check(now, limit, window));
129        }
130
131        let mut window_state = WindowState {
132            started_at: now,
133            count: 0,
134        };
135        let decision = window_state.check(now, limit, window);
136        state
137            .windows
138            .insert(identity.as_str().to_owned(), window_state);
139        Ok(decision)
140    }
141}
142
143#[derive(Clone)]
144struct WindowState {
145    started_at: Instant,
146    count: u64,
147}
148
149impl WindowState {
150    fn check(&mut self, now: Instant, limit: u64, window: Duration) -> RateLimitDecision {
151        if now.duration_since(self.started_at) >= window {
152            self.started_at = now;
153            self.count = 0;
154        }
155
156        let allowed = self.count < limit;
157        if allowed {
158            self.count += 1;
159        }
160        let remaining = limit.saturating_sub(self.count);
161        let reset_after = window.saturating_sub(now.duration_since(self.started_at));
162        RateLimitDecision {
163            allowed,
164            limit,
165            remaining,
166            reset_after,
167        }
168    }
169}
170
171/// Typed config for production-shaped rate limiting.
172#[derive(Clone)]
173pub struct RateLimitConfig {
174    limit: u64,
175    window: Duration,
176    store: Arc<dyn RateLimitStore>,
177    identity: IdentityFn,
178    fail_open: bool,
179}
180
181impl RateLimitConfig {
182    /// Creates a rate-limit config with an explicit store.
183    pub fn new(limit: u64, window: Duration, store: impl RateLimitStore) -> Self {
184        Self {
185            limit,
186            window,
187            store: Arc::new(store),
188            identity: Arc::new(|_parts| Some(RequestIdentity::new("anonymous"))),
189            fail_open: true,
190        }
191    }
192
193    /// Replaces the identity extractor.
194    pub fn identity(mut self, extractor: impl IdentityExtractor) -> Self {
195        self.identity = Arc::new(move |parts| extractor.extract(parts));
196        self
197    }
198
199    /// Allows requests when the backing store fails.
200    pub fn fail_open(mut self) -> Self {
201        self.fail_open = true;
202        self
203    }
204
205    /// Rejects requests when the backing store fails.
206    pub fn fail_closed(mut self) -> Self {
207        self.fail_open = false;
208        self
209    }
210
211    /// Creates a Tower layer from this config.
212    pub fn layer(self) -> RateLimitLayer {
213        RateLimitLayer { config: self }
214    }
215}
216
217/// Tower layer that applies configured rate limiting.
218#[derive(Clone)]
219pub struct RateLimitLayer {
220    config: RateLimitConfig,
221}
222
223impl<S> Layer<S> for RateLimitLayer {
224    type Service = RateLimitService<S>;
225
226    fn layer(&self, inner: S) -> Self::Service {
227        RateLimitService {
228            inner,
229            config: self.config.clone(),
230        }
231    }
232}
233
234/// Service produced by [`RateLimitLayer`].
235#[derive(Clone)]
236pub struct RateLimitService<S> {
237    inner: S,
238    config: RateLimitConfig,
239}
240
241impl<S> Service<Request> for RateLimitService<S>
242where
243    S: Service<Request, Response = Response<Body>> + Send + 'static,
244    S::Future: Send + 'static,
245    S::Error: Send + 'static,
246{
247    type Response = Response<Body>;
248    type Error = S::Error;
249    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
250
251    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
252        self.inner.poll_ready(cx)
253    }
254
255    fn call(&mut self, request: Request) -> Self::Future {
256        let config = &self.config;
257        let (parts, body) = request.into_parts();
258        let identity =
259            (config.identity)(&parts).unwrap_or_else(|| RequestIdentity::new("anonymous"));
260        let decision = config
261            .store
262            .check(&identity, config.limit, config.window)
263            .unwrap_or(RateLimitDecision {
264                allowed: config.fail_open,
265                limit: config.limit,
266                remaining: if config.fail_open { config.limit } else { 0 },
267                reset_after: config.window,
268            });
269
270        if !decision.allowed {
271            return Box::pin(async move { Ok(rate_limited_response(decision)) });
272        }
273
274        let future = self.inner.call(Request::from_parts(parts, body));
275        Box::pin(async move {
276            let mut response = future.await?;
277            insert_rate_limit_headers(response.headers_mut(), &decision);
278            Ok(response)
279        })
280    }
281}
282
283fn rate_limited_response(decision: RateLimitDecision) -> Response<Body> {
284    let mut response = Response::new(Body::from("rate limit exceeded"));
285    *response.status_mut() = StatusCode::TOO_MANY_REQUESTS;
286    insert_rate_limit_headers(response.headers_mut(), &decision);
287    response.headers_mut().insert(
288        http::header::RETRY_AFTER,
289        HeaderValue::from(decision.reset_after.as_secs().max(1)),
290    );
291    response
292}
293
294fn insert_rate_limit_headers(headers: &mut http::HeaderMap, decision: &RateLimitDecision) {
295    headers.insert("ratelimit-limit", HeaderValue::from(decision.limit));
296    headers.insert("ratelimit-remaining", HeaderValue::from(decision.remaining));
297    headers.insert(
298        "ratelimit-reset",
299        HeaderValue::from(decision.reset_after.as_secs().max(1)),
300    );
301}
302
303#[cfg(test)]
304mod tests {
305    use std::{
306        convert::Infallible,
307        sync::{
308            Weak,
309            atomic::{AtomicUsize, Ordering},
310        },
311    };
312
313    use tower::{ServiceExt, service_fn};
314
315    use super::*;
316
317    struct StoreCloneProbe {
318        observed_strong_count: Arc<AtomicUsize>,
319        owner: Mutex<Option<Weak<dyn RateLimitStore>>>,
320    }
321
322    impl RateLimitStore for StoreCloneProbe {
323        fn check(
324            &self,
325            _identity: &RequestIdentity,
326            limit: u64,
327            window: Duration,
328        ) -> Result<RateLimitDecision, RateLimitError> {
329            let owner = self.owner.lock().unwrap();
330            self.observed_strong_count
331                .store(owner.as_ref().unwrap().strong_count(), Ordering::SeqCst);
332            Ok(RateLimitDecision {
333                allowed: true,
334                limit,
335                remaining: limit.saturating_sub(1),
336                reset_after: window,
337            })
338        }
339    }
340
341    #[tokio::test]
342    async fn service_borrows_the_config_store_during_synchronous_policy_work() {
343        let observed_strong_count = Arc::new(AtomicUsize::new(0));
344        let probe = Arc::new(StoreCloneProbe {
345            observed_strong_count: Arc::clone(&observed_strong_count),
346            owner: Mutex::new(None),
347        });
348        let store: Arc<dyn RateLimitStore> = probe.clone();
349        *probe.owner.lock().unwrap() = Some(Arc::downgrade(&store));
350        drop(probe);
351
352        let service = RateLimitService {
353            inner: service_fn(|_request: Request| async {
354                Ok::<_, Infallible>(Response::new(Body::empty()))
355            }),
356            config: RateLimitConfig {
357                limit: 10,
358                window: Duration::from_secs(60),
359                store,
360                identity: Arc::new(|_parts| Some(RequestIdentity::new("test"))),
361                fail_open: true,
362            },
363        };
364
365        let response = service.oneshot(Request::new(Body::empty())).await.unwrap();
366
367        assert_eq!(response.status(), StatusCode::OK);
368        assert_eq!(observed_strong_count.load(Ordering::SeqCst), 1);
369    }
370}
nidus_http/middleware/rate_limit.rs

nidus_http/middleware/
rate_limit.rs
