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