move rate limiting for existing crate publishes to the app

Author: pietroalbini

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;
-		}
 	}
 }

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,7 +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.rate_limiter))?;
+        let krate = persist.create_or_update(&conn, user.id)?;
 
         let owners = krate.owners(&conn)?;
         if user.rights(req.app(), &owners)? < Rights::Publish {

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 {

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::rate_limiter::{LimitedAction, RateLimiter};
 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<&RateLimiter>,
-    ) -> 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, LimitedAction::PublishNew, conn)?;
-                }
                 return Ok(krate);
             }
 

src/rate_limiter.rs

@@ -9,6 +9,7 @@ use std::time::Duration;
 crate::pg_enum! {
     pub enum LimitedAction {
         PublishNew = 0,
+        PublishExisting = 1,
     }
 }
 
@@ -17,13 +18,15 @@ impl LimitedAction {
     pub fn default_rate_seconds(&self) -> u64 {
         match self {
             LimitedAction::PublishNew => 60 * 10,
+            LimitedAction::PublishExisting => 60,
         }
     }
 
     /// How many requests a user can make before the rate limit goes into effect.
     pub fn default_burst(&self) -> i32 {
         match self {
             LimitedAction::PublishNew => 5,
+            LimitedAction::PublishExisting => 30,
         }
     }
 
@@ -32,13 +35,17 @@ impl LimitedAction {
             LimitedAction::PublishNew => {
                 "You have published too many new crates in a short period of time."
             }
+            LimitedAction::PublishExisting => {
+                "You have published too many versions of existing crates in a short period of time."
+            }
         }
     }
 
     /// Key used to identify this action in environment variables. See `src/config.rs`.
     pub fn env_var_key(&self) -> &'static str {
         match self {
             LimitedAction::PublishNew => "PUBLISH_NEW",
+            LimitedAction::PublishExisting => "PUBLISH_EXISTING",
         }
     }
 }
@@ -350,6 +357,49 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn two_actions_dont_interfere_with_each_other() -> QueryResult<()> {
+        let conn = pg_connection();
+        let now = now();
+        let user_id = new_user(&conn, "user")?;
+
+        let mut config = HashMap::new();
+        config.insert(
+            LimitedAction::PublishNew,
+            RateLimiterConfig {
+                rate: Duration::from_secs(1),
+                burst: 10,
+            },
+        );
+        config.insert(
+            LimitedAction::PublishExisting,
+            RateLimiterConfig {
+                rate: Duration::from_secs(1),
+                burst: 20,
+            },
+        );
+        let rate = RateLimiter::new(config);
+
+        let refill_time = now + chrono::Duration::seconds(4);
+        assert_eq!(
+            10,
+            rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
+                .tokens,
+        );
+        assert_eq!(
+            9,
+            rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
+                .tokens,
+        );
+        assert_eq!(
+            20,
+            rate.take_token(user_id, LimitedAction::PublishExisting, refill_time, &conn)?
+                .tokens,
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
         let conn = pg_connection();
@@ -375,6 +425,47 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn override_is_different_for_each_action() -> QueryResult<()> {
+        let conn = pg_connection();
+        let now = now();
+        let user_id = new_user(&conn, "user")?;
+
+        let mut config = HashMap::new();
+        for action in &[LimitedAction::PublishNew, LimitedAction::PublishExisting] {
+            config.insert(
+                *action,
+                RateLimiterConfig {
+                    rate: Duration::from_secs(1),
+                    burst: 10,
+                },
+            );
+        }
+        let rate = RateLimiter::new(config);
+
+        diesel::insert_into(publish_rate_overrides::table)
+            .values((
+                publish_rate_overrides::user_id.eq(user_id),
+                publish_rate_overrides::action.eq(LimitedAction::PublishNew),
+                publish_rate_overrides::burst.eq(20),
+            ))
+            .execute(&conn)?;
+
+        let refill_time = now + chrono::Duration::seconds(4);
+        assert_eq!(
+            20,
+            rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
+                .tokens,
+        );
+        assert_eq!(
+            10,
+            rate.take_token(user_id, LimitedAction::PublishExisting, refill_time, &conn)?
+                .tokens,
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn overrides_can_expire() -> QueryResult<()> {
         let conn = pg_connection();

src/tasks/update_downloads.rs

@@ -98,7 +98,7 @@ mod test {
             name: "foo",
             ..Default::default()
         }
-        .create_or_update(conn, user_id, None)
+        .create_or_update(conn, user_id)
         .unwrap();
         let version = NewVersion::new(
             krate.id,

src/tests/builders/krate.rs

@@ -114,9 +114,7 @@ impl<'a> CrateBuilder<'a> {
     pub fn build(mut self, connection: &PgConnection) -> AppResult<Crate> {
         use diesel::{insert_into, select, update};
 
-        let mut krate = self
-            .krate
-            .create_or_update(connection, self.owner_id, None)?;
+        let mut krate = self.krate.create_or_update(connection, self.owner_id)?;
 
         // Since we are using `NewCrate`, we can't set all the
         // crate properties in a single DB call.