cratestack_axum/ratelimit/rest_ops_filter.rs
1//! Rate-limit filter for REST transport: check if a route is exempt
2//! from rate limiting based on its `RouteTransportDescriptor.rate_limited_by_default`.
3//!
4//! REST resolves op identity very differently from RPC (`rpc_ops_filter`):
5//! there is no single `/rpc/{op_id}` path segment to read, so this filter
6//! keys off [`axum::extract::MatchedPath`] — the route *pattern* the
7//! request matched (e.g. `/widgets/{id}`), not the concrete request path
8//! (e.g. `/widgets/42`). `RouteTransportDescriptor::path` is emitted in
9//! that same `{param}` shape (see `cratestack-macros/src/transport/rest.rs`),
10//! so the two compare directly with no path-param parsing needed.
11//!
12//! # `Router::layer` and `Router::route_layer` both work
13//!
14//! Unlike some middleware, `MatchedPath` is populated for either mount
15//! method — axum's `Router::layer` applies the middleware to each
16//! endpoint's `Route` individually (same as `route_layer`), it just
17//! *also* wraps the router's fallback (404) service, which `route_layer`
18//! does not. Practically: with `route_layer`, a request that matches no
19//! route skips this filter (and the rate limiter) entirely; with `layer`,
20//! it still runs the filter, finds no match, and fails closed (rate-limits
21//! the 404). Either is safe — pick based on whether 404s should count
22//! against the budget.
23use axum::extract::{MatchedPath, Request};
24use cratestack_core::RouteTransportDescriptor;
25
26/// Build a rate-limit filter function for REST schemas.
27///
28/// Returns a function that:
29/// - Reads the matched route pattern via [`MatchedPath`] (populated by
30/// axum under both `Router::layer` and `Router::route_layer` — see
31/// module docs).
32/// - Looks up the route (matched pattern + HTTP method) in the provided
33/// descriptors.
34/// - Returns `false` (exempt) if `rate_limited_by_default` is false.
35/// - Returns `true` (apply rate limit) if the route participates, or if
36/// lookup fails for any reason.
37///
38/// Fails closed: if `MatchedPath` is absent (the request matched no
39/// route — a 404) or the route isn't found in `routes` (a schema/router
40/// mismatch), the request is rate-limited. This prevents accidental
41/// exemptions from missing data or misconfiguration.
42pub fn build_rest_ops_filter(
43 routes: &'static [RouteTransportDescriptor],
44) -> impl Fn(&Request) -> bool + Send + Sync {
45 move |req: &Request| {
46 let Some(matched) = req.extensions().get::<MatchedPath>() else {
47 // No matched path: the request hit no route (a 404). Fail closed.
48 return true;
49 };
50 let path = matched.as_str();
51 let method = req.method().as_str();
52
53 match routes
54 .iter()
55 .find(|route| route.method == method && route.path == path)
56 {
57 Some(route) => route.rate_limited_by_default,
58 None => true,
59 }
60 }
61}
62
63#[cfg(test)]
64mod tests {
65 use axum::Router;
66 use axum::body::Body;
67 use axum::extract::ConnectInfo;
68 use axum::http::{Request as HttpRequest, StatusCode};
69 use axum::routing::{get, post};
70 use cratestack_core::{RouteTransportCapabilities, RouteTransportDescriptor};
71 use tower::ServiceExt;
72
73 use super::build_rest_ops_filter;
74 use crate::ratelimit::{InMemoryRateLimitStore, RateLimitConfig, RateLimitLayer};
75
76 /// Every request in this module that actually reaches the rate-limit
77 /// store (i.e. isn't exempted by the filter under test) needs a
78 /// verifiable caller identity — cratestack#416 made the default key fn
79 /// refuse requests with neither an `Authorization` header nor a
80 /// `ConnectInfo<SocketAddr>` peer, and `oneshot` never populates
81 /// `ConnectInfo` on its own.
82 fn with_peer(mut req: HttpRequest<Body>) -> HttpRequest<Body> {
83 let peer: std::net::SocketAddr = "192.0.2.50:1".parse().unwrap();
84 req.extensions_mut().insert(ConnectInfo(peer));
85 req
86 }
87
88 const CAPS: RouteTransportCapabilities = RouteTransportCapabilities {
89 request_types: &[],
90 response_types: &[],
91 default_response_type: "",
92 supports_sequence_response: false,
93 };
94
95 const ROUTES: &[RouteTransportDescriptor] = &[
96 RouteTransportDescriptor {
97 name: "createPayment",
98 method: "POST",
99 path: "/$procs/createPayment",
100 capabilities: CAPS,
101 idempotent_by_default: false,
102 rate_limited_by_default: false,
103 },
104 RouteTransportDescriptor {
105 name: "Widget",
106 method: "GET",
107 path: "/widgets/{id}",
108 capabilities: CAPS,
109 idempotent_by_default: true,
110 rate_limited_by_default: true,
111 },
112 ];
113
114 async fn ok() -> &'static str {
115 "ok"
116 }
117
118 fn app() -> Router {
119 Router::new()
120 .route("/$procs/createPayment", post(ok))
121 .route("/widgets/{id}", get(ok))
122 .route_layer(
123 RateLimitLayer::new(
124 std::sync::Arc::new(InMemoryRateLimitStore::default()),
125 RateLimitConfig::new(1, 0.001),
126 )
127 .with_should_rate_limit_fn(build_rest_ops_filter(ROUTES)),
128 )
129 }
130
131 /// AC1/AC2 mirrored for REST: an exempt route survives past its
132 /// burst, an un-annotated route (with a path param, proving
133 /// `MatchedPath` — not the concrete request path — is what's
134 /// compared) is throttled.
135 #[tokio::test]
136 async fn route_layer_exempts_annotated_route_and_throttles_others() {
137 let router = app();
138
139 for _ in 0..3 {
140 let resp = router
141 .clone()
142 .oneshot(
143 HttpRequest::post("/$procs/createPayment")
144 .body(Body::empty())
145 .unwrap(),
146 )
147 .await
148 .unwrap();
149 assert_eq!(
150 resp.status(),
151 StatusCode::OK,
152 "@no_rate_limit route should never be throttled"
153 );
154 }
155
156 let first = router
157 .clone()
158 .oneshot(with_peer(
159 HttpRequest::get("/widgets/42").body(Body::empty()).unwrap(),
160 ))
161 .await
162 .unwrap();
163 assert_eq!(first.status(), StatusCode::OK, "first request within burst");
164
165 let second = router
166 .clone()
167 .oneshot(with_peer(
168 HttpRequest::get("/widgets/7").body(Body::empty()).unwrap(),
169 ))
170 .await
171 .unwrap();
172 assert_eq!(
173 second.status(),
174 StatusCode::TOO_MANY_REQUESTS,
175 "un-annotated route is throttled regardless of the concrete id in the path, \
176 proving the comparison is against the matched route pattern"
177 );
178 }
179
180 /// `Router::layer` (unlike `route_layer`) also wraps the fallback
181 /// (404) service, so a request that matches no route still runs
182 /// through the filter, finds no `MatchedPath`, and fails closed —
183 /// consuming budget even for a 404. This proves the filter never
184 /// fails open on an unmatched path.
185 #[tokio::test]
186 async fn plain_layer_fails_closed_and_throttles_unmatched_paths() {
187 let router = Router::new()
188 .route("/$procs/createPayment", post(ok))
189 .layer(
190 RateLimitLayer::new(
191 std::sync::Arc::new(InMemoryRateLimitStore::default()),
192 RateLimitConfig::new(1, 0.001),
193 )
194 .with_should_rate_limit_fn(build_rest_ops_filter(ROUTES)),
195 );
196
197 let first = router
198 .clone()
199 .oneshot(with_peer(
200 HttpRequest::get("/does/not/exist")
201 .body(Body::empty())
202 .unwrap(),
203 ))
204 .await
205 .unwrap();
206 assert_eq!(
207 first.status(),
208 StatusCode::NOT_FOUND,
209 "first 404 within burst still reaches the fallback"
210 );
211
212 let second = router
213 .clone()
214 .oneshot(with_peer(
215 HttpRequest::get("/does/not/exist")
216 .body(Body::empty())
217 .unwrap(),
218 ))
219 .await
220 .unwrap();
221 assert_eq!(
222 second.status(),
223 StatusCode::TOO_MANY_REQUESTS,
224 "Router::layer wraps the fallback too, so an unmatched path has no \
225 MatchedPath, fails closed, and is throttled just like any other route"
226 );
227 }
228
229 /// `Router::route_layer`'s counterpart: it does *not* wrap the
230 /// fallback, so unmatched paths bypass the rate limiter (and this
231 /// filter) entirely — every 404 stays a 404, never a 429, and never
232 /// consumes budget.
233 #[tokio::test]
234 async fn route_layer_does_not_throttle_unmatched_paths() {
235 let router = Router::new()
236 .route("/$procs/createPayment", post(ok))
237 .route_layer(
238 RateLimitLayer::new(
239 std::sync::Arc::new(InMemoryRateLimitStore::default()),
240 RateLimitConfig::new(1, 0.001),
241 )
242 .with_should_rate_limit_fn(build_rest_ops_filter(ROUTES)),
243 );
244
245 for _ in 0..5 {
246 let resp = router
247 .clone()
248 .oneshot(
249 HttpRequest::get("/does/not/exist")
250 .body(Body::empty())
251 .unwrap(),
252 )
253 .await
254 .unwrap();
255 assert_eq!(
256 resp.status(),
257 StatusCode::NOT_FOUND,
258 "route_layer skips the fallback entirely, so repeated unmatched \
259 requests are always 404, never 429"
260 );
261 }
262 }
263}