Skip to content

Commit

Permalink
rate limit yanking and unyanking
Browse files Browse the repository at this point in the history
  • Loading branch information
pietroalbini committed Jul 26, 2023
1 parent 4358690 commit aeb1356
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/controllers/version/yank.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ use crate::controllers::cargo_prelude::*;
use crate::models::token::EndpointScope;
use crate::models::Rights;
use crate::models::{insert_version_owner_action, VersionAction};
use crate::rate_limiter::LimitedAction;
use crate::schema::versions;

/// Handles the `DELETE /crates/:crate_id/:version/yank` route.
Expand Down Expand Up @@ -58,6 +59,10 @@ fn modify_yank(
.for_crate(crate_name)
.check(req, conn)?;

state
.rate_limiter
.check_rate_limit(auth.user_id(), LimitedAction::YankUnyank, conn)?;

let (version, krate) = version_and_crate(conn, crate_name, version)?;
let api_token_id = auth.api_token_id();
let user = auth.user();
Expand Down
7 changes: 7 additions & 0 deletions src/rate_limiter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ pg_enum! {
pub enum LimitedAction {
PublishNew = 0,
PublishUpdate = 1,
YankUnyank = 2,
}
}

Expand All @@ -21,20 +22,23 @@ impl LimitedAction {
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 => 30,
}
}

pub fn env_var_key(&self) -> &'static str {
match self {
LimitedAction::PublishNew => "PUBLISH_NEW",
LimitedAction::PublishUpdate => "PUBLISH_EXISTING",
LimitedAction::YankUnyank => "YANK_UNYANK",
}
}

Expand All @@ -46,6 +50,9 @@ impl LimitedAction {
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"
}
}
}
}
Expand Down
40 changes: 40 additions & 0 deletions src/tests/krate/yanking.rs
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -60,6 +62,44 @@ 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");
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();
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();
check_yanked(&app, false);
}

#[test]
fn yank_max_version() {
let (_, anon, _, token) = TestApp::full().with_token();
Expand Down

0 comments on commit aeb1356

Please sign in to comment.