Auto merge of #3790 - pietroalbini:pa-rate-limit-updates, r=jtgeibel

Move new versions rate limits to the application This PR refactors the rate limiting code to support more kinds of rate limits, and moves the new versions rate limit from nginx to the application. See #3780 for the rationale of this change. * The rate limiting for new publishes of existing crates is now tracked per-user instead of per-IP, with the ability to add overrides similar to the rate limiting for new crates. The limit is still the usual 1 version/min with a burst of 30, but in practice this **could lower** the rate limiting, as before the limit was between 30 and 60 depending how lucky users were with the load balancing. * `PublishRateLimit` is now called `RateLimiter`, and for every operation it requires a `LimitedAction`. This allows different rate limits to exist for `LimitedAction::PublishNew` and `LimitedAction::PublishExisting`, and adding new kinds of rate limits only requires adding a new enum variant. * The environment variables to override the default rate limits changed to `RATE_LIMITING_{KIND}_RATE_SECONDS` and `RATE_LIMITING_{KIND}_BURST`. There are two things that will be done in followup PRs: * Remove the default for the `action` column. * Rename the `publish_limit_buckets` to `rate_limit_buckets` and `publish_rate_overrides` to `rate_limit_overrides`
rust-lang · Aug 11, 2021 · 1097171 · 1097171
2 parents 6be0846 + 1a5b0e8
commit 1097171
Show file tree

Hide file tree

Showing 25 changed files with 1,212 additions and 531 deletions.
diff --git a/config/nginx.conf.erb b/config/nginx.conf.erb
@@ -176,8 +176,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;
 	}
@@ -261,12 +259,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;
-		}
 	}
 }
diff --git a/migrations/2021-07-18-125813_add_rate_limit_action/down.sql b/migrations/2021-07-18-125813_add_rate_limit_action/down.sql
@@ -0,0 +1,9 @@
+DELETE FROM publish_limit_buckets WHERE action != 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_limit_buckets DROP COLUMN action;
+
+DELETE FROM publish_rate_overrides WHERE action != 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_rate_overrides DROP COLUMN action;
diff --git a/migrations/2021-07-18-125813_add_rate_limit_action/up.sql b/migrations/2021-07-18-125813_add_rate_limit_action/up.sql
@@ -0,0 +1,7 @@
+ALTER TABLE publish_limit_buckets ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id, action);
+
+ALTER TABLE publish_rate_overrides ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id, action);
diff --git a/src/app.rs b/src/app.rs
@@ -8,6 +8,7 @@ use crate::downloads_counter::DownloadsCounter;
 use crate::email::Emails;
 use crate::github::GitHubClient;
 use crate::metrics::{InstanceMetrics, ServiceMetrics};
+use crate::rate_limiter::RateLimiter;
 use diesel::r2d2;
 use oauth2::basic::BasicClient;
 use reqwest::blocking::Client;
@@ -43,6 +44,8 @@ pub struct App {
     /// Metrics related to this specific instance of the service
     pub instance_metrics: InstanceMetrics,
 
+    pub rate_limiter: RateLimiter,
+
     /// A configured client for outgoing HTTP requests
     ///
     /// In production this shares a single connection pool across requests.  In tests
@@ -165,6 +168,7 @@ impl App {
             read_only_replica_database: replica_database,
             github,
             github_oauth,
+            rate_limiter: RateLimiter::new(config.rate_limiter.clone()),
             config,
             downloads_counter: DownloadsCounter::new(),
             emails: Emails::from_environment(),

diff --git a/src/config.rs b/src/config.rs
@@ -1,5 +1,7 @@
-use crate::publish_rate_limit::PublishRateLimit;
+use crate::rate_limiter::{LimitedAction, RateLimiterConfig};
 use crate::{env, env_optional, uploaders::Uploader, Env};
+use std::collections::HashMap;
+use std::time::Duration;
 
 mod base;
 mod database_pools;
@@ -16,7 +18,6 @@ pub struct Server {
     pub gh_base_url: String,
     pub max_upload_size: u64,
     pub max_unpack_size: u64,
-    pub publish_rate_limit: PublishRateLimit,
     pub blocked_traffic: Vec<(String, Vec<String>)>,
     pub max_allowed_page_offset: u32,
     pub page_offset_ua_blocklist: Vec<String>,
@@ -27,6 +28,7 @@ pub struct Server {
     pub metrics_authorization_token: Option<String>,
     pub use_test_database_pool: bool,
     pub instance_metrics_log_every_seconds: Option<u64>,
+    pub rate_limiter: HashMap<LimitedAction, RateLimiterConfig>,
 }
 
 impl Default for Server {
@@ -64,12 +66,34 @@ impl Default for Server {
             .split(',')
             .map(ToString::to_string)
             .collect();
+
         let page_offset_ua_blocklist = match env_optional::<String>("WEB_PAGE_OFFSET_UA_BLOCKLIST")
         {
             None => vec![],
             Some(s) if s.is_empty() => vec![],
             Some(s) => s.split(',').map(String::from).collect(),
         };
+
+        // 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 {
+            rate_limiter.insert(
+                *action,
+                RateLimiterConfig {
+                    rate: Duration::from_secs(
+                        env_optional(&format!(
+                            "RATE_LIMITER_{}_RATE_SECONDS",
+                            action.env_var_key()
+                        ))
+                        .unwrap_or_else(|| action.default_rate_seconds()),
+                    ),
+                    burst: env_optional(&format!("RATE_LIMITER_{}_BURST", action.env_var_key()))
+                        .unwrap_or_else(|| action.default_burst()),
+                },
+            );
+        }
+
         Server {
             db: DatabasePools::full_from_environment(),
             base: Base::from_environment(),
@@ -79,7 +103,6 @@ impl Default for Server {
             gh_base_url: "https://api.github.com".to_string(),
             max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
             max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed
-            publish_rate_limit: Default::default(),
             blocked_traffic: blocked_traffic(),
             max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),
             page_offset_ua_blocklist,
@@ -96,6 +119,7 @@ impl Default for Server {
             metrics_authorization_token: dotenv::var("METRICS_AUTHORIZATION_TOKEN").ok(),
             use_test_database_pool: false,
             instance_metrics_log_every_seconds: env_optional("INSTANCE_METRICS_LOG_EVERY_SECONDS"),
+            rate_limiter,
         }
     }
 }

diff --git a/src/controllers/krate/publish.rs b/src/controllers/krate/publish.rs
@@ -11,6 +11,7 @@ use crate::models::{
     NewVersion, Rights, VersionAction,
 };
 
+use crate::rate_limiter::LimitedAction;
 use crate::render;
 use crate::schema::*;
 use crate::util::errors::{cargo_err, AppResult};
@@ -72,6 +73,16 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
         ))
     })?;
 
+    // Check the rate limits on the endpoint. Publishing a new crate uses a different (stricter)
+    // limiting pool than publishing a new version of an existing crate to prevent mass squatting.
+    let limited_action = if Crate::by_name(&new_crate.name).execute(&*conn)? == 0 {
+        LimitedAction::PublishNew
+    } else {
+        LimitedAction::PublishExisting
+    };
+    app.rate_limiter
+        .check_rate_limit(user.id, limited_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(|| {
@@ -107,8 +118,7 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
         };
 
         let license_file = new_crate.license_file.as_deref();
-        let krate =
-            persist.create_or_update(&conn, user.id, Some(&app.config.publish_rate_limit))?;
+        let krate = persist.create_or_update(&conn, user.id)?;
 
         let owners = krate.owners(&conn)?;
         if user.rights(req.app(), &owners)? < Rights::Publish {

diff --git a/src/downloads_counter.rs b/src/downloads_counter.rs
@@ -438,7 +438,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 {

diff --git a/src/lib.rs b/src/lib.rs
@@ -44,12 +44,13 @@ pub mod git;
 pub mod github;
 pub mod metrics;
 pub mod middleware;
-mod publish_rate_limit;
+pub mod rate_limiter;
 pub mod render;
 pub mod schema;
 pub mod tasks;
 mod test_util;
 pub mod uploaders;
+#[macro_use]
 pub mod util;
 
 pub mod controllers;

diff --git a/src/models/action.rs b/src/models/action.rs
@@ -1,23 +1,14 @@
-use chrono::NaiveDateTime;
-use diesel::prelude::*;
-use diesel::{
-    deserialize::{self, FromSql},
-    pg::Pg,
-    serialize::{self, Output, ToSql},
-    sql_types::Integer,
-};
-use std::io::Write;
-
 use crate::models::{ApiToken, User, Version};
 use crate::schema::*;
+use chrono::NaiveDateTime;
+use diesel::prelude::*;
 
-#[derive(Debug, Clone, Copy, PartialEq, FromSqlRow, AsExpression)]
-#[repr(i32)]
-#[sql_type = "Integer"]
-pub enum VersionAction {
-    Publish = 0,
-    Yank = 1,
-    Unyank = 2,
+pg_enum! {
+    pub enum VersionAction {
+        Publish = 0,
+        Yank = 1,
+        Unyank = 2,
+    }
 }
 
 impl From<VersionAction> for &'static str {
@@ -38,23 +29,6 @@ impl From<VersionAction> for String {
     }
 }
 
-impl FromSql<Integer, Pg> for VersionAction {
-    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
-        match <i32 as FromSql<Integer, Pg>>::from_sql(bytes)? {
-            0 => Ok(VersionAction::Publish),
-            1 => Ok(VersionAction::Yank),
-            2 => Ok(VersionAction::Unyank),
-            n => Err(format!("unknown version action: {}", n).into()),
-        }
-    }
-}
-
-impl ToSql<Integer, Pg> for VersionAction {
-    fn to_sql<W: Write>(&self, out: &mut Output<'_, W, Pg>) -> serialize::Result {
-        ToSql::<Integer, Pg>::to_sql(&(*self as i32), out)
-    }
-}
-
 #[derive(Debug, Clone, Copy, Queryable, Identifiable, Associations)]
 #[belongs_to(Version)]
 #[belongs_to(User, foreign_key = "user_id")]

diff --git a/src/models/dependency.rs b/src/models/dependency.rs
@@ -1,7 +1,3 @@
-use diesel::deserialize::{self, FromSql};
-use diesel::pg::Pg;
-use diesel::sql_types::Integer;
-
 use crate::models::{Crate, Version};
 use crate::schema::*;
 
@@ -32,23 +28,10 @@ pub struct ReverseDependency {
     pub name: String,
 }
 
-#[derive(Copy, Clone, Serialize, Deserialize, Debug, FromSqlRow)]
-#[serde(rename_all = "lowercase")]
-#[repr(u32)]
-pub enum DependencyKind {
-    Normal = 0,
-    Build = 1,
-    Dev = 2,
-    // if you add a kind here, be sure to update `from_row` below.
-}
-
-impl FromSql<Integer, Pg> for DependencyKind {
-    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
-        match <i32 as FromSql<Integer, Pg>>::from_sql(bytes)? {
-            0 => Ok(DependencyKind::Normal),
-            1 => Ok(DependencyKind::Build),
-            2 => Ok(DependencyKind::Dev),
-            n => Err(format!("unknown kind: {}", n).into()),
-        }
+pg_enum! {
+    pub enum DependencyKind {
+        Normal = 0,
+        Build = 1,
+        Dev = 2,
     }
 }
diff --git a/src/models/krate.rs b/src/models/krate.rs
@@ -15,7 +15,6 @@ use crate::models::{
 use crate::util::errors::{cargo_err, AppResult};
 
 use crate::models::helpers::with_count::*;
-use crate::publish_rate_limit::PublishRateLimit;
 use crate::schema::*;
 
 #[derive(Debug, Queryable, Identifiable, Associations, Clone, Copy)]
@@ -93,12 +92,7 @@ pub struct NewCrate<'a> {
 }
 
 impl<'a> NewCrate<'a> {
-    pub fn create_or_update(
-        self,
-        conn: &PgConnection,
-        uploader: i32,
-        rate_limit: Option<&PublishRateLimit>,
-    ) -> AppResult<Crate> {
+    pub fn create_or_update(self, conn: &PgConnection, uploader: i32) -> AppResult<Crate> {
         use diesel::update;
 
         self.validate()?;
@@ -108,9 +102,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);
             }