feat(api): add full API rate limits

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

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;

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]
+}

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

src/version.rs

@@ -1 +1 @@
-pub const VERSION: &str = "0.5.3-alpha.16";
+pub const VERSION: &str = "0.5.3";