Merge pull request #6875 from ferrous-systems/pa-rate-limits-part2

Add application-level rate limits for crate updates, yanking and unyanking

Author: Turbo87

config/nginx.conf.erb

@@ -185,8 +185,6 @@ http {
 	client_body_timeout 30;
 	client_max_body_size 50m;
 
-	limit_req_zone $remote_addr zone=publish:10m rate=1r/m;
-
 	upstream app_server {
 		server localhost:8888 fail_timeout=0;
 	}
@@ -262,12 +260,5 @@ http {
 				proxy_pass http://app_server;
 			}
 		<% end %>
-
-		location ~ ^/api/v./crates/new$ {
-			proxy_pass http://app_server;
-
-			limit_req zone=publish burst=30 nodelay;
-			limit_req_status 429;
-		}
 	}
 }

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

@@ -19,6 +19,7 @@ use crate::models::{
 
 use crate::middleware::log_request::RequestLogExt;
 use crate::models::token::EndpointScope;
+use crate::rate_limiter::LimitedAction;
 use crate::schema::*;
 use crate::util::errors::{cargo_err, internal, AppResult};
 use crate::util::Maximums;
@@ -101,6 +102,14 @@ pub async fn publish(app: AppState, req: BytesRequest) -> AppResult<Json<GoodCra
             ))
         })?;
 
+        // Use a different rate limit whether this is a new or an existing crate.
+        let rate_limit_action = match existing_crate {
+            Some(_) => LimitedAction::PublishUpdate,
+            None => LimitedAction::PublishNew,
+        };
+        app.rate_limiter
+            .check_rate_limit(user.id, rate_limit_action, conn)?;
+
         // Create a transaction on the database, if there are no errors,
         // commit the transactions to record a new or updated crate.
         conn.transaction(|conn| {
@@ -137,7 +146,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)?;
 
             let owners = krate.owners(conn)?;
             if user.rights(&app, &owners)? < Rights::Publish {

src/controllers/version/yank.rs

@@ -8,6 +8,7 @@ use crate::controllers::cargo_prelude::*;
 use crate::models::token::EndpointScope;
 use crate::models::Rights;
 use crate::models::{insert_version_owner_action, VersionAction};
+use crate::rate_limiter::LimitedAction;
 use crate::schema::versions;
 
 /// Handles the `DELETE /crates/:crate_id/:version/yank` route.
@@ -58,6 +59,10 @@ fn modify_yank(
         .for_crate(crate_name)
         .check(req, conn)?;
 
+    state
+        .rate_limiter
+        .check_rate_limit(auth.user_id(), LimitedAction::YankUnyank, conn)?;
+
     let (version, krate) = version_and_crate(conn, crate_name, version)?;
     let api_token_id = auth.api_token_id();
     let user = auth.user();

src/downloads_counter.rs

@@ -439,7 +439,7 @@ mod tests {
                 name: "foo",
                 ..NewCrate::default()
             }
-            .create_or_update(conn, user.id, None)
+            .create_or_update(conn, user.id)
             .expect("failed to create crate");
 
             Self {

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;
 pub mod sql;
 pub mod ssh;

src/models/krate.rs

@@ -17,7 +17,6 @@ use crate::models::{
 use crate::util::errors::{cargo_err, AppResult};
 
 use crate::models::helpers::with_count::*;
-use crate::rate_limiter::RateLimiter;
 use crate::schema::*;
 use crate::sql::canon_crate_name;
 
@@ -103,12 +102,7 @@ pub struct NewCrate<'a> {
 }
 
 impl<'a> NewCrate<'a> {
-    pub fn create_or_update(
-        self,
-        conn: &mut PgConnection,
-        uploader: i32,
-        rate_limit: Option<&RateLimiter>,
-    ) -> AppResult<Crate> {
+    pub fn create_or_update(self, conn: &mut PgConnection, uploader: i32) -> AppResult<Crate> {
         use diesel::update;
 
         self.validate()?;
@@ -118,9 +112,6 @@ impl<'a> NewCrate<'a> {
             // To avoid race conditions, we try to insert
             // 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)?;
-                }
                 return Ok(krate);
             }