feat(api): add full API rate limits

2026-07-14 05:26:59 +00:00 · 2022-04-14 17:09:16 +01:00
parent 7ea5146e58
commit d2f0fb3414
9 changed files with 399 additions and 52 deletions
--- 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"
--- 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"
--- 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
--- 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
--- 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
--- 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;
--- a/src/util/ratelimiter.rs
+++ 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<u64, Entry> = 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<Ratelimiter, u128> {
+        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<Self, Self::Error> {
+        let ratelimiter = request
+            .local_cache_async(async {
+                use rocket::outcome::Outcome;
+                let identifier = if let Outcome::Success(session) =
+                    request.guard::<rauth::entities::Session>().await
+                {
+                    session.id.expect("`id` on User")
+                } else if let Outcome::Success(addr) = request.guard::<SocketAddr>().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<RequestHeaderInput> {
+        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::<Ratelimiter>().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::<Ratelimiter>().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<Self, Self::Error> {
+        Outcome::Success(match request.guard::<Ratelimiter>().await {
+            Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter),
+            Outcome::Failure((_, retry_after)) => RatelimitInformation::Failure { retry_after },
+            _ => unreachable!(),
+        })
+    }
+}
+
+#[get("/ratelimit")]
+fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
+    Json(info)
+}
+
+pub fn routes() -> Vec<rocket::Route> {
+    routes![ratelimit_info]
+}
--- 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);

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