rate_limit

Author: epompeii

lib/bencher_config/src/config_tx.rs

@@ -13,7 +13,7 @@ use bencher_json::{
 use bencher_rbac::init_rbac;
 use bencher_schema::context::{ApiContext, Database, DbConnection};
 #[cfg(feature = "plus")]
-use bencher_schema::model::server::QueryServer;
+use bencher_schema::{context::RateLimit, model::server::QueryServer};
 use bencher_token::TokenKey;
 #[cfg(feature = "plus")]
 use diesel::connection::SimpleConnection;
@@ -52,6 +52,8 @@ pub enum ConfigTxError {
     DatabaseConnection(String, diesel::ConnectionError),
     #[error("Failed to parse data store: {0}")]
     DataStore(bencher_schema::context::DataStoreError),
+    #[error("Failed to configure rate limits: {0}")]
+    RateLimit(Box<bencher_schema::context::RateLimitError>),
     #[error("Failed to register endpoint: {0}")]
     Register(dropshot::ApiDescriptionRegisterError),
     #[error("Failed to create server: {0}")]
@@ -104,7 +106,8 @@ impl ConfigTx {
             restart_tx,
             #[cfg(feature = "plus")]
             plus,
-        )?;
+        )
+        .await?;
         debug!(log, "Configuring TLS");
         let tls = server.tls.take().map(|json_tls| match json_tls {
             JsonTls::AsFile {
@@ -158,7 +161,7 @@ impl ConfigTx {
     }
 }
 
-fn into_context(
+async fn into_context(
     log: &Logger,
     console: JsonConsole,
     security: JsonSecurity,
@@ -190,17 +193,20 @@ fn into_context(
         None
     };
 
+    let database = Database {
+        path: json_database.file,
+        connection: Arc::new(tokio::sync::Mutex::new(database_connection)),
+        data_store,
+    };
+
     info!(&log, "Loading secret key");
     let token_key = TokenKey::new(
         security.issuer.unwrap_or_else(|| console_url.to_string()),
         &security.secret_key,
     );
 
     #[cfg(feature = "plus")]
-    let rate_limit = plus
-        .as_ref()
-        .and_then(|plus| plus.rate_limit.map(Into::into))
-        .unwrap_or_default();
+    let rate_limit = plus.as_ref().and_then(|plus| plus.rate_limit);
 
     info!(&log, "Configuring Bencher Plus");
     #[cfg(feature = "plus")]
@@ -215,17 +221,25 @@ fn into_context(
     #[cfg(feature = "plus")]
     let is_bencher_cloud = bencher_json::is_bencher_cloud(&console_url) && biller.is_some();
 
+    #[cfg(feature = "plus")]
+    let rate_limit = RateLimit::new(
+        log,
+        &database.connection,
+        &licensor,
+        is_bencher_cloud,
+        rate_limit,
+    )
+    .await
+    .map_err(Box::new)
+    .map_err(ConfigTxError::RateLimit)?;
+
     debug!(&log, "Creating API context");
     Ok(ApiContext {
         console_url,
         token_key,
         rbac: init_rbac().map_err(ConfigTxError::Polar)?.into(),
         messenger: smtp.into(),
-        database: Database {
-            path: json_database.file,
-            connection: Arc::new(tokio::sync::Mutex::new(database_connection)),
-            data_store,
-        },
+        database,
         restart_tx,
         #[cfg(feature = "plus")]
         rate_limit,

lib/bencher_schema/src/context/mod.rs

@@ -27,7 +27,7 @@ pub use indexer::{IndexError, Indexer};
 pub use messenger::ServerStatsBody;
 pub use messenger::{Body, ButtonBody, Email, Message, Messenger, NewUserBody};
 #[cfg(feature = "plus")]
-pub use rate_limit::RateLimit;
+pub use rate_limit::{RateLimit, RateLimitError};
 pub use rbac::{Rbac, RbacError};
 #[cfg(feature = "plus")]
 pub use stats::StatsSettings;

lib/bencher_schema/src/context/rate_limit.rs

@@ -1,6 +1,19 @@
 use std::time::Duration;
 
-use bencher_json::{system::config::JsonRateLimit, DateTime};
+use bencher_json::{system::config::JsonRateLimit, DateTime, PlanLevel};
+use bencher_license::Licensor;
+use slog::Logger;
+
+use crate::{
+    error::BencherResource,
+    model::{
+        organization::{plan::LicenseUsage, QueryOrganization},
+        project::{branch::QueryBranch, threshold::QueryThreshold, QueryProject},
+        user::QueryUser,
+    },
+};
+
+use super::DbConnection;
 
 const DAY: Duration = Duration::from_secs(24 * 60 * 60);
 const UNCLAIMED_RATE_LIMIT: u32 = u8::MAX as u32;
@@ -12,6 +25,46 @@ pub struct RateLimit {
     pub claimed: u32,
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum RateLimitError {
+    #[error("User ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = user.uuid)]
+    User {
+        user: QueryUser,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+    #[error("Organization ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = organization.uuid)]
+    Organization {
+        organization: QueryOrganization,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+    #[error("Unclaimed project ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage or claim the project: https://bencher.dev/auth/signup?claim={uuid}", uuid = project.uuid)]
+    UnclaimedProject {
+        project: QueryProject,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+    #[error("Claimed project ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = project.uuid)]
+    ClaimedProject {
+        project: QueryProject,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+    #[error("Branch ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = branch.uuid)]
+    Branch {
+        branch: QueryBranch,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+    #[error("Threshold ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = threshold.uuid)]
+    Threshold {
+        threshold: QueryThreshold,
+        resource: BencherResource,
+        rate_limit: u32,
+    },
+}
+
 impl From<JsonRateLimit> for RateLimit {
     fn from(json: JsonRateLimit) -> Self {
         let JsonRateLimit {
@@ -38,6 +91,36 @@ impl Default for RateLimit {
 }
 
 impl RateLimit {
+    pub async fn new(
+        log: &Logger,
+        conn: &tokio::sync::Mutex<DbConnection>,
+        licensor: &Licensor,
+        is_bencher_cloud: bool,
+        rate_limit: Option<JsonRateLimit>,
+    ) -> Result<Self, RateLimitError> {
+        let Some(rate_limit) = rate_limit else {
+            return Ok(Self::default());
+        };
+
+        if !is_bencher_cloud {
+            match LicenseUsage::get_for_server(conn, licensor, Some(PlanLevel::Team)).await {
+                Ok(license_usages) if license_usages.is_empty() => {
+                    slog::warn!(log, "Custom rate limits are set, but there is no valid Bencher Plus license key! This is a violation of the Bencher License: https://bencher.dev/legal/license");
+                    slog::warn!(
+                        log,
+                        "Please purchase a license key: https://bencher.dev/pricing"
+                    );
+                },
+                Ok(_) => {},
+                Err(e) => {
+                    slog::error!(log, "Failed to check license for custom rate limits: {e}");
+                },
+            }
+        }
+
+        Ok(rate_limit.into())
+    }
+
     pub fn window(&self) -> (DateTime, DateTime) {
         let end_time = chrono::Utc::now();
         let start_time = end_time - self.window;

lib/bencher_schema/src/macros/rate_limit.rs

@@ -1,52 +1,3 @@
-use crate::{
-    error::BencherResource,
-    model::{
-        organization::QueryOrganization,
-        project::{branch::QueryBranch, threshold::QueryThreshold, QueryProject},
-        user::QueryUser,
-    },
-};
-
-#[derive(Debug, thiserror::Error)]
-pub enum RateLimitError {
-    #[error("User ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = user.uuid)]
-    User {
-        user: QueryUser,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-    #[error("Organization ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = organization.uuid)]
-    Organization {
-        organization: QueryOrganization,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-    #[error("Unclaimed project ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage or claim the project: https://bencher.dev/auth/signup?claim={uuid}", uuid = project.uuid)]
-    UnclaimedProject {
-        project: QueryProject,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-    #[error("Claimed project ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = project.uuid)]
-    ClaimedProject {
-        project: QueryProject,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-    #[error("Branch ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = branch.uuid)]
-    Branch {
-        branch: QueryBranch,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-    #[error("Threshold ({uuid}) has exceeded the daily rate limit ({rate_limit}) for {resource} creation. Please, reduce your daily usage.", uuid = threshold.uuid)]
-    Threshold {
-        threshold: QueryThreshold,
-        resource: BencherResource,
-        rate_limit: u32,
-    },
-}
-
 #[macro_export]
 macro_rules! fn_rate_limit {
     ($table:ident, $resource:ident) => {
@@ -80,7 +31,7 @@ macro_rules! fn_rate_limit {
                     if creation_count >= context.rate_limit.unclaimed =>
                 {
                     Err($crate::error::too_many_requests(
-                        $crate::macros::rate_limit::RateLimitError::UnclaimedProject {
+                        $crate::context::RateLimitError::UnclaimedProject {
                             project: query_project,
                             resource: $crate::error::BencherResource::$resource,
                             rate_limit: context.rate_limit.unclaimed,
@@ -91,7 +42,7 @@ macro_rules! fn_rate_limit {
                     if creation_count >= context.rate_limit.claimed =>
                 {
                     Err($crate::error::too_many_requests(
-                        $crate::macros::rate_limit::RateLimitError::ClaimedProject {
+                        $crate::context::RateLimitError::ClaimedProject {
                             project: query_project,
                             resource: $crate::error::BencherResource::$resource,
                             rate_limit: context.rate_limit.claimed,

lib/bencher_schema/src/model/organization/mod.rs

@@ -306,7 +306,7 @@ pub struct InsertOrganization {
 impl InsertOrganization {
     #[cfg(feature = "plus")]
     pub async fn rate_limit(context: &ApiContext, query_user: &QueryUser) -> Result<(), HttpError> {
-        use crate::macros::rate_limit::RateLimitError;
+        use crate::context::RateLimitError;
 
         let resource = BencherResource::Organization;
         let (start_time, end_time) = context.rate_limit.window();

lib/bencher_schema/src/model/project/branch/head.rs

@@ -206,7 +206,7 @@ impl InsertHead {
         context: &ApiContext,
         query_branch: &QueryBranch,
     ) -> Result<(), HttpError> {
-        use crate::{error::BencherResource, macros::rate_limit::RateLimitError};
+        use crate::{context::RateLimitError, error::BencherResource};
 
         let resource = BencherResource::Head;
         let (start_time, end_time) = context.rate_limit.window();

lib/bencher_schema/src/model/project/mod.rs

@@ -503,7 +503,7 @@ impl InsertProject {
         context: &ApiContext,
         query_organization: &QueryOrganization,
     ) -> Result<(), HttpError> {
-        use crate::macros::rate_limit::RateLimitError;
+        use crate::context::RateLimitError;
 
         let resource = BencherResource::Project;
         let (start_time, end_time) = context.rate_limit.window();

lib/bencher_schema/src/model/project/threshold/model.rs

@@ -128,7 +128,7 @@ impl InsertModel {
         context: &crate::ApiContext,
         query_threshold: &QueryThreshold,
     ) -> Result<(), HttpError> {
-        use crate::{conn_lock, error::issue_error, macros::rate_limit::RateLimitError};
+        use crate::{conn_lock, context::RateLimitError, error::issue_error};
 
         let resource = BencherResource::Model;
         let (start_time, end_time) = context.rate_limit.window();

lib/bencher_schema/src/model/user/token.rs

@@ -98,7 +98,7 @@ impl InsertToken {
         context: &crate::ApiContext,
         query_user: &QueryUser,
     ) -> Result<(), HttpError> {
-        use crate::{conn_lock, macros::rate_limit::RateLimitError};
+        use crate::{conn_lock, context::RateLimitError};
 
         let resource = BencherResource::Token;
         let (start_time, end_time) = context.rate_limit.window();