cratestack_axum/ratelimit/
layer.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use axum::extract::Request;
5use cratestack_core::CratestackError;
6use tower::Layer;
7
8use super::budget::RateLimitBucketBudget;
9use super::budget::warn::BudgetWarnings;
10use super::config::RateLimitConfig;
11use super::key_fn::{default_key_fn, default_should_rate_limit_fn};
12use super::policy::{DEFAULT_STORE_TIMEOUT, StoreErrorPolicy, StoreErrorWarnings};
13use super::scope::{KeyDerivation, UnverifiedAuthPolicy};
14use super::service::RateLimitService;
15use super::store::RateLimitStore;
16
17pub(super) type KeyFn =
18    Arc<dyn Fn(&Request) -> Result<KeyDerivation, CratestackError> + Send + Sync>;
19
20#[derive(Clone)]
21pub struct RateLimitLayer {
22    store: Arc<dyn RateLimitStore>,
23    config: RateLimitConfig,
24    key_fn: Option<KeyFn>,
25    should_rate_limit_fn: Arc<dyn Fn(&Request) -> bool + Send + Sync>,
26    store_error_policy: StoreErrorPolicy,
27    store_timeout: Duration,
28    bucket_budget: Option<RateLimitBucketBudget>,
29    unverified_auth_policy: UnverifiedAuthPolicy,
30    warnings: Arc<StoreErrorWarnings>,
31    budget_warnings: Arc<BudgetWarnings>,
32}
33
34impl RateLimitLayer {
35    pub fn new(store: Arc<dyn RateLimitStore>, config: RateLimitConfig) -> Self {
36        Self {
37            store,
38            config,
39            key_fn: None,
40            should_rate_limit_fn: Arc::new(default_should_rate_limit_fn),
41            store_error_policy: StoreErrorPolicy::default(),
42            store_timeout: DEFAULT_STORE_TIMEOUT,
43            bucket_budget: Some(RateLimitBucketBudget::default()),
44            unverified_auth_policy: UnverifiedAuthPolicy::default(),
45            warnings: Arc::new(StoreErrorWarnings::default()),
46            budget_warnings: Arc::new(BudgetWarnings::default()),
47        }
48    }
49
50    /// Choose what happens when the backing store itself fails, as
51    /// opposed to when a caller is genuinely over budget. Defaults to
52    /// [`StoreErrorPolicy::Allow`], which serves through **transport-class
53    /// failures only** — see that type's docs for the distinction, why a
54    /// reachable-but-refusing store stays closed regardless, and why key
55    /// derivation deliberately does not follow suit.
56    pub fn with_store_error_policy(mut self, policy: StoreErrorPolicy) -> Self {
57        self.store_error_policy = policy;
58        self
59    }
60
61    /// Ceiling on how long one store lookup may take before the layer
62    /// gives up and applies [`StoreErrorPolicy`] to a synthetic
63    /// transport-class error. Defaults to [`DEFAULT_STORE_TIMEOUT`].
64    ///
65    /// This is ONE budget for the whole lookup, including any retry the
66    /// backend performs internally — the point is to bound what the
67    /// caller waits, and a per-attempt budget silently doubles when a
68    /// store retries. Without it, "degrade to unlimited" degrades only
69    /// after the driver's own reconnect cycle finishes, which was
70    /// measured at nineteen seconds per request against a real outage.
71    pub fn with_store_timeout(mut self, timeout: Duration) -> Self {
72        self.store_timeout = timeout;
73        self
74    }
75
76    /// Tune how many distinct buckets one scope may create
77    /// (cratestack#871). Defaults to [`RateLimitBucketBudget::default`].
78    pub fn with_bucket_budget(mut self, budget: RateLimitBucketBudget) -> Self {
79        self.bucket_budget = Some(budget);
80        self
81    }
82
83    /// Let an unverified `Authorization` header mint buckets without any
84    /// cardinality bound — the pre-cratestack#871 behaviour.
85    ///
86    /// Only correct when something else already bounds the keyspace (an
87    /// authenticating proxy in front, a `with_key_fn` that keys on
88    /// verified material, mTLS). Otherwise this restores the measured
89    /// amplification primitive: one store key per request, attacker-chosen.
90    pub fn without_bucket_budget(mut self) -> Self {
91        self.bucket_budget = None;
92        self
93    }
94
95    /// What the default key function does with an `Authorization` header
96    /// nothing has verified. See [`UnverifiedAuthPolicy`].
97    pub fn with_unverified_auth_policy(mut self, policy: UnverifiedAuthPolicy) -> Self {
98        self.unverified_auth_policy = policy;
99        self
100    }
101
102    /// Override how the layer derives the bucket key. The supplied closure
103    /// is infallible by design — opting out of the default's fail-closed
104    /// behavior is the caller's explicit choice, including any deliberate
105    /// shared bucket.
106    ///
107    /// An override carries **no** bucket budget: the layer has no basis to
108    /// invent a scope or a fallback for a key whose derivation it cannot
109    /// see. A consumer whose key function reads caller-supplied material
110    /// owns bounding it, exactly as it owns the fail-closed decision.
111    pub fn with_key_fn(mut self, f: impl Fn(&Request) -> String + Send + Sync + 'static) -> Self {
112        self.key_fn = Some(Arc::new(move |req| Ok(KeyDerivation::unbudgeted(f(req)))));
113        self
114    }
115
116    pub fn with_should_rate_limit_fn(
117        mut self,
118        f: impl Fn(&Request) -> bool + Send + Sync + 'static,
119    ) -> Self {
120        self.should_rate_limit_fn = Arc::new(f);
121        self
122    }
123
124    /// Test seam: the warning counters this layer shares with every
125    /// service it builds. Lets a test assert that the per-request path
126    /// actually *called* `consume::report`, which deleting outright used
127    /// to leave every test green (cratestack#871 review, should-fix 3).
128    pub(super) fn _budget_warnings(&self) -> &BudgetWarnings {
129        &self.budget_warnings
130    }
131
132    /// Bind the layer's configuration into the closure `RateLimitService`
133    /// calls, so the per-request path never has to branch on "default or
134    /// override" again.
135    fn resolved_key_fn(&self) -> KeyFn {
136        if let Some(key_fn) = &self.key_fn {
137            return key_fn.clone();
138        }
139        let budget = self.bucket_budget;
140        let policy = self.unverified_auth_policy;
141        let warnings = self.budget_warnings.clone();
142        Arc::new(move |req| match budget {
143            Some(budget) => default_key_fn(req, budget, policy, &warnings),
144            // `without_bucket_budget()`: derive exactly as before, then
145            // drop the budget rather than skipping derivation, so the
146            // key SHAPE (and therefore every existing bucket) is
147            // untouched by the opt-out.
148            None => default_key_fn(req, RateLimitBucketBudget::default(), policy, &warnings)
149                .map(|derivation| KeyDerivation::unbudgeted(derivation.key)),
150        })
151    }
152}
153
154impl<S> Layer<S> for RateLimitLayer {
155    type Service = RateLimitService<S>;
156
157    fn layer(&self, inner: S) -> Self::Service {
158        RateLimitService {
159            inner,
160            store: self.store.clone(),
161            config: self.config,
162            key_fn: self.resolved_key_fn(),
163            should_rate_limit_fn: self.should_rate_limit_fn.clone(),
164            store_error_policy: self.store_error_policy,
165            store_timeout: self.store_timeout,
166            warnings: self.warnings.clone(),
167            budget_warnings: self.budget_warnings.clone(),
168        }
169    }
170}
cratestack_axum/ratelimit/layer.rs

cratestack_axum/ratelimit/
layer.rs
