rocket_governor/
header.rs

1//! The headers used when a [RocketGovernor](super::RocketGovernor) guarded
2//! path responds with [`TooManyRequests`](http::Status::TooManyRequests).  
3//!
4//! Depending on setup some headers are also set on responding successful
5//! to let the client know about request limitations in the near future.
6//!
7//! There is an [RFC Draft](https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers).
8// TODO: Check RFC Draft for publication
9
10use rocket::http;
11
12/// HTTP headers used for rate-limiting.
13pub enum Header {
14    /// Standard header for status 429 Too Many Requests
15    /// ([RFC 6585](https://tools.ietf.org/html/rfc6585#section-4)).
16    /// This should indicate a client for how long it should wait in seconds
17    /// for retry.
18    RetryAfter(u64),
19
20    /// Custom header for reporting problems with rate limiter.
21    XRateLimitError(&'static str),
22
23    /// Header provides information about limitation of the route.
24    XRateLimitLimit(u64),
25
26    /// Header provides information about how many requests are left for the
27    /// endpoint.
28    XRateLimitRemaining(u64),
29
30    /// Header provides the time in seconds when a request to the route is not
31    /// rate limited and the rate limiter bucket is full again.
32    XRateLimitReset(u64),
33}
34
35/// Standard header for status 429 Too Many Requests
36/// ([RFC 6585](https://tools.ietf.org/html/rfc6585#section-4)).
37/// This should indicate a client for how long it should wait in seconds for
38/// retry.
39pub const RETRY_AFTER: &str = "retry-after";
40
41/// Custom header for reporting problems with rate limiter.
42pub const X_RATELIMIT_ERROR: &str = "x-ratelimit-error";
43
44// TODO: https://github.com/kolbma/rocket-governor/issues/2
45
46// TODO: Check ratelimit-headers draft for publication
47// https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers
48
49/// Header provides information about limitation of the route.
50pub const X_RATELIMIT_LIMIT: &str = "x-ratelimit-limit";
51
52/// Header provides information about how many requests are left for the
53/// endpoint.
54pub const X_RATELIMIT_REMAINING: &str = "x-ratelimit-remaining";
55
56/// Header provides the time in seconds when a request to the route is not
57/// rate limited and the rate limiter bucket is full again.
58pub const X_RATELIMIT_RESET: &str = "x-ratelimit-reset";
59
60#[doc(hidden)]
61impl From<Header> for http::Header<'_> {
62    fn from(header: Header) -> Self {
63        match header {
64            Header::RetryAfter(sec) => http::Header::new(RETRY_AFTER, sec.to_string()),
65            Header::XRateLimitError(err) => http::Header::new(X_RATELIMIT_ERROR, err),
66            Header::XRateLimitLimit(limit) => {
67                http::Header::new(X_RATELIMIT_LIMIT, limit.to_string())
68            }
69            Header::XRateLimitRemaining(remaining) => {
70                http::Header::new(X_RATELIMIT_REMAINING, remaining.to_string())
71            }
72            Header::XRateLimitReset(sec) => http::Header::new(X_RATELIMIT_RESET, sec.to_string()),
73        }
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::Header;
80    use super::{
81        RETRY_AFTER, X_RATELIMIT_ERROR, X_RATELIMIT_LIMIT, X_RATELIMIT_REMAINING, X_RATELIMIT_RESET,
82    };
83    use rocket::http;
84    use std::str::FromStr;
85
86    #[test]
87    fn test_header() {
88        let h: http::Header = Header::RetryAfter(10).into();
89        assert_eq!(RETRY_AFTER, h.name());
90        assert_eq!(10, u64::from_str(h.value()).unwrap());
91
92        let h: http::Header = Header::XRateLimitError("some error").into();
93        assert_eq!(X_RATELIMIT_ERROR, h.name());
94        assert_eq!("some error", h.value());
95
96        let h: http::Header = Header::XRateLimitLimit(100).into();
97        assert_eq!(X_RATELIMIT_LIMIT, h.name());
98        assert_eq!(100, u64::from_str(h.value()).unwrap());
99
100        let h: http::Header = Header::XRateLimitRemaining(1).into();
101        assert_eq!(X_RATELIMIT_REMAINING, h.name());
102        assert_eq!(1, u64::from_str(h.value()).unwrap());
103
104        let h: http::Header = Header::XRateLimitReset(5).into();
105        assert_eq!(X_RATELIMIT_RESET, h.name());
106        assert_eq!(5, u64::from_str(h.value()).unwrap());
107    }
108}