diff --git a/config/nginx.conf.erb b/config/nginx.conf.erb index 5afcefc59a0..3f169d2b362 100644 --- a/config/nginx.conf.erb +++ b/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; - } } } diff --git a/src/app.rs b/src/app.rs index dcbac4a8ac6..f1129388678 100644 --- a/src/app.rs +++ b/src/app.rs @@ -10,6 +10,7 @@ use crate::downloads_counter::DownloadsCounter; use crate::email::Emails; use crate::github::{GitHubClient, RealGitHubClient}; use crate::metrics::{InstanceMetrics, ServiceMetrics}; +use crate::rate_limiter::RateLimiter; use crate::storage::Storage; use axum::extract::{FromRef, FromRequestParts, State}; use diesel::r2d2; @@ -68,6 +69,9 @@ pub struct App { /// In-flight request counters for the `balance_capacity` middleware. pub balance_capacity: BalanceCapacityState, + + /// Rate limit select actions. + pub rate_limiter: RateLimiter, } impl App { @@ -178,6 +182,7 @@ impl App { http_client, fastboot_client, balance_capacity: Default::default(), + rate_limiter: RateLimiter::new(config.rate_limiter.clone()), config, } } diff --git a/src/config/server.rs b/src/config/server.rs index 4e2d083112a..df35d6dbaf0 100644 --- a/src/config/server.rs +++ b/src/config/server.rs @@ -2,7 +2,7 @@ use anyhow::{anyhow, Context}; use ipnetwork::IpNetwork; use oauth2::{ClientId, ClientSecret}; -use crate::rate_limiter::RateLimiter; +use crate::rate_limiter::{LimitedAction, RateLimiterConfig}; use crate::{env, env_optional, Env}; use super::base::Base; @@ -10,7 +10,7 @@ use super::database_pools::DatabasePools; use crate::config::balance_capacity::BalanceCapacityConfig; use crate::storage::StorageConfig; use http::HeaderValue; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::net::IpAddr; use std::time::Duration; @@ -30,7 +30,7 @@ pub struct Server { pub gh_client_secret: ClientSecret, pub max_upload_size: u64, pub max_unpack_size: u64, - pub rate_limiter: RateLimiter, + pub rate_limiter: HashMap, pub new_version_rate_limit: Option, pub blocked_traffic: Vec<(String, Vec)>, pub max_allowed_page_offset: u32, @@ -140,6 +140,24 @@ impl Default for Server { .map(|s| s.parse().expect("SERVER_THREADS was not a valid number")) .ok(); + // Dynamically load the configuration for all the rate limiting actions. See + // `src/rate_limiter.rs` for their definition. + let mut rate_limiter = HashMap::new(); + for action in LimitedAction::VARIANTS { + let env_var_key = action.env_var_key(); + rate_limiter.insert( + *action, + RateLimiterConfig { + rate: Duration::from_secs( + env_optional(&format!("RATE_LIMITER_{env_var_key}_RATE_SECONDS")) + .unwrap_or_else(|| action.default_rate_seconds()), + ), + burst: env_optional(&format!("RATE_LIMITER_{env_var_key}_BURST")) + .unwrap_or_else(|| action.default_burst()), + }, + ); + } + Server { db: DatabasePools::full_from_environment(&base), storage: StorageConfig::from_environment(), @@ -153,7 +171,7 @@ impl Default for Server { gh_client_secret: ClientSecret::new(env("GH_CLIENT_SECRET")), max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed - rate_limiter: Default::default(), + rate_limiter, 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), diff --git a/src/controllers/krate/publish.rs b/src/controllers/krate/publish.rs index 41ab0f6f34c..a0354ce69cb 100644 --- a/src/controllers/krate/publish.rs +++ b/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 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 { } impl<'a> NewCrate<'a> { - pub fn create_or_update( - self, - conn: &mut PgConnection, - uploader: i32, - rate_limit: Option<&RateLimiter>, - ) -> AppResult { + pub fn create_or_update(self, conn: &mut PgConnection, uploader: i32) -> AppResult { 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, conn)?; - } return Ok(krate); } diff --git a/src/rate_limiter.rs b/src/rate_limiter.rs index e85deb9792a..f408b3d67a0 100644 --- a/src/rate_limiter.rs +++ b/src/rate_limiter.rs @@ -1,52 +1,93 @@ +use crate::schema::{publish_limit_buckets, publish_rate_overrides}; +use crate::sql::{date_part, floor, greatest, interval_part, least, pg_enum}; +use crate::util::errors::{AppResult, TooManyRequests}; use chrono::{NaiveDateTime, Utc}; -use diesel::data_types::PgInterval; +use diesel::dsl::IntervalDsl; use diesel::prelude::*; use diesel::sql_types::Interval; +use std::borrow::Cow; +use std::collections::HashMap; use std::time::Duration; -use crate::schema::{publish_limit_buckets, publish_rate_overrides}; -use crate::sql::{date_part, floor, greatest, interval_part, least, pg_enum}; -use crate::util::errors::{AppResult, TooManyRequests}; - pg_enum! { pub enum LimitedAction { PublishNew = 0, + PublishUpdate = 1, + YankUnyank = 2, + } +} + +impl LimitedAction { + pub fn default_rate_seconds(&self) -> u64 { + match self { + LimitedAction::PublishNew => 60 * 60, + LimitedAction::PublishUpdate => 60, + LimitedAction::YankUnyank => 60, + } + } + + pub fn default_burst(&self) -> i32 { + match self { + LimitedAction::PublishNew => 5, + LimitedAction::PublishUpdate => 30, + LimitedAction::YankUnyank => 100, + } + } + + pub fn env_var_key(&self) -> &'static str { + match self { + LimitedAction::PublishNew => "PUBLISH_NEW", + LimitedAction::PublishUpdate => "PUBLISH_UPDATE", + LimitedAction::YankUnyank => "YANK_UNYANK", + } + } + + pub fn error_messagge(&self) -> &'static str { + match self { + 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" + } + LimitedAction::YankUnyank => { + "You have yanked or unyanked too many versions in a short period of time" + } + } } } #[derive(Debug, Clone, Copy)] -pub struct RateLimiter { +pub struct RateLimiterConfig { pub rate: Duration, pub burst: i32, } -impl Default for RateLimiter { - fn default() -> Self { - let minutes = dotenvy::var("WEB_NEW_PKG_RATE_LIMIT_RATE_MINUTES") - .unwrap_or_default() - .parse() - .ok() - .unwrap_or(10); - let burst = dotenvy::var("WEB_NEW_PKG_RATE_LIMIT_BURST") - .unwrap_or_default() - .parse() - .ok() - .unwrap_or(5); - Self { - rate: Duration::from_secs(60) * minutes, - burst, - } - } +#[derive(Debug)] +pub struct RateLimiter { + config: HashMap, } impl RateLimiter { - pub fn check_rate_limit(&self, uploader: i32, conn: &mut PgConnection) -> AppResult<()> { - let bucket = self.take_token(uploader, Utc::now().naive_utc(), conn)?; + pub fn new(config: HashMap) -> Self { + Self { config } + } + + pub fn check_rate_limit( + &self, + uploader: i32, + performed_action: LimitedAction, + conn: &mut PgConnection, + ) -> AppResult<()> { + let bucket = self.take_token(uploader, performed_action, 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(), + action: performed_action, + retry_after: bucket.last_refill + + chrono::Duration::from_std(self.config_for_action(performed_action).rate) + .unwrap(), })) } } @@ -62,12 +103,14 @@ impl RateLimiter { fn take_token( &self, uploader: i32, + performed_action: LimitedAction, now: NaiveDateTime, conn: &mut PgConnection, ) -> QueryResult { use self::publish_limit_buckets::dsl::*; - let performed_action = LimitedAction::PublishNew; + let config = self.config_for_action(performed_action); + let refill_rate = (config.rate.as_millis() as i64).milliseconds(); let burst: i32 = publish_rate_overrides::table .find((uploader, performed_action)) @@ -79,14 +122,14 @@ impl RateLimiter { .select(publish_rate_overrides::burst) .first(conn) .optional()? - .unwrap_or(self.burst); + .unwrap_or(config.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()), + / interval_part("epoch", refill_rate), ); diesel::insert_into(publish_limit_buckets) @@ -100,15 +143,20 @@ impl RateLimiter { .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), + last_refill.eq(last_refill + refill_rate.into_sql::() * tokens_to_add), )) .get_result(conn) } - fn refill_rate(&self) -> PgInterval { - use diesel::dsl::*; - (self.rate.as_millis() as i64).milliseconds() + fn config_for_action(&self, action: LimitedAction) -> Cow<'_, RateLimiterConfig> { + // The wrapper returns the default config for the action when not configured. + match self.config.get(&action) { + Some(config) => Cow::Borrowed(config), + None => Cow::Owned(RateLimiterConfig { + rate: Duration::from_secs(action.default_rate_seconds()), + burst: action.default_burst(), + }), + } } } @@ -133,11 +181,18 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; - let bucket = rate.take_token(new_user(conn, "user1")?, now, conn)?; + action: LimitedAction::PublishNew, + } + .create(); + let bucket = rate.take_token( + new_user(conn, "user1")?, + LimitedAction::PublishNew, + now, + conn, + )?; let expected = Bucket { user_id: bucket.user_id, tokens: 10, @@ -146,11 +201,18 @@ mod tests { }; assert_eq!(expected, bucket); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_millis(50), burst: 20, - }; - let bucket = rate.take_token(new_user(conn, "user2")?, now, conn)?; + action: LimitedAction::PublishNew, + } + .create(); + let bucket = rate.take_token( + new_user(conn, "user2")?, + LimitedAction::PublishNew, + now, + conn, + )?; let expected = Bucket { user_id: bucket.user_id, tokens: 20, @@ -166,12 +228,14 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); let user_id = new_user_bucket(conn, 5, now)?.user_id; - let bucket = rate.take_token(user_id, now, conn)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; let expected = Bucket { user_id, tokens: 4, @@ -187,13 +251,15 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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 bucket = rate.take_token(user_id, LimitedAction::PublishNew, refill_time, conn)?; let expected = Bucket { user_id, tokens: 6, @@ -213,13 +279,15 @@ mod tests { NaiveDateTime::parse_from_str("2019-03-19T21:11:24.620401", "%Y-%m-%dT%H:%M:%S%.f") .unwrap(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_millis(100), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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 bucket = rate.take_token(user_id, LimitedAction::PublishNew, refill_time, conn)?; let expected = Bucket { user_id, tokens: 7, @@ -235,12 +303,19 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_millis(100), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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 bucket = rate.take_token( + user_id, + LimitedAction::PublishNew, + now + chrono::Duration::milliseconds(250), + conn, + )?; let expected_refill_time = now + chrono::Duration::milliseconds(200); let expected = Bucket { user_id, @@ -257,12 +332,14 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); let user_id = new_user_bucket(conn, 1, now)?.user_id; - let bucket = rate.take_token(user_id, now, conn)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; let expected = Bucket { user_id, tokens: 0, @@ -271,7 +348,7 @@ mod tests { }; assert_eq!(expected, bucket); - let bucket = rate.take_token(user_id, now, conn)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; assert_eq!(expected, bucket); Ok(()) } @@ -281,13 +358,15 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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 bucket = rate.take_token(user_id, LimitedAction::PublishNew, refill_time, conn)?; let expected = Bucket { user_id, tokens: 1, @@ -304,13 +383,15 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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 bucket = rate.take_token(user_id, LimitedAction::PublishNew, refill_time, conn)?; let expected = Bucket { user_id, tokens: 10, @@ -322,27 +403,73 @@ mod tests { Ok(()) } + #[test] + fn two_actions_dont_interfere_with_each_other() -> QueryResult<()> { + let conn = &mut pg_connection(); + let now = now(); + + let mut config = HashMap::new(); + config.insert( + LimitedAction::PublishNew, + RateLimiterConfig { + rate: Duration::from_secs(1), + burst: 10, + }, + ); + config.insert( + LimitedAction::YankUnyank, + RateLimiterConfig { + rate: Duration::from_secs(1), + burst: 20, + }, + ); + let rate = RateLimiter::new(config); + + let user_id = new_user(conn, "user")?; + + assert_eq!( + 10, + rate.take_token(user_id, LimitedAction::PublishNew, now, conn)? + .tokens + ); + assert_eq!( + 9, + rate.take_token(user_id, LimitedAction::PublishNew, now, conn)? + .tokens + ); + assert_eq!( + 20, + rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)? + .tokens + ); + + Ok(()) + } + #[test] fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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::action.eq(LimitedAction::PublishNew), 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)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; + let other_bucket = rate.take_token(other_user_id, LimitedAction::PublishNew, now, conn)?; assert_eq!(20, bucket.tokens); assert_eq!(10, other_bucket.tokens); @@ -354,23 +481,26 @@ mod tests { let conn = &mut pg_connection(); let now = now(); - let rate = RateLimiter { + let rate = SampleRateLimiter { rate: Duration::from_secs(1), burst: 10, - }; + action: LimitedAction::PublishNew, + } + .create(); 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::action.eq(LimitedAction::PublishNew), publish_rate_overrides::burst.eq(20), publish_rate_overrides::expires_at.eq(now + chrono::Duration::days(30)), )) .execute(conn)?; - let bucket = rate.take_token(user_id, now, conn)?; - let other_bucket = rate.take_token(other_user_id, now, conn)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; + let other_bucket = rate.take_token(other_user_id, LimitedAction::PublishNew, now, conn)?; assert_eq!(20, bucket.tokens); assert_eq!(10, other_bucket.tokens); @@ -381,8 +511,8 @@ mod tests { .filter(publish_rate_overrides::user_id.eq(user_id)) .execute(conn)?; - let bucket = rate.take_token(user_id, now, conn)?; - let other_bucket = rate.take_token(other_user_id, now, conn)?; + let bucket = rate.take_token(user_id, LimitedAction::PublishNew, now, conn)?; + let other_bucket = rate.take_token(other_user_id, LimitedAction::PublishNew, now, conn)?; // The number of tokens of user_id is 10 and not 9 because when the new burst limit is // lower than the amount of available tokens, the number of available tokens is reset to @@ -393,6 +523,46 @@ mod tests { Ok(()) } + #[test] + fn override_is_different_for_each_action() -> QueryResult<()> { + let conn = &mut pg_connection(); + let now = now(); + let user_id = new_user(conn, "user")?; + + let mut config = HashMap::new(); + for action in [LimitedAction::PublishNew, LimitedAction::YankUnyank] { + config.insert( + action, + RateLimiterConfig { + rate: Duration::from_secs(1), + burst: 10, + }, + ); + } + let rate = RateLimiter::new(config); + + diesel::insert_into(publish_rate_overrides::table) + .values(( + publish_rate_overrides::user_id.eq(user_id), + publish_rate_overrides::action.eq(LimitedAction::PublishNew), + publish_rate_overrides::burst.eq(20), + )) + .execute(conn)?; + + assert_eq!( + 20, + rate.take_token(user_id, LimitedAction::PublishNew, now, conn)? + .tokens, + ); + assert_eq!( + 10, + rate.take_token(user_id, LimitedAction::YankUnyank, now, conn)? + .tokens, + ); + + Ok(()) + } + fn new_user(conn: &mut PgConnection, gh_login: &str) -> QueryResult { use crate::models::NewUser; @@ -419,6 +589,26 @@ mod tests { .get_result(conn) } + struct SampleRateLimiter { + rate: Duration, + burst: i32, + action: LimitedAction, + } + + impl SampleRateLimiter { + fn create(self) -> RateLimiter { + let mut config = HashMap::new(); + config.insert( + self.action, + RateLimiterConfig { + rate: self.rate, + burst: self.burst, + }, + ); + RateLimiter::new(config) + } + } + /// 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 diff --git a/src/sql.rs b/src/sql.rs index e3bbc3bc47f..0e050362d6d 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -20,7 +20,7 @@ macro_rules! pg_enum { $($item:ident = $int:expr,)* } ) => { - #[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, FromSqlRow, AsExpression)] + #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, FromSqlRow, AsExpression)] #[diesel(sql_type = diesel::sql_types::Integer)] #[serde(rename_all = "snake_case")] #[repr(i32)] @@ -28,6 +28,10 @@ macro_rules! pg_enum { $($item = $int,)* } + impl $name { + $vis const VARIANTS: &[$name] = &[$($name::$item),*]; + } + impl diesel::deserialize::FromSql for $name { fn from_sql(bytes: diesel::pg::PgValue<'_>) -> diesel::deserialize::Result { match >::from_sql(bytes)? { diff --git a/src/tests/builders/krate.rs b/src/tests/builders/krate.rs index 9249fd6466c..0d21a49f1e4 100644 --- a/src/tests/builders/krate.rs +++ b/src/tests/builders/krate.rs @@ -114,9 +114,7 @@ impl<'a> CrateBuilder<'a> { pub fn build(mut self, connection: &mut PgConnection) -> AppResult { 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. diff --git a/src/tests/krate/publish.rs b/src/tests/krate/publish.rs index cce099c10dd..63ef5154386 100644 --- a/src/tests/krate/publish.rs +++ b/src/tests/krate/publish.rs @@ -6,6 +6,7 @@ use crates_io::controllers::krate::publish::{ missing_metadata_error_message, MISSING_RIGHTS_ERROR_MESSAGE, }; use crates_io::models::krate::MAX_NAME_LENGTH; +use crates_io::rate_limiter::LimitedAction; use crates_io::schema::{api_tokens, emails, versions_published_by}; use crates_io::views::GoodCrate; use crates_io_tarball::TarballBuilder; @@ -1004,7 +1005,7 @@ fn tarball_bigger_than_max_upload_size() { #[test] fn publish_new_crate_rate_limited() { let (app, anon, _, token) = TestApp::full() - .with_publish_rate_limit(Duration::from_millis(500), 1) + .with_rate_limit(LimitedAction::PublishNew, Duration::from_millis(500), 1) .with_token(); // Upload a new crate @@ -1015,8 +1016,9 @@ fn publish_new_crate_rate_limited() { // Uploading a second crate is limited let crate_to_publish = PublishBuilder::new("rate_limited2", "1.0.0"); - 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(LimitedAction::PublishNew); assert_eq!(app.stored_files().len(), 2); @@ -1036,9 +1038,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_publish_rate_limit(Duration::from_millis(500), 1) + .with_rate_limit(LimitedAction::PublishNew, Duration::from_secs(60 * 60), 1) .with_token(); // Upload a new crate @@ -1049,6 +1051,69 @@ 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", "1.0.0"); + 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", "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", "1.0.2"); + token + .publish_crate(crate_to_publish) + .assert_rate_limited(LimitedAction::PublishUpdate); + + // 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", "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", "1.0.0"); + token.publish_crate(crate_to_publish).good(); + + // Upload a second new crate + let crate_to_publish = PublishBuilder::new("rate_limited2", "1.0.0"); + token.publish_crate(crate_to_publish).good(); +} + #[test] fn features_version_2() { let (app, _, user, token) = TestApp::full().with_token(); diff --git a/src/tests/krate/yanking.rs b/src/tests/krate/yanking.rs index 6319d295ef9..6f222a336d7 100644 --- a/src/tests/krate/yanking.rs +++ b/src/tests/krate/yanking.rs @@ -1,6 +1,8 @@ use crate::builders::PublishBuilder; use crate::routes::crates::versions::yank_unyank::YankRequestHelper; use crate::util::{RequestHelper, TestApp}; +use crates_io::rate_limiter::LimitedAction; +use std::time::Duration; #[test] #[allow(unknown_lints, clippy::bool_assert_comparison)] // for claim::assert_some_eq! with bool @@ -60,6 +62,48 @@ fn yank_works_as_intended() { assert!(!json.version.yanked); } +#[test] +fn test_yank_rate_limiting() { + #[track_caller] + fn check_yanked(app: &TestApp, is_yanked: bool) { + let crates = app.crates_from_index_head("yankable"); + assert_eq!(crates.len(), 1); + assert_some_eq!(crates[0].yanked, is_yanked); + } + + let (app, _, _, token) = TestApp::full() + .with_rate_limit(LimitedAction::YankUnyank, Duration::from_millis(500), 1) + .with_token(); + + // Upload a new crate + let crate_to_publish = PublishBuilder::new("yankable", "1.0.0"); + token.publish_crate(crate_to_publish).good(); + check_yanked(&app, false); + + // Yank the crate + token.yank("yankable", "1.0.0").good(); + check_yanked(&app, true); + + // Try unyanking the crate, will get rate limited + token + .unyank("yankable", "1.0.0") + .assert_rate_limited(LimitedAction::YankUnyank); + check_yanked(&app, true); + + // Let the rate limit refill. + std::thread::sleep(Duration::from_millis(500)); + + // Unyanking now works. + token.unyank("yankable", "1.0.0").good(); + check_yanked(&app, false); + + // Yanking again will trigger the rate limit. + token + .yank("yankable", "1.0.0") + .assert_rate_limited(LimitedAction::YankUnyank); + check_yanked(&app, false); +} + #[test] fn yank_max_version() { let (_, anon, _, token) = TestApp::full().with_token(); diff --git a/src/tests/util/response.rs b/src/tests/util/response.rs index 513327e3490..d69bef618bc 100644 --- a/src/tests/util/response.rs +++ b/src/tests/util/response.rs @@ -2,6 +2,7 @@ use serde_json::Value; use std::marker::PhantomData; use std::ops::Deref; +use crates_io::rate_limiter::LimitedAction; use http::{header, StatusCode}; /// A type providing helper methods for working with responses @@ -57,6 +58,29 @@ impl Response { .ends_with(target)); self } + + /// Assert that the status code is 429 and that the body matches a rate limit. + #[track_caller] + pub fn assert_rate_limited(self, action: LimitedAction) { + #[derive(serde::Deserialize)] + #[serde(deny_unknown_fields)] + struct ErrorResponse { + errors: Vec, + } + #[derive(serde::Deserialize)] + #[serde(deny_unknown_fields)] + struct ErrorDetails { + detail: String, + } + + assert_eq!(StatusCode::TOO_MANY_REQUESTS, self.status()); + + let expected_message_start = + format!("{}. Please try again after ", action.error_messagge()); + let error: ErrorResponse = json(self.response); + assert_eq!(1, error.errors.len()); + assert!(error.errors[0].detail.starts_with(&expected_message_start)); + } } impl Response<()> { diff --git a/src/tests/util/test_app.rs b/src/tests/util/test_app.rs index b3ae270dbfb..5b16703421f 100644 --- a/src/tests/util/test_app.rs +++ b/src/tests/util/test_app.rs @@ -11,6 +11,7 @@ use std::{rc::Rc, sync::Arc, time::Duration}; use crate::util::github::{MockGitHubClient, MOCK_GITHUB_DATA}; use anyhow::Context; use crates_io::models::token::{CrateScope, EndpointScope}; +use crates_io::rate_limiter::{LimitedAction, RateLimiterConfig}; use crates_io::swirl::Runner; use diesel::PgConnection; use futures_util::TryStreamExt; @@ -329,10 +330,11 @@ impl TestAppBuilder { self } - pub fn with_publish_rate_limit(self, rate: Duration, burst: i32) -> Self { + pub fn with_rate_limit(self, action: LimitedAction, rate: Duration, burst: i32) -> Self { self.with_config(|config| { - config.rate_limiter.rate = rate; - config.rate_limiter.burst = burst; + config + .rate_limiter + .insert(action, RateLimiterConfig { rate, burst }); }) } diff --git a/src/util/errors/json.rs b/src/util/errors/json.rs index 7d584f04373..b64f16ada41 100644 --- a/src/util/errors/json.rs +++ b/src/util/errors/json.rs @@ -4,6 +4,7 @@ use std::fmt; use super::{AppError, BoxedAppError, InternalAppErrorStatic}; +use crate::rate_limiter::LimitedAction; use chrono::NaiveDateTime; use http::{header, StatusCode}; @@ -74,6 +75,7 @@ pub(super) struct ServerError(pub(super) String); pub(crate) struct ServiceUnavailable(pub(super) String); #[derive(Debug)] pub(crate) struct TooManyRequests { + pub action: LimitedAction, pub retry_after: NaiveDateTime, } @@ -131,9 +133,9 @@ impl AppError for TooManyRequests { let retry_after = self.retry_after.format(HTTP_DATE_FORMAT); let detail = format!( - "You have published too many crates in a \ - short period of time. Please try again after {retry_after} or email \ - help@crates.io to have your limit increased." + "{}. Please try again after {retry_after} or email \ + help@crates.io to have your limit increased.", + self.action.error_messagge() ); let mut response = json_error(&detail, StatusCode::TOO_MANY_REQUESTS); response.headers_mut().insert( diff --git a/src/worker/update_downloads.rs b/src/worker/update_downloads.rs index b90e38b329c..a03e5ab8b95 100644 --- a/src/worker/update_downloads.rs +++ b/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,