axum_limit/
snapshot.rs

1use crate::time::{millis_to_secs, saturating_sub_ms};
2use http::header::{HeaderName, HeaderValue, RETRY_AFTER};
3use http::HeaderMap;
4
5/// Result of a rate limit check, including metadata for standard response headers.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub struct RateLimitSnapshot {
8    /// Whether the request is allowed.
9    pub allowed: bool,
10    /// Configured maximum requests per period.
11    pub limit: usize,
12    /// Remaining requests in the current window or bucket.
13    pub remaining: usize,
14    /// UTC timestamp in milliseconds when the limit fully resets.
15    pub reset_at_ms: u64,
16}
17
18impl RateLimitSnapshot {
19    /// Returns how long to wait before retrying from `now_ms`.
20    pub fn retry_after_ms(&self, now_ms: u64) -> u64 {
21        saturating_sub_ms(self.reset_at_ms, now_ms)
22    }
23
24    /// Returns `Retry-After` in whole seconds, at least 1 when limited.
25    pub fn retry_after_secs(&self, now_ms: u64) -> u64 {
26        self.retry_after_ms(now_ms).div_ceil(1000).max(1)
27    }
28
29    /// Builds standard rate limit headers from this snapshot.
30    pub fn to_headers(&self, now_ms: u64) -> HeaderMap {
31        let mut headers = HeaderMap::new();
32        insert_header(
33            &mut headers,
34            HeaderName::from_static("x-ratelimit-limit"),
35            self.limit,
36        );
37        insert_header(
38            &mut headers,
39            HeaderName::from_static("x-ratelimit-remaining"),
40            self.remaining,
41        );
42        insert_header(
43            &mut headers,
44            HeaderName::from_static("x-ratelimit-reset"),
45            millis_to_secs(self.reset_at_ms),
46        );
47
48        if !self.allowed {
49            insert_header(&mut headers, RETRY_AFTER, self.retry_after_secs(now_ms));
50        }
51
52        headers
53    }
54}
55
56/// Extension type stored in request parts after a successful limit check.
57#[derive(Debug, Clone, Copy)]
58pub struct RateLimitInfo(pub RateLimitSnapshot);
59
60fn insert_header(headers: &mut HeaderMap, name: HeaderName, value: impl ToString) {
61    if let Ok(value) = HeaderValue::from_str(&value.to_string()) {
62        headers.insert(name, value);
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    use crate::time::now_ms;
71
72    #[test]
73    fn headers_include_limit_metadata() {
74        let now_ms = now_ms();
75        let snapshot = RateLimitSnapshot {
76            allowed: false,
77            limit: 10,
78            remaining: 0,
79            reset_at_ms: now_ms + 30_000,
80        };
81
82        let headers = snapshot.to_headers(now_ms);
83        assert_eq!(
84            headers
85                .get("x-ratelimit-limit")
86                .and_then(|v| v.to_str().ok()),
87            Some("10")
88        );
89        assert_eq!(
90            headers
91                .get("x-ratelimit-remaining")
92                .and_then(|v| v.to_str().ok()),
93            Some("0")
94        );
95        assert!(headers.get(RETRY_AFTER).is_some());
96    }
97}
axum_limit/snapshot.rs

axum_limit/
snapshot.rs
