sentry/transports/
ratelimit.rs

1use httpdate::parse_http_date;
2use std::time::{Duration, SystemTime};
3
4use crate::protocol::EnvelopeItem;
5use crate::Envelope;
6
7/// A Utility that helps with rate limiting sentry requests.
8#[derive(Debug, Default)]
9pub struct RateLimiter {
10    global: Option<SystemTime>,
11    error: Option<SystemTime>,
12    session: Option<SystemTime>,
13    transaction: Option<SystemTime>,
14    attachment: Option<SystemTime>,
15}
16
17impl RateLimiter {
18    /// Create a new RateLimiter.
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    /// Updates the RateLimiter with information from a `Retry-After` header.
24    pub fn update_from_retry_after(&mut self, header: &str) {
25        let new_time = if let Ok(value) = header.parse::<f64>() {
26            SystemTime::now() + Duration::from_secs(value.ceil() as u64)
27        } else if let Ok(value) = parse_http_date(header) {
28            value
29        } else {
30            SystemTime::now() + Duration::from_secs(60)
31        };
32
33        self.global = Some(new_time);
34    }
35
36    /// Updates the RateLimiter with information from a `X-Sentry-Rate-Limits` header.
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                    "attachment" => self.attachment = new_time,
59                    _ => {}
60                }
61            }
62            Some(())
63        };
64
65        for group in header.split(',') {
66            parse_group(group.trim());
67        }
68    }
69
70    /// Updates the RateLimiter in response to a `429` status code.
71    pub fn update_from_429(&mut self) {
72        self.global = Some(SystemTime::now() + Duration::from_secs(60));
73    }
74
75    /// Query the RateLimiter if a certain category of event is currently rate limited.
76    ///
77    /// If the given category is rate limited, it will return the remaining
78    /// [`Duration`] for which it is.
79    pub fn is_disabled(&self, category: RateLimitingCategory) -> Option<Duration> {
80        if let Some(ts) = self.global {
81            let time_left = ts.duration_since(SystemTime::now()).ok();
82            if time_left.is_some() {
83                return time_left;
84            }
85        }
86        let time_left = match category {
87            RateLimitingCategory::Any => self.global,
88            RateLimitingCategory::Error => self.error,
89            RateLimitingCategory::Session => self.session,
90            RateLimitingCategory::Transaction => self.transaction,
91            RateLimitingCategory::Attachment => self.attachment,
92        }?;
93        time_left.duration_since(SystemTime::now()).ok()
94    }
95
96    /// Query the RateLimiter for a certain category of event.
97    ///
98    /// Returns `true` if the category is *not* rate limited and should be sent.
99    pub fn is_enabled(&self, category: RateLimitingCategory) -> bool {
100        self.is_disabled(category).is_none()
101    }
102
103    /// Filters the [`Envelope`] according to the current rate limits.
104    ///
105    /// Returns [`None`] if all the envelope items were filtered out.
106    pub fn filter_envelope(&self, envelope: Envelope) -> Option<Envelope> {
107        envelope.filter(|item| {
108            self.is_enabled(match item {
109                EnvelopeItem::Event(_) => RateLimitingCategory::Error,
110                EnvelopeItem::SessionUpdate(_) | EnvelopeItem::SessionAggregates(_) => {
111                    RateLimitingCategory::Session
112                }
113                EnvelopeItem::Transaction(_) => RateLimitingCategory::Transaction,
114                EnvelopeItem::Attachment(_) => RateLimitingCategory::Attachment,
115                _ => RateLimitingCategory::Any,
116            })
117        })
118    }
119}
120
121/// The Category of payload that a Rate Limit refers to.
122#[non_exhaustive]
123pub enum RateLimitingCategory {
124    /// Rate Limit for any kind of payload.
125    Any,
126    /// Rate Limit pertaining to Errors.
127    Error,
128    /// Rate Limit pertaining to Sessions.
129    Session,
130    /// Rate Limit pertaining to Transactions.
131    Transaction,
132    /// Rate Limit pertaining to Attachments.
133    Attachment,
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn test_sentry_header() {
142        let mut rl = RateLimiter::new();
143        rl.update_from_sentry_header("120:error:project:reason, 60:session:foo");
144
145        assert!(rl.is_disabled(RateLimitingCategory::Error).unwrap() <= Duration::from_secs(120));
146        assert!(rl.is_disabled(RateLimitingCategory::Session).unwrap() <= Duration::from_secs(60));
147        assert!(rl.is_disabled(RateLimitingCategory::Transaction).is_none());
148        assert!(rl.is_disabled(RateLimitingCategory::Any).is_none());
149
150        rl.update_from_sentry_header(
151            r#"
152                30::bar,
153                120:invalid:invalid,
154                4711:foo;bar;baz;security:project
155            "#,
156        );
157
158        assert!(
159            rl.is_disabled(RateLimitingCategory::Transaction).unwrap() <= Duration::from_secs(30)
160        );
161        assert!(rl.is_disabled(RateLimitingCategory::Any).unwrap() <= Duration::from_secs(30));
162    }
163
164    #[test]
165    fn test_retry_after() {
166        let mut rl = RateLimiter::new();
167        rl.update_from_retry_after("60");
168
169        assert!(rl.is_disabled(RateLimitingCategory::Error).unwrap() <= Duration::from_secs(60));
170        assert!(rl.is_disabled(RateLimitingCategory::Session).unwrap() <= Duration::from_secs(60));
171        assert!(
172            rl.is_disabled(RateLimitingCategory::Transaction).unwrap() <= Duration::from_secs(60)
173        );
174        assert!(rl.is_disabled(RateLimitingCategory::Any).unwrap() <= Duration::from_secs(60));
175    }
176}