Skip to main content

lychee_lib/ratelimit/
headers.rs

1//! Handle rate limiting headers.
2//! Note that we might want to replace this module with
3//! <https://github.com/mre/rate-limits> at some point in the future.
4
5use http::HeaderValue;
6use std::time::{Duration, SystemTime};
7use thiserror::Error;
8
9#[derive(Debug, Error, PartialEq, Eq)]
10pub(crate) enum RetryAfterParseError {
11    #[error("Unable to parse value '{0}'")]
12    ValueError(String),
13
14    #[error("Header value contains invalid chars")]
15    HeaderValueError,
16}
17
18/// Parse the "Retry-After" header as specified per
19/// [RFC 7231 section 7.1.3](https://www.rfc-editor.org/rfc/rfc7231#section-7.1.3)
20pub(crate) fn parse_retry_after(value: &HeaderValue) -> Result<Duration, RetryAfterParseError> {
21    let value = value
22        .to_str()
23        .map_err(|_| RetryAfterParseError::HeaderValueError)?;
24
25    // RFC 7231: Retry-After = HTTP-date / delay-seconds
26    value.parse::<u64>().map(Duration::from_secs).or_else(|_| {
27        httpdate::parse_http_date(value)
28            .map(|s| {
29                s.duration_since(SystemTime::now())
30                    // if date is in the past, we can use ZERO
31                    .unwrap_or(Duration::ZERO)
32            })
33            .map_err(|_| RetryAfterParseError::ValueError(value.into()))
34    })
35}
36
37/// Parse the common "X-RateLimit" header fields.
38/// Unfortunately, this is not standardised yet, but there is an
39/// [IETF draft](https://datatracker.ietf.org/doc/draft-ietf-httpapi-ratelimit-headers/).
40pub(crate) fn parse_common_rate_limit_header_fields(
41    headers: &http::HeaderMap,
42) -> (Option<usize>, Option<usize>) {
43    let remaining = self::parse_header_value(
44        headers,
45        &[
46            "x-ratelimit-remaining",
47            "x-rate-limit-remaining",
48            "ratelimit-remaining",
49        ],
50    );
51
52    let limit = self::parse_header_value(
53        headers,
54        &["x-ratelimit-limit", "x-rate-limit-limit", "ratelimit-limit"],
55    );
56
57    (remaining, limit)
58}
59
60/// Helper method to parse numeric header values from common rate limit headers
61fn parse_header_value(headers: &http::HeaderMap, header_names: &[&str]) -> Option<usize> {
62    for header_name in header_names {
63        if let Some(value) = headers.get(*header_name)
64            && let Ok(value_str) = value.to_str()
65            && let Ok(number) = value_str.parse()
66        {
67            return Some(number);
68        }
69    }
70    None
71}
72
73#[cfg(test)]
74mod tests {
75    use std::time::Duration;
76
77    use http::HeaderValue;
78
79    use crate::ratelimit::headers::{RetryAfterParseError, parse_retry_after};
80
81    #[test]
82    fn test_retry_after() {
83        assert_eq!(parse_retry_after(&value("1")), Ok(Duration::from_secs(1)));
84        assert_eq!(
85            parse_retry_after(&value("-1")),
86            Err(RetryAfterParseError::ValueError("-1".into()))
87        );
88
89        assert_eq!(
90            parse_retry_after(&value("Fri, 15 May 2015 15:34:21 GMT")),
91            Ok(Duration::ZERO)
92        );
93
94        let result = parse_retry_after(&value("Fri, 15 May 4099 15:34:21 GMT"));
95        let is_in_future = matches!(result, Ok(d) if d.as_secs() > 0);
96        assert!(is_in_future);
97    }
98
99    fn value(v: &str) -> HeaderValue {
100        HeaderValue::from_str(v).unwrap()
101    }
102}