diff --git a/RustConfig b/RustConfig index 82de0875c07..f7cf220b751 100644 --- a/RustConfig +++ b/RustConfig @@ -1 +1 @@ -VERSION=1.32.0 +VERSION=1.33.0 diff --git a/migrations/2019-03-18-233900_create_publish_limit_buckets/down.sql b/migrations/2019-03-18-233900_create_publish_limit_buckets/down.sql new file mode 100644 index 00000000000..4a6bdefead2 --- /dev/null +++ b/migrations/2019-03-18-233900_create_publish_limit_buckets/down.sql @@ -0,0 +1 @@ +DROP TABLE publish_limit_buckets; diff --git a/migrations/2019-03-18-233900_create_publish_limit_buckets/up.sql b/migrations/2019-03-18-233900_create_publish_limit_buckets/up.sql new file mode 100644 index 00000000000..af4feb4217c --- /dev/null +++ b/migrations/2019-03-18-233900_create_publish_limit_buckets/up.sql @@ -0,0 +1,5 @@ +CREATE TABLE publish_limit_buckets( + user_id INTEGER PRIMARY KEY NOT NULL REFERENCES users, + tokens INTEGER NOT NULL, + last_refill TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP +); diff --git a/migrations/2019-04-04-192902_create_publish_rate_overrides/down.sql b/migrations/2019-04-04-192902_create_publish_rate_overrides/down.sql new file mode 100644 index 00000000000..9a19fcc2547 --- /dev/null +++ b/migrations/2019-04-04-192902_create_publish_rate_overrides/down.sql @@ -0,0 +1 @@ +DROP TABLE publish_rate_overrides; diff --git a/migrations/2019-04-04-192902_create_publish_rate_overrides/up.sql b/migrations/2019-04-04-192902_create_publish_rate_overrides/up.sql new file mode 100644 index 00000000000..2136089670b --- /dev/null +++ b/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 +); diff --git a/src/bin/update-downloads.rs b/src/bin/update-downloads.rs index 0f16cb60855..3ada76a7106 100644 --- a/src/bin/update-downloads.rs +++ b/src/bin/update-downloads.rs @@ -107,7 +107,7 @@ mod test { name: "foo", ..Default::default() } - .create_or_update(conn, None, user_id) + .create_or_update(conn, None, user_id, None) .unwrap(); let version = NewVersion::new( krate.id, diff --git a/src/config.rs b/src/config.rs index 44b030370aa..b5162f34908 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,3 +1,4 @@ +use crate::publish_rate_limit::PublishRateLimit; use crate::{env, uploaders::Uploader, Env, Replica}; use std::path::PathBuf; use url::Url; @@ -16,6 +17,7 @@ pub struct Config { pub max_unpack_size: u64, pub mirror: Replica, pub api_protocol: String, + pub publish_rate_limit: PublishRateLimit, } impl Default for Config { @@ -132,6 +134,7 @@ impl Default for Config { max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed mirror, api_protocol, + publish_rate_limit: Default::default(), } } } diff --git a/src/controllers/krate/publish.rs b/src/controllers/krate/publish.rs index 1192ffb29a0..d9be8618a1e 100644 --- a/src/controllers/krate/publish.rs +++ b/src/controllers/krate/publish.rs @@ -84,7 +84,12 @@ pub fn publish(req: &mut dyn Request) -> CargoResult { }; let license_file = new_crate.license_file.as_ref().map(|s| &**s); - let krate = persist.create_or_update(&conn, license_file, user.id)?; + let krate = persist.create_or_update( + &conn, + license_file, + user.id, + Some(&app.config.publish_rate_limit), + )?; let owners = krate.owners(&conn)?; if user.rights(req.app(), &owners)? < Rights::Publish { diff --git a/src/lib.rs b/src/lib.rs index fca561e4cb0..5b5ca0a94a9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,8 +37,10 @@ pub mod email; pub mod git; pub mod github; pub mod middleware; +mod publish_rate_limit; pub mod render; pub mod schema; +mod test_util; pub mod uploaders; pub mod util; diff --git a/src/models/category.rs b/src/models/category.rs index 44b220b8361..f07ecc4437c 100644 --- a/src/models/category.rs +++ b/src/models/category.rs @@ -183,12 +183,11 @@ impl<'a> NewCategory<'a> { #[cfg(test)] mod tests { use super::*; + use crate::test_util::pg_connection_no_transaction; use diesel::connection::SimpleConnection; fn pg_connection() -> PgConnection { - let database_url = - dotenv::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests"); - let conn = PgConnection::establish(&database_url).unwrap(); + let conn = pg_connection_no_transaction(); // These tests deadlock if run concurrently conn.batch_execute("BEGIN; LOCK categories IN ACCESS EXCLUSIVE MODE") .unwrap(); diff --git a/src/models/krate.rs b/src/models/krate.rs index 91aff10c31d..327665108be 100644 --- a/src/models/krate.rs +++ b/src/models/krate.rs @@ -15,6 +15,7 @@ use crate::models::{ use crate::views::{EncodableCrate, EncodableCrateLinks}; use crate::models::helpers::with_count::*; +use crate::publish_rate_limit::PublishRateLimit; use crate::schema::*; /// Hosts in this list are known to not be hosting documentation, @@ -105,6 +106,7 @@ impl<'a> NewCrate<'a> { conn: &PgConnection, license_file: Option<&'a str>, uploader: i32, + rate_limit: Option<&PublishRateLimit>, ) -> CargoResult { use diesel::update; @@ -115,6 +117,9 @@ 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, conn)?; + } return Ok(krate); } diff --git a/src/models/user.rs b/src/models/user.rs index f40dbded535..2f22c20a19e 100644 --- a/src/models/user.rs +++ b/src/models/user.rs @@ -21,7 +21,7 @@ pub struct User { pub gh_id: i32, } -#[derive(Insertable, Debug)] +#[derive(Insertable, Debug, Default)] #[table_name = "users"] pub struct NewUser<'a> { pub gh_id: i32, diff --git a/src/publish_rate_limit.rs b/src/publish_rate_limit.rs new file mode 100644 index 00000000000..d21d47ab574 --- /dev/null +++ b/src/publish_rate_limit.rs @@ -0,0 +1,356 @@ +use chrono::{NaiveDateTime, Utc}; +use diesel::data_types::PgInterval; +use diesel::prelude::*; +use std::time::Duration; + +use crate::schema::{publish_limit_buckets, publish_rate_overrides}; +use crate::util::errors::{CargoResult, TooManyRequests}; + +#[derive(Debug, Clone, Copy)] +pub struct PublishRateLimit { + pub rate: Duration, + pub burst: i32, +} + +impl Default for PublishRateLimit { + fn default() -> Self { + Self { + rate: Duration::from_secs(60) * 10, + burst: 30, + } + } +} + +#[derive(Queryable, Insertable, Debug, PartialEq, Clone, Copy)] +#[table_name = "publish_limit_buckets"] +#[allow(dead_code)] // Most fields only read in tests +struct Bucket { + user_id: i32, + tokens: i32, + last_refill: NaiveDateTime, +} + +impl PublishRateLimit { + pub fn check_rate_limit(&self, uploader: i32, conn: &PgConnection) -> CargoResult<()> { + let bucket = self.take_token(uploader, Utc::now().naive_utc(), conn)?; + if bucket.tokens >= 1 { + Ok(()) + } else { + Err(Box::new(TooManyRequests { + retry_after: bucket.last_refill + chrono::Duration::from_std(self.rate).unwrap(), + })) + } + } + + /// Refill a user's bucket as needed, take a token from it, + /// and returns the result. + /// + /// The number of tokens remaining will always be between 0 and self.burst. + /// If the number is 0, the request should be rejected, as the user doesn't + /// have a token to take. Technically a "full" bucket would have + /// `self.burst + 1` tokens in it, but that value would never be returned + /// since we only refill buckets when trying to take a token from it. + fn take_token( + &self, + uploader: i32, + now: NaiveDateTime, + conn: &PgConnection, + ) -> CargoResult { + use self::publish_limit_buckets::dsl::*; + use diesel::sql_types::{Double, Interval, Text, Timestamp}; + + sql_function!(fn date_part(x: Text, y: Timestamp) -> Double); + sql_function! { + #[sql_name = "date_part"] + fn interval_part(x: Text, y: Interval) -> Double; + } + sql_function!(fn floor(x: Double) -> Integer); + sql_function!(fn greatest(x: T, y: T) -> T); + sql_function!(fn least(x: T, y: T) -> T); + + let burst = publish_rate_overrides::table + .find(uploader) + .select(publish_rate_overrides::burst) + .first::(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. + let tokens_to_add = floor( + (date_part("epoch", now) - date_part("epoch", last_refill)) + / interval_part("epoch", self.refill_rate()), + ); + + diesel::insert_into(publish_limit_buckets) + .values((user_id.eq(uploader), tokens.eq(burst), last_refill.eq(now))) + .on_conflict(user_id) + .do_update() + .set(( + tokens.eq(least(burst, greatest(0, tokens - 1) + tokens_to_add)), + last_refill + .eq(last_refill + self.refill_rate().into_sql::() * tokens_to_add), + )) + .get_result(conn) + .map_err(Into::into) + } + + fn refill_rate(&self) -> PgInterval { + use diesel::dsl::*; + (self.rate.as_millis() as i64).milliseconds() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_util::*; + + #[test] + fn take_token_with_no_bucket_creates_new_one() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let bucket = rate.take_token(new_user(&conn, "user1")?, now, &conn)?; + let expected = Bucket { + user_id: bucket.user_id, + tokens: 10, + last_refill: now, + }; + assert_eq!(expected, bucket); + + let rate = PublishRateLimit { + rate: Duration::from_millis(50), + burst: 20, + }; + let bucket = rate.take_token(new_user(&conn, "user2")?, now, &conn)?; + let expected = Bucket { + user_id: bucket.user_id, + tokens: 20, + last_refill: now, + }; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn take_token_with_existing_bucket_modifies_existing_bucket() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 5, now)?.user_id; + let bucket = rate.take_token(user_id, now, &conn)?; + let expected = Bucket { + user_id, + tokens: 4, + last_refill: now, + }; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn take_token_after_delay_refills() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 5, now)?.user_id; + let refill_time = now + chrono::Duration::seconds(2); + let bucket = rate.take_token(user_id, refill_time, &conn)?; + let expected = Bucket { + user_id, + tokens: 6, + last_refill: refill_time, + }; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn refill_subsecond_rate() -> CargoResult<()> { + let conn = pg_connection(); + // Subsecond rates have floating point rounding issues, so use a known + // timestamp that rounds fine + let now = + NaiveDateTime::parse_from_str("2019-03-19T21:11:24.620401", "%Y-%m-%dT%H:%M:%S%.f")?; + + let rate = PublishRateLimit { + rate: Duration::from_millis(100), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 5, now)?.user_id; + let refill_time = now + chrono::Duration::milliseconds(300); + let bucket = rate.take_token(user_id, refill_time, &conn)?; + let expected = Bucket { + user_id, + tokens: 7, + last_refill: refill_time, + }; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn last_refill_always_advanced_by_multiple_of_rate() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_millis(100), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 5, now)?.user_id; + let bucket = rate.take_token(user_id, now + chrono::Duration::milliseconds(250), &conn)?; + let expected_refill_time = now + chrono::Duration::milliseconds(200); + let expected = Bucket { + user_id, + tokens: 6, + last_refill: expected_refill_time, + }; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn zero_tokens_returned_when_user_has_no_tokens_left() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 1, now)?.user_id; + let bucket = rate.take_token(user_id, now, &conn)?; + let expected = Bucket { + user_id, + tokens: 0, + last_refill: now, + }; + assert_eq!(expected, bucket); + + let bucket = rate.take_token(user_id, now, &conn)?; + assert_eq!(expected, bucket); + Ok(()) + } + + #[test] + fn a_user_with_no_tokens_gets_a_token_after_exactly_rate() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 0, now)?.user_id; + let refill_time = now + chrono::Duration::seconds(1); + let bucket = rate.take_token(user_id, refill_time, &conn)?; + let expected = Bucket { + user_id, + tokens: 1, + last_refill: refill_time, + }; + assert_eq!(expected, bucket); + + Ok(()) + } + + #[test] + fn tokens_never_refill_past_burst() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + let rate = PublishRateLimit { + rate: Duration::from_secs(1), + burst: 10, + }; + let user_id = new_user_bucket(&conn, 8, now)?.user_id; + let refill_time = now + chrono::Duration::seconds(4); + let bucket = rate.take_token(user_id, refill_time, &conn)?; + let expected = Bucket { + user_id, + tokens: 10, + last_refill: refill_time, + }; + assert_eq!(expected, bucket); + + Ok(()) + } + + #[test] + fn override_is_used_instead_of_global_burst_if_present() -> CargoResult<()> { + let conn = pg_connection(); + let now = now(); + + 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 { + use crate::models::NewUser; + + let user = NewUser { + gh_login, + ..NewUser::default() + } + .create_or_update(conn)?; + Ok(user.id) + } + + fn new_user_bucket( + conn: &PgConnection, + tokens: i32, + now: NaiveDateTime, + ) -> CargoResult { + diesel::insert_into(publish_limit_buckets::table) + .values(Bucket { + user_id: new_user(conn, "new_user")?, + tokens, + last_refill: now, + }) + .get_result(conn) + .map_err(Into::into) + } + + /// Strips ns precision from `Utc::now`. PostgreSQL only has microsecond + /// precision, but some platforms (notably Linux) provide nanosecond + /// precision, meaning that round tripping through the database would + /// change the value. + fn now() -> NaiveDateTime { + let now = Utc::now().naive_utc(); + let nanos = now.timestamp_subsec_nanos(); + now - chrono::Duration::nanoseconds(nanos.into()) + } +} diff --git a/src/schema.patch b/src/schema.patch index 1a65430ec4c..212cfe21933 100644 --- a/src/schema.patch +++ b/src/schema.patch @@ -55,7 +55,7 @@ index df884e4..18e08cd 100644 /// Representation of the `reserved_crate_names` table. /// -@@ -881,21 +901,23 @@ table! { +@@ -881,23 +901,25 @@ table! { } joinable!(api_tokens -> users (user_id)); @@ -73,6 +73,8 @@ index df884e4..18e08cd 100644 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)); @@ -80,12 +82,13 @@ index df884e4..18e08cd 100644 joinable!(version_downloads -> versions (version_id)); joinable!(versions -> crates (crate_id)); -@@ -913,12 +935,13 @@ allow_tables_to_appear_in_same_query!( - dependencies, +@@ -913,13 +935,14 @@ allow_tables_to_appear_in_same_query!( emails, follows, keywords, metadata, + publish_limit_buckets, + publish_rate_overrides, readme_renderings, + recent_crate_downloads, reserved_crate_names, diff --git a/src/schema.rs b/src/schema.rs index 934c3d2e424..6fd4bfb22d7 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -586,6 +586,58 @@ table! { } } +table! { + use diesel::sql_types::*; + use diesel_full_text_search::{TsVector as Tsvector}; + + /// Representation of the `publish_limit_buckets` table. + /// + /// (Automatically generated by Diesel.) + publish_limit_buckets (user_id) { + /// The `user_id` column of the `publish_limit_buckets` table. + /// + /// Its SQL type is `Int4`. + /// + /// (Automatically generated by Diesel.) + user_id -> Int4, + /// The `tokens` column of the `publish_limit_buckets` table. + /// + /// Its SQL type is `Int4`. + /// + /// (Automatically generated by Diesel.) + tokens -> Int4, + /// The `last_refill` column of the `publish_limit_buckets` table. + /// + /// Its SQL type is `Timestamp`. + /// + /// (Automatically generated by Diesel.) + last_refill -> Timestamp, + } +} + +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}; @@ -928,6 +980,8 @@ joinable!(dependencies -> versions (version_id)); 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)); @@ -952,6 +1006,8 @@ allow_tables_to_appear_in_same_query!( follows, keywords, metadata, + publish_limit_buckets, + publish_rate_overrides, readme_renderings, recent_crate_downloads, reserved_crate_names, diff --git a/src/test_util.rs b/src/test_util.rs new file mode 100644 index 00000000000..0cae4fe771e --- /dev/null +++ b/src/test_util.rs @@ -0,0 +1,15 @@ +#![cfg(test)] + +use diesel::prelude::*; + +pub fn pg_connection_no_transaction() -> PgConnection { + let database_url = + dotenv::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL must be set to run tests"); + PgConnection::establish(&database_url).unwrap() +} + +pub fn pg_connection() -> PgConnection { + let conn = pg_connection_no_transaction(); + conn.begin_test_transaction().unwrap(); + conn +} diff --git a/src/tests/all.rs b/src/tests/all.rs index a420a94ae7c..c535e4cc384 100644 --- a/src/tests/all.rs +++ b/src/tests/all.rs @@ -114,22 +114,11 @@ pub struct OkBool { ok: bool, } -/// Initialize the app and a proxy that can record and playback outgoing HTTP requests -fn app_with_proxy() -> ( - record::Bomb, - Arc, - conduit_middleware::MiddlewareBuilder, -) { - let (proxy, bomb) = record::proxy(); - let (app, handler) = init_app(Some(proxy)); - (bomb, app, handler) -} - fn app() -> (Arc, conduit_middleware::MiddlewareBuilder) { - init_app(None) + build_app(simple_config(), None) } -fn init_app(proxy: Option) -> (Arc, conduit_middleware::MiddlewareBuilder) { +fn simple_config() -> Config { let uploader = Uploader::S3 { bucket: s3::Bucket::new( String::from("alexcrichton-test"), @@ -143,7 +132,7 @@ fn init_app(proxy: Option) -> (Arc, conduit_middleware::MiddlewareB cdn: None, }; - let config = Config { + Config { uploader, session_key: "test this has to be over 32 bytes long".to_string(), git_repo_checkout: git::checkout(), @@ -158,8 +147,14 @@ fn init_app(proxy: Option) -> (Arc, conduit_middleware::MiddlewareB // When testing we route all API traffic over HTTP so we can // sniff/record it, but everywhere else we use https api_protocol: String::from("http"), - }; + publish_rate_limit: Default::default(), + } +} +fn build_app( + config: Config, + proxy: Option, +) -> (Arc, conduit_middleware::MiddlewareBuilder) { let client = if let Some(proxy) = proxy { let mut builder = Client::builder(); builder = builder diff --git a/src/tests/builders.rs b/src/tests/builders.rs index 1e41fb02898..819545f49fb 100644 --- a/src/tests/builders.rs +++ b/src/tests/builders.rs @@ -234,7 +234,7 @@ impl<'a> CrateBuilder<'a> { let mut krate = self .krate - .create_or_update(connection, None, self.owner_id)?; + .create_or_update(connection, None, self.owner_id, None)?; // Since we are using `NewCrate`, we can't set all the // crate properties in a single DB call. diff --git a/src/tests/http-data/krate_publish_new_crate_rate_limited b/src/tests/http-data/krate_publish_new_crate_rate_limited new file mode 100644 index 00000000000..ef6d3225b66 --- /dev/null +++ b/src/tests/http-data/krate_publish_new_crate_rate_limited @@ -0,0 +1,144 @@ +[ + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited1/rate_limited1-1.0.0.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + }, + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited2/rate_limited2-1.0.0.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + } +] diff --git a/src/tests/http-data/krate_publish_rate_limit_doesnt_affect_existing_crates b/src/tests/http-data/krate_publish_rate_limit_doesnt_affect_existing_crates new file mode 100644 index 00000000000..b0a60893f39 --- /dev/null +++ b/src/tests/http-data/krate_publish_rate_limit_doesnt_affect_existing_crates @@ -0,0 +1,144 @@ +[ + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited1/rate_limited1-1.0.0.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + }, + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited1/rate_limited1-1.0.1.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + } +] diff --git a/src/tests/http-data/krate_user_never_gets_more_than_max_tokens_added b/src/tests/http-data/krate_user_never_gets_more_than_max_tokens_added new file mode 100644 index 00000000000..ef6d3225b66 --- /dev/null +++ b/src/tests/http-data/krate_user_never_gets_more_than_max_tokens_added @@ -0,0 +1,144 @@ +[ + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited1/rate_limited1-1.0.0.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + }, + { + "request": { + "uri": "http://alexcrichton-test.s3.amazonaws.com/crates/rate_limited2/rate_limited2-1.0.0.crate", + "method": "PUT", + "headers": [ + [ + "content-length", + "35" + ], + [ + "content-type", + "application/x-tar" + ], + [ + "accept-encoding", + "gzip" + ], + [ + "authorization", + "AWS AKIAICL5IWUZYWWKA7JA:3P5wvArAHvV7o8atB0gUc0RTCCc=" + ], + [ + "accept", + "*/*" + ], + [ + "date", + "Fri, 15 Sep 2017 07:53:06 -0700" + ], + [ + "user-agent", + "reqwest/0.9.1" + ], + [ + "host", + "alexcrichton-test.s3.amazonaws.com" + ] + ], + "body": "H4sIAAAAAAAA/+3AAQEAAACCIP+vbkhQwKsBLq+17wAEAAA=" + }, + "response": { + "status": 200, + "headers": [ + [ + "x-amz-id-2", + "9ZLFTsEwh2iNu+BlWzaTZ85mFA7pxZgsGhQCj3Qi67LqT/iB5eiCOJQPYww2BEkoivJbr0mWruo=" + ], + [ + "date", + "Fri, 15 Sep 2017 14:53:07 GMT" + ], + [ + "content-length", + "0" + ], + [ + "x-amz-request-id", + "FD643F2FC49A7DF3" + ], + [ + "Server", + "AmazonS3" + ], + [ + "ETag", + "\"f9016ad360cebb4fe2e6e96e5949f022\"" + ] + ], + "body": "" + } + } +] diff --git a/src/tests/krate.rs b/src/tests/krate.rs index 0d8650cddf8..84c623a7ea5 100644 --- a/src/tests/krate.rs +++ b/src/tests/krate.rs @@ -14,6 +14,8 @@ use cargo_registry::{ use std::{ collections::HashMap, io::{self, prelude::*}, + thread, + time::Duration, }; use chrono::Utc; @@ -2110,3 +2112,46 @@ fn new_krate_tarball_with_hard_links() { json.errors ); } + +#[test] +fn publish_new_crate_rate_limited() { + let (app, anon, _, token) = TestApp::full() + .with_publish_rate_limit(Duration::from_millis(500), 1) + .with_token(); + + // Upload a new crate + let crate_to_publish = PublishBuilder::new("rate_limited1"); + token.enqueue_publish(crate_to_publish).good(); + + // Uploading a second crate is limited + let crate_to_publish = PublishBuilder::new("rate_limited2"); + token.enqueue_publish(crate_to_publish).assert_status(429); + app.run_pending_background_jobs(); + + anon.get::<()>("/api/v1/crates/rate_limited2") + .assert_status(404); + + // Wait for the limit to be up + thread::sleep(Duration::from_millis(500)); + + let crate_to_publish = PublishBuilder::new("rate_limited2"); + token.enqueue_publish(crate_to_publish).good(); + + let json = anon.show_crate("rate_limited2"); + assert_eq!(json.krate.max_version, "1.0.0"); +} + +#[test] +fn publish_rate_limit_doesnt_affect_existing_crates() { + let (app, _, _, token) = TestApp::full() + .with_publish_rate_limit(Duration::from_millis(500), 1) + .with_token(); + + // Upload a new crate + let crate_to_publish = PublishBuilder::new("rate_limited1"); + token.enqueue_publish(crate_to_publish).good(); + + let new_version = PublishBuilder::new("rate_limited1").version("1.0.1"); + token.enqueue_publish(new_version).good(); + app.run_pending_background_jobs(); +} diff --git a/src/tests/util.rs b/src/tests/util.rs index f56cea921a0..8cd15f98bfe 100644 --- a/src/tests/util.rs +++ b/src/tests/util.rs @@ -20,15 +20,14 @@ //! to the underlying database model value (`User` and `ApiToken` respectively). use crate::{ - app, app_with_proxy, builders::PublishBuilder, record, CrateList, CrateResponse, GoodCrate, - OkBool, VersionResponse, + builders::PublishBuilder, record, CrateList, CrateResponse, GoodCrate, OkBool, VersionResponse, }; use cargo_registry::{ background_jobs::Environment, db::DieselPool, middleware::current_user::AuthenticationSource, models::{ApiToken, User}, - App, + App, Config, }; use diesel::PgConnection; use std::{rc::Rc, sync::Arc, time::Duration}; @@ -88,71 +87,29 @@ impl Drop for TestAppInner { } /// A representation of the app and its database transaction +#[derive(Clone)] pub struct TestApp(Rc); impl TestApp { /// Initialize an application with an `Uploader` that panics pub fn init() -> TestAppBuilder { - let (app, middle) = app(); - let inner = Rc::new(TestAppInner { - app, - _bomb: None, - middle, + TestAppBuilder { + config: crate::simple_config(), + proxy: None, + bomb: None, index: None, - runner: None, - }); - TestAppBuilder(TestApp(inner)) + build_job_runner: false, + } } /// Initialize the app and a proxy that can record and playback outgoing HTTP requests pub fn with_proxy() -> TestAppBuilder { - let (bomb, app, middle) = app_with_proxy(); - let inner = Rc::new(TestAppInner { - app, - _bomb: Some(bomb), - middle, - index: None, - runner: None, - }); - TestAppBuilder(TestApp(inner)) + Self::init().with_proxy() } /// Initialize a full application, with a proxy, index, and background worker pub fn full() -> TestAppBuilder { - use crate::git; - - let (bomb, app, middle) = app_with_proxy(); - git::init(); - - let thread_local_path = git::bare(); - let index = UpstreamRepository::open_bare(thread_local_path).unwrap(); - - let index_clone = - WorkerRepository::open(&app.config.index_location).expect("Could not clone index"); - let connection_pool = app.diesel_database.clone(); - let environment = Environment::new( - index_clone, - None, - connection_pool.clone(), - app.config.uploader.clone(), - app.http_client().clone(), - ); - - let runner = Runner::builder(connection_pool, environment) - // We only have 1 connection in tests, so trying to run more than - // 1 job concurrently will just block - .thread_count(1) - .job_start_timeout(Duration::from_secs(1)) - .build(); - - let inner = Rc::new(TestAppInner { - app, - _bomb: Some(bomb), - middle, - index: Some(index), - runner: Some(runner), - }); - TestAppBuilder(TestApp(inner)) + Self::with_proxy().with_git_index().with_job_runner() } /// Obtain the database connection and pass it to the closure @@ -236,15 +193,63 @@ impl TestApp { } } -pub struct TestAppBuilder(TestApp); +pub struct TestAppBuilder { + config: Config, + proxy: Option, + bomb: Option, + index: Option, + build_job_runner: bool, +} impl TestAppBuilder { /// Create a `TestApp` with an empty database pub fn empty(self) -> (TestApp, MockAnonymousUser) { + let (app, middle) = crate::build_app(self.config, self.proxy); + + let runner = if self.build_job_runner { + let connection_pool = app.diesel_database.clone(); + let index = + WorkerRepository::open(&app.config.index_location).expect("Could not clone index"); + let environment = Environment::new( + index, + None, + connection_pool.clone(), + app.config.uploader.clone(), + app.http_client().clone(), + ); + + Some( + Runner::builder(connection_pool, environment) + // We only have 1 connection in tests, so trying to run more than + // 1 job concurrently will just block + .thread_count(1) + .job_start_timeout(Duration::from_secs(1)) + .build(), + ) + } else { + None + }; + + let test_app_inner = TestAppInner { + app, + _bomb: self.bomb, + middle, + index: self.index, + runner, + }; + let test_app = TestApp(Rc::new(test_app_inner)); let anon = MockAnonymousUser { - app: TestApp(Rc::clone(&(self.0).0)), + app: test_app.clone(), }; - (self.0, anon) + (test_app, anon) + } + + /// Create a proxy for use with this app + pub fn with_proxy(mut self) -> Self { + let (proxy, bomb) = record::proxy(); + self.proxy = Some(proxy); + self.bomb = Some(bomb); + self } // Create a `TestApp` with a database including a default user @@ -261,6 +266,27 @@ impl TestAppBuilder { let token = user.db_new_token("bar"); (app, anon, user, token) } + + pub fn with_publish_rate_limit(mut self, rate: Duration, burst: i32) -> Self { + self.config.publish_rate_limit.rate = rate; + self.config.publish_rate_limit.burst = burst; + self + } + + pub fn with_git_index(mut self) -> Self { + use crate::git; + + git::init(); + + let thread_local_path = git::bare(); + self.index = Some(UpstreamRepository::open_bare(thread_local_path).unwrap()); + self + } + + pub fn with_job_runner(mut self) -> Self { + self.build_job_runner = true; + self + } } /// A colleciton of helper methods for the 3 authentication types diff --git a/src/util/errors.rs b/src/util/errors.rs index 9903c0ba463..481b97881a4 100644 --- a/src/util/errors.rs +++ b/src/util/errors.rs @@ -2,6 +2,7 @@ use std::any::{Any, TypeId}; use std::error::Error; use std::fmt; +use chrono::NaiveDateTime; use conduit::Response; use diesel::result::Error as DieselError; @@ -383,3 +384,45 @@ impl fmt::Display for ReadOnlyMode { "Tried to write in read only mode".fmt(f) } } + +#[derive(Debug, Clone, Copy)] +pub struct TooManyRequests { + pub retry_after: NaiveDateTime, +} + +impl CargoError for TooManyRequests { + fn description(&self) -> &str { + "too many requests" + } + + fn response(&self) -> Option { + const HTTP_DATE_FORMAT: &str = "%a, %d %b %Y %H:%M:%S GMT"; + 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 \ + help@crates.io 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 { + true + } +} + +impl fmt::Display for TooManyRequests { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + "Too many requests".fmt(f) + } +}