update rate limiter to support more than one limited action

Author: pietroalbini

src/app.rs

@@ -10,6 +10,7 @@ use crate::downloads_counter::DownloadsCounter;
 use crate::email::Emails;
 use crate::github::{GitHubClient, RealGitHubClient};
 use crate::metrics::{InstanceMetrics, ServiceMetrics};
+use crate::rate_limiter::RateLimiter;
 use crate::storage::Storage;
 use axum::extract::{FromRef, FromRequestParts, State};
 use diesel::r2d2;
@@ -68,6 +69,9 @@ pub struct App {
 
     /// In-flight request counters for the `balance_capacity` middleware.
     pub balance_capacity: BalanceCapacityState,
+
+    /// Rate limit select actions.
+    pub rate_limiter: RateLimiter,
 }
 
 impl App {
@@ -178,6 +182,7 @@ impl App {
             http_client,
             fastboot_client,
             balance_capacity: Default::default(),
+            rate_limiter: RateLimiter::new(config.rate_limiter.clone()),
             config,
         }
     }

src/config/server.rs

@@ -2,15 +2,15 @@ use anyhow::{anyhow, Context};
 use ipnetwork::IpNetwork;
 use oauth2::{ClientId, ClientSecret};
 
-use crate::rate_limiter::RateLimiter;
+use crate::rate_limiter::{LimitedAction, RateLimiterConfig};
 use crate::{env, env_optional, Env};
 
 use super::base::Base;
 use super::database_pools::DatabasePools;
 use crate::config::balance_capacity::BalanceCapacityConfig;
 use crate::storage::StorageConfig;
 use http::HeaderValue;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::net::IpAddr;
 use std::time::Duration;
 
@@ -30,7 +30,7 @@ pub struct Server {
     pub gh_client_secret: ClientSecret,
     pub max_upload_size: u64,
     pub max_unpack_size: u64,
-    pub rate_limiter: RateLimiter,
+    pub rate_limiter: HashMap<LimitedAction, RateLimiterConfig>,
     pub new_version_rate_limit: Option<u32>,
     pub blocked_traffic: Vec<(String, Vec<String>)>,
     pub max_allowed_page_offset: u32,
@@ -140,6 +140,24 @@ impl Default for Server {
             .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
             .ok();
 
+        // Dynamically load the configuration for all the rate limiting actions. See
+        // `src/rate_limiter.rs` for their definition.
+        let mut rate_limiter = HashMap::new();
+        for action in LimitedAction::VARIANTS {
+            let env_var_key = action.env_var_key();
+            rate_limiter.insert(
+                *action,
+                RateLimiterConfig {
+                    rate: Duration::from_secs(
+                        env_optional(&format!("RATE_LIMITER_{env_var_key}_RATE_SECONDS"))
+                            .unwrap_or_else(|| action.default_rate_seconds()),
+                    ),
+                    burst: env_optional(&format!("RATE_LIMITER_{env_var_key}_BURST"))
+                        .unwrap_or_else(|| action.default_burst()),
+                },
+            );
+        }
+
         Server {
             db: DatabasePools::full_from_environment(&base),
             storage: StorageConfig::from_environment(),
@@ -153,7 +171,7 @@ impl Default for Server {
             gh_client_secret: ClientSecret::new(env("GH_CLIENT_SECRET")),
             max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
             max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed
-            rate_limiter: Default::default(),
+            rate_limiter,
             new_version_rate_limit: env_optional("MAX_NEW_VERSIONS_DAILY"),
             blocked_traffic: blocked_traffic(),
             max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),

src/controllers/krate/publish.rs

@@ -137,7 +137,7 @@ pub async fn publish(app: AppState, req: BytesRequest) -> AppResult<Json<GoodCra
             };
 
             let license_file = new_crate.license_file.as_deref();
-            let krate = persist.create_or_update(conn, user.id, Some(&app.config.rate_limiter))?;
+            let krate = persist.create_or_update(conn, user.id, Some(&app.rate_limiter))?;
 
             let owners = krate.owners(conn)?;
             if user.rights(&app, &owners)? < Rights::Publish {

src/lib.rs

@@ -46,7 +46,7 @@ pub mod github;
 pub mod headers;
 pub mod metrics;
 pub mod middleware;
-mod rate_limiter;
+pub mod rate_limiter;
 pub mod schema;
 #[macro_use]
 pub mod sql;

src/models/krate.rs

@@ -17,7 +17,7 @@ use crate::models::{
 use crate::util::errors::{cargo_err, AppResult};
 
 use crate::models::helpers::with_count::*;
-use crate::rate_limiter::RateLimiter;
+use crate::rate_limiter::{LimitedAction, RateLimiter};
 use crate::schema::*;
 use crate::sql::canon_crate_name;
 
@@ -119,7 +119,7 @@ impl<'a> NewCrate<'a> {
             // first so we know whether to add an owner
             if let Some(krate) = self.save_new_crate(conn, uploader)? {
                 if let Some(rate_limit) = rate_limit {
-                    rate_limit.check_rate_limit(uploader, conn)?;
+                    rate_limit.check_rate_limit(uploader, LimitedAction::PublishNew, conn)?;
                 }
                 return Ok(krate);
             }