Allow per-user rate limit overrides

Since this is primarily an anti-spam mechanism, if a legitimate user
really is just trying to publish 50 crates at once or something we want
to be able to give them that ability. There's no UI for this, as it's
not expected to be used often (if ever), so it'll require someone with
prod access to add the override for now.

I expect the bar for setting this to be pretty low. Mostly "you are
clearly actually trying to upload a bunch of real crates and are not a
spammer".

Author: sgrif

migrations/2019-04-04-192902_create_publish_rate_overrides/down.sql

@@ -0,0 +1 @@
+DROP TABLE publish_rate_overrides;

migrations/2019-04-04-192902_create_publish_rate_overrides/up.sql

@@ -0,0 +1,4 @@
+CREATE TABLE publish_rate_overrides (
+  user_id INTEGER PRIMARY KEY REFERENCES users,
+  burst INTEGER NOT NULL
+);

src/publish_rate_limit.rs

@@ -3,7 +3,7 @@ use diesel::data_types::PgInterval;
 use diesel::prelude::*;
 use std::time::Duration;
 
-use crate::schema::publish_limit_buckets;
+use crate::schema::{publish_limit_buckets, publish_rate_overrides};
 use crate::util::errors::{CargoResult, TooManyRequests};
 
 #[derive(Debug, Clone, Copy)]
@@ -68,6 +68,13 @@ impl PublishRateLimit {
         sql_function!(fn greatest<T>(x: T, y: T) -> T);
         sql_function!(fn least<T>(x: T, y: T) -> T);
 
+        let burst = publish_rate_overrides::table
+            .find(uploader)
+            .select(publish_rate_overrides::burst)
+            .first::<i32>(conn)
+            .optional()?
+            .unwrap_or(self.burst);
+
         // Interval division is poorly defined in general (what is 1 month / 30 days?)
         // However, for the intervals we're dealing with, it is always well
         // defined, so we convert to an f64 of seconds to represent this.
@@ -79,13 +86,13 @@ impl PublishRateLimit {
         diesel::insert_into(publish_limit_buckets)
             .values((
                 user_id.eq(uploader),
-                tokens.eq(self.burst),
+                tokens.eq(burst),
                 last_refill.eq(now),
             ))
             .on_conflict(user_id)
             .do_update()
             .set((
-                tokens.eq(least(self.burst, greatest(0, tokens - 1) + tokens_to_add)),
+                tokens.eq(least(burst, greatest(0, tokens - 1) + tokens_to_add)),
                 last_refill
                     .eq(last_refill + self.refill_rate().into_sql::<Interval>() * tokens_to_add),
             ))
@@ -288,6 +295,33 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn override_is_used_instead_of_global_burst_if_present() -> CargoResult<()> {
+        let conn = pg_connection();
+        let now = Utc::now().naive_utc();
+
+        let rate = PublishRateLimit {
+            rate: Duration::from_secs(1),
+            burst: 10,
+        };
+        let user_id = new_user(&conn, "user1")?;
+        let other_user_id = new_user(&conn, "user2")?;
+
+        diesel::insert_into(publish_rate_overrides::table)
+            .values((
+                publish_rate_overrides::user_id.eq(user_id),
+                publish_rate_overrides::burst.eq(20),
+            ))
+            .execute(&conn)?;
+
+        let bucket = rate.take_token(user_id, now, &conn)?;
+        let other_bucket = rate.take_token(other_user_id, now, &conn)?;
+
+        assert_eq!(20, bucket.tokens);
+        assert_eq!(10, other_bucket.tokens);
+        Ok(())
+    }
+
     fn new_user(conn: &PgConnection, gh_login: &str) -> CargoResult<i32> {
         use crate::models::NewUser;
 

src/schema.patch

@@ -55,7 +55,7 @@ index df884e4..18e08cd 100644
 
      /// Representation of the `reserved_crate_names` table.
      ///
-@@ -881,22 +901,24 @@ table! {
+@@ -881,23 +901,25 @@ table! {
  }
 
  joinable!(api_tokens -> users (user_id));
@@ -74,6 +74,7 @@ index df884e4..18e08cd 100644
  joinable!(follows -> crates (crate_id));
  joinable!(follows -> users (user_id));
  joinable!(publish_limit_buckets -> users (user_id));
+ joinable!(publish_rate_overrides -> users (user_id));
  joinable!(readme_renderings -> versions (version_id));
 +joinable!(recent_crate_downloads -> crates (crate_id));
  joinable!(version_authors -> users (user_id));
@@ -82,12 +83,12 @@ index df884e4..18e08cd 100644
  joinable!(versions -> crates (crate_id));
 
 @@ -913,13 +935,14 @@ allow_tables_to_appear_in_same_query!(
-     dependencies,
      emails,
      follows,
      keywords,
      metadata,
      publish_limit_buckets,
+     publish_rate_overrides,
      readme_renderings,
 +    recent_crate_downloads,
      reserved_crate_names,

src/schema.rs

@@ -621,6 +621,29 @@ table! {
     }
 }
 
+table! {
+    use diesel::sql_types::*;
+    use diesel_full_text_search::{TsVector as Tsvector};
+
+    /// Representation of the `publish_rate_overrides` table.
+    ///
+    /// (Automatically generated by Diesel.)
+    publish_rate_overrides (user_id) {
+        /// The `user_id` column of the `publish_rate_overrides` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        user_id -> Int4,
+        /// The `burst` column of the `publish_rate_overrides` table.
+        ///
+        /// Its SQL type is `Int4`.
+        ///
+        /// (Automatically generated by Diesel.)
+        burst -> Int4,
+    }
+}
+
 table! {
     use diesel::sql_types::*;
     use diesel_full_text_search::{TsVector as Tsvector};
@@ -964,6 +987,7 @@ joinable!(emails -> users (user_id));
 joinable!(follows -> crates (crate_id));
 joinable!(follows -> users (user_id));
 joinable!(publish_limit_buckets -> users (user_id));
+joinable!(publish_rate_overrides -> users (user_id));
 joinable!(readme_renderings -> versions (version_id));
 joinable!(recent_crate_downloads -> crates (crate_id));
 joinable!(version_authors -> users (user_id));
@@ -989,6 +1013,7 @@ allow_tables_to_appear_in_same_query!(
     keywords,
     metadata,
     publish_limit_buckets,
+    publish_rate_overrides,
     readme_renderings,
     recent_crate_downloads,
     reserved_crate_names,

src/tests/krate.rs

@@ -2112,7 +2112,7 @@ fn publish_new_crate_rate_limited() {
 
     // Uploading a second crate is limited
     let crate_to_publish = PublishBuilder::new("rate_limited2");
-    token.publish(crate_to_publish).assert_status(529);
+    token.publish(crate_to_publish).assert_status(429);
 
     anon.get::<()>("/api/v1/crates/rate_limited2")
         .assert_status(404);
@@ -2141,7 +2141,3 @@ fn publish_rate_limit_doesnt_affect_existing_crates() {
         .version("1.0.1");
     token.publish(new_version).good();
 }
-
-#[test]
-#[ignore]
-fn publish_rate_limit_per_user_override() {}

src/util/errors.rs

@@ -1,5 +1,4 @@
 use std::any::{Any, TypeId};
-use std::collections::HashMap;
 use std::error::Error;
 use std::fmt;
 
@@ -397,19 +396,19 @@ impl CargoError for TooManyRequests {
     }
 
     fn response(&self) -> Option<Response> {
-        use std::io::Cursor;
-
         const HTTP_DATE_FORMAT: &str = "%a, %d %b %Y %H:%M:%S GMT";
-        let mut headers = HashMap::new();
-        headers.insert(
-            "Retry-After".into(),
-            vec![self.retry_after.format(HTTP_DATE_FORMAT).to_string()],
-        );
-        Some(Response {
-            status: (529, "TOO MANY REQUESTS"),
-            headers,
-            body: Box::new(Cursor::new(Vec::new())),
-        })
+        let retry_after = self.retry_after.format(HTTP_DATE_FORMAT);
+
+        let mut response = json_response(&Bad {
+            errors: vec![StringError {
+                detail: format!("You have published too many crates in a \
+                    short period of time. Please try again after {} or email \
+                    [email protected] to have your limit increased.", retry_after),
+            }],
+        });
+        response.status = (429, "TOO MANY REQUESTS");
+        response.headers.insert("Retry-After".into(), vec![retry_after.to_string()]);
+        Some(response)
     }
 
     fn human(&self) -> bool {