From 2f452b15b15fb0bffbbb4b1b7e87073558961428 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 7 Sep 2021 15:48:06 +0100 Subject: [PATCH 1/6] add ratelimits --- Cargo.lock | 88 +------------ Cargo.toml | 1 - src/main.rs | 6 +- src/routes/channels/message_send.rs | 6 +- src/routes/root.rs | 7 +- src/util/ratelimit.rs | 188 +++++++++++++++++++++++++--- src/util/result.rs | 6 +- 7 files changed, 190 insertions(+), 112 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 85243f11..709bf03a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,12 +17,6 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee2a4ec343196209d6594e19543ae87a39f96d5534d7174822a3ad825dd6ed7e" -[[package]] -name = "ahash" -version = "0.3.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8fd72866655d1904d6b0997d0b07ba561047d070fbe29de039031c641b61217" - [[package]] name = "aho-corasick" version = "0.7.15" @@ -611,16 +605,6 @@ dependencies = [ "syn", ] -[[package]] -name = "dashmap" -version = "4.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e77a43b28d0668df09411cb0bc9a8c2adc40f9a048afe863e05fd43251e8e39c" -dependencies = [ - "cfg-if 1.0.0", - "num_cpus", -] - [[package]] name = "derivative" version = "2.2.0" @@ -942,12 +926,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "futures-timer" -version = "3.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e64b03909df88034c26dc1547e8970b91f98bdb65165d6a4e9110d94263dbb2c" - [[package]] name = "futures-util" version = "0.3.12" @@ -1047,23 +1025,6 @@ dependencies = [ "web-sys", ] -[[package]] -name = "governor" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06c5d2f987ee8f6dff3fa1a352058dc59b990e447e4c7846aa7d804971314f7b" -dependencies = [ - "dashmap", - "futures", - "futures-timer", - "no-std-compat", - "nonzero_ext", - "parking_lot", - "quanta", - "rand 0.8.3", - "smallvec", -] - [[package]] name = "h2" version = "0.2.7" @@ -1117,16 +1078,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "hashbrown" -version = "0.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e91b62f79061a0bc2e046024cb7ba44b08419ed238ecbd9adbd787434b9e8c25" -dependencies = [ - "ahash", - "autocfg 1.0.1", -] - [[package]] name = "hashbrown" version = "0.9.1" @@ -1376,7 +1327,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4fb1fa934250de4de8aef298d81c729a7d33d8c239daa3a7575e6b92bfc7313b" dependencies = [ "autocfg 1.0.1", - "hashbrown 0.9.1", + "hashbrown", "serde", ] @@ -1813,15 +1764,6 @@ dependencies = [ "libc", ] -[[package]] -name = "no-std-compat" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" -dependencies = [ - "hashbrown 0.8.2", -] - [[package]] name = "nom" version = "6.1.0" @@ -1833,12 +1775,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "nonzero_ext" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44a1290799eababa63ea60af0cbc3f03363e328e58f32fb0294798ed3e85f444" - [[package]] name = "ntapi" version = "0.3.6" @@ -2273,16 +2209,6 @@ dependencies = [ "yansi", ] -[[package]] -name = "quanta" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d98dc777a7a39b76b1a26ae9d3f691f4c1bc0455090aa0b64dfa8cb7fc34c135" -dependencies = [ - "libc", - "winapi 0.3.9", -] - [[package]] name = "quick-error" version = "1.2.3" @@ -2730,7 +2656,6 @@ dependencies = [ "reqwest 0.11.4", "rmp-serde", "rocket", - "rocket-governor", "rocket_cors", "serde", "serde_json", @@ -2815,17 +2740,6 @@ dependencies = [ "yansi", ] -[[package]] -name = "rocket-governor" -version = "0.0.1-rc.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a7603d1d66d1c4c41997f1f80fb624c16fd8a063a0f562d6e511e7eec4d26d2" -dependencies = [ - "governor", - "lazy_static", - "rocket", -] - [[package]] name = "rocket_codegen" version = "0.5.0-rc.1" diff --git a/Cargo.toml b/Cargo.toml index c9be9731..bf335e5a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,7 +52,6 @@ rauth = { git = "https://gitlab.insrt.uk/insert/rauth", rev = "19bef300ecc269c02 hive_pubsub = { git = "https://gitlab.insrt.uk/insert/hive", rev = "a89826df2b30166220e68a6ed01a58b751456604", features = ["mongo"] } # web -rocket-governor = { version = "0.0.1-rc.9", features = ["logger"] } async-tungstenite = { version = "0.10.0", features = ["async-std-runtime"] } rocket = { version = "0.5.0-rc.1", default-features = false, features = ["json"] } mongodb = { version = "1.2.2", features = ["tokio-runtime"], default-features = false } diff --git a/src/main.rs b/src/main.rs index d0756e0c..63e0c7a8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -31,11 +31,12 @@ use rauth::{ }; use std::str::FromStr; use rocket_cors::AllowedOrigins; -use rocket::catchers; +// use rocket::catchers; use util::variables::{ APP_URL, HCAPTCHA_KEY, INVITE_ONLY, PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME, USE_EMAIL, USE_HCAPTCHA, }; +use crate::util::ratelimit::RatelimitState; #[async_std::main] async fn main() { @@ -146,8 +147,9 @@ Reset your password here: {{url}}", .mount("/auth", rauth::routes::routes()) .manage(auth) .manage(cors.clone()) + .manage(RatelimitState::new()) .attach(cors) - .register("/", catchers![rocket_governor::rocket_governor_catcher]) + // .register("/", catchers![rocket_governor::rocket_governor_catcher]) .launch() .await .unwrap(); diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs index b49d4698..22d8b9a7 100644 --- a/src/routes/channels/message_send.rs +++ b/src/routes/channels/message_send.rs @@ -1,7 +1,7 @@ use std::collections::HashSet; use crate::database::*; -use crate::util::ratelimit::RateLimited; +use crate::util::ratelimit::{Ratelimiter, RatelimitResponse}; use crate::util::result::{Error, Result}; use mongodb::{bson::doc, options::FindOneOptions}; @@ -34,7 +34,7 @@ lazy_static! { } #[post("//messages", data = "")] -pub async fn message_send(_r: RateLimited<'_>, user: User, target: Ref, message: Json) -> Result { +pub async fn message_send(_r: Ratelimiter, user: User, target: Ref, message: Json) -> Result> { let message = message.into_inner(); message .validate() @@ -147,5 +147,5 @@ pub async fn message_send(_r: RateLimited<'_>, user: User, target: Ref, message: msg.clone().publish(&target, perm.get_embed_links()).await?; - Ok(json!(msg)) + Ok(RatelimitResponse(json!(msg))) } diff --git a/src/routes/root.rs b/src/routes/root.rs index 1b0a15d4..c15536f1 100644 --- a/src/routes/root.rs +++ b/src/routes/root.rs @@ -1,14 +1,13 @@ -use crate::util::{ratelimit::RateLimitGuard, variables::{ +use crate::util::{ratelimit::Ratelimiter, variables::{ APP_URL, AUTUMN_URL, EXTERNAL_WS_URL, HCAPTCHA_SITEKEY, INVITE_ONLY, JANUARY_URL, USE_AUTUMN, USE_EMAIL, USE_HCAPTCHA, USE_JANUARY, USE_VOSO, VAPID_PUBLIC_KEY, VOSO_URL, VOSO_WS_HOST, }}; use mongodb::bson::doc; use rocket::{http::Status, serde::json::Value}; -use rocket_governor::RocketGovernor; #[get("/")] -pub async fn root(_limitguard: RocketGovernor<'_, RateLimitGuard>) -> Value { +pub async fn root() -> Value { json!({ "revolt": crate::version::VERSION, "features": { @@ -39,6 +38,6 @@ pub async fn root(_limitguard: RocketGovernor<'_, RateLimitGuard>) -> Value { } #[get("/ping")] -pub async fn ping(_limitguard: RocketGovernor<'_, RateLimitGuard>) -> Status { +pub async fn ping(_limitguard: Ratelimiter) -> Status { Status::Ok } diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs index f951b452..3fcbc905 100644 --- a/src/util/ratelimit.rs +++ b/src/util/ratelimit.rs @@ -1,22 +1,182 @@ -use rocket_governor::{Method, Quota, RocketGovernable, RocketGovernor, NonZeroU32}; -use phf::phf_map; +use rocket::{async_trait, http::Status, request::{Outcome, FromRequest}, response}; +use std::{collections::{HashMap, hash_map::DefaultHasher}, time}; +use rauth::auth::Session; +use crate::util::result::Error; +use std::hash::{Hash, Hasher}; +use std::sync::{Arc, Mutex}; -pub struct RateLimitGuard; +#[inline] +fn now() -> f64 { + // will this ever actually panic? + time::SystemTime::now().duration_since(time::SystemTime::UNIX_EPOCH).expect("???").as_secs_f64() +} -static ROUTE_QUOTAS: phf::Map<&'static str, &'static Quota> = phf_map! { - "message_send" => &Quota::per_second(NonZeroU32::new(1u32).unwrap()) - .allow_burst(NonZeroU32::new(10u32).unwrap()) -}; +struct Ratelimit { + pub rate: u8, + pub per: u8, + window: f64, + tokens: u8, + pub last: f64, +} -impl<'r> RocketGovernable<'r> for RateLimitGuard { - fn quota(_method: Method, route: &str) -> Quota { - if let Some(q) = ROUTE_QUOTAS.get(route) { - **q +impl Ratelimit { + pub fn new(rate: u8, per: u8) -> Self { + Self { + rate, + per, + window: 0f64, + tokens: rate, + last: 0.0 + } + } + + fn get_tokens(&self, current: Option) -> u8 { + let current = current.unwrap_or(now()); + if current > (self.window + self.per as f64) { + self.rate } else { - Quota::per_second(Self::nonzero(1u32)) - .allow_burst(Self::nonzero(5u32)) + self.tokens + } + } + + fn update_ratelimit(&mut self) -> Option { + let current = now(); + self.last = current; + + self.tokens = self.get_tokens(Some(current)); + + if self.tokens == self.rate { + self.window = current + } + + if self.tokens == 0 { + return Some(self.per as f64 - (current - self.window)) + } + + self.tokens -= 1; + + if self.tokens == 0 { + self.window = current + } + + None + } +} + +impl Clone for Ratelimit { + fn clone(&self) -> Ratelimit { + Ratelimit::new(self.rate, self.per) + } +} + +#[derive(Clone)] +struct RatelimitMapping { + cache: HashMap, + cooldown: Ratelimit +} + +impl RatelimitMapping { + pub fn new(rate: u8, per: u8) -> Self { + RatelimitMapping { + cache: HashMap::new(), + cooldown: Ratelimit::new(rate, per) + } + } + + fn verify_cache(&mut self) { + let current = now(); + self.cache.retain(|_, v| current < v.last + v.per as f64); + } + + pub fn get_bucket(&mut self, key: u64) -> &mut Ratelimit { + self.verify_cache(); + self.cache.entry(key).or_insert(self.cooldown.clone()) + } +} + +#[derive(Copy, Clone)] +pub struct Ratelimiter { + bucket: u64, + limit: u8, + remaining: u8, + reset: f64 +} + +pub struct RatelimitState(Arc>>); + +impl RatelimitState { + pub fn new() -> Self { + let mut hashmap = HashMap::new(); + hashmap.insert("message_send", RatelimitMapping::new(10, 10)); + hashmap.insert("root", RatelimitMapping::new(10, 10)); + let arc = Arc::new(Mutex::new(hashmap)); + RatelimitState(arc) + } +} + +#[async_trait] +impl<'r> FromRequest<'r> for Ratelimiter { + type Error = Error; + + async fn from_request(request: &'r rocket::Request<'_>) -> Outcome { + let res = request.local_cache_async(async { + if let Some(route) = request.route() { + if let Some(route_name) = &route.name { + let session = request.guard::().await.unwrap(); + let state = request.guard::<&rocket::State>().await.unwrap().inner(); + let arc = Arc::clone(&state.0); + let mut mutex = arc.lock().unwrap(); + println!("{}", route_name); + let mapping = mutex.get_mut(route_name.as_ref()).unwrap(); + + // create a unique key tied to the user id and route they use + let key = format!("{}:{}:{}", session.user_id, route.method.as_str(), route.uri.as_str()); + let mut hasher = DefaultHasher::new(); + + key.hash(&mut hasher); + let hashed = hasher.finish(); + + let bucket = mapping.get_bucket(hashed); + let ret = bucket.update_ratelimit(); + + if let Some(retry_after) = ret { + Err(Error::TooManyRequests { retry_after }) + } else { + Ok(Ratelimiter { + bucket: hashed, + limit: bucket.rate, + remaining: bucket.get_tokens(None), + reset: bucket.window + (bucket.per as f64) + }) + } + } else { + todo!() + } + } else { + todo!() + } + }).await; + + match res { + Ok(rl) => Outcome::Success((*rl).clone()), + Err(e) => Outcome::Failure((Status::TooManyRequests, (*e).clone())) } } } -pub type RateLimited<'a> = RocketGovernor<'a, RateLimitGuard>; +pub struct RatelimitResponse(pub R); + +impl<'r, 'o: 'r, R: response::Responder<'r, 'o>> response::Responder<'r, 'o> for RatelimitResponse { + fn respond_to(self, request: &'r rocket::Request<'_>) -> response::Result<'o> { + let res: &Result = request.local_cache(|| unreachable!()); + let ratelimiter = res.as_ref().unwrap(); + + Ok(response::Builder::new(rocket::Response::new()) + .raw_header("x-ratelimit-bucket", ratelimiter.bucket.to_string()) + .raw_header("x-ratelimit-limit", ratelimiter.limit.to_string()) + .raw_header("x-ratelimit-remaining", ratelimiter.remaining.to_string()) + .raw_header("x-ratelimit-reset", ratelimiter.reset.to_string()) + .merge(self.0.respond_to(request)?) + .finalize()) + } +} diff --git a/src/util/result.rs b/src/util/result.rs index 6279a598..678a4b29 100644 --- a/src/util/result.rs +++ b/src/util/result.rs @@ -6,7 +6,7 @@ use serde::Serialize; use std::io::Cursor; use validator::ValidationErrors; -#[derive(Serialize, Debug)] +#[derive(Serialize, Debug, Clone)] #[serde(tag = "type")] pub enum Error { LabelMe, @@ -66,6 +66,9 @@ pub enum Error { VosoUnavailable, NotFound, NoEffect, + TooManyRequests { + retry_after: f64 + } } pub struct EmptyResponse; @@ -127,6 +130,7 @@ impl<'r> Responder<'r, 'static> for Error { Error::VosoUnavailable => Status::BadRequest, Error::NotFound => Status::NotFound, Error::NoEffect => Status::Ok, + Error::TooManyRequests { .. } => Status::TooManyRequests, }; // Serialize the error data structure into JSON. From fcb59b249355565ea453ab8e5b3ea8b8c79c05c2 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 7 Sep 2021 15:50:56 +0100 Subject: [PATCH 2/6] remove comments --- src/main.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/main.rs b/src/main.rs index 63e0c7a8..1578b9cc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -31,7 +31,6 @@ use rauth::{ }; use std::str::FromStr; use rocket_cors::AllowedOrigins; -// use rocket::catchers; use util::variables::{ APP_URL, HCAPTCHA_KEY, INVITE_ONLY, PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME, USE_EMAIL, USE_HCAPTCHA, @@ -149,7 +148,6 @@ Reset your password here: {{url}}", .manage(cors.clone()) .manage(RatelimitState::new()) .attach(cors) - // .register("/", catchers![rocket_governor::rocket_governor_catcher]) .launch() .await .unwrap(); From fc5d14d4aa22ea88457717024178ad035a48b355 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 7 Sep 2021 15:52:33 +0100 Subject: [PATCH 3/6] remove line --- src/util/ratelimit.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs index 3fcbc905..f69bbdd0 100644 --- a/src/util/ratelimit.rs +++ b/src/util/ratelimit.rs @@ -108,7 +108,6 @@ impl RatelimitState { pub fn new() -> Self { let mut hashmap = HashMap::new(); hashmap.insert("message_send", RatelimitMapping::new(10, 10)); - hashmap.insert("root", RatelimitMapping::new(10, 10)); let arc = Arc::new(Mutex::new(hashmap)); RatelimitState(arc) } From b99807cf2a79bf157bee7ddd3d7203a32957bff0 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Tue, 7 Sep 2021 15:57:38 +0100 Subject: [PATCH 4/6] use unwrap --- src/util/ratelimit.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs index f69bbdd0..71e5ef46 100644 --- a/src/util/ratelimit.rs +++ b/src/util/ratelimit.rs @@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex}; #[inline] fn now() -> f64 { // will this ever actually panic? - time::SystemTime::now().duration_since(time::SystemTime::UNIX_EPOCH).expect("???").as_secs_f64() + time::SystemTime::now().duration_since(time::SystemTime::UNIX_EPOCH).unwrap().as_secs_f64() } struct Ratelimit { From 46c4de8c68cae371edc72b1422b58ebad5eddf2c Mon Sep 17 00:00:00 2001 From: Zomatree Date: Wed, 15 Sep 2021 20:55:02 +0100 Subject: [PATCH 5/6] cleanup --- src/util/ratelimit.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs index 71e5ef46..16d7f744 100644 --- a/src/util/ratelimit.rs +++ b/src/util/ratelimit.rs @@ -116,7 +116,7 @@ impl RatelimitState { #[async_trait] impl<'r> FromRequest<'r> for Ratelimiter { type Error = Error; - + async fn from_request(request: &'r rocket::Request<'_>) -> Outcome { let res = request.local_cache_async(async { if let Some(route) = request.route() { @@ -125,7 +125,6 @@ impl<'r> FromRequest<'r> for Ratelimiter { let state = request.guard::<&rocket::State>().await.unwrap().inner(); let arc = Arc::clone(&state.0); let mut mutex = arc.lock().unwrap(); - println!("{}", route_name); let mapping = mutex.get_mut(route_name.as_ref()).unwrap(); // create a unique key tied to the user id and route they use @@ -149,16 +148,16 @@ impl<'r> FromRequest<'r> for Ratelimiter { }) } } else { - todo!() + unreachable!() } } else { - todo!() + unreachable!() } }).await; match res { - Ok(rl) => Outcome::Success((*rl).clone()), - Err(e) => Outcome::Failure((Status::TooManyRequests, (*e).clone())) + Ok(rl) => Outcome::Success(*rl), + Err(e) => Outcome::Failure((Status::TooManyRequests, e.clone())) } } } From 6e70a39fce7c547eb286a7e911260e663bdbcb43 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Wed, 15 Sep 2021 21:37:42 +0100 Subject: [PATCH 6/6] fix rauth --- src/util/ratelimit.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs index 16d7f744..8af34a4c 100644 --- a/src/util/ratelimit.rs +++ b/src/util/ratelimit.rs @@ -1,6 +1,6 @@ use rocket::{async_trait, http::Status, request::{Outcome, FromRequest}, response}; use std::{collections::{HashMap, hash_map::DefaultHasher}, time}; -use rauth::auth::Session; +use rauth::entities::Session; use crate::util::result::Error; use std::hash::{Hash, Hasher}; use std::sync::{Arc, Mutex};