move check_daily_limit inside PushRateLimit

Author: Marco Napetti

src/config.rs

@@ -14,7 +14,6 @@ use std::time::Duration;
 
 const DEFAULT_VERSION_ID_CACHE_SIZE: u64 = 10_000;
 const DEFAULT_VERSION_ID_CACHE_TTL: u64 = 5 * 60; // 5 minutes
-const DEFAULT_MAX_NEW_VERSIONS_DAILY: i64 = 10;
 
 pub struct Server {
     pub base: Base,
@@ -42,7 +41,6 @@ pub struct Server {
     pub blocked_routes: HashSet<String>,
     pub version_id_cache_size: u64,
     pub version_id_cache_ttl: Duration,
-    pub max_new_versions_daily: i64,
 }
 
 impl Default for Server {
@@ -78,7 +76,6 @@ impl Default for Server {
     ///   endpoint even with a healthy database pool.
     /// - `BLOCKED_ROUTES`: A comma separated list of HTTP route patterns that are manually blocked
     ///   by an operator (e.g. `/crates/:crate_id/:version/download`).
-    /// - `MAX_NEW_VERSIONS_DAILY`: Maximum number of new crate versions creatable in a 24 hours span.
     ///
     /// # Panics
     ///
@@ -148,8 +145,6 @@ impl Default for Server {
             version_id_cache_ttl: Duration::from_secs(
                 env_optional("VERSION_ID_CACHE_TTL").unwrap_or(DEFAULT_VERSION_ID_CACHE_TTL),
             ),
-            max_new_versions_daily: env_optional("MAX_NEW_VERSIONS_DAILY")
-                .unwrap_or(DEFAULT_MAX_NEW_VERSIONS_DAILY),
         }
     }
 }

src/controllers/krate/publish.rs

@@ -128,6 +128,10 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
             )));
         }
 
+        app.config
+            .publish_rate_limit
+            .check_daily_limit(krate.id, &conn)?;
+
         // Length of the .crate tarball, which appears after the metadata in the request body.
         // TODO: Not sure why we're using the total content length (metadata + .crate file length)
         // to compare against the max upload size... investigate that and perhaps change to use
@@ -173,11 +177,7 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
             hex_cksum.clone(),
             links.clone(),
         )?
-        .save(
-            &conn,
-            &verified_email_address,
-            Some(app.config.max_new_versions_daily),
-        )?;
+        .save(&conn, &verified_email_address)?;
 
         insert_version_owner_action(
             &conn,

src/downloads_counter.rs

@@ -461,7 +461,7 @@ mod tests {
                 None,
             )
             .expect("failed to create version")
-            .save(conn, "[email protected]", None)
+            .save(conn, "[email protected]")
             .expect("failed to save version");
 
             self.next_version += 1;

src/lib.rs

@@ -46,7 +46,7 @@ pub mod email;
 pub mod github;
 pub mod metrics;
 pub mod middleware;
-mod publish_rate_limit;
+pub mod publish_rate_limit;
 pub mod schema;
 pub mod sql;
 mod test_util;

src/models/version.rs

@@ -155,12 +155,7 @@ impl NewVersion {
         Ok(new_version)
     }
 
-    pub fn save(
-        &self,
-        conn: &PgConnection,
-        published_by_email: &str,
-        daily_limit: Option<i64>,
-    ) -> AppResult<Version> {
+    pub fn save(&self, conn: &PgConnection, published_by_email: &str) -> AppResult<Version> {
         use crate::schema::versions::dsl::*;
         use diesel::dsl::exists;
         use diesel::{insert_into, select};
@@ -177,10 +172,6 @@ impl NewVersion {
                 )));
             }
 
-            if let Some(limit) = daily_limit {
-                self.rate_limit(conn, limit)?;
-            }
-
             let version: Version = insert_into(versions).values(self).get_result(conn)?;
 
             insert_into(versions_published_by::table)
@@ -204,26 +195,6 @@ impl NewVersion {
         }
         Ok(())
     }
-
-    fn rate_limit(&self, conn: &PgConnection, daily_limit: i64) -> AppResult<()> {
-        use crate::schema::versions::dsl::*;
-        use diesel::dsl::{count_star, now, IntervalDsl};
-
-        let today: i64 = versions
-            .filter(crate_id.eq(self.crate_id))
-            .filter(created_at.gt(now - 24.hours()))
-            .select(count_star())
-            .first(conn)
-            .optional()?
-            .unwrap_or_default();
-        if today >= daily_limit {
-            Err(cargo_err(
-                "You have published too many crates in the last 24 hours",
-            ))
-        } else {
-            Ok(())
-        }
-    }
 }
 
 fn validate_license_expr(s: &str) -> AppResult<()> {

src/publish_rate_limit.rs

@@ -6,12 +6,13 @@ use std::time::Duration;
 
 use crate::schema::{publish_limit_buckets, publish_rate_overrides};
 use crate::sql::{date_part, floor, greatest, interval_part, least};
-use crate::util::errors::{AppResult, TooManyRequests};
+use crate::util::errors::{cargo_err, AppResult, TooManyRequests};
 
 #[derive(Debug, Clone, Copy)]
 pub struct PublishRateLimit {
     pub rate: Duration,
     pub burst: i32,
+    pub daily: Option<i64>,
 }
 
 impl Default for PublishRateLimit {
@@ -26,9 +27,14 @@ impl Default for PublishRateLimit {
             .parse()
             .ok()
             .unwrap_or(5);
+        let daily = dotenv::var("MAX_NEW_VERSIONS_DAILY")
+            .unwrap_or_default()
+            .parse()
+            .ok();
         Self {
             rate: Duration::from_secs(60) * minutes,
             burst,
+            daily,
         }
     }
 }
@@ -106,6 +112,27 @@ impl PublishRateLimit {
         use diesel::dsl::*;
         (self.rate.as_millis() as i64).milliseconds()
     }
+
+    pub fn check_daily_limit(&self, krate_id: i32, conn: &PgConnection) -> AppResult<()> {
+        use crate::schema::versions::dsl::*;
+        use diesel::dsl::{count_star, now, IntervalDsl};
+
+        if let Some(daily_limit) = self.daily {
+            let today: i64 = versions
+                .filter(crate_id.eq(krate_id))
+                .filter(created_at.gt(now - 24.hours()))
+                .select(count_star())
+                .first(conn)
+                .optional()?
+                .unwrap_or_default();
+            if today >= daily_limit {
+                return Err(cargo_err(
+                    "You have published too many crates in the last 24 hours",
+                ));
+            }
+        }
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -122,6 +149,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let bucket = rate.take_token(new_user(&conn, "user1")?, now, &conn)?;
         let expected = Bucket {
@@ -134,6 +162,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_millis(50),
             burst: 20,
+            daily: None,
         };
         let bucket = rate.take_token(new_user(&conn, "user2")?, now, &conn)?;
         let expected = Bucket {
@@ -153,6 +182,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 5, now)?.user_id;
         let bucket = rate.take_token(user_id, now, &conn)?;
@@ -173,6 +203,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 5, now)?.user_id;
         let refill_time = now + chrono::Duration::seconds(2);
@@ -198,6 +229,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_millis(100),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 5, now)?.user_id;
         let refill_time = now + chrono::Duration::milliseconds(300);
@@ -219,6 +251,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_millis(100),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 5, now)?.user_id;
         let bucket = rate.take_token(user_id, now + chrono::Duration::milliseconds(250), &conn)?;
@@ -240,6 +273,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 1, now)?.user_id;
         let bucket = rate.take_token(user_id, now, &conn)?;
@@ -263,6 +297,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 0, now)?.user_id;
         let refill_time = now + chrono::Duration::seconds(1);
@@ -285,6 +320,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user_bucket(&conn, 8, now)?.user_id;
         let refill_time = now + chrono::Duration::seconds(4);
@@ -307,6 +343,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user(&conn, "user1")?;
         let other_user_id = new_user(&conn, "user2")?;
@@ -334,6 +371,7 @@ mod tests {
         let rate = PublishRateLimit {
             rate: Duration::from_secs(1),
             burst: 10,
+            daily: None,
         };
         let user_id = new_user(&conn, "user1")?;
         let other_user_id = new_user(&conn, "user2")?;

src/tests/builders/version.rs

@@ -104,7 +104,7 @@ impl<'a> VersionBuilder<'a> {
             self.checksum,
             self.links,
         )?
-        .save(connection, "[email protected]", None)?;
+        .save(connection, "[email protected]")?;
 
         if self.yanked {
             vers = update(&vers)