jmug/stoatchat
1
0
Fork 1
mirror of https://github.com/stoatchat/stoatchat.git synced 2026-07-14 05:26:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

add ratelimits

Browse Source
This commit is contained in:
Zomatree
2021-09-07 15:48:06 +01:00
parent b0d2e99588
commit 2f452b15b1
7 changed files with 190 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
Cargo.lock generated
View File

@@ -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"

1
Cargo.toml
View File

@@ -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 }

6
src/main.rs
View File

@@ -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();

6
src/routes/channels/message_send.rs
View File

@@ -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("/<target>/messages", data = "<message>")]
pub async fn message_send(_r: RateLimited<'_>, user: User, target: Ref, message: Json<Data>) -> Result<Value> {
pub async fn message_send(_r: Ratelimiter, user: User, target: Ref, message: Json<Data>) -> Result<RatelimitResponse<Value>> {
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)))
}

7
src/routes/root.rs
View File

@@ -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
}

188
src/util/ratelimit.rs
View File

@@ -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<f64>) -> 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<f64> {
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<u64, Ratelimit>,
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<Mutex<HashMap<&'static str, RatelimitMapping>>>);
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<Self, Self::Error> {
let res = request.local_cache_async(async {
if let Some(route) = request.route() {
if let Some(route_name) = &route.name {
let session = request.guard::<Session>().await.unwrap();
let state = request.guard::<&rocket::State<RatelimitState>>().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<R>(pub R);
impl<'r, 'o: 'r, R: response::Responder<'r, 'o>> response::Responder<'r, 'o> for RatelimitResponse<R> {
fn respond_to(self, request: &'r rocket::Request<'_>) -> response::Result<'o> {
let res: &Result<Ratelimiter, Error> = 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())
}
}

6
src/util/result.rs
View File

@@ -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.