forked from jmug/stoatchat
Merge pull request #79 from Zomatree/ratelimits
This commit is contained in:
Paul Makles
GitHub
156
Cargo.lock
generated
156
Cargo.lock
generated
@@ -27,12 +27,6 @@ version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
|
||||
|
||||
[[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.18"
|
||||
@@ -121,7 +115,7 @@ dependencies = [
|
||||
"parking",
|
||||
"polling",
|
||||
"slab",
|
||||
"socket2 0.4.1",
|
||||
"socket2 0.4.2",
|
||||
"waker-fn",
|
||||
"winapi 0.3.9",
|
||||
]
|
||||
@@ -665,16 +659,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"
|
||||
@@ -994,12 +978,6 @@ version = "0.3.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d3d00f4eddb73e498a54394f228cd55853bdf059259e8e7bc6e69d408892e99"
|
||||
|
||||
[[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.17"
|
||||
@@ -1100,23 +1078,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.4",
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.2.7"
|
||||
@@ -1170,16 +1131,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.11.2"
|
||||
@@ -1344,7 +1295,7 @@ dependencies = [
|
||||
"httpdate 1.0.1",
|
||||
"itoa",
|
||||
"pin-project-lite 0.2.7",
|
||||
"socket2 0.4.1",
|
||||
"socket2 0.4.2",
|
||||
"tokio 1.11.0",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
@@ -1429,7 +1380,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bc633605454125dec4b66843673f01c7df2b89479b32e0ed634e43a91cff62a5"
|
||||
dependencies = [
|
||||
"autocfg 1.0.1",
|
||||
"hashbrown 0.11.2",
|
||||
"hashbrown",
|
||||
"serde",
|
||||
]
|
||||
|
||||
@@ -1492,9 +1443,9 @@ checksum = "b71991ff56294aa922b450139ee08b3bfc70982c6b2c7562771375cf73542dd4"
|
||||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.54"
|
||||
version = "0.3.55"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1866b355d9c878e5e607473cbe3f63282c0b7aad2db1dbebf55076c686918254"
|
||||
checksum = "7cc9ffccd38c451a86bf13657df244e9c3f37493cce8e5e21e940963777acc84"
|
||||
dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
@@ -1552,9 +1503,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.101"
|
||||
version = "0.2.102"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3cb00336871be5ed2c8ed44b60ae9959dc5b9f08539422ed43f09e34ecaeba21"
|
||||
checksum = "a2a5ac8f984bfcf3a823267e5fde638acc3325f6496633a5da6bb6eb2171e103"
|
||||
|
||||
[[package]]
|
||||
name = "linked-hash-map"
|
||||
@@ -1857,15 +1808,6 @@ dependencies = [
|
||||
"memoffset",
|
||||
]
|
||||
|
||||
[[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.2"
|
||||
@@ -1878,12 +1820,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"
|
||||
@@ -2235,9 +2171,9 @@ checksum = "ac74c624d6b2d21f425f752262f42188365d7b8ff1aff74c82e45136510a4857"
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro-crate"
|
||||
version = "1.0.0"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41fdbd1df62156fbc5945f4762632564d7d038153091c3fcf1067f6aef7cff92"
|
||||
checksum = "1ebace6889caf889b4d3f76becee12e90353f2b8c7d875534a71e5742f8f6f83"
|
||||
dependencies = [
|
||||
"thiserror",
|
||||
"toml",
|
||||
@@ -2301,16 +2237,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"
|
||||
@@ -2757,7 +2683,6 @@ dependencies = [
|
||||
"reqwest 0.11.4",
|
||||
"rmp-serde",
|
||||
"rocket",
|
||||
"rocket-governor",
|
||||
"rocket_cors",
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -2840,17 +2765,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"
|
||||
@@ -3091,9 +3005,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.67"
|
||||
version = "1.0.68"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a7f9e390c27c3c0ce8bc5d725f6e4d30a29d26659494aa4b17535f7522c5c950"
|
||||
checksum = "0f690853975602e1bfe1ccbf50504d67174e3bcf340f23b5ea9992e0587a52d8"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"itoa",
|
||||
@@ -3169,9 +3083,9 @@ checksum = "2579985fda508104f7587689507983eadd6a6e84dd35d6d115361f530916fa0d"
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.9.6"
|
||||
version = "0.9.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9204c41a1597a8c5af23c82d1c921cb01ec0a4c59e07a9c7306062829a3903f3"
|
||||
checksum = "b69f9a4c9740d74c5baa3fd2e547f9525fa8088a8a958e0ca2409a514e33f5fa"
|
||||
dependencies = [
|
||||
"block-buffer 0.9.0",
|
||||
"cfg-if 1.0.0",
|
||||
@@ -3220,9 +3134,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "socket2"
|
||||
version = "0.4.1"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "765f090f0e423d2b55843402a07915add955e7d60657db13707a159727326cad"
|
||||
checksum = "5dc90fe6c7be1a323296982db1836d1ea9e47b6839496dde9a541bc496df3516"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"winapi 0.3.9",
|
||||
@@ -3473,9 +3387,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec"
|
||||
version = "1.3.1"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "848a1e1181b9f6753b5e96a092749e29b11d19ede67dfbbd6c7dc7e0f49b5338"
|
||||
checksum = "5241dd6f21443a3606b432718b166d3cedc962fd4b8bea54a8bc7f514ebda986"
|
||||
dependencies = [
|
||||
"tinyvec_macros",
|
||||
]
|
||||
@@ -3658,9 +3572,9 @@ checksum = "360dfd1d6d30e05fda32ace2c8c70e9c0a9da713275777f5a4dbb8a1893930c6"
|
||||
|
||||
[[package]]
|
||||
name = "tracing"
|
||||
version = "0.1.26"
|
||||
version = "0.1.27"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09adeb8c97449311ccd28a427f96fb563e7fd31aabf994189879d9da2394b89d"
|
||||
checksum = "c2ba9ab62b7d6497a8638dfda5e5c4fb3b2d5a7fca4118f2b96151c8ef1a437e"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.0",
|
||||
"log",
|
||||
@@ -3670,9 +3584,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tracing-core"
|
||||
version = "0.1.19"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ca517f43f0fb96e0c3072ed5c275fe5eece87e8cb52f4a77b69226d3b1c9df8"
|
||||
checksum = "46125608c26121c81b0c6d693eab5a420e416da7e43c426d2e8f7df8da8a3acf"
|
||||
dependencies = [
|
||||
"lazy_static",
|
||||
]
|
||||
@@ -4023,9 +3937,9 @@ checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f"
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.77"
|
||||
version = "0.2.78"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e68338db6becec24d3c7977b5bf8a48be992c934b5d07177e3931f5dc9b076c"
|
||||
checksum = "632f73e236b219150ea279196e54e610f5dbafa5d61786303d4da54f84e47fce"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.0",
|
||||
"serde",
|
||||
@@ -4035,9 +3949,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-backend"
|
||||
version = "0.2.77"
|
||||
version = "0.2.78"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f34c405b4f0658583dba0c1c7c9b694f3cac32655db463b56c254a1c75269523"
|
||||
checksum = "a317bf8f9fba2476b4b2c85ef4c4af8ff39c3c7f0cdfeed4f82c34a880aa837b"
|
||||
dependencies = [
|
||||
"bumpalo",
|
||||
"lazy_static",
|
||||
@@ -4050,9 +3964,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-futures"
|
||||
version = "0.4.27"
|
||||
version = "0.4.28"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a87d738d4abc4cf22f6eb142f5b9a81301331ee3c767f2fef2fda4e325492060"
|
||||
checksum = "8e8d7523cb1f2a4c96c1317ca690031b714a51cc14e05f712446691f413f5d39"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.0",
|
||||
"js-sys",
|
||||
@@ -4062,9 +3976,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro"
|
||||
version = "0.2.77"
|
||||
version = "0.2.78"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b9d5a6580be83b19dc570a8f9c324251687ab2184e57086f71625feb57ec77c8"
|
||||
checksum = "d56146e7c495528bf6587663bea13a8eb588d39b36b679d83972e1a2dbbdacf9"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"wasm-bindgen-macro-support",
|
||||
@@ -4072,9 +3986,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro-support"
|
||||
version = "0.2.77"
|
||||
version = "0.2.78"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e3775a030dc6f5a0afd8a84981a21cc92a781eb429acef9ecce476d0c9113e92"
|
||||
checksum = "7803e0eea25835f8abdc585cd3021b3deb11543c6fe226dcd30b228857c5c5ab"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -4085,9 +3999,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-shared"
|
||||
version = "0.2.77"
|
||||
version = "0.2.78"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c279e376c7a8e8752a8f1eaa35b7b0bee6bb9fb0cdacfa97cc3f1f289c87e2b4"
|
||||
checksum = "0237232789cf037d5480773fe568aac745bfe2afbc11a863e97901780a6b47cc"
|
||||
|
||||
[[package]]
|
||||
name = "web-push"
|
||||
@@ -4116,9 +4030,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.54"
|
||||
version = "0.3.55"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0a84d70d1ec7d2da2d26a5bd78f4bca1b8c3254805363ce743b7a05bc30d195a"
|
||||
checksum = "38eb105f1c59d9eaa6b5cdc92b859d85b926e82cb2e0945cd0c9259faa6fe9fb"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
|
||||
@@ -52,7 +52,6 @@ rauth = { git = "https://github.com/insertish/rauth", rev = "ad2f5ad1657f4b20c8e
|
||||
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 }
|
||||
|
||||
@@ -27,13 +27,13 @@ use rauth::{
|
||||
config::{Captcha, Config, EmailVerification, SMTPSettings, Template, Templates},
|
||||
logic::Auth,
|
||||
};
|
||||
use rocket::catchers;
|
||||
use rocket_cors::AllowedOrigins;
|
||||
use std::str::FromStr;
|
||||
use rocket_cors::AllowedOrigins;
|
||||
use util::variables::{
|
||||
APP_URL, HCAPTCHA_KEY, INVITE_ONLY, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME,
|
||||
USE_EMAIL, USE_HCAPTCHA,
|
||||
};
|
||||
use crate::util::ratelimit::RatelimitState;
|
||||
|
||||
#[async_std::main]
|
||||
async fn main() {
|
||||
@@ -132,8 +132,8 @@ async fn launch_web() {
|
||||
.mount("/auth/session", rauth::web::session::routes())
|
||||
.manage(auth)
|
||||
.manage(cors.clone())
|
||||
.manage(RatelimitState::new())
|
||||
.attach(cors)
|
||||
.register("/", catchers![rocket_governor::rocket_governor_catcher])
|
||||
.launch()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -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};
|
||||
@@ -35,7 +35,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()
|
||||
@@ -149,5 +149,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)))
|
||||
}
|
||||
|
||||
@@ -1,18 +1,13 @@
|
||||
use crate::util::{
|
||||
ratelimit::RateLimitGuard,
|
||||
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 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": {
|
||||
@@ -43,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
|
||||
}
|
||||
|
||||
@@ -1,22 +1,180 @@
|
||||
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::entities::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).unwrap().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));
|
||||
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();
|
||||
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 {
|
||||
unreachable!()
|
||||
}
|
||||
} else {
|
||||
unreachable!()
|
||||
}
|
||||
}).await;
|
||||
|
||||
match res {
|
||||
Ok(rl) => Outcome::Success(*rl),
|
||||
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,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,
|
||||
@@ -69,6 +69,9 @@ pub enum Error {
|
||||
VosoUnavailable,
|
||||
NotFound,
|
||||
NoEffect,
|
||||
TooManyRequests {
|
||||
retry_after: f64
|
||||
}
|
||||
}
|
||||
|
||||
pub struct EmptyResponse;
|
||||
@@ -131,6 +134,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.
|
||||
|
||||
Reference in New Issue
Block a user