aster_forge_middleware/axum/
rate_limit.rs

1//! Axum rate limits with explicit peer/key and rejection-response boundaries.
2pub use crate::shared::rate_limit::{NormalizedStringRateLimiter, RateLimitRejection};
3use axum::{
4    extract::{Request, State},
5    middleware::Next,
6    response::Response,
7};
8use governor::{RateLimiter, clock::DefaultClock, state::keyed::DefaultKeyedStateStore};
9use ipnet::IpNet;
10use std::{
11    net::{IpAddr, Ipv4Addr},
12    num::{NonZeroU32, NonZeroU64},
13    sync::Arc,
14};
15
16type IpLimiter = RateLimiter<IpAddr, DefaultKeyedStateStore<IpAddr>, DefaultClock>;
17type PeerResolver = dyn Fn(&Request) -> Option<IpAddr> + Send + Sync;
18type RejectionMapper = dyn Fn(RateLimitRejection) -> Response + Send + Sync;
19type KeyResolver = dyn Fn(&Request) -> Result<String, Box<Response>> + Send + Sync;
20
21/// State for `from_fn_with_state(config, ip_rate_limit)`. Clones share the same quota buckets.
22#[derive(Clone)]
23pub struct IpRateLimitConfig {
24    enabled: bool,
25    trusted: Vec<IpNet>,
26    limiter: Arc<IpLimiter>,
27    peer: Arc<PeerResolver>,
28    rejection: Arc<RejectionMapper>,
29}
30
31impl IpRateLimitConfig {
32    /// Builds a shared IP limiter; the product supplies its 429 envelope and Retry-After policy.
33    pub fn new<F>(
34        enabled: bool,
35        seconds_per_request: NonZeroU64,
36        burst_size: NonZeroU32,
37        trusted_proxies: &[String],
38        rejection: F,
39    ) -> Self
40    where
41        F: Fn(RateLimitRejection) -> Response + Send + Sync + 'static,
42    {
43        Self {
44            enabled,
45            trusted: aster_forge_utils::net::parse_trusted_proxies(trusted_proxies),
46            limiter: Arc::new(RateLimiter::keyed(
47                crate::shared::rate_limit::rate_limit_quota(seconds_per_request, burst_size),
48            )),
49            peer: Arc::new(super::client_ip::direct_peer),
50            rejection: Arc::new(rejection),
51        }
52    }
53    /// Overrides direct-peer extraction for an alternate trusted transport extension.
54    /// This resolver must never infer a peer from unverified request headers.
55    #[must_use]
56    pub fn peer_resolver<P>(mut self, peer: P) -> Self
57    where
58        P: Fn(&Request) -> Option<IpAddr> + Send + Sync + 'static,
59    {
60        self.peer = Arc::new(peer);
61        self
62    }
63    /// Removes expired buckets. Products may invoke this from their maintenance runtime.
64    pub fn retain_recent(&self) {
65        self.limiter.retain_recent();
66    }
67}
68
69/// Applies trusted-proxy IP limits. Missing peers/UDS share a localhost bucket and ignore headers.
70/// Products should disable IP limits or use keyed limits on UDS deployments.
71pub async fn ip_rate_limit(
72    State(config): State<IpRateLimitConfig>,
73    request: Request,
74    next: Next,
75) -> Response {
76    if config.enabled {
77        let key = (config.peer)(&request).map_or(IpAddr::V4(Ipv4Addr::LOCALHOST), |peer| {
78            super::client_ip::real_ip_from_trusted_headers(request.headers(), peer, &config.trusted)
79        });
80        if let Err(negative) = config.limiter.check_key(&key) {
81            return (config.rejection)(RateLimitRejection::from_not_until(&negative));
82        }
83    }
84    next.run(request).await
85}
86
87/// State for keyed limits. Product code extracts the key; Forge trims and lowercases it.
88#[derive(Clone)]
89pub struct KeyedRateLimitConfig {
90    limiter: NormalizedStringRateLimiter,
91    key: Arc<KeyResolver>,
92    rejection: Arc<RejectionMapper>,
93}
94impl KeyedRateLimitConfig {
95    /// Supplies a shared limiter, product key extraction, and product rejection response.
96    pub fn new<K, F>(limiter: NormalizedStringRateLimiter, key: K, rejection: F) -> Self
97    where
98        K: Fn(&Request) -> Result<String, Box<Response>> + Send + Sync + 'static,
99        F: Fn(RateLimitRejection) -> Response + Send + Sync + 'static,
100    {
101        Self {
102            limiter,
103            key: Arc::new(key),
104            rejection: Arc::new(rejection),
105        }
106    }
107}
108/// Applies a normalized string quota, preserving product key-extraction errors.
109pub async fn keyed_rate_limit(
110    State(config): State<KeyedRateLimitConfig>,
111    request: Request,
112    next: Next,
113) -> Response {
114    if config.limiter.enabled() {
115        let key = match (config.key)(&request) {
116            Ok(key) => key,
117            Err(response) => return *response,
118        };
119        if let Some(rejection) = config.limiter.check(&key) {
120            return (config.rejection)(rejection);
121        }
122    }
123    next.run(request).await
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use axum::{
130        Json, Router,
131        body::{Body, to_bytes},
132        extract::ConnectInfo,
133        middleware,
134        response::IntoResponse,
135        routing::get,
136    };
137    use http::{Request as HttpRequest, StatusCode};
138    use std::net::SocketAddr;
139    use tower::ServiceExt;
140
141    fn reject(rejection: RateLimitRejection) -> Response {
142        let wait = rejection.retry_after_seconds();
143        let mut response = (
144            StatusCode::TOO_MANY_REQUESTS,
145            Json(serde_json::json!({"code": "product_limit", "retry_after": wait})),
146        )
147            .into_response();
148        response
149            .headers_mut()
150            .insert("retry-after", wait.to_string().parse().unwrap());
151        response
152    }
153    fn ip_app(enabled: bool, burst: u32, trusted: &[String]) -> Router {
154        let config = IpRateLimitConfig::new(
155            enabled,
156            NonZeroU64::new(60).unwrap(),
157            NonZeroU32::new(burst).unwrap(),
158            trusted,
159            reject,
160        );
161        Router::new()
162            .route("/", get(|| async { "called" }))
163            .layer(middleware::from_fn_with_state(config, ip_rate_limit))
164    }
165    fn request(peer: Option<&str>, forwarded: &str) -> HttpRequest<Body> {
166        let mut request = HttpRequest::builder()
167            .header("x-forwarded-for", forwarded)
168            .body(Body::empty())
169            .unwrap();
170        if let Some(peer) = peer {
171            request
172                .extensions_mut()
173                .insert(ConnectInfo(peer.parse::<SocketAddr>().unwrap()));
174        }
175        request
176    }
177
178    #[tokio::test]
179    async fn ip_limits_isolate_clients_and_preserve_custom_429_retry_after() {
180        let router = ip_app(true, 1, &["10.0.0.0/8".into()]);
181        assert_eq!(
182            router
183                .clone()
184                .oneshot(request(Some("10.0.0.5:1234"), "203.0.113.1"))
185                .await
186                .unwrap()
187                .status(),
188            StatusCode::OK
189        );
190        assert_eq!(
191            router
192                .clone()
193                .oneshot(request(Some("10.0.0.5:1234"), "203.0.113.2"))
194                .await
195                .unwrap()
196                .status(),
197            StatusCode::OK
198        );
199        let response = router
200            .oneshot(request(Some("10.0.0.5:1234"), "203.0.113.1"))
201            .await
202            .unwrap();
203        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
204        let delay: u64 = response.headers()["retry-after"]
205            .to_str()
206            .unwrap()
207            .parse()
208            .unwrap();
209        assert!((1..=60).contains(&delay));
210        let value: serde_json::Value =
211            serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap();
212        assert_eq!(value["code"], "product_limit");
213        assert_eq!(value["retry_after"], delay);
214    }
215
216    #[tokio::test]
217    async fn untrusted_and_missing_peers_cannot_rotate_forwarded_headers_to_bypass_limits() {
218        for peer in [Some("198.51.100.1:1234"), None] {
219            let router = ip_app(true, 1, &["127.0.0.1".into()]);
220            assert_eq!(
221                router
222                    .clone()
223                    .oneshot(request(peer, "203.0.113.1"))
224                    .await
225                    .unwrap()
226                    .status(),
227                StatusCode::OK
228            );
229            assert_eq!(
230                router
231                    .oneshot(request(peer, "203.0.113.2"))
232                    .await
233                    .unwrap()
234                    .status(),
235                StatusCode::TOO_MANY_REQUESTS
236            );
237        }
238    }
239
240    #[tokio::test]
241    async fn disabled_ip_limit_never_extracts_the_peer() {
242        let config = IpRateLimitConfig::new(
243            false,
244            NonZeroU64::new(60).unwrap(),
245            NonZeroU32::new(1).unwrap(),
246            &[],
247            reject,
248        )
249        .peer_resolver(|_| panic!("disabled limiter must not invoke peer resolver"));
250        let router = Router::new()
251            .route("/", get(|| async { "called" }))
252            .layer(middleware::from_fn_with_state(config, ip_rate_limit));
253        for _ in 0..3 {
254            assert_eq!(
255                router
256                    .clone()
257                    .oneshot(request(None, "anything"))
258                    .await
259                    .unwrap()
260                    .status(),
261                StatusCode::OK
262            );
263        }
264    }
265
266    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
267    async fn concurrent_requests_share_exactly_one_burst_across_router_clones() {
268        let router = ip_app(true, 4, &[]);
269        let mut jobs = Vec::new();
270        for _ in 0..32 {
271            let router = router.clone();
272            jobs.push(tokio::spawn(async move {
273                router
274                    .oneshot(request(Some("203.0.113.1:1234"), ""))
275                    .await
276                    .unwrap()
277                    .status()
278            }));
279        }
280        let mut allowed = 0;
281        for job in jobs {
282            match job.await.unwrap() {
283                StatusCode::OK => allowed += 1,
284                StatusCode::TOO_MANY_REQUESTS => (),
285                status => panic!("unexpected status {status}"),
286            }
287        }
288        assert_eq!(allowed, 4);
289    }
290
291    fn keyed_app(enabled: bool) -> Router {
292        let limiter = NormalizedStringRateLimiter::new(
293            enabled,
294            NonZeroU64::new(60).unwrap(),
295            NonZeroU32::new(1).unwrap(),
296        );
297        let config = KeyedRateLimitConfig::new(
298            limiter,
299            |request| {
300                request
301                    .headers()
302                    .get("x-key")
303                    .and_then(|value| value.to_str().ok())
304                    .map(str::to_string)
305                    .ok_or_else(|| {
306                        Box::new(
307                            (StatusCode::UNPROCESSABLE_ENTITY, "missing product key")
308                                .into_response(),
309                        )
310                    })
311            },
312            reject,
313        );
314        Router::new()
315            .route("/", get(|| async { "called" }))
316            .layer(middleware::from_fn_with_state(config, keyed_rate_limit))
317    }
318    #[tokio::test]
319    async fn string_keys_are_normalized_isolated_and_product_extraction_errors_are_retained() {
320        let router = keyed_app(true);
321        for (key, expected) in [
322            ("User@Example.com", StatusCode::OK),
323            ("other@example.com", StatusCode::OK),
324            (" user@example.com ", StatusCode::TOO_MANY_REQUESTS),
325        ] {
326            let request = HttpRequest::builder()
327                .header("x-key", key)
328                .body(Body::empty())
329                .unwrap();
330            assert_eq!(
331                router.clone().oneshot(request).await.unwrap().status(),
332                expected
333            );
334        }
335        let response = router
336            .oneshot(HttpRequest::new(Body::empty()))
337            .await
338            .unwrap();
339        assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
340        assert_eq!(
341            to_bytes(response.into_body(), 4096).await.unwrap().as_ref(),
342            b"missing product key"
343        );
344    }
345    #[tokio::test]
346    async fn disabled_keyed_limit_bypasses_product_key_extraction() {
347        let router = keyed_app(false);
348        for _ in 0..3 {
349            assert_eq!(
350                router
351                    .clone()
352                    .oneshot(HttpRequest::new(Body::empty()))
353                    .await
354                    .unwrap()
355                    .status(),
356                StatusCode::OK
357            );
358        }
359    }
360}
aster_forge_middleware/axum/rate_limit.rs

aster_forge_middleware/axum/
rate_limit.rs
