rate_limit_config

Author: epompeii

lib/api_auth/src/github.rs

@@ -56,10 +56,11 @@ async fn post_inner(
     // If not on Bencher Cloud, then at least one organization must have a valid Bencher Plus license
     if !context.is_bencher_cloud
         && LicenseUsage::get_for_server(
-            conn_lock!(context),
+            &context.database.connection,
             &context.licensor,
             Some(PlanLevel::Enterprise),
-        )?
+        )
+        .await?
         .is_empty()
     {
         return Err(payment_required_error(

lib/api_organizations/src/projects.rs

@@ -219,7 +219,7 @@ async fn post_inner(
         QueryProject::is_visibility_public(visibility)?;
         #[cfg(feature = "plus")]
         PlanKind::check_for_organization(
-            conn_lock!(context),
+            context,
             context.biller.as_ref(),
             &context.licensor,
             &query_organization,

lib/api_projects/src/projects.rs

@@ -284,7 +284,7 @@ async fn patch_inner(
         QueryProject::is_visibility_public(visibility)?;
         #[cfg(feature = "plus")]
         PlanKind::check_for_project(
-            conn_lock!(context),
+            context,
             context.biller.as_ref(),
             &context.licensor,
             &query_project,

lib/bencher_config/src/config_tx.rs

@@ -196,6 +196,12 @@ fn into_context(
         &security.secret_key,
     );
 
+    #[cfg(feature = "plus")]
+    let rate_limit = plus
+        .as_ref()
+        .and_then(|plus| plus.rate_limit.map(Into::into))
+        .unwrap_or_default();
+
     info!(&log, "Configuring Bencher Plus");
     #[cfg(feature = "plus")]
     let Plus {
@@ -222,6 +228,8 @@ fn into_context(
         },
         restart_tx,
         #[cfg(feature = "plus")]
+        rate_limit,
+        #[cfg(feature = "plus")]
         github,
         #[cfg(feature = "plus")]
         stats,

lib/bencher_json/src/system/config/mod.rs

@@ -22,6 +22,7 @@ pub use plus::{
         JsonCloud,
     },
     litestream::{JsonLitestream, JsonReplica},
+    rate_limit::JsonRateLimit,
     stats::JsonStats,
     JsonPlus,
 };

lib/bencher_json/src/system/config/plus/mod.rs

@@ -8,16 +8,20 @@ use serde::{Deserialize, Serialize};
 pub mod cloud;
 pub mod github;
 pub mod litestream;
+pub mod rate_limit;
 pub mod stats;
 
 pub use cloud::JsonCloud;
 pub use github::JsonGitHub;
 pub use litestream::JsonLitestream;
+pub use rate_limit::JsonRateLimit;
 pub use stats::JsonStats;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[cfg_attr(feature = "schema", derive(JsonSchema))]
 pub struct JsonPlus {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub rate_limit: Option<JsonRateLimit>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub github: Option<JsonGitHub>,
     #[serde(alias = "disaster_recovery", skip_serializing_if = "Option::is_none")]

lib/bencher_json/src/system/config/plus/rate_limit.rs

@@ -0,0 +1,11 @@
+#[cfg(feature = "schema")]
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+#[cfg_attr(feature = "schema", derive(JsonSchema))]
+pub struct JsonRateLimit {
+    pub window: Option<u32>,
+    pub unclaimed: Option<u32>,
+    pub claimed: Option<u32>,
+}

lib/bencher_json/src/system/config/plus/stats.rs

@@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[cfg_attr(feature = "schema", derive(JsonSchema))]
 pub struct JsonStats {
-    // Number of seconds from midnight
+    /// Number of seconds from midnight
     pub offset: Option<u32>,
+    /// Enable stats collection
     pub enabled: Option<bool>,
 }

lib/bencher_schema/src/context/mod.rs

@@ -14,6 +14,8 @@ use crate::model::project::QueryProject;
 mod database;
 mod indexer;
 mod messenger;
+#[cfg(feature = "plus")]
+mod rate_limit;
 mod rbac;
 #[cfg(feature = "plus")]
 mod stats;
@@ -24,6 +26,8 @@ pub use indexer::{IndexError, Indexer};
 #[cfg(feature = "plus")]
 pub use messenger::ServerStatsBody;
 pub use messenger::{Body, ButtonBody, Email, Message, Messenger, NewUserBody};
+#[cfg(feature = "plus")]
+pub use rate_limit::RateLimit;
 pub use rbac::{Rbac, RbacError};
 #[cfg(feature = "plus")]
 pub use stats::StatsSettings;
@@ -36,6 +40,8 @@ pub struct ApiContext {
     pub database: Database,
     pub restart_tx: Sender<()>,
     #[cfg(feature = "plus")]
+    pub rate_limit: RateLimit,
+    #[cfg(feature = "plus")]
     pub github: Option<GitHub>,
     #[cfg(feature = "plus")]
     pub indexer: Option<Indexer>,

lib/bencher_schema/src/context/rate_limit.rs

@@ -0,0 +1,46 @@
+use std::time::Duration;
+
+use bencher_json::{system::config::JsonRateLimit, DateTime};
+
+const DAY: Duration = Duration::from_secs(24 * 60 * 60);
+const UNCLAIMED_RATE_LIMIT: u32 = u8::MAX as u32;
+const CLAIMED_RATE_LIMIT: u32 = u16::MAX as u32;
+
+pub struct RateLimit {
+    pub window: Duration,
+    pub unclaimed: u32,
+    pub claimed: u32,
+}
+
+impl From<JsonRateLimit> for RateLimit {
+    fn from(json: JsonRateLimit) -> Self {
+        let JsonRateLimit {
+            window,
+            unclaimed,
+            claimed,
+        } = json;
+        Self {
+            window: window.map(u64::from).map_or(DAY, Duration::from_secs),
+            unclaimed: unclaimed.unwrap_or(UNCLAIMED_RATE_LIMIT),
+            claimed: claimed.unwrap_or(CLAIMED_RATE_LIMIT),
+        }
+    }
+}
+
+impl Default for RateLimit {
+    fn default() -> Self {
+        Self {
+            window: DAY,
+            unclaimed: UNCLAIMED_RATE_LIMIT,
+            claimed: CLAIMED_RATE_LIMIT,
+        }
+    }
+}
+
+impl RateLimit {
+    pub fn window(&self) -> (DateTime, DateTime) {
+        let end_time = chrono::Utc::now();
+        let start_time = end_time - self.window;
+        (start_time.into(), end_time.into())
+    }
+}