1 use httpdate::parse_http_date;
2 use std::time::{Duration, SystemTime};
3 
4 // TODO: maybe move this someplace where we can filter an `Envelope`s items.
5 
6 /// A Utility that helps with rate limiting sentry requests.
7 #[derive(Debug, Default)]
8 pub struct RateLimiter {
9     global: Option<SystemTime>,
10     error: Option<SystemTime>,
11     session: Option<SystemTime>,
12     transaction: Option<SystemTime>,
13 }
14 
15 impl RateLimiter {
16     /// Create a new RateLimiter.
new() -> Self17     pub fn new() -> Self {
18         Self::default()
19     }
20 
21     /// Updates the RateLimiter with information from a `Retry-After` header.
update_from_retry_after(&mut self, header: &str)22     pub fn update_from_retry_after(&mut self, header: &str) {
23         let new_time = if let Ok(value) = header.parse::<f64>() {
24             Some(SystemTime::now() + Duration::from_secs(value.ceil() as u64))
25         } else if let Ok(value) = parse_http_date(header) {
26             Some(value)
27         } else {
28             None
29         };
30 
31         if new_time.is_some() {
32             self.global = new_time;
33         }
34     }
35 
36     /// Updates the RateLimiter with information from a `X-Sentry-Rate-Limits` header.
update_from_sentry_header(&mut self, header: &str)37     pub fn update_from_sentry_header(&mut self, header: &str) {
38         // <rate-limit> = (<group>,)+
39         // <group> = <time>:(<category>;)+:<scope>(:<reason>)?
40 
41         let mut parse_group = |group: &str| {
42             let mut splits = group.split(':');
43             let seconds = splits.next()?.parse::<f64>().ok()?;
44             let categories = splits.next()?;
45             let _scope = splits.next()?;
46 
47             let new_time = Some(SystemTime::now() + Duration::from_secs(seconds.ceil() as u64));
48 
49             if categories.is_empty() {
50                 self.global = new_time;
51             }
52 
53             for category in categories.split(';') {
54                 match category {
55                     "error" => self.error = new_time,
56                     "session" => self.session = new_time,
57                     "transaction" => self.transaction = new_time,
58                     _ => {}
59                 }
60             }
61             Some(())
62         };
63 
64         for group in header.split(',') {
65             parse_group(group.trim());
66         }
67     }
68 
69     /// Query the RateLimiter for a certain category of event.
is_disabled(&self, category: RateLimitingCategory) -> Option<Duration>70     pub fn is_disabled(&self, category: RateLimitingCategory) -> Option<Duration> {
71         if let Some(ts) = self.global {
72             let time_left = ts.duration_since(SystemTime::now()).ok();
73             if time_left.is_some() {
74                 return time_left;
75             }
76         }
77         let time_left = match category {
78             RateLimitingCategory::Any => self.global,
79             RateLimitingCategory::Error => self.error,
80             RateLimitingCategory::Session => self.session,
81             RateLimitingCategory::Transaction => self.transaction,
82         }?;
83         time_left.duration_since(SystemTime::now()).ok()
84     }
85 }
86 
87 /// The Category of payload that a Rate Limit refers to.
88 #[non_exhaustive]
89 #[allow(dead_code)]
90 pub enum RateLimitingCategory {
91     /// Rate Limit for any kind of payload.
92     Any,
93     /// Rate Limit pertaining to Errors.
94     Error,
95     /// Rate Limit pertaining to Sessions.
96     Session,
97     /// Rate Limit pertaining to Transactions.
98     Transaction,
99 }
100 
101 #[cfg(test)]
102 mod tests {
103     use super::*;
104 
105     #[test]
test_sentry_header()106     fn test_sentry_header() {
107         let mut rl = RateLimiter::new();
108         rl.update_from_sentry_header("120:error:project:reason, 60:session:foo");
109 
110         assert!(rl.is_disabled(RateLimitingCategory::Error).unwrap() <= Duration::from_secs(120));
111         assert!(rl.is_disabled(RateLimitingCategory::Session).unwrap() <= Duration::from_secs(60));
112         assert!(rl.is_disabled(RateLimitingCategory::Transaction).is_none());
113         assert!(rl.is_disabled(RateLimitingCategory::Any).is_none());
114 
115         rl.update_from_sentry_header(
116             r#"
117                 30::bar,
118                 120:invalid:invalid,
119                 4711:foo;bar;baz;security:project
120             "#,
121         );
122 
123         assert!(
124             rl.is_disabled(RateLimitingCategory::Transaction).unwrap() <= Duration::from_secs(30)
125         );
126         assert!(rl.is_disabled(RateLimitingCategory::Any).unwrap() <= Duration::from_secs(30));
127     }
128 
129     #[test]
test_retry_after()130     fn test_retry_after() {
131         let mut rl = RateLimiter::new();
132         rl.update_from_retry_after("60");
133 
134         assert!(rl.is_disabled(RateLimitingCategory::Error).unwrap() <= Duration::from_secs(60));
135         assert!(rl.is_disabled(RateLimitingCategory::Session).unwrap() <= Duration::from_secs(60));
136         assert!(
137             rl.is_disabled(RateLimitingCategory::Transaction).unwrap() <= Duration::from_secs(60)
138         );
139         assert!(rl.is_disabled(RateLimitingCategory::Any).unwrap() <= Duration::from_secs(60));
140     }
141 }
142