Skip to content

Commit 1a5b0e8

Browse files
pietroalbinipietroalbini
committed
move rate limiting for existing crate publishes to the app
1 parent da22a27 commit 1a5b0e8

12 files changed

+583
-27
lines changed

‎config/nginx.conf.erb

-9
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,6 @@ http {
176176
client_body_timeout 30;
177177
client_max_body_size 50m;
178178

179-
limit_req_zone $remote_addr zone=publish:10m rate=1r/m;
180-
181179
upstream app_server {
182180
server localhost:8888 fail_timeout=0;
183181
}
@@ -261,12 +259,5 @@ http {
261259
proxy_pass http://app_server;
262260
}
263261
<% end %>
264-
265-
location ~ ^/api/v./crates/new$ {
266-
proxy_pass http://app_server;
267-
268-
limit_req zone=publish burst=30 nodelay;
269-
limit_req_status 429;
270-
}
271262
}
272263
}

‎src/controllers/krate/publish.rs

+12-1
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use crate::models::{
1111
NewVersion, Rights, VersionAction,
1212
};
1313

14+
use crate::rate_limiter::LimitedAction;
1415
use crate::render;
1516
use crate::schema::*;
1617
use crate::util::errors::{cargo_err, AppResult};
@@ -72,6 +73,16 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
7273
))
7374
})?;
7475

76+
// Check the rate limits on the endpoint. Publishing a new crate uses a different (stricter)
77+
// limiting pool than publishing a new version of an existing crate to prevent mass squatting.
78+
let limited_action = if Crate::by_name(&new_crate.name).execute(&*conn)? == 0 {
79+
LimitedAction::PublishNew
80+
} else {
81+
LimitedAction::PublishExisting
82+
};
83+
app.rate_limiter
84+
.check_rate_limit(user.id, limited_action, &conn)?;
85+
7586
// Create a transaction on the database, if there are no errors,
7687
// commit the transactions to record a new or updated crate.
7788
conn.transaction(|| {
@@ -107,7 +118,7 @@ pub fn publish(req: &mut dyn RequestExt) -> EndpointResult {
107118
};
108119

109120
let license_file = new_crate.license_file.as_deref();
110-
let krate = persist.create_or_update(&conn, user.id, Some(&app.rate_limiter))?;
121+
let krate = persist.create_or_update(&conn, user.id)?;
111122

112123
let owners = krate.owners(&conn)?;
113124
if user.rights(req.app(), &owners)? < Rights::Publish {

‎src/downloads_counter.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ mod tests {
438438
name: "foo",
439439
..NewCrate::default()
440440
}
441-
.create_or_update(conn, user.id, None)
441+
.create_or_update(conn, user.id)
442442
.expect("failed to create crate");
443443

444444
Self {

‎src/models/krate.rs

+1-10
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ use crate::models::{
1515
use crate::util::errors::{cargo_err, AppResult};
1616

1717
use crate::models::helpers::with_count::*;
18-
use crate::rate_limiter::{LimitedAction, RateLimiter};
1918
use crate::schema::*;
2019

2120
#[derive(Debug, Queryable, Identifiable, Associations, Clone, Copy)]
@@ -93,12 +92,7 @@ pub struct NewCrate<'a> {
9392
}
9493

9594
impl<'a> NewCrate<'a> {
96-
pub fn create_or_update(
97-
self,
98-
conn: &PgConnection,
99-
uploader: i32,
100-
rate_limit: Option<&RateLimiter>,
101-
) -> AppResult<Crate> {
95+
pub fn create_or_update(self, conn: &PgConnection, uploader: i32) -> AppResult<Crate> {
10296
use diesel::update;
10397

10498
self.validate()?;
@@ -108,9 +102,6 @@ impl<'a> NewCrate<'a> {
108102
// To avoid race conditions, we try to insert
109103
// first so we know whether to add an owner
110104
if let Some(krate) = self.save_new_crate(conn, uploader)? {
111-
if let Some(rate_limit) = rate_limit {
112-
rate_limit.check_rate_limit(uploader, LimitedAction::PublishNew, conn)?;
113-
}
114105
return Ok(krate);
115106
}
116107

‎src/rate_limiter.rs

+91
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use std::time::Duration;
99
crate::pg_enum! {
1010
pub enum LimitedAction {
1111
PublishNew = 0,
12+
PublishExisting = 1,
1213
}
1314
}
1415

@@ -17,13 +18,15 @@ impl LimitedAction {
1718
pub fn default_rate_seconds(&self) -> u64 {
1819
match self {
1920
LimitedAction::PublishNew => 60 * 10,
21+
LimitedAction::PublishExisting => 60,
2022
}
2123
}
2224

2325
/// How many requests a user can make before the rate limit goes into effect.
2426
pub fn default_burst(&self) -> i32 {
2527
match self {
2628
LimitedAction::PublishNew => 5,
29+
LimitedAction::PublishExisting => 30,
2730
}
2831
}
2932

@@ -32,13 +35,17 @@ impl LimitedAction {
3235
LimitedAction::PublishNew => {
3336
"You have published too many new crates in a short period of time."
3437
}
38+
LimitedAction::PublishExisting => {
39+
"You have published too many versions of existing crates in a short period of time."
40+
}
3541
}
3642
}
3743

3844
/// Key used to identify this action in environment variables. See `src/config.rs`.
3945
pub fn env_var_key(&self) -> &'static str {
4046
match self {
4147
LimitedAction::PublishNew => "PUBLISH_NEW",
48+
LimitedAction::PublishExisting => "PUBLISH_EXISTING",
4249
}
4350
}
4451
}
@@ -350,6 +357,49 @@ mod tests {
350357
Ok(())
351358
}
352359

360+
#[test]
361+
fn two_actions_dont_interfere_with_each_other() -> QueryResult<()> {
362+
let conn = pg_connection();
363+
let now = now();
364+
let user_id = new_user(&conn, "user")?;
365+
366+
let mut config = HashMap::new();
367+
config.insert(
368+
LimitedAction::PublishNew,
369+
RateLimiterConfig {
370+
rate: Duration::from_secs(1),
371+
burst: 10,
372+
},
373+
);
374+
config.insert(
375+
LimitedAction::PublishExisting,
376+
RateLimiterConfig {
377+
rate: Duration::from_secs(1),
378+
burst: 20,
379+
},
380+
);
381+
let rate = RateLimiter::new(config);
382+
383+
let refill_time = now + chrono::Duration::seconds(4);
384+
assert_eq!(
385+
10,
386+
rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
387+
.tokens,
388+
);
389+
assert_eq!(
390+
9,
391+
rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
392+
.tokens,
393+
);
394+
assert_eq!(
395+
20,
396+
rate.take_token(user_id, LimitedAction::PublishExisting, refill_time, &conn)?
397+
.tokens,
398+
);
399+
400+
Ok(())
401+
}
402+
353403
#[test]
354404
fn override_is_used_instead_of_global_burst_if_present() -> QueryResult<()> {
355405
let conn = pg_connection();
@@ -375,6 +425,47 @@ mod tests {
375425
Ok(())
376426
}
377427

428+
#[test]
429+
fn override_is_different_for_each_action() -> QueryResult<()> {
430+
let conn = pg_connection();
431+
let now = now();
432+
let user_id = new_user(&conn, "user")?;
433+
434+
let mut config = HashMap::new();
435+
for action in &[LimitedAction::PublishNew, LimitedAction::PublishExisting] {
436+
config.insert(
437+
*action,
438+
RateLimiterConfig {
439+
rate: Duration::from_secs(1),
440+
burst: 10,
441+
},
442+
);
443+
}
444+
let rate = RateLimiter::new(config);
445+
446+
diesel::insert_into(publish_rate_overrides::table)
447+
.values((
448+
publish_rate_overrides::user_id.eq(user_id),
449+
publish_rate_overrides::action.eq(LimitedAction::PublishNew),
450+
publish_rate_overrides::burst.eq(20),
451+
))
452+
.execute(&conn)?;
453+
454+
let refill_time = now + chrono::Duration::seconds(4);
455+
assert_eq!(
456+
20,
457+
rate.take_token(user_id, LimitedAction::PublishNew, refill_time, &conn)?
458+
.tokens,
459+
);
460+
assert_eq!(
461+
10,
462+
rate.take_token(user_id, LimitedAction::PublishExisting, refill_time, &conn)?
463+
.tokens,
464+
);
465+
466+
Ok(())
467+
}
468+
378469
#[test]
379470
fn overrides_can_expire() -> QueryResult<()> {
380471
let conn = pg_connection();

‎src/tasks/update_downloads.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ mod test {
9898
name: "foo",
9999
..Default::default()
100100
}
101-
.create_or_update(conn, user_id, None)
101+
.create_or_update(conn, user_id)
102102
.unwrap();
103103
let version = NewVersion::new(
104104
krate.id,

‎src/tests/builders/krate.rs

+1-3
Original file line numberDiff line numberDiff line change
@@ -114,9 +114,7 @@ impl<'a> CrateBuilder<'a> {
114114
pub fn build(mut self, connection: &PgConnection) -> AppResult<Crate> {
115115
use diesel::{insert_into, select, update};
116116

117-
let mut krate = self
118-
.krate
119-
.create_or_update(connection, self.owner_id, None)?;
117+
let mut krate = self.krate.create_or_update(connection, self.owner_id)?;
120118

121119
// Since we are using `NewCrate`, we can't set all the
122120
// crate properties in a single DB call.

0 commit comments

Comments
 (0)