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 cratestack_exec::OpExecutor;
14use tower::Service;
15
16use super::budget::warn::BudgetWarnings;
17use super::config::RateLimitConfig;
18use super::layer::{KeyFn, OpResolver};
19use super::policy::{StoreErrorPolicy, StoreErrorWarnings};
20
21#[derive(Clone)]
22pub struct RateLimitService<S> {
23    pub(super) inner: S,
24    /// Owns the decision (ADR 0015 slice 2): whether an op is limited at
25    /// all, and the store call. Everything else here is transport.
26    pub(super) executor: OpExecutor,
27    /// Still read for the `X-RateLimit-Limit` header on allowed responses.
28    pub(super) config: RateLimitConfig,
29    pub(super) key_fn: KeyFn,
30    pub(super) op_resolver: OpResolver,
31    pub(super) store_error_policy: StoreErrorPolicy,
32    pub(super) store_timeout: Duration,
33    pub(super) warnings: Arc<StoreErrorWarnings>,
34    pub(super) budget_warnings: Arc<BudgetWarnings>,
35}
36
37impl<S> Service<Request> for RateLimitService<S>
38where
39    S: Service<Request, Response = Response, Error = std::convert::Infallible>
40        + Clone
41        + Send
42        + 'static,
43    S::Future: Send + 'static,
44{
45    type Response = Response;
46    type Error = std::convert::Infallible;
47    type Future =
48        std::pin::Pin<Box<dyn std::future::Future<Output = Result<Response, Self::Error>> + Send>>;
49
50    fn poll_ready(
51        &mut self,
52        cx: &mut std::task::Context<'_>,
53    ) -> std::task::Poll<Result<(), Self::Error>> {
54        self.inner.poll_ready(cx)
55    }
56
57    fn call(&mut self, req: Request) -> Self::Future {
58        let op = (self.op_resolver)(&req);
59        let should_rate_limit = self.executor.rate_limit_applies(&op);
60        // Clone the whole service, not just `inner`: the async body needs
61        // the executor, the key fn and both warning budgets, and cloning
62        // once is cheaper than seven `Arc::clone`s at the call site.
63        let mut service = self.clone();
64        Box::pin(async move {
65            // If the operation is exempt from rate limiting, skip the check
66            // entirely — including key derivation. An exempt route must not
67            // be refused just because the default key fn can't verify the
68            // caller's identity; only routes that actually need a bucket
69            // pay that cost.
70            if !should_rate_limit {
71                return service.inner.call(req).await;
72            }
73            Ok(super::consume::run(service, req, op).await)
74        })
75    }
76}
cratestack_axum/ratelimit/service.rs

cratestack_axum/ratelimit/
service.rs
