diff --git a/Cargo.lock b/Cargo.lock index d8b4fc82..9a4ddb28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -278,7 +278,7 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3410529e8288c463bedb5930f82833bc0c90e5d2fe639a56582a4d09220b281" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", ] [[package]] @@ -306,9 +306,9 @@ checksum = "1d49d90015b3c36167a20fe2810c5cd875ad504b39cff3d4eae7977e6b7c1cb2" [[package]] name = "autocfg" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdb031dd78e28731d87d56cc8ffef4a8f36ca26c38fe2de700543e627f8a464a" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" [[package]] name = "backtrace" @@ -764,6 +764,17 @@ dependencies = [ "syn 1.0.76", ] +[[package]] +name = "dashmap" +version = "5.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c8858831f7781322e539ea39e72449c46b059638250c14344fec8d0aa6e539c" +dependencies = [ + "cfg-if 1.0.0", + "num_cpus", + "parking_lot 0.12.0", +] + [[package]] name = "data-encoding" version = "2.3.2" @@ -1504,7 +1515,7 @@ version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc633605454125dec4b66843673f01c7df2b89479b32e0ed634e43a91cff62a5" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", "hashbrown", "serde", ] @@ -1655,10 +1666,11 @@ dependencies = [ [[package]] name = "lock_api" -version = "0.4.5" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712a4d093c9976e24e7dbca41db895dabcbac38eb5f4045393d17a95bdfb1109" +checksum = "327fa5b6a6940e4699ec49a9beae1ea4845c6bab9314e4f84ac68742139d8c53" dependencies = [ + "autocfg 1.1.0", "scopeguard", ] @@ -1750,7 +1762,7 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59accc507f1338036a0477ef61afdae33cde60840f4dfe481319ce3ad116ddf9" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", ] [[package]] @@ -1766,7 +1778,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a92518e98c078586bc6c934028adcca4c92a53d6a958196de835170a01d84e4b" dependencies = [ "adler", - "autocfg 1.0.1", + "autocfg 1.1.0", ] [[package]] @@ -2039,7 +2051,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2cc698a63b549a70bc047073d2949cce27cd1c7b0a4a862d08a8031bc2801db" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", "num-traits", ] @@ -2049,14 +2061,14 @@ version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", ] [[package]] name = "num_cpus" -version = "1.13.0" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05499f3756671c15885fee9034446956fff3f243d6077b91e5767df161f766b3" +checksum = "19e64526ebdee182341572e50e9ad03965aa510cd94427a4549448f285e957a1" dependencies = [ "hermit-abi", "libc", @@ -2147,7 +2159,7 @@ version = "0.9.66" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1996d2d305e561b70d1ee0c53f1542833f4e1ac6ce9a6708b6ff2738ca67dc82" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", "cc", "libc", "pkg-config", @@ -2210,7 +2222,17 @@ checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" dependencies = [ "instant", "lock_api", - "parking_lot_core", + "parking_lot_core 0.8.5", +] + +[[package]] +name = "parking_lot" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87f5ec2493a61ac0506c0f4199f99070cbe83857b0337006a30f3e6719b8ef58" +dependencies = [ + "lock_api", + "parking_lot_core 0.9.2", ] [[package]] @@ -2227,6 +2249,19 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "parking_lot_core" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "995f667a6c822200b0433ac218e05582f0e2efa1b922a3fd2fbaadc5f87bab37" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "redox_syscall", + "smallvec", + "windows-sys", +] + [[package]] name = "pbkdf2" version = "0.8.0" @@ -2511,7 +2546,7 @@ dependencies = [ "fnv", "lazy_static", "libc", - "parking_lot", + "parking_lot 0.11.2", "procfs", "regex", "thiserror", @@ -2557,7 +2592,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f" dependencies = [ "log", - "parking_lot", + "parking_lot 0.11.2", "scheduled-thread-pool", ] @@ -2974,6 +3009,7 @@ dependencies = [ "bitfield", "chrono", "ctrlc", + "dashmap", "dotenv", "env_logger", "futures", @@ -3013,6 +3049,7 @@ dependencies = [ [[package]] name = "revolt-quark" version = "0.1.0" +source = "git+https://github.com/revoltchat/quark?rev=d65f49168c186b15a8465f16cbeccebac9b2bb7b#d65f49168c186b15a8465f16cbeccebac9b2bb7b" dependencies = [ "async-recursion", "async-std", @@ -3118,7 +3155,7 @@ dependencies = [ "memchr", "multer", "num_cpus", - "parking_lot", + "parking_lot 0.11.2", "pin-project-lite 0.2.7", "rand 0.8.4", "ref-cast", @@ -3182,7 +3219,7 @@ dependencies = [ "log", "memchr", "mime", - "parking_lot", + "parking_lot 0.11.2", "pear", "percent-encoding", "pin-project-lite 0.2.7", @@ -3319,7 +3356,7 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7" dependencies = [ - "parking_lot", + "parking_lot 0.11.2", ] [[package]] @@ -3898,7 +3935,7 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4efe6fc2395938c8155973d7be49fe8d03a843726e285e100a8a383cc0154ce" dependencies = [ - "autocfg 1.0.1", + "autocfg 1.1.0", "bytes 1.1.0", "libc", "memchr", @@ -4151,7 +4188,7 @@ dependencies = [ "lazy_static", "log", "lru-cache", - "parking_lot", + "parking_lot 0.11.2", "resolv-conf", "smallvec", "thiserror", @@ -4639,6 +4676,49 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-sys" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5acdd78cb4ba54c0045ac14f62d8f94a03d10047904ae2a40afa1e99d8f70825" +dependencies = [ + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_msvc" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17cffbe740121affb56fad0fc0e421804adf0ae00891205213b5cecd30db881d" + +[[package]] +name = "windows_i686_gnu" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2564fde759adb79129d9b4f54be42b32c89970c18ebf93124ca8870a498688ed" + +[[package]] +name = "windows_i686_msvc" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cd9d32ba70453522332c14d38814bceeb747d80b3958676007acadd7e166956" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfce6deae227ee8d356d19effc141a509cc503dfd1f850622ec4b0f84428e1f4" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d19538ccc21819d01deaf88d6a17eae6596a12e9aafdbb97916fb49896d89de9" + [[package]] name = "winreg" version = "0.6.2" diff --git a/Cargo.toml b/Cargo.toml index 9a70a7e4..860ce85c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ lru = "0.7.0" url = "2.2.2" log = "0.4.11" dotenv = "0.15.0" +dashmap = "5.2.0" linkify = "0.6.0" once_cell = "1.4.1" env_logger = "0.7.1" diff --git a/docker-compose.yml b/docker-compose.yml index 499db3c7..f2b9fcf4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -14,5 +14,4 @@ services: - REVOLT_UNSAFE_NO_CAPTCHA=1 ports: - "8000:8000" - - "9000:9000" restart: unless-stopped diff --git a/set_version.sh b/set_version.sh index df2f4415..fad28050 100755 --- a/set_version.sh +++ b/set_version.sh @@ -1,3 +1,3 @@ #!/bin/bash -export version=0.5.3-alpha.16 +export version=0.5.3 echo "pub const VERSION: &str = \"${version}\";" > src/version.rs diff --git a/src/main.rs b/src/main.rs index 2e68acea..fb36d705 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,16 +8,11 @@ extern crate serde_json; extern crate lazy_static; extern crate ctrlc; -//pub mod database; -//pub mod notifications; pub mod routes; -//pub mod redis; pub mod util; pub mod version; //pub mod task_queue; -use async_std::task; -use futures::join; use log::info; use rauth::{ config::{Captcha, Config, EmailVerification, SMTPSettings, Template, Templates}, @@ -30,7 +25,6 @@ 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() { @@ -38,33 +32,19 @@ async fn main() { env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")); info!( - "Starting REVOLT server [version {}].", + "Starting Revolt server [version {}].", crate::version::VERSION ); - /*util::variables::preflight_checks(); - database::connect().await; - redis::connect().await; - notifications::hive::init_hive().await; - task_queue::start_queues().await;*/ + util::variables::preflight_checks(); + #[cfg(debug_assertions)] ctrlc::set_handler(move || { // Force ungraceful exit to avoid hang. std::process::exit(0); }) .expect("Error setting Ctrl-C handler"); - let web_task = task::spawn(launch_web()); - //let hive_task = task::spawn_local(notifications::hive::listen()); - - join!( - web_task, - //hive_task, - //notifications::websocket::launch_server() - ); -} - -async fn launch_web() { let cors = rocket_cors::CorsOptions { allowed_origins: AllowedOrigins::All, allowed_methods: [ @@ -136,6 +116,7 @@ async fn launch_web() { let rocket = rocket::build(); routes::mount(rocket) .mount("/", rocket_cors::catch_all_options_routes()) + .mount("/", util::ratelimiter::routes()) .mount( "/swagger/", rocket_okapi::swagger_ui::make_swagger_ui(&rocket_okapi::swagger_ui::SwaggerUIConfig { @@ -147,6 +128,7 @@ async fn launch_web() { .manage(db) .manage(cors.clone()) //.manage(RatelimitState::new()) + .attach(util::ratelimiter::RatelimitFairing) .attach(cors) .launch() .await diff --git a/src/util/mod.rs b/src/util/mod.rs index 89e1680f..3e327cab 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,4 +1,4 @@ -pub mod variables; -// pub mod ratelimit; pub mod idempotency; +pub mod ratelimiter; pub mod regex; +pub mod variables; diff --git a/src/util/ratelimiter.rs b/src/util/ratelimiter.rs new file mode 100644 index 00000000..8159d77d --- /dev/null +++ b/src/util/ratelimiter.rs @@ -0,0 +1,288 @@ +//! Pulled from lightspeed-tv/backend. +//! +//! This will be replaced again in the near future since +//! I don't want duplication between two different projects. + +use std::collections::hash_map::DefaultHasher; +use std::hash::Hasher; +use std::net::SocketAddr; +use std::ops::Add; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use rocket::fairing::{Fairing, Info, Kind}; +use rocket::http::uri::Origin; +use rocket::http::{Method, Status}; +use rocket::request::{FromRequest, Outcome}; +use rocket::serde::json::Json; +use rocket::{Data, Request, Response}; + +use rocket_okapi::gen::OpenApiGenerator; +use rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput}; + +use serde::Serialize; + +use dashmap::DashMap; + +use log::info; + +/// Ratelimit Bucket +#[derive(Clone, Copy)] +struct Entry { + used: u8, + reset: u128, +} + +lazy_static! { + static ref MAP: DashMap = DashMap::new(); +} + +/// Get the current time from Unix Epoch as a Duration +fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time went backwards...") +} + +impl Entry { + /// Find bucket by its key + pub fn from(key: u64) -> Entry { + MAP.get(&key).map(|x| *x).unwrap_or_else(|| Entry { + used: 0, + reset: now().add(Duration::from_secs(10)).as_millis(), + }) + } + + /// Deduct one unit from the bucket and save + pub fn deduct(mut self, key: u64) { + let current_time = now().as_millis(); + if current_time > self.reset { + self.used = 1; + self.reset = now().add(Duration::from_secs(10)).as_millis(); + } else { + self.used += 1; + } + + MAP.insert(key, self); + } + + /// Get remaining units in the bucket + pub fn get_remaining(&self, limit: u8) -> u8 { + if now().as_millis() > self.reset { + limit + } else { + limit - self.used + } + } + + /// Get how long bucket has until reset + pub fn left_until_reset(&self) -> u128 { + let current_time = now().as_millis(); + if current_time > self.reset { + 0 + } else { + self.reset - current_time + } + } +} + +/// Ratelimit Guard +#[derive(Serialize, Clone, Copy, Debug)] +#[allow(dead_code)] +pub struct Ratelimiter { + key: u64, + limit: u8, + remaining: u8, + reset: u128, +} + +/// Find bucket from given request +/// +/// Optionally, include a resource id to hash against. +fn resolve_bucket<'r>(request: &'r rocket::Request<'_>) -> (&'r str, Option<&'r str>) { + if let Some(segment) = request.routed_segment(0) { + let resource = request.routed_segment(1); + match (segment, resource) { + ("users", _) => ("users", None), + ("bots", _) => ("bots", None), + ("channels", Some(id)) => { + if request.method() == Method::Post { + if let Some("messages") = request.routed_segment(2) { + return ("messaging", Some(id)); + } + } + + ("channels", Some(id)) + } + ("servers", Some(id)) => ("servers", Some(id)), + ("auth", _) => ("auth", None), + ("swagger", _) => ("swagger", None), + _ => ("any", None), + } + } else { + ("any", None) + } +} + +/// Resolve per-bucket limits +fn resolve_bucket_limit(bucket: &str) -> u8 { + match bucket { + "users" => 20, + "bots" => 10, + "messaging" => 10, + "channels" => 15, + "servers" => 5, + "auth" => 3, + "swagger" => 100, + _ => 20, + } +} + +impl Ratelimiter { + /// Generate guard from identifier and target bucket + pub fn from( + identifier: &str, + (bucket, resource): (&str, Option<&str>), + ) -> Result { + let mut key = DefaultHasher::new(); + key.write(identifier.as_bytes()); + key.write(bucket.as_bytes()); + + if let Some(id) = resource { + key.write(id.as_bytes()); + } + + let key = key.finish(); + let limit = resolve_bucket_limit(bucket); + let entry = Entry::from(key); + + let remaining = entry.get_remaining(limit); + let reset = entry.left_until_reset(); + + if remaining > 0 { + entry.deduct(key); + Ok(Ratelimiter { + key, + limit, + remaining: remaining - 1, + reset, + }) + } else { + Err(reset) + } + } +} + +#[async_trait] +impl<'r> FromRequest<'r> for Ratelimiter { + type Error = u128; + + async fn from_request(request: &'r rocket::Request<'_>) -> Outcome { + let ratelimiter = request + .local_cache_async(async { + use rocket::outcome::Outcome; + let identifier = if let Outcome::Success(session) = + request.guard::().await + { + session.id.expect("`id` on User") + } else if let Outcome::Success(addr) = request.guard::().await { + addr.ip().to_string() + } else { + panic!("No identifier!"); + }; + + Ratelimiter::from(&identifier, resolve_bucket(request)) + }) + .await; + + match ratelimiter { + Ok(ratelimiter) => Outcome::Success(*ratelimiter), + Err(retry_after) => Outcome::Failure((Status::TooManyRequests, *retry_after)), + } + } +} + +impl<'r> OpenApiFromRequest<'r> for Ratelimiter { + fn from_request_input( + _gen: &mut OpenApiGenerator, + _name: String, + _required: bool, + ) -> rocket_okapi::Result { + Ok(RequestHeaderInput::None) + } +} + +/// Attach ratelimiter to the Rocket application +pub struct RatelimitFairing; + +#[rocket::async_trait] +impl Fairing for RatelimitFairing { + fn info(&self) -> Info { + Info { + name: "Ratelimiter", + kind: Kind::Request | Kind::Response, + } + } + + async fn on_request(&self, request: &mut Request<'_>, _: &mut Data<'_>) { + use rocket::outcome::Outcome; + if let Outcome::Failure(_) = request.guard::().await { + info!( + "User rate-limited on route {}! (IP = {:?})", + request.uri(), + request.remote().map(|x| x.ip()) + ); + + request.set_uri(Origin::parse("/ratelimit").unwrap()) + } + } + + async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) { + use rocket::outcome::Outcome; + match request.guard::().await { + Outcome::Success(ratelimiter) => { + let Ratelimiter { + key, + limit, + remaining, + reset, + } = ratelimiter; + + response.set_raw_header("X-RateLimit-Limit", limit.to_string()); + response.set_raw_header("X-RateLimit-Bucket", key.to_string()); + response.set_raw_header("X-RateLimit-Remaining", remaining.to_string()); + response.set_raw_header("X-RateLimit-Reset-After", reset.to_string()); + } + Outcome::Failure(_) => {} + Outcome::Forward(_) => unreachable!(), + } + } +} + +#[derive(Serialize)] +#[serde(untagged)] +pub enum RatelimitInformation { + Success(Ratelimiter), + Failure { retry_after: u128 }, +} + +#[async_trait] +impl<'r> FromRequest<'r> for RatelimitInformation { + type Error = u128; + + async fn from_request(request: &'r rocket::Request<'_>) -> Outcome { + Outcome::Success(match request.guard::().await { + Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter), + Outcome::Failure((_, retry_after)) => RatelimitInformation::Failure { retry_after }, + _ => unreachable!(), + }) + } +} + +#[get("/ratelimit")] +fn ratelimit_info(info: RatelimitInformation) -> Json { + Json(info) +} + +pub fn routes() -> Vec { + routes![ratelimit_info] +} diff --git a/src/util/variables.rs b/src/util/variables.rs index ad5d8121..fca7f006 100644 --- a/src/util/variables.rs +++ b/src/util/variables.rs @@ -5,8 +5,6 @@ use log::warn; lazy_static! { // Application Settings - pub static ref MONGO_URI: String = - env::var("REVOLT_MONGO_URI").expect("Missing REVOLT_MONGO_URI environment variable."); pub static ref REDIS_URI: String = env::var("REVOLT_REDIS_URI").expect("Missing REVOLT_REDIS_URI environment variable."); pub static ref WS_HOST: String = @@ -76,7 +74,6 @@ lazy_static! { pub fn preflight_checks() { format!("url = {}", *APP_URL); - format!("mongo = {}", *MONGO_URI); format!("public = {}", *PUBLIC_URL); format!("external = {}", *EXTERNAL_WS_URL); diff --git a/src/version.rs b/src/version.rs index 842167f1..b6b78bba 100644 --- a/src/version.rs +++ b/src/version.rs @@ -1 +1 @@ -pub const VERSION: &str = "0.5.3-alpha.16"; +pub const VERSION: &str = "0.5.3";