Simplify daily version limit implementation

The `PublishRateLimit` struct is somewhat closely tied to the current system of limiting new crate creation. Instead of adding the new daily versions limit to it, we can just keep the code in the controller without having to touch any of the `PublishRateLimit` constructor calls.

Author: Turbo87

src/config.rs

@@ -25,6 +25,7 @@ pub struct Server {
     pub max_upload_size: u64,
     pub max_unpack_size: u64,
     pub publish_rate_limit: PublishRateLimit,
+    pub new_version_rate_limit: Option<u32>,
     pub blocked_traffic: Vec<(String, Vec<String>)>,
     pub max_allowed_page_offset: u32,
     pub page_offset_ua_blocklist: Vec<String>,
@@ -118,6 +119,7 @@ impl Default for Server {
             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(),
+            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),
             page_offset_ua_blocklist,

src/controllers/krate/publish.rs

@@ -128,9 +128,14 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
             )));
         }
 
-        app.config
-            .publish_rate_limit
-            .check_daily_limit(krate.id, &conn)?;
+        if let Some(daily_version_limit) = app.config.new_version_rate_limit {
+            let published_today = count_versions_published_today(krate.id, &conn)?;
+            if published_today >= daily_version_limit as i64 {
+                return Err(cargo_err(
+                    "You have published too many versions of this crate in the last 24 hours",
+                ));
+            }
+        }
 
         // 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)
@@ -263,6 +268,19 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
     })
 }
 
+/// Counts the number of versions for `krate_id` that were published within
+/// the last 24 hours.
+fn count_versions_published_today(krate_id: i32, conn: &PgConnection) -> QueryResult<i64> {
+    use crate::schema::versions::dsl::*;
+    use diesel::dsl::{now, IntervalDsl};
+
+    versions
+        .filter(crate_id.eq(krate_id))
+        .filter(created_at.gt(now - 24.hours()))
+        .count()
+        .get_result(conn)
+}
+
 /// Used by the `krate::new` function.
 ///
 /// This function parses the JSON headers to interpret the data and validates

src/lib.rs

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

src/publish_rate_limit.rs

@@ -6,13 +6,12 @@ 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::{cargo_err, AppResult, TooManyRequests};
+use crate::util::errors::{AppResult, TooManyRequests};
 
 #[derive(Debug, Clone, Copy)]
 pub struct PublishRateLimit {
     pub rate: Duration,
     pub burst: i32,
-    pub daily: Option<i64>,
 }
 
 impl Default for PublishRateLimit {
@@ -27,14 +26,9 @@ 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,
         }
     }
 }
@@ -112,27 +106,6 @@ 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 versions of this crate in the last 24 hours",
-                ));
-            }
-        }
-        Ok(())
-    }
 }
 
 #[cfg(test)]
@@ -149,7 +122,6 @@ 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 {
@@ -162,7 +134,6 @@ 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 {
@@ -182,7 +153,6 @@ 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)?;
@@ -203,7 +173,6 @@ 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);
@@ -229,7 +198,6 @@ 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);
@@ -251,7 +219,6 @@ 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)?;
@@ -273,7 +240,6 @@ 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)?;
@@ -297,7 +263,6 @@ 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);
@@ -320,7 +285,6 @@ 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);
@@ -343,7 +307,6 @@ 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")?;
@@ -371,7 +334,6 @@ 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/util/test_app.rs

@@ -2,7 +2,6 @@ use super::{MockAnonymousUser, MockCookieUser, MockTokenUser};
 use crate::record;
 use crate::util::{chaosproxy::ChaosProxy, fresh_schema::FreshSchema};
 use cargo_registry::config::{self, DbPoolConfig};
-use cargo_registry::publish_rate_limit::PublishRateLimit;
 use cargo_registry::{background_jobs::Environment, db::DieselPool, App, Emails};
 use cargo_registry_index::testing::UpstreamIndex;
 use cargo_registry_index::{Credentials, Repository as WorkerRepository, RepositoryConfig};
@@ -341,10 +340,8 @@ fn simple_config() -> config::Server {
         gh_base_url: "http://api.github.com".to_string(),
         max_upload_size: 3000,
         max_unpack_size: 2000,
-        publish_rate_limit: PublishRateLimit {
-            daily: Some(10),
-            ..Default::default()
-        },
+        publish_rate_limit: Default::default(),
+        new_version_rate_limit: Some(10),
         blocked_traffic: Default::default(),
         max_allowed_page_offset: 200,
         page_offset_ua_blocklist: vec![],

src/tests/version.rs

@@ -180,9 +180,10 @@ fn version_size() {
 
 #[test]
 fn daily_limit() {
-    let (_, _, user) = TestApp::full().with_user();
+    let (app, _, user) = TestApp::full().with_user();
 
-    for version in 1..=10 {
+    let max_daily_versions = app.as_inner().config.new_version_rate_limit.unwrap();
+    for version in 1..=max_daily_versions {
         let crate_to_publish =
             PublishBuilder::new("foo_daily_limit").version(&format!("0.0.{}", version));
         user.publish_crate(crate_to_publish).good();