Skip to main content

api_bones/
ratelimit.rs

1//! Rate limit metadata types.
2//!
3//! [`RateLimitInfo`] carries the structured data normally surfaced through
4//! `X-RateLimit-*` HTTP response headers, making it easy to include quota
5//! information in both successful responses and 429 error bodies.
6//!
7//! # Example
8//!
9//! ```rust
10//! use api_bones::ratelimit::RateLimitInfo;
11//!
12//! let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(60);
13//! assert!(info.is_exceeded());
14//! assert_eq!(info.retry_after, Some(60));
15//! ```
16
17#[cfg(feature = "serde")]
18use serde::{Deserialize, Serialize};
19
20// ---------------------------------------------------------------------------
21// RateLimitInfo
22// ---------------------------------------------------------------------------
23
24/// Structured rate-limit metadata matching `X-RateLimit-*` headers.
25///
26/// | Field           | HTTP header              | Meaning                               |
27/// |-----------------|--------------------------|---------------------------------------|
28/// | `limit`         | `X-RateLimit-Limit`      | Max requests allowed in the window    |
29/// | `remaining`     | `X-RateLimit-Remaining`  | Requests still available              |
30/// | `reset`         | `X-RateLimit-Reset`      | Unix timestamp when the window resets |
31/// | `retry_after`   | `Retry-After`            | Seconds to wait before retrying (429) |
32#[derive(Debug, Clone, PartialEq, Eq)]
33#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
34#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
35#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
36#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
37#[cfg_attr(feature = "proptest", derive(proptest_derive::Arbitrary))]
38pub struct RateLimitInfo {
39    /// Maximum number of requests allowed in the current window.
40    pub limit: u64,
41
42    /// Number of requests remaining in the current window.
43    pub remaining: u64,
44
45    /// Unix timestamp (seconds) at which the current window resets.
46    pub reset: u64,
47
48    /// Seconds the client should wait before retrying (present on 429 responses).
49    #[cfg_attr(
50        feature = "serde",
51        serde(default, skip_serializing_if = "Option::is_none")
52    )]
53    pub retry_after: Option<u64>,
54}
55
56impl RateLimitInfo {
57    /// Create a new `RateLimitInfo`.
58    ///
59    /// # Examples
60    ///
61    /// ```rust
62    /// use api_bones::ratelimit::RateLimitInfo;
63    ///
64    /// let info = RateLimitInfo::new(100, 50, 1_700_000_000);
65    /// assert_eq!(info.limit, 100);
66    /// assert_eq!(info.remaining, 50);
67    /// assert!(info.retry_after.is_none());
68    /// ```
69    #[must_use]
70    pub fn new(limit: u64, remaining: u64, reset: u64) -> Self {
71        Self {
72            limit,
73            remaining,
74            reset,
75            retry_after: None,
76        }
77    }
78
79    /// Set the `retry_after` hint (builder-style).
80    ///
81    /// # Examples
82    ///
83    /// ```rust
84    /// use api_bones::ratelimit::RateLimitInfo;
85    ///
86    /// let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(60);
87    /// assert_eq!(info.retry_after, Some(60));
88    /// ```
89    #[must_use]
90    pub fn retry_after(mut self, seconds: u64) -> Self {
91        self.retry_after = Some(seconds);
92        self
93    }
94
95    /// Return `true` when no requests remain in the current window.
96    ///
97    /// # Examples
98    ///
99    /// ```rust
100    /// use api_bones::ratelimit::RateLimitInfo;
101    ///
102    /// let exceeded = RateLimitInfo::new(100, 0, 1_700_000_000);
103    /// assert!(exceeded.is_exceeded());
104    ///
105    /// let available = RateLimitInfo::new(100, 50, 1_700_000_000);
106    /// assert!(!available.is_exceeded());
107    /// ```
108    #[must_use]
109    pub fn is_exceeded(&self) -> bool {
110        self.remaining == 0
111    }
112}
113
114// ---------------------------------------------------------------------------
115// Axum: header extraction / injection
116// ---------------------------------------------------------------------------
117
118#[cfg(feature = "http")]
119mod http_impl {
120    use super::RateLimitInfo;
121    #[cfg(not(feature = "std"))]
122    use alloc::string::ToString;
123    use http::{HeaderMap, HeaderValue};
124
125    /// Header name constants.
126    pub const HEADER_LIMIT: &str = "x-ratelimit-limit";
127    pub const HEADER_REMAINING: &str = "x-ratelimit-remaining";
128    pub const HEADER_RESET: &str = "x-ratelimit-reset";
129    pub const HEADER_RETRY_AFTER: &str = "retry-after";
130
131    impl RateLimitInfo {
132        /// Inject rate-limit headers into a [`HeaderMap`].
133        ///
134        /// Inserts `X-RateLimit-Limit`, `X-RateLimit-Remaining`, and
135        /// `X-RateLimit-Reset`.  Also inserts `Retry-After` when
136        /// [`retry_after`](RateLimitInfo::retry_after) is set.
137        pub fn inject_headers(&self, headers: &mut HeaderMap) {
138            // These values are u64 formatted to ASCII digits — infallible.
139            headers.insert(
140                HEADER_LIMIT,
141                HeaderValue::from_str(&self.limit.to_string())
142                    .expect("u64 decimal is always a valid header value"),
143            );
144            headers.insert(
145                HEADER_REMAINING,
146                HeaderValue::from_str(&self.remaining.to_string())
147                    .expect("u64 decimal is always a valid header value"),
148            );
149            headers.insert(
150                HEADER_RESET,
151                HeaderValue::from_str(&self.reset.to_string())
152                    .expect("u64 decimal is always a valid header value"),
153            );
154            if let Some(secs) = self.retry_after {
155                headers.insert(
156                    HEADER_RETRY_AFTER,
157                    HeaderValue::from_str(&secs.to_string())
158                        .expect("u64 decimal is always a valid header value"),
159                );
160            }
161        }
162
163        /// Extract a `RateLimitInfo` from a [`HeaderMap`].
164        ///
165        /// Returns `None` if any of the three required headers
166        /// (`X-RateLimit-Limit`, `X-RateLimit-Remaining`, `X-RateLimit-Reset`)
167        /// are missing or cannot be parsed as `u64`.
168        #[must_use]
169        pub fn from_headers(headers: &HeaderMap) -> Option<Self> {
170            let parse = |name| -> Option<u64> {
171                headers
172                    .get(name)
173                    .and_then(|v| v.to_str().ok())
174                    .and_then(|s| s.parse().ok())
175            };
176
177            let limit = parse(HEADER_LIMIT)?;
178            let remaining = parse(HEADER_REMAINING)?;
179            let reset = parse(HEADER_RESET)?;
180            let retry_after = parse(HEADER_RETRY_AFTER);
181
182            Some(Self {
183                limit,
184                remaining,
185                reset,
186                retry_after,
187            })
188        }
189    }
190}
191
192// ---------------------------------------------------------------------------
193// Tests
194// ---------------------------------------------------------------------------
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199
200    // -----------------------------------------------------------------------
201    // Construction
202    // -----------------------------------------------------------------------
203
204    #[test]
205    fn new_sets_fields() {
206        let info = RateLimitInfo::new(100, 42, 1_700_000_000);
207        assert_eq!(info.limit, 100);
208        assert_eq!(info.remaining, 42);
209        assert_eq!(info.reset, 1_700_000_000);
210        assert!(info.retry_after.is_none());
211    }
212
213    #[test]
214    fn retry_after_builder() {
215        let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(60);
216        assert_eq!(info.retry_after, Some(60));
217    }
218
219    // -----------------------------------------------------------------------
220    // is_exceeded
221    // -----------------------------------------------------------------------
222
223    #[test]
224    fn is_exceeded_when_remaining_zero() {
225        let info = RateLimitInfo::new(100, 0, 1_700_000_000);
226        assert!(info.is_exceeded());
227    }
228
229    #[test]
230    fn is_not_exceeded_when_remaining_nonzero() {
231        let info = RateLimitInfo::new(100, 1, 1_700_000_000);
232        assert!(!info.is_exceeded());
233    }
234
235    // -----------------------------------------------------------------------
236    // Serde round-trips
237    // -----------------------------------------------------------------------
238
239    #[cfg(feature = "serde")]
240    #[test]
241    fn serde_round_trip_without_retry_after() {
242        let info = RateLimitInfo::new(100, 50, 1_700_000_000);
243        let json = serde_json::to_value(&info).unwrap();
244        assert_eq!(json["limit"], 100);
245        assert_eq!(json["remaining"], 50);
246        assert_eq!(json["reset"], 1_700_000_000_u64);
247        assert!(json.get("retry_after").is_none());
248        let back: RateLimitInfo = serde_json::from_value(json).unwrap();
249        assert_eq!(back, info);
250    }
251
252    #[cfg(feature = "serde")]
253    #[test]
254    fn serde_round_trip_with_retry_after() {
255        let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(30);
256        let json = serde_json::to_value(&info).unwrap();
257        assert_eq!(json["retry_after"], 30);
258        let back: RateLimitInfo = serde_json::from_value(json).unwrap();
259        assert_eq!(back, info);
260    }
261
262    #[cfg(feature = "serde")]
263    #[test]
264    fn serde_omits_retry_after_when_none() {
265        let info = RateLimitInfo::new(10, 5, 999);
266        let json = serde_json::to_value(&info).unwrap();
267        assert!(json.get("retry_after").is_none());
268    }
269
270    // -----------------------------------------------------------------------
271    // Axum header injection / extraction
272    // -----------------------------------------------------------------------
273
274    #[cfg(feature = "http")]
275    mod http_tests {
276        use super::*;
277        use http::HeaderMap;
278
279        #[test]
280        fn inject_and_extract_without_retry_after() {
281            let info = RateLimitInfo::new(200, 150, 1_700_000_000);
282            let mut headers = HeaderMap::new();
283            info.inject_headers(&mut headers);
284
285            let extracted = RateLimitInfo::from_headers(&headers).unwrap();
286            assert_eq!(extracted.limit, 200);
287            assert_eq!(extracted.remaining, 150);
288            assert_eq!(extracted.reset, 1_700_000_000);
289            assert!(extracted.retry_after.is_none());
290        }
291
292        #[test]
293        fn inject_and_extract_with_retry_after() {
294            let info = RateLimitInfo::new(100, 0, 1_700_000_000).retry_after(45);
295            let mut headers = HeaderMap::new();
296            info.inject_headers(&mut headers);
297
298            let extracted = RateLimitInfo::from_headers(&headers).unwrap();
299            assert_eq!(extracted.retry_after, Some(45));
300        }
301
302        #[test]
303        fn from_headers_returns_none_on_missing_required_header() {
304            let headers = HeaderMap::new();
305            assert!(RateLimitInfo::from_headers(&headers).is_none());
306        }
307
308        #[test]
309        fn from_headers_returns_none_on_invalid_value() {
310            let mut headers = HeaderMap::new();
311            headers.insert("x-ratelimit-limit", "not-a-number".parse().unwrap());
312            headers.insert("x-ratelimit-remaining", "5".parse().unwrap());
313            headers.insert("x-ratelimit-reset", "999".parse().unwrap());
314            assert!(RateLimitInfo::from_headers(&headers).is_none());
315        }
316    }
317}