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