move existing crate update rate limit to the application

Author: pietroalbini

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

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

src/rate_limiter.rs

@@ -12,25 +12,29 @@ use std::time::Duration;
 crate::pg_enum! {
     pub enum LimitedAction {
         PublishNew = 0,
+        PublishUpdate = 1,
     }
 }
 
 impl LimitedAction {
     pub fn default_rate_seconds(&self) -> u64 {
         match self {
             LimitedAction::PublishNew => 60 * 60,
+            LimitedAction::PublishUpdate => 60,
         }
     }
 
     pub fn default_burst(&self) -> i32 {
         match self {
             LimitedAction::PublishNew => 5,
+            LimitedAction::PublishUpdate => 30,
         }
     }
 
     pub fn env_var_key(&self) -> &'static str {
         match self {
             LimitedAction::PublishNew => "PUBLISH_NEW",
+            LimitedAction::PublishUpdate => "PUBLISH_EXISTING",
         }
     }
 
@@ -39,6 +43,9 @@ impl LimitedAction {
             LimitedAction::PublishNew => {
                 "You have published too many new crates in a short period of time"
             }
+            LimitedAction::PublishUpdate => {
+                "You have published too many updates to existing crates in a short period of time"
+            }
         }
     }
 }

src/tests/builders/krate.rs

@@ -114,9 +114,7 @@ impl<'a> CrateBuilder<'a> {
     pub fn build(mut self, connection: &mut 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.

src/tests/krate/publish.rs

@@ -1028,8 +1028,7 @@ fn publish_new_crate_rate_limited() {
 
     // Uploading a second crate is limited
     let crate_to_publish = PublishBuilder::new("rate_limited2");
-    let response = token.publish_crate(crate_to_publish);
-    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
+    token.publish_crate(crate_to_publish).assert_rate_limited();
 
     assert_eq!(app.stored_files().len(), 2);
 
@@ -1049,9 +1048,9 @@ fn publish_new_crate_rate_limited() {
 }
 
 #[test]
-fn publish_rate_limit_doesnt_affect_existing_crates() {
+fn publish_new_crate_rate_limit_doesnt_affect_existing_crates() {
     let (_, _, _, token) = TestApp::full()
-        .with_rate_limit(LimitedAction::PublishNew, Duration::from_millis(500), 1)
+        .with_rate_limit(LimitedAction::PublishNew, Duration::from_secs(60 * 60), 1)
         .with_token();
 
     // Upload a new crate
@@ -1062,6 +1061,67 @@ fn publish_rate_limit_doesnt_affect_existing_crates() {
     token.publish_crate(new_version).good();
 }
 
+#[test]
+fn publish_existing_crate_rate_limited() {
+    let (app, anon, _, token) = TestApp::full()
+        .with_rate_limit(LimitedAction::PublishUpdate, Duration::from_millis(500), 1)
+        .with_token();
+
+    // Upload a new crate
+    let crate_to_publish = PublishBuilder::new("rate_limited1");
+    token.publish_crate(crate_to_publish).good();
+
+    let json = anon.show_crate("rate_limited1");
+    assert_eq!(json.krate.max_version, "1.0.0");
+    assert_eq!(app.stored_files().len(), 2);
+
+    // Uploading the first update to the crate works
+    let crate_to_publish = PublishBuilder::new("rate_limited1").version("1.0.1");
+    token.publish_crate(crate_to_publish).good();
+
+    let json = anon.show_crate("rate_limited1");
+    assert_eq!(json.krate.max_version, "1.0.1");
+    assert_eq!(app.stored_files().len(), 3);
+
+    // Uploading the second update to the crate is rate limited
+    let crate_to_publish = PublishBuilder::new("rate_limited1").version("1.0.2");
+    token.publish_crate(crate_to_publish).assert_rate_limited();
+
+    // Check that  version 1.0.2 was not published
+    let json = anon.show_crate("rate_limited1");
+    assert_eq!(json.krate.max_version, "1.0.1");
+    assert_eq!(app.stored_files().len(), 3);
+
+    // Wait for the limit to be up
+    thread::sleep(Duration::from_millis(500));
+
+    let crate_to_publish = PublishBuilder::new("rate_limited1").version("1.0.2");
+    token.publish_crate(crate_to_publish).good();
+
+    let json = anon.show_crate("rate_limited1");
+    assert_eq!(json.krate.max_version, "1.0.2");
+    assert_eq!(app.stored_files().len(), 4);
+}
+
+#[test]
+fn publish_existing_crate_rate_limit_doesnt_affect_new_crates() {
+    let (_, _, _, token) = TestApp::full()
+        .with_rate_limit(
+            LimitedAction::PublishUpdate,
+            Duration::from_secs(60 * 60),
+            1,
+        )
+        .with_token();
+
+    // Upload a new crate
+    let crate_to_publish = PublishBuilder::new("rate_limited1");
+    token.publish_crate(crate_to_publish).good();
+
+    // Upload a second new crate
+    let crate_to_publish = PublishBuilder::new("rate_limited2");
+    token.publish_crate(crate_to_publish).good();
+}
+
 #[test]
 fn features_version_2() {
     let (app, _, user, token) = TestApp::full().with_token();

src/tests/util/response.rs

@@ -57,6 +57,12 @@ impl<T> Response<T> {
             .ends_with(target));
         self
     }
+
+    /// Assert that the status code is 429
+    #[track_caller]
+    pub fn assert_rate_limited(&self) {
+        assert_eq!(StatusCode::TOO_MANY_REQUESTS, self.status());
+    }
 }
 
 impl Response<()> {

src/worker/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,