cratestack_axum/ratelimit/
service.rs

1//! The `tower::Service` half of [`super::RateLimitLayer`].
2//!
3//! Split from `layer.rs` for the workspace's 200-line ceiling when
4//! cratestack#871 added the bucket-budget knobs to the builder. The
5//! per-request body itself lives one module further out, in
6//! `super::consume`.
7
8use std::sync::Arc;
9use std::time::Duration;
10
11use axum::extract::Request;
12use axum::response::Response;
13use tower::Service;
14
15use super::budget::warn::BudgetWarnings;
16use super::config::RateLimitConfig;
17use super::layer::KeyFn;
18use super::policy::{StoreErrorPolicy, StoreErrorWarnings};
19use super::store::RateLimitStore;
20
21#[derive(Clone)]
22pub struct RateLimitService<S> {
23    pub(super) inner: S,
24    pub(super) store: Arc<dyn RateLimitStore>,
25    pub(super) config: RateLimitConfig,
26    pub(super) key_fn: KeyFn,
27    pub(super) should_rate_limit_fn: Arc<dyn Fn(&Request) -> bool + Send + Sync>,
28    pub(super) store_error_policy: StoreErrorPolicy,
29    pub(super) store_timeout: Duration,
30    pub(super) warnings: Arc<StoreErrorWarnings>,
31    pub(super) budget_warnings: Arc<BudgetWarnings>,
32}
33
34impl<S> Service<Request> for RateLimitService<S>
35where
36    S: Service<Request, Response = Response, Error = std::convert::Infallible>
37        + Clone
38        + Send
39        + 'static,
40    S::Future: Send + 'static,
41{
42    type Response = Response;
43    type Error = std::convert::Infallible;
44    type Future =
45        std::pin::Pin<Box<dyn std::future::Future<Output = Result<Response, Self::Error>> + Send>>;
46
47    fn poll_ready(
48        &mut self,
49        cx: &mut std::task::Context<'_>,
50    ) -> std::task::Poll<Result<(), Self::Error>> {
51        self.inner.poll_ready(cx)
52    }
53
54    fn call(&mut self, req: Request) -> Self::Future {
55        let should_rate_limit = (self.should_rate_limit_fn)(&req);
56        // Clone the whole service, not just `inner`: the async body needs
57        // the store, the key fn and both warning budgets, and cloning
58        // once is cheaper than seven `Arc::clone`s at the call site.
59        let mut service = self.clone();
60        Box::pin(async move {
61            // If the operation is exempt from rate limiting, skip the check
62            // entirely — including key derivation. An exempt route must not
63            // be refused just because the default key fn can't verify the
64            // caller's identity; only routes that actually need a bucket
65            // pay that cost.
66            if !should_rate_limit {
67                return service.inner.call(req).await;
68            }
69            Ok(super::consume::run(service, req).await)
70        })
71    }
72}
cratestack_axum/ratelimit/service.rs

cratestack_axum/ratelimit/
service.rs
