1const HEADER_RATE_LIMIT: &str = "X-RateLimit-Limit";
7const HEADER_RATE_REMAINING: &str = "X-RateLimit-Remaining";
8const HEADER_RATE_RESET: &str = "X-RateLimit-Reset";
9const HEADER_RATE_RETRY: &str = "Retry-After";
10
11use reqwest::header::HeaderMap;
12use std::str::FromStr;
13use std::time::{Duration, SystemTime, UNIX_EPOCH};
14
15#[derive(Clone, Default, Debug)]
22pub struct Rate {
23 pub limit: u32,
24 pub remaining: u32,
25 pub reset: u64,
26 pub retry: u64,
27}
28
29impl Rate {
30 pub(crate) fn from_headers(headers: &HeaderMap) -> Option<Self> {
31 let limit = fetch_header(headers, HEADER_RATE_LIMIT)?;
32 let remaining = fetch_header(headers, HEADER_RATE_REMAINING)?;
33 let reset_secs: u64 = fetch_header(headers, HEADER_RATE_RESET)?;
34 let retry_secs: u64 = fetch_header(headers, HEADER_RATE_RETRY)?;
35
36 let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?;
37 let reset = now + Duration::from_secs(reset_secs);
38 let retry = now + Duration::from_secs(retry_secs);
39
40 Some(Self {
41 limit,
42 remaining,
43 reset: reset.as_secs(),
44 retry: retry.as_secs(),
45 })
46 }
47}
48
49fn fetch_header<T>(headers: &HeaderMap, header: &str) -> Option<T>
50where
51 T: FromStr,
52{
53 headers.get(header)?.to_str().ok()?.parse::<T>().ok()
54}
55
56#[cfg(test)]
58mod tests {
59 use super::*;
60 use httpmock::Method::GET;
61 use httpmock::MockServer;
62
63 #[tokio::test]
64 async fn test_rate_from_headers() {
65 let server = MockServer::start();
66 let mock = server.mock(|when, then| {
67 when.method(GET).path("/test");
68 then.status(200)
69 .header(HEADER_RATE_LIMIT, "100")
70 .header(HEADER_RATE_REMAINING, "50")
71 .header(HEADER_RATE_RESET, "1000")
72 .header(HEADER_RATE_RETRY, "1000");
73 });
74
75 let client = reqwest::Client::new();
76 let response = client.get(server.url("/test")).send().await.unwrap();
77 let rate = Rate::from_headers(response.headers()).unwrap();
78
79 assert_eq!(rate.limit, 100);
80 assert_eq!(rate.remaining, 50);
81 let now = SystemTime::now().duration_since(UNIX_EPOCH).ok().unwrap();
84 let reset = now + Duration::from_secs(1000);
85 assert_eq!(rate.reset, reset.as_secs());
86 let retry = now + Duration::from_secs(1000);
87 assert_eq!(rate.retry, retry.as_secs());
88 mock.assert();
89 }
90
91 #[tokio::test]
92 async fn test_no_headers() {
93 let server = MockServer::start();
94 let mock = server.mock(|when, then| {
95 when.method(GET).path("/test");
96 then.status(200);
97 });
98
99 let client = reqwest::Client::new();
100 let response = client.get(server.url("/test")).send().await.unwrap();
101 let rate = Rate::from_headers(response.headers());
102 assert!(rate.is_none());
103 mock.assert();
104 }
105}