From 6a451bcf194d5dbead44ea6beb8bcf1879056734 Mon Sep 17 00:00:00 2001 From: Jan Michael Auer Date: Wed, 3 Jun 2020 16:48:47 +0200 Subject: [PATCH] ref(quotas): Split up quotas code into submodules --- relay-quotas/Cargo.toml | 2 +- relay-quotas/src/legacy.rs | 2 +- relay-quotas/src/lib.rs | 19 +- relay-quotas/src/quota.rs | 704 +++++++++++++++ relay-quotas/src/{types.rs => rate_limit.rs} | 833 +++--------------- .../src/{rate_limiter.rs => redis.rs} | 25 +- relay-server/Cargo.toml | 2 +- relay-server/src/actors/events.rs | 10 +- 8 files changed, 849 insertions(+), 748 deletions(-) create mode 100644 relay-quotas/src/quota.rs rename relay-quotas/src/{types.rs => rate_limit.rs} (53%) rename relay-quotas/src/{rate_limiter.rs => redis.rs} (96%) diff --git a/relay-quotas/Cargo.toml b/relay-quotas/Cargo.toml index d13a98a2d1..9e97c27438 100644 --- a/relay-quotas/Cargo.toml +++ b/relay-quotas/Cargo.toml @@ -11,7 +11,7 @@ publish = false [features] default = [] -rate-limiter = [ +redis = [ "failure", "log", "relay-redis/impl", diff --git a/relay-quotas/src/legacy.rs b/relay-quotas/src/legacy.rs index 9b22e139fb..456cd36f83 100644 --- a/relay-quotas/src/legacy.rs +++ b/relay-quotas/src/legacy.rs @@ -8,7 +8,7 @@ use serde::{ser::SerializeSeq, Deserialize, Deserializer, Serialize, Serializer}; use smallvec::smallvec; -use crate::types::{DataCategory, Quota, QuotaScope, ReasonCode}; +use crate::quota::{DataCategory, Quota, QuotaScope, ReasonCode}; /// Legacy format of the `Quota` type. #[derive(Deserialize, Serialize)] diff --git a/relay-quotas/src/lib.rs b/relay-quotas/src/lib.rs index 33052d1b45..597c37574c 100644 --- a/relay-quotas/src/lib.rs +++ b/relay-quotas/src/lib.rs @@ -2,13 +2,20 @@ #![warn(missing_docs)] -mod types; -pub use self::types::*; +/// The default timeout to apply when a scope is fully rejected. This +/// typically happens for disabled keys, projects, or organizations. +const REJECT_ALL_SECS: u64 = 60; + +mod quota; +mod rate_limit; + +pub use self::quota::*; +pub use self::rate_limit::*; #[cfg(feature = "legacy")] pub mod legacy; -#[cfg(feature = "rate-limiter")] -mod rate_limiter; -#[cfg(feature = "rate-limiter")] -pub use self::rate_limiter::*; +#[cfg(feature = "redis")] +mod redis; +#[cfg(feature = "redis")] +pub use self::redis::*; diff --git a/relay-quotas/src/quota.rs b/relay-quotas/src/quota.rs new file mode 100644 index 0000000000..d7e8efcb71 --- /dev/null +++ b/relay-quotas/src/quota.rs @@ -0,0 +1,704 @@ +use std::fmt; +use std::str::FromStr; + +use serde::{Deserialize, Serialize}; +use smallvec::SmallVec; + +use relay_common::ProjectId; + +/// Data scoping information. +/// +/// This structure holds information of all scopes required for attributing an item to quotas. +#[derive(Clone, Debug)] +pub struct Scoping { + /// The organization id. + pub organization_id: u64, + + /// The project id. + pub project_id: ProjectId, + + /// The DSN public key. + pub public_key: String, + + /// The public key's internal id. + pub key_id: Option, +} + +impl Scoping { + /// Returns an `ItemScoping` for this scope. + /// + /// The item scoping will contain a reference to this scope and the information passed to this + /// function. This is a cheap operation to allow rate limiting for an individual item. + pub fn item(&self, category: DataCategory) -> ItemScoping<'_> { + ItemScoping { + category, + scoping: self, + } + } +} + +/// Data categorization and scoping information. +/// +/// `ItemScoping` is always attached to a `Scope` and references it internally. It is a cheap, +/// copyable type intended for the use with `RateLimits` and `RateLimiter`. It implements +/// `Deref` and `AsRef` for ease of use. +#[derive(Clone, Copy, Debug)] +pub struct ItemScoping<'a> { + /// The data category of the item. + pub category: DataCategory, + + /// Scoping of the data. + pub scoping: &'a Scoping, +} + +impl AsRef for ItemScoping<'_> { + fn as_ref(&self) -> &Scoping { + &self.scoping + } +} + +impl std::ops::Deref for ItemScoping<'_> { + type Target = Scoping; + + fn deref(&self) -> &Self::Target { + &self.scoping + } +} + +impl ItemScoping<'_> { + /// Returns the identifier of the given scope. + pub fn scope_id(&self, scope: QuotaScope) -> Option { + match scope { + QuotaScope::Organization => Some(self.organization_id), + QuotaScope::Project => Some(self.project_id.value()), + QuotaScope::Key => self.key_id, + QuotaScope::Unknown => None, + } + } + + /// Checks whether the category matches any of the quota's categories. + pub(crate) fn matches_categories(&self, categories: &DataCategories) -> bool { + // An empty list of categories means that this quota matches all categories. Note that we + // skip `Unknown` categories silently. If the list of categories only contains `Unknown`s, + // we do **not** match, since apparently the quota is meant for some data this Relay does + // not support yet. + categories.is_empty() || categories.iter().any(|cat| *cat == self.category) + } +} + +/// Classifies the type of data that is being ingested. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum DataCategory { + /// Events with an `event_type` not explicitly listed below. + Default, + /// Error events. + Error, + /// Transaction events. + Transaction, + /// Events with an event type of `csp`, `hpkp`, `expectct` and `expectstaple`. + Security, + /// An attachment. Quantity is the size of the attachment in bytes. + Attachment, + /// Session updates. Quantity is the number of updates in the batch. + Session, + /// Any other data category not known by this Relay. + #[serde(other)] + Unknown, +} + +impl DataCategory { + /// Returns the data category corresponding to the given name. + pub fn from_name(string: &str) -> Self { + match string { + "default" => Self::Default, + "error" => Self::Error, + "transaction" => Self::Transaction, + "security" => Self::Security, + "attachment" => Self::Attachment, + "session" => Self::Session, + _ => Self::Unknown, + } + } + + /// Returns the canonical name of this data category. + pub fn name(self) -> &'static str { + match self { + Self::Default => "default", + Self::Error => "error", + Self::Transaction => "transaction", + Self::Security => "security", + Self::Attachment => "attachment", + Self::Session => "session", + Self::Unknown => "unknown", + } + } +} + +impl fmt::Display for DataCategory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.name()) + } +} + +impl FromStr for DataCategory { + type Err = (); + + fn from_str(string: &str) -> Result { + Ok(Self::from_name(string)) + } +} + +/// An efficient container for data categories that avoids allocations. +/// +/// `DataCategories` is to be treated like a set. +pub type DataCategories = SmallVec<[DataCategory; 8]>; + +/// The scope that a quota applies to. +/// +/// Except for the `Unknown` variant, this type directly translates to the variants of +/// `RateLimitScope` which are used by rate limits. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum QuotaScope { + /// The organization that this project belongs to. + /// + /// This is the top-level scope. + Organization, + + /// The project. + /// + /// This is a sub-scope of `Organization`. + Project, + + /// A project key, which corresponds to a DSN entry. + /// + /// This is a sub-scope of `Project`. + Key, + + /// Any other scope that is not known by this Relay. + #[serde(other)] + Unknown, +} + +impl QuotaScope { + /// Returns the quota scope corresponding to the given name. + pub fn from_name(string: &str) -> Self { + match string { + "organization" => Self::Organization, + "project" => Self::Project, + "key" => Self::Key, + _ => Self::Unknown, + } + } + + /// Returns the canonical name of this scope. + pub fn name(self) -> &'static str { + match self { + Self::Key => "key", + Self::Project => "project", + Self::Organization => "organization", + Self::Unknown => "unknown", + } + } +} + +impl fmt::Display for QuotaScope { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.name()) + } +} + +impl FromStr for QuotaScope { + type Err = (); + + fn from_str(string: &str) -> Result { + Ok(Self::from_name(string)) + } +} + +fn default_scope() -> QuotaScope { + QuotaScope::Organization +} + +/// A machine readable, freeform reason code for rate limits. +#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] +pub struct ReasonCode(String); + +impl ReasonCode { + /// Creates a new reason code from a string. + /// + /// This method is only to be used by tests. Reason codes should only be deserialized from + /// quotas, but never constructed manually. + #[cfg(test)] + pub fn new>(code: S) -> Self { + Self(code.into()) + } + + /// Returns the string representation of this reason code. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Configuration for a data ingestion quota (rate limiting). +/// +/// Sentry applies multiple quotas to incoming data before accepting it, some of which can be +/// configured by the customer. Each piece of data (such as event, attachment) will be counted +/// against all quotas that it matches with based on the `category`. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct Quota { + /// The unique identifier for counting this quota. Required, except for quotas with a `limit` of + /// `0`, since they are statically enforced. + #[serde(default)] + pub id: Option, + + /// A set of data categories that this quota applies to. If missing or empty, this quota + /// applies to all data. + #[serde(default = "DataCategories::new")] + pub categories: DataCategories, + + /// A scope for this quota. This quota is enforced separately within each instance of this scope + /// (e.g. for each project key separately). Defaults to `QuotaScope::Organization`. + #[serde(default = "default_scope")] + pub scope: QuotaScope, + + /// Identifier of the scope to apply to. If set, then this quota will only apply to the + /// specified scope instance (e.g. a project key). Requires `scope` to be set explicitly. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scope_id: Option, + + /// Maxmimum number of matching events allowed. Can be `0` to reject all events, `None` for an + /// unlimited counted quota, or a positive number for enforcement. Requires `window` if the + /// limit is not `0`. + #[serde(default)] + pub limit: Option, + + /// The time window in seconds to enforce this quota in. Required in all cases except `limit=0`, + /// since those quotas are not measured. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub window: Option, + + /// A machine readable reason returned when this quota is exceeded. Required in all cases except + /// `limit=None`, since unlimited quotas can never be exceeded. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reason_code: Option, +} + +impl Quota { + /// Checks whether this quota's scope matches the given item scoping. + /// + /// This quota matches, if: + /// - there is no `scope_id` constraint + /// - the `scope_id` constraint is not numeric + /// - the scope identifier matches the one from ascoping and the scope is known + fn matches_scope(&self, scoping: ItemScoping<'_>) -> bool { + // Check for a scope identifier constraint. If there is no constraint, this means that the + // quota matches any scope. In case the scope is unknown, it will be coerced to the most + // specific scope later. + let scope_id = match self.scope_id { + Some(ref scope_id) => scope_id, + None => return true, + }; + + // Check if the scope identifier in the quota is parseable. If not, this means we cannot + // fulfill the constraint, so the quota does not match. + let parsed = match scope_id.parse::() { + Ok(parsed) => parsed, + Err(_) => return false, + }; + + // At this stage, require that the scope is known since we have to fulfill the constraint. + scoping.scope_id(self.scope) == Some(parsed) + } + + /// Checks whether the quota's constraints match the current item. + pub fn matches(&self, scoping: ItemScoping<'_>) -> bool { + self.matches_scope(scoping) && scoping.matches_categories(&self.categories) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use smallvec::smallvec; + + #[test] + fn test_parse_quota_reject_all() { + let json = r#"{ + "limit": 0, + "reasonCode": "not_yet" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: None, + categories: [], + scope: organization, + limit: Some(0), + reasonCode: Some(ReasonCode("not_yet")), + ) + "###); + } + + #[test] + fn test_parse_quota_reject_transactions() { + let json = r#"{ + "limit": 0, + "categories": ["transaction"], + "reasonCode": "not_yet" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: None, + categories: [ + transaction, + ], + scope: organization, + limit: Some(0), + reasonCode: Some(ReasonCode("not_yet")), + ) + "###); + } + + #[test] + fn test_parse_quota_limited() { + let json = r#"{ + "id": "o", + "limit": 4711, + "window": 42, + "reasonCode": "not_so_fast" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: Some("o"), + categories: [], + scope: organization, + limit: Some(4711), + window: Some(42), + reasonCode: Some(ReasonCode("not_so_fast")), + ) + "###); + } + + #[test] + fn test_parse_quota_project() { + let json = r#"{ + "id": "p", + "scope": "project", + "scopeId": "1", + "limit": 4711, + "window": 42, + "reasonCode": "not_so_fast" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: Some("p"), + categories: [], + scope: project, + scopeId: Some("1"), + limit: Some(4711), + window: Some(42), + reasonCode: Some(ReasonCode("not_so_fast")), + ) + "###); + } + + #[test] + fn test_parse_quota_key() { + let json = r#"{ + "id": "k", + "scope": "key", + "scopeId": "1", + "limit": 4711, + "window": 42, + "reasonCode": "not_so_fast" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: Some("k"), + categories: [], + scope: key, + scopeId: Some("1"), + limit: Some(4711), + window: Some(42), + reasonCode: Some(ReasonCode("not_so_fast")), + ) + "###); + } + + #[test] + fn test_parse_quota_unknown_variants() { + let json = r#"{ + "id": "f", + "categories": ["future"], + "scope": "future", + "scopeId": "1", + "limit": 4711, + "window": 42, + "reasonCode": "not_so_fast" + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: Some("f"), + categories: [ + unknown, + ], + scope: unknown, + scopeId: Some("1"), + limit: Some(4711), + window: Some(42), + reasonCode: Some(ReasonCode("not_so_fast")), + ) + "###); + } + + #[test] + fn test_parse_quota_unlimited() { + let json = r#"{ + "id": "o", + "window": 42 + }"#; + + let quota = serde_json::from_str::(json).expect("parse quota"); + + insta::assert_ron_snapshot!(quota, @r###" + Quota( + id: Some("o"), + categories: [], + scope: organization, + limit: None, + window: Some(42), + ) + "###); + } + + #[test] + fn test_quota_matches_no_categories() { + let quota = Quota { + id: None, + categories: DataCategories::new(), + scope: QuotaScope::Organization, + scope_id: None, + limit: None, + window: None, + reason_code: None, + }; + + assert!(quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_unknown_category() { + let quota = Quota { + id: None, + categories: smallvec![DataCategory::Unknown], + scope: QuotaScope::Organization, + scope_id: None, + limit: None, + window: None, + reason_code: None, + }; + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_multiple_categores() { + let quota = Quota { + id: None, + categories: smallvec![DataCategory::Unknown, DataCategory::Error], + scope: QuotaScope::Organization, + scope_id: None, + limit: None, + window: None, + reason_code: None, + }; + + assert!(quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Transaction, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_no_invalid_scope() { + let quota = Quota { + id: None, + categories: DataCategories::new(), + scope: QuotaScope::Organization, + scope_id: Some("not_a_number".to_owned()), + limit: None, + window: None, + reason_code: None, + }; + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_organization_scope() { + let quota = Quota { + id: None, + categories: DataCategories::new(), + scope: QuotaScope::Organization, + scope_id: Some("42".to_owned()), + limit: None, + window: None, + reason_code: None, + }; + + assert!(quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 0, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_project_scope() { + let quota = Quota { + id: None, + categories: DataCategories::new(), + scope: QuotaScope::Project, + scope_id: Some("21".to_owned()), + limit: None, + window: None, + reason_code: None, + }; + + assert!(quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(0), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + } + + #[test] + fn test_quota_matches_key_scope() { + let quota = Quota { + id: None, + categories: DataCategories::new(), + scope: QuotaScope::Key, + scope_id: Some("17".to_owned()), + limit: None, + window: None, + reason_code: None, + }; + + assert!(quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(17), + } + })); + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(0), + } + })); + + assert!(!quota.matches(ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: None, + } + })); + } +} diff --git a/relay-quotas/src/types.rs b/relay-quotas/src/rate_limit.rs similarity index 53% rename from relay-quotas/src/types.rs rename to relay-quotas/src/rate_limit.rs index e329886877..2afb043dc5 100644 --- a/relay-quotas/src/types.rs +++ b/relay-quotas/src/rate_limit.rs @@ -2,323 +2,10 @@ use std::fmt; use std::str::FromStr; use std::time::{Duration, Instant}; -use serde::{Deserialize, Serialize}; -use smallvec::SmallVec; - use relay_common::ProjectId; -/// Data scoping information. -/// -/// This structure holds information of all scopes required for attributing an item to quotas. -#[derive(Clone, Debug)] -pub struct Scoping { - /// The organization id. - pub organization_id: u64, - - /// The project id. - pub project_id: ProjectId, - - /// The DSN public key. - pub public_key: String, - - /// The public key's internal id. - pub key_id: Option, -} - -impl Scoping { - /// Returns an `ItemScoping` for this scope. - /// - /// The item scoping will contain a reference to this scope and the information passed to this - /// function. This is a cheap operation to allow rate limiting for an individual item. - pub fn item(&self, category: DataCategory) -> ItemScoping<'_> { - ItemScoping { - category, - scoping: self, - } - } -} - -/// Data categorization and scoping information. -/// -/// `ItemScoping` is always attached to a `Scope` and references it internally. It is a cheap, -/// copyable type intended for the use with `RateLimits` and `RateLimiter`. It implements -/// `Deref` and `AsRef` for ease of use. -#[derive(Clone, Copy, Debug)] -pub struct ItemScoping<'a> { - /// The data category of the item. - pub category: DataCategory, - - /// Scoping of the data. - pub scoping: &'a Scoping, -} - -impl AsRef for ItemScoping<'_> { - fn as_ref(&self) -> &Scoping { - &self.scoping - } -} - -impl std::ops::Deref for ItemScoping<'_> { - type Target = Scoping; - - fn deref(&self) -> &Self::Target { - &self.scoping - } -} - -impl ItemScoping<'_> { - /// Returns the identifier of the given scope. - pub fn scope_id(&self, scope: QuotaScope) -> Option { - match scope { - QuotaScope::Organization => Some(self.organization_id), - QuotaScope::Project => Some(self.project_id.value()), - QuotaScope::Key => self.key_id, - QuotaScope::Unknown => None, - } - } - - /// Checks whether the category matches any of the quota's categories. - fn matches_categories(&self, categories: &DataCategories) -> bool { - // An empty list of categories means that this quota matches all categories. Note that we - // skip `Unknown` categories silently. If the list of categories only contains `Unknown`s, - // we do **not** match, since apparently the quota is meant for some data this Relay does - // not support yet. - categories.is_empty() || categories.iter().any(|cat| *cat == self.category) - } -} - -/// Classifies the type of data that is being ingested. -#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)] -#[serde(rename_all = "lowercase")] -pub enum DataCategory { - /// Events with an `event_type` not explicitly listed below. - Default, - /// Error events. - Error, - /// Transaction events. - Transaction, - /// Events with an event type of `csp`, `hpkp`, `expectct` and `expectstaple`. - Security, - /// An attachment. Quantity is the size of the attachment in bytes. - Attachment, - /// Session updates. Quantity is the number of updates in the batch. - Session, - /// Any other data category not known by this Relay. - #[serde(other)] - Unknown, -} - -impl DataCategory { - /// Returns the data category corresponding to the given name. - pub fn from_name(string: &str) -> Self { - match string { - "default" => Self::Default, - "error" => Self::Error, - "transaction" => Self::Transaction, - "security" => Self::Security, - "attachment" => Self::Attachment, - "session" => Self::Session, - _ => Self::Unknown, - } - } - - /// Returns the canonical name of this data category. - pub fn name(self) -> &'static str { - match self { - Self::Default => "default", - Self::Error => "error", - Self::Transaction => "transaction", - Self::Security => "security", - Self::Attachment => "attachment", - Self::Session => "session", - Self::Unknown => "unknown", - } - } -} - -impl fmt::Display for DataCategory { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.name()) - } -} - -impl FromStr for DataCategory { - type Err = (); - - fn from_str(string: &str) -> Result { - Ok(Self::from_name(string)) - } -} - -/// An efficient container for data categories that avoids allocations. -/// -/// `DataCategories` is to be treated like a set. -pub type DataCategories = SmallVec<[DataCategory; 8]>; - -/// The scope that a quota applies to. -/// -/// Except for the `Unknown` variant, this type directly translates to the variants of -/// `RateLimitScope` which are used by rate limits. -#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)] -#[serde(rename_all = "lowercase")] -pub enum QuotaScope { - /// The organization that this project belongs to. - /// - /// This is the top-level scope. - Organization, - - /// The project. - /// - /// This is a sub-scope of `Organization`. - Project, - - /// A project key, which corresponds to a DSN entry. - /// - /// This is a sub-scope of `Project`. - Key, - - /// Any other scope that is not known by this Relay. - #[serde(other)] - Unknown, -} - -impl QuotaScope { - /// Returns the quota scope corresponding to the given name. - pub fn from_name(string: &str) -> Self { - match string { - "organization" => Self::Organization, - "project" => Self::Project, - "key" => Self::Key, - _ => Self::Unknown, - } - } - - /// Returns the canonical name of this scope. - pub fn name(self) -> &'static str { - match self { - Self::Key => "key", - Self::Project => "project", - Self::Organization => "organization", - Self::Unknown => "unknown", - } - } -} - -impl fmt::Display for QuotaScope { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.name()) - } -} - -impl FromStr for QuotaScope { - type Err = (); - - fn from_str(string: &str) -> Result { - Ok(Self::from_name(string)) - } -} - -fn default_scope() -> QuotaScope { - QuotaScope::Organization -} - -/// A machine readable, freeform reason code for rate limits. -#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)] -pub struct ReasonCode(String); - -impl ReasonCode { - /// Creates a new reason code from a string. - /// - /// This method is only to be used by tests. Reason codes should only be deserialized from - /// quotas, but never constructed manually. - #[cfg(test)] - pub fn new>(code: S) -> Self { - Self(code.into()) - } - - /// Returns the string representation of this reason code. - pub fn as_str(&self) -> &str { - &self.0 - } -} - -/// Configuration for a data ingestion quota (rate limiting). -/// -/// Sentry applies multiple quotas to incoming data before accepting it, some of which can be -/// configured by the customer. Each piece of data (such as event, attachment) will be counted -/// against all quotas that it matches with based on the `category`. -#[derive(Clone, Debug, Deserialize, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct Quota { - /// The unique identifier for counting this quota. Required, except for quotas with a `limit` of - /// `0`, since they are statically enforced. - #[serde(default)] - pub id: Option, - - /// A set of data categories that this quota applies to. If missing or empty, this quota - /// applies to all data. - #[serde(default = "DataCategories::new")] - pub categories: DataCategories, - - /// A scope for this quota. This quota is enforced separately within each instance of this scope - /// (e.g. for each project key separately). Defaults to `QuotaScope::Organization`. - #[serde(default = "default_scope")] - pub scope: QuotaScope, - - /// Identifier of the scope to apply to. If set, then this quota will only apply to the - /// specified scope instance (e.g. a project key). Requires `scope` to be set explicitly. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub scope_id: Option, - - /// Maxmimum number of matching events allowed. Can be `0` to reject all events, `None` for an - /// unlimited counted quota, or a positive number for enforcement. Requires `window` if the - /// limit is not `0`. - #[serde(default)] - pub limit: Option, - - /// The time window in seconds to enforce this quota in. Required in all cases except `limit=0`, - /// since those quotas are not measured. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub window: Option, - - /// A machine readable reason returned when this quota is exceeded. Required in all cases except - /// `limit=None`, since unlimited quotas can never be exceeded. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub reason_code: Option, -} - -impl Quota { - /// Checks whether this quota's scope matches the given item scoping. - /// - /// This quota matches, if: - /// - there is no `scope_id` constraint - /// - the `scope_id` constraint is not numeric - /// - the scope identifier matches the one from ascoping and the scope is known - fn matches_scope(&self, scoping: ItemScoping<'_>) -> bool { - // Check for a scope identifier constraint. If there is no constraint, this means that the - // quota matches any scope. In case the scope is unknown, it will be coerced to the most - // specific scope later. - let scope_id = match self.scope_id { - Some(ref scope_id) => scope_id, - None => return true, - }; - - // Check if the scope identifier in the quota is parseable. If not, this means we cannot - // fulfill the constraint, so the quota does not match. - let parsed = match scope_id.parse::() { - Ok(parsed) => parsed, - Err(_) => return false, - }; - - // At this stage, require that the scope is known since we have to fulfill the constraint. - scoping.scope_id(self.scope) == Some(parsed) - } - - /// Checks whether the quota's constraints match the current item. - pub fn matches(&self, scoping: ItemScoping<'_>) -> bool { - self.matches_scope(scoping) && scoping.matches_categories(&self.categories) - } -} +use crate::quota::{DataCategories, ItemScoping, Quota, QuotaScope, ReasonCode, Scoping}; +use crate::REJECT_ALL_SECS; /// A monotonic expiration marker for `RateLimit`s. /// @@ -413,7 +100,7 @@ impl FromStr for RetryAfter { /// information about the scope instance. That is, the specific identifiers of the individual scopes /// that a rate limit applied to. #[derive(Clone, Debug, Eq, PartialEq)] -#[cfg_attr(test, derive(Serialize))] +#[cfg_attr(test, derive(serde::Serialize))] pub enum RateLimitScope { /// An organization with identifier. Organization(u64), @@ -447,7 +134,7 @@ impl RateLimitScope { /// A bounded rate limit. #[derive(Clone, Debug, PartialEq)] -#[cfg_attr(test, derive(Serialize))] +#[cfg_attr(test, derive(serde::Serialize))] pub struct RateLimit { /// A set of data categories that this quota applies to. If missing or empty, this rate limit /// applies to all data. @@ -495,7 +182,7 @@ impl RateLimit { /// can be iterated over using `iter`. Additionally, rate limits can be checked for items by /// invoking `check` with the respective `ItemScoping`. #[derive(Clone, Debug, Default)] -#[cfg_attr(test, derive(Serialize))] +#[cfg_attr(test, derive(serde::Serialize))] pub struct RateLimits { limits: Vec, } @@ -546,24 +233,41 @@ impl RateLimits { self.iter().any(|limit| !limit.retry_after.expired()) } - /// Checks whether any rate limits apply to the given scoping returning the matched limits. + /// Removes expired rate limits from this instance. + pub fn clean_expired(&mut self) { + self.limits.retain(|limit| !limit.retry_after.expired()); + } + + /// Checks whether any rate limits apply to the given scoping. /// /// If no limits match, then the returned `RateLimits` instance evalutes `is_ok`. Otherwise, it /// contains rate limits that match the given scoping. - pub fn check(&mut self, scoping: ItemScoping<'_>) -> Self { + pub fn check(&self, scoping: ItemScoping<'_>) -> Self { + self.check_with_quotas(&[], scoping) + } + + /// Checks whether any rate limits apply to the given scoping. + /// + /// This is similar to `check`. Additionally, it checks for quotas with a static limit `0`, and + /// rejects items even if there is no active rate limit in this instance. + /// + /// If no limits or quotas match, then the returned `RateLimits` instance evalutes `is_ok`. + /// Otherwise, it contains rate limits that match the given scoping. + pub fn check_with_quotas(&self, quotas: &[Quota], scoping: ItemScoping<'_>) -> Self { let mut applied_limits = Self::new(); - self.limits.retain(|limit| { - if limit.retry_after.expired() { - return false; + for quota in quotas { + if quota.limit == Some(0) && quota.matches(scoping) { + let retry_after = RetryAfter::from_secs(REJECT_ALL_SECS); + applied_limits.add(RateLimit::from_quota(quota, &*scoping, retry_after)); } + } + for limit in &self.limits { if limit.matches(scoping) { applied_limits.add(limit.clone()); } - - true - }); + } applied_limits } @@ -637,386 +341,9 @@ impl<'a> IntoIterator for &'a RateLimits { #[cfg(test)] mod tests { use super::*; + use crate::quota::DataCategory; use smallvec::smallvec; - #[test] - fn test_parse_quota_reject_all() { - let json = r#"{ - "limit": 0, - "reasonCode": "not_yet" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: None, - categories: [], - scope: organization, - limit: Some(0), - reasonCode: Some(ReasonCode("not_yet")), - ) - "###); - } - - #[test] - fn test_parse_quota_reject_transactions() { - let json = r#"{ - "limit": 0, - "categories": ["transaction"], - "reasonCode": "not_yet" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: None, - categories: [ - transaction, - ], - scope: organization, - limit: Some(0), - reasonCode: Some(ReasonCode("not_yet")), - ) - "###); - } - - #[test] - fn test_parse_quota_limited() { - let json = r#"{ - "id": "o", - "limit": 4711, - "window": 42, - "reasonCode": "not_so_fast" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: Some("o"), - categories: [], - scope: organization, - limit: Some(4711), - window: Some(42), - reasonCode: Some(ReasonCode("not_so_fast")), - ) - "###); - } - - #[test] - fn test_parse_quota_project() { - let json = r#"{ - "id": "p", - "scope": "project", - "scopeId": "1", - "limit": 4711, - "window": 42, - "reasonCode": "not_so_fast" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: Some("p"), - categories: [], - scope: project, - scopeId: Some("1"), - limit: Some(4711), - window: Some(42), - reasonCode: Some(ReasonCode("not_so_fast")), - ) - "###); - } - - #[test] - fn test_parse_quota_key() { - let json = r#"{ - "id": "k", - "scope": "key", - "scopeId": "1", - "limit": 4711, - "window": 42, - "reasonCode": "not_so_fast" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: Some("k"), - categories: [], - scope: key, - scopeId: Some("1"), - limit: Some(4711), - window: Some(42), - reasonCode: Some(ReasonCode("not_so_fast")), - ) - "###); - } - - #[test] - fn test_parse_quota_unknown_variants() { - let json = r#"{ - "id": "f", - "categories": ["future"], - "scope": "future", - "scopeId": "1", - "limit": 4711, - "window": 42, - "reasonCode": "not_so_fast" - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: Some("f"), - categories: [ - unknown, - ], - scope: unknown, - scopeId: Some("1"), - limit: Some(4711), - window: Some(42), - reasonCode: Some(ReasonCode("not_so_fast")), - ) - "###); - } - - #[test] - fn test_parse_quota_unlimited() { - let json = r#"{ - "id": "o", - "window": 42 - }"#; - - let quota = serde_json::from_str::(json).expect("parse quota"); - - insta::assert_ron_snapshot!(quota, @r###" - Quota( - id: Some("o"), - categories: [], - scope: organization, - limit: None, - window: Some(42), - ) - "###); - } - - #[test] - fn test_quota_matches_no_categories() { - let quota = Quota { - id: None, - categories: DataCategories::new(), - scope: QuotaScope::Organization, - scope_id: None, - limit: None, - window: None, - reason_code: None, - }; - - assert!(quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_unknown_category() { - let quota = Quota { - id: None, - categories: smallvec![DataCategory::Unknown], - scope: QuotaScope::Organization, - scope_id: None, - limit: None, - window: None, - reason_code: None, - }; - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_multiple_categores() { - let quota = Quota { - id: None, - categories: smallvec![DataCategory::Unknown, DataCategory::Error], - scope: QuotaScope::Organization, - scope_id: None, - limit: None, - window: None, - reason_code: None, - }; - - assert!(quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Transaction, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_no_invalid_scope() { - let quota = Quota { - id: None, - categories: DataCategories::new(), - scope: QuotaScope::Organization, - scope_id: Some("not_a_number".to_owned()), - limit: None, - window: None, - reason_code: None, - }; - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_organization_scope() { - let quota = Quota { - id: None, - categories: DataCategories::new(), - scope: QuotaScope::Organization, - scope_id: Some("42".to_owned()), - limit: None, - window: None, - reason_code: None, - }; - - assert!(quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 0, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_project_scope() { - let quota = Quota { - id: None, - categories: DataCategories::new(), - scope: QuotaScope::Project, - scope_id: Some("21".to_owned()), - limit: None, - window: None, - reason_code: None, - }; - - assert!(quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(0), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - } - - #[test] - fn test_quota_matches_key_scope() { - let quota = Quota { - id: None, - categories: DataCategories::new(), - scope: QuotaScope::Key, - scope_id: Some("17".to_owned()), - limit: None, - window: None, - reason_code: None, - }; - - assert!(quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(17), - } - })); - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: Some(0), - } - })); - - assert!(!quota.matches(ItemScoping { - category: DataCategory::Error, - scoping: &Scoping { - organization_id: 42, - project_id: ProjectId::new(21), - public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), - key_id: None, - } - })); - } - #[test] fn test_parse_retry_after() { // positive float @@ -1322,7 +649,7 @@ mod tests { } #[test] - fn test_rate_limits_check() { + fn test_rate_limits_clean_expired() { let mut rate_limits = RateLimits::new(); // Active error limit @@ -1341,6 +668,40 @@ mod tests { retry_after: RetryAfter::from_secs(0), }); + // Sanity check before running `clean_expired` + assert_eq!(rate_limits.iter().count(), 2); + + rate_limits.clean_expired(); + + // Check that the expired limit has been removed + insta::assert_ron_snapshot!(rate_limits, @r###" + RateLimits( + limits: [ + RateLimit( + categories: [ + error, + ], + scope: Organization(42), + reason_code: None, + retry_after: RetryAfter(1), + ), + ], + ) + "###); + } + + #[test] + fn test_rate_limits_check() { + let mut rate_limits = RateLimits::new(); + + // Active error limit + rate_limits.add(RateLimit { + categories: smallvec![DataCategory::Error], + scope: RateLimitScope::Organization(42), + reason_code: None, + retry_after: RetryAfter::from_secs(1), + }); + // Active transaction limit rate_limits.add(RateLimit { categories: smallvec![DataCategory::Transaction], @@ -1349,9 +710,6 @@ mod tests { retry_after: RetryAfter::from_secs(1), }); - // Sanity check before running `check` - assert_eq!(rate_limits.iter().count(), 3); - let applied_limits = rate_limits.check(ItemScoping { category: DataCategory::Error, scoping: &Scoping { @@ -1377,9 +735,51 @@ mod tests { ], ) "###); + } - // Check that the expired limit has been removed - insta::assert_ron_snapshot!(rate_limits, @r###" + #[test] + fn test_rate_limits_check_quotas() { + let mut rate_limits = RateLimits::new(); + + // Active error limit + rate_limits.add(RateLimit { + categories: smallvec![DataCategory::Error], + scope: RateLimitScope::Organization(42), + reason_code: None, + retry_after: RetryAfter::from_secs(1), + }); + + // Active transaction limit + rate_limits.add(RateLimit { + categories: smallvec![DataCategory::Transaction], + scope: RateLimitScope::Organization(42), + reason_code: None, + retry_after: RetryAfter::from_secs(1), + }); + + let item_scoping = ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(21), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: None, + }, + }; + + let quotas = &[Quota { + id: None, + categories: smallvec![DataCategory::Error], + scope: QuotaScope::Organization, + scope_id: Some("42".to_owned()), + limit: Some(0), + window: None, + reason_code: Some(ReasonCode::new("zero")), + }]; + + let applied_limits = rate_limits.check_with_quotas(quotas, item_scoping); + + insta::assert_ron_snapshot!(applied_limits, @r###" RateLimits( limits: [ RateLimit( @@ -1387,21 +787,12 @@ mod tests { error, ], scope: Organization(42), - reason_code: None, - retry_after: RetryAfter(1), - ), - RateLimit( - categories: [ - transaction, - ], - scope: Organization(42), - reason_code: None, - retry_after: RetryAfter(1), + reason_code: Some(ReasonCode("zero")), + retry_after: RetryAfter(60), ), ], ) "###); - assert_eq!(rate_limits.iter().count(), 2); } #[test] diff --git a/relay-quotas/src/rate_limiter.rs b/relay-quotas/src/redis.rs similarity index 96% rename from relay-quotas/src/rate_limiter.rs rename to relay-quotas/src/redis.rs index 0abb5b1319..d5da70f269 100644 --- a/relay-quotas/src/rate_limiter.rs +++ b/relay-quotas/src/redis.rs @@ -7,18 +7,16 @@ use relay_common::UnixTimestamp; use relay_redis::{redis::Script, RedisError, RedisPool}; use sentry::protocol::value; -use crate::types::{ItemScoping, Quota, QuotaScope, RateLimit, RateLimits, RetryAfter}; +use crate::quota::{ItemScoping, Quota, QuotaScope}; +use crate::rate_limit::{RateLimit, RateLimits, RetryAfter}; +use crate::REJECT_ALL_SECS; /// The `grace` period allows accomodating for clock drift in TTL /// calculation since the clock on the Redis instance used to store quota /// metrics may not be in sync with the computer running this code. const GRACE: u64 = 60; -/// The default timeout to apply when a scope is fully rejected. This -/// typically happens for disabled keys, projects, or organizations. -const REJECT_ALL_SECS: u64 = 60; - -/// An error returned by `RateLimiter`. +/// An error returned by `RedisRateLimiter`. #[derive(Debug, Fail)] pub enum RateLimitingError { /// Failed to communicate with Redis. @@ -133,18 +131,18 @@ impl std::ops::Deref for RedisQuota<'_> { /// quotas allow to specify the data categories they apply to, for example error events or /// attachments. For more information on quota parameters, see `QuotaConfig`. /// -/// Requires the `rate-limiter` feature. +/// Requires the `redis` feature. #[derive(Clone)] -pub struct RateLimiter { +pub struct RedisRateLimiter { pool: RedisPool, script: Arc