Skip to main content

rate_limits/retryafter/
mod.rs

1//! Retry-After header parsing
2//!
3//! See <https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After>
4use std::str::FromStr;
5
6use headers::HeaderValue;
7use time::{Date, format_description::well_known::Rfc2822};
8
9use crate::{
10    casesensitive_headermap::CaseSensitiveHeaderMap,
11    reset_time::{ResetTime, ResetTimeKind},
12};
13
14use super::error::{Error, Result};
15
16/// HTTP rate limits as parsed from header values
17#[derive(Copy, Clone, Debug, PartialEq)]
18pub struct RateLimit {
19    /// Time at which the rate limit will be reset
20    pub reset: ResetTime,
21}
22
23impl RateLimit {
24    /// Rate limit implementation based on `Retry-After` header value
25    ///
26    /// See <https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After>
27    pub fn new<T: Into<CaseSensitiveHeaderMap>>(headers: T) -> std::result::Result<Self, Error> {
28        let headers = headers.into();
29        let reset = match Self::get_retry_after_header(&headers) {
30            Some(retry_after) => {
31                if Date::parse(retry_after.to_str()?, &Rfc2822).is_ok() {
32                    ResetTime::new(retry_after, ResetTimeKind::ImfFixdate)?
33                } else {
34                    ResetTime::new(retry_after, ResetTimeKind::Seconds)?
35                }
36            }
37            None => return Err(Error::MissingRetryAfter),
38        };
39
40        Ok(RateLimit { reset })
41    }
42
43    /// Get the Retry-After header value
44    ///
45    /// This does not need to be case sensitive because the header name is
46    /// not ambiguous.
47    fn get_retry_after_header(header_map: &CaseSensitiveHeaderMap) -> Option<&HeaderValue> {
48        header_map
49            .get("Retry-After")
50            .or_else(|| header_map.get("retry-after"))
51    }
52
53    /// Get the time at which the rate limit will be reset
54    #[must_use]
55    pub const fn reset(&self) -> ResetTime {
56        self.reset
57    }
58}
59
60impl FromStr for RateLimit {
61    type Err = Error;
62
63    fn from_str(map: &str) -> Result<Self> {
64        RateLimit::new(CaseSensitiveHeaderMap::from_str(map)?)
65    }
66}
67
68#[cfg(test)]
69mod tests {
70
71    use super::*;
72    use indoc::indoc;
73    use time::macros::datetime;
74
75    #[test]
76    fn parse_retry_after_seconds() {
77        let map = CaseSensitiveHeaderMap::from_str("Retry-After: 30").unwrap();
78        let retry = RateLimit::get_retry_after_header(&map).unwrap();
79
80        assert_eq!("30", retry);
81    }
82
83    #[test]
84    fn retry_after_seconds() {
85        let headers = indoc! {"
86            Retry-After: 19
87        "};
88
89        let rate = RateLimit::from_str(headers).unwrap();
90        assert_eq!(rate.reset(), ResetTime::Seconds(19));
91    }
92
93    #[test]
94    fn retry_after_seconds_case_sensitive() {
95        let headers = indoc! {"
96            retry-after: 19
97        "};
98
99        let rate = RateLimit::from_str(headers).unwrap();
100        assert_eq!(rate.reset(), ResetTime::Seconds(19));
101    }
102
103    #[test]
104    fn retry_after_imf_fixdate() {
105        let headers = indoc! {"
106            Retry-After: Fri, 31 Dec 1999 23:59:59 GMT
107        "};
108
109        let rate = RateLimit::from_str(headers).unwrap();
110        assert_eq!(
111            rate.reset(),
112            ResetTime::DateTime(datetime!(1999-12-31 23:59:59 UTC))
113        );
114    }
115}