refactor: move ratelimits to a generic system for all web servers

Cargo.lock

@@ -962,6 +962,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
 dependencies = [
  "async-trait",
  "axum-core",
+ "axum-macros",
  "bytes 1.10.1",
  "futures-util",
  "http 1.3.1",
@@ -6259,6 +6260,7 @@ dependencies = [
  "revolt-config",
  "revolt-database",
  "revolt-files",
+ "revolt-ratelimits",
  "revolt-result",
  "revolt_clamav-client",
  "serde",
@@ -6415,6 +6417,7 @@ dependencies = [
  "revolt-models",
  "revolt-permissions",
  "revolt-presence",
+ "revolt-ratelimits",
  "revolt-result",
  "revolt_rocket_okapi",
  "rocket",
@@ -6563,6 +6566,22 @@ dependencies = [
  "web-push",
 ]
 
+[[package]]
+name = "revolt-ratelimits"
+version = "0.8.8"
+dependencies = [
+ "async-trait",
+ "authifier",
+ "axum",
+ "dashmap",
+ "log",
+ "revolt-database",
+ "revolt-result",
+ "revolt_rocket_okapi",
+ "rocket",
+ "serde",
+]
+
 [[package]]
 name = "revolt-result"
 version = "0.8.8"

crates/core/database/src/models/users/axum.rs

@@ -1,14 +1,21 @@
-use axum::{extract::FromRequestParts, http::request::Parts};
+use axum::{
+    extract::{FromRef, FromRequestParts},
+    http::request::Parts,
+};
 
 use revolt_result::{create_error, Error, Result};
 
 use crate::{Database, User};
 
 #[async_trait::async_trait]
-impl FromRequestParts<Database> for User {
+impl<S: Send + Sync> FromRequestParts<S> for User
+where
+    Database: FromRef<S>,
+{
     type Rejection = Error;
 
-    async fn from_request_parts(parts: &mut Parts, db: &Database) -> Result<User> {
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<User> {
+        let db = Database::from_ref(state);
         if let Some(Ok(bot_token)) = parts.headers.get("x-bot-token").map(|v| v.to_str()) {
             let bot = db.fetch_bot_by_token(bot_token).await?;
             db.fetch_user(&bot.id).await

crates/core/ratelimits/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "revolt-ratelimits"
+version = "0.8.8"
+edition = "2024"
+
+[features]
+rocket = ["dep:rocket", "dep:revolt_rocket_okapi", "revolt-database/rocket-impl"]
+axum = ["dep:axum", "revolt-database/axum-impl"]
+
+default = ["rocket", "axum"]
+
+[dependencies]
+revolt-database = { version = "0.8.8", path = "../database"}
+revolt-result = { version = "0.8.8", path = "../result" }
+
+rocket = { version = "0.5.1", optional = true }
+revolt_rocket_okapi = { version = "0.10.0", optional = true }
+
+axum = { version = "0.7.5", optional = true, features = ["macros"] }
+
+serde = { version = "1", features = ["derive"] }
+authifier = { version = "1.0.15" }
+dashmap = "5.2.0"
+async-trait = "0.1.81"
+log = "0.4"

crates/core/ratelimits/src/axum.rs

@@ -0,0 +1,193 @@
+use std::net::SocketAddr;
+
+use async_trait::async_trait;
+use axum::{
+    Json, RequestPartsExt, Router,
+    body::Body,
+    extract::{ConnectInfo, FromRef, FromRequestParts, State},
+    http::{HeaderValue, Request, StatusCode, request::Parts},
+    middleware::Next,
+    response::{IntoResponse, Response},
+    routing::get,
+};
+use revolt_database::{Database, User};
+
+use crate::ratelimiter::{RatelimitInformation, Ratelimiter, RequestKind};
+
+#[derive(Clone, Copy)]
+pub struct AxumRequestKind;
+
+impl RequestKind for AxumRequestKind {
+    type R<'a> = Parts;
+}
+
+pub type RatelimitStorage = crate::ratelimiter::RatelimitStorage<AxumRequestKind>;
+
+fn to_ip(parts: &Parts) -> String {
+    parts
+        .extensions
+        .get::<ConnectInfo<SocketAddr>>()
+        .map(|info| info.ip().to_string())
+        .unwrap_or_default()
+}
+
+fn to_real_ip(parts: &Parts) -> String {
+    if let Ok(true) = std::env::var("TRUST_CLOUDFLARE").map(|x| x == "1") {
+        parts
+            .headers
+            .get("CF-Connecting-IP")
+            .map(|x| x.to_str().unwrap().to_string())
+            .unwrap_or_else(|| to_ip(parts))
+    } else {
+        to_ip(parts)
+    }
+}
+
+#[async_trait]
+impl<S: Send + Sync> FromRequestParts<S> for Ratelimiter
+where
+    Database: FromRef<S>,
+    RatelimitStorage: FromRef<S>,
+{
+    type Rejection = Json<Ratelimiter>;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        if parts
+            .extensions
+            .get::<Result<Ratelimiter, Json<Ratelimiter>>>()
+            .is_none()
+        {
+            let storage = RatelimitStorage::from_ref(state);
+
+            let identifier = if let Ok(user) = parts.extract_with_state::<User, _>(state).await {
+                user.id
+            } else {
+                to_real_ip(parts)
+            };
+
+            let (bucket, resource) = storage.resolver.resolve_bucket(parts);
+            let limit = storage.resolver.resolve_bucket_limit(bucket);
+
+            let ratelimiter =
+                Ratelimiter::from(&storage.map, &identifier, limit, (bucket, resource));
+
+            parts.extensions.insert(ratelimiter.map_err(Json));
+        };
+
+        *parts
+            .extensions
+            .get::<Result<Ratelimiter, Json<Ratelimiter>>>()
+            .unwrap()
+    }
+}
+
+#[async_trait]
+impl<S: Send + Sync> FromRequestParts<S> for RatelimitInformation
+where
+    Database: FromRef<S>,
+    RatelimitStorage: FromRef<S>,
+{
+    type Rejection = Json<RatelimitInformation>;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        if parts
+            .extensions
+            .get::<Result<Ratelimiter, Json<Ratelimiter>>>()
+            .is_none()
+        {
+            let ratelimiter = parts.extract_with_state::<Ratelimiter, S>(state).await;
+
+            parts.extensions.insert(ratelimiter);
+        };
+
+        let ratelimiter = *parts
+            .extensions
+            .get::<Result<Ratelimiter, Json<Ratelimiter>>>()
+            .unwrap();
+
+        match ratelimiter {
+            Ok(ratelimter) => Ok(RatelimitInformation::Success(ratelimter)),
+            Err(ratelimiter) => Err(Json(RatelimitInformation::Failure {
+                retry_after: ratelimiter.reset,
+            })),
+        }
+    }
+}
+
+pub async fn ratelimit_middleware(
+    State(database): State<Database>,
+    State(ratelimit_storage): State<RatelimitStorage>,
+    request: Request<Body>,
+    next: Next,
+) -> Response {
+    #[derive(axum::extract::FromRef)]
+    struct TempState {
+        database: Database,
+        ratelimit_storage: RatelimitStorage,
+    }
+
+    let state = TempState {
+        database,
+        ratelimit_storage,
+    };
+
+    let (mut parts, body) = request.into_parts();
+
+    let res = Ratelimiter::from_request_parts(&mut parts, &state).await;
+
+    let (Ok(ratelimiter) | Err(Json(ratelimiter))) = &res;
+
+    let mut response = if res.is_ok() {
+        let request = Request::from_parts(parts, body);
+
+        next.run(request).await
+    } else {
+        let ratelimit_info = RatelimitInformation::from_request_parts(&mut parts, &state).await;
+
+        ratelimit_info.map(Json).into_response()
+    };
+
+    let Ratelimiter {
+        key,
+        limit,
+        remaining,
+        reset,
+    } = ratelimiter;
+
+    let headers = response.headers_mut();
+
+    headers.insert(
+        "X-RateLimit-Limit",
+        HeaderValue::from_str(&limit.to_string()).unwrap(),
+    );
+    headers.insert(
+        "X-RateLimit-Bucket",
+        HeaderValue::from_str(&key.to_string()).unwrap(),
+    );
+    headers.insert(
+        "X-RateLimit-Remaining",
+        HeaderValue::from_str(&remaining.to_string()).unwrap(),
+    );
+    headers.insert(
+        "X-RateLimit-Reset-After",
+        HeaderValue::from_str(&reset.to_string()).unwrap(),
+    );
+
+    if res.is_err() {
+        *response.status_mut() = StatusCode::TOO_MANY_REQUESTS;
+    };
+
+    response
+}
+
+async fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
+    Json(info)
+}
+
+pub fn routes<S: Clone + Send + Sync + 'static>() -> Router<S>
+where
+    Database: FromRef<S>,
+    RatelimitStorage: FromRef<S>,
+{
+    Router::new().route("/ratelimit", get(ratelimit_info))
+}

crates/core/ratelimits/src/lib.rs

@@ -0,0 +1,7 @@
+pub mod ratelimiter;
+
+#[cfg(feature = "rocket")]
+pub mod rocket;
+
+#[cfg(feature = "axum")]
+pub mod axum;

crates/core/ratelimits/src/ratelimiter.rs

@@ -0,0 +1,145 @@
+use std::collections::hash_map::DefaultHasher;
+use std::hash::Hasher;
+use std::ops::Add;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use serde::Serialize;
+
+use dashmap::DashMap;
+
+pub trait RequestKind {
+    type R<'a>;
+}
+
+pub trait RatelimitResolver<R>: Send + Sync {
+    fn resolve_bucket<'a>(&self, request: &'a R) -> (&'a str, Option<&'a str>);
+    fn resolve_bucket_limit(&self, bucket: &str) -> u32;
+}
+
+#[derive(Clone)]
+pub struct RatelimitStorage<K: RequestKind> {
+    pub resolver: Arc<dyn for<'a> RatelimitResolver<K::R<'a>>>,
+    pub map: Arc<DashMap<u64, Entry>>,
+}
+
+impl<K: RequestKind> RatelimitStorage<K> {
+    pub fn new<R: for<'a> RatelimitResolver<K::R<'a>> + 'static>(resolver: R) -> Self {
+        Self {
+            resolver: Arc::new(resolver),
+            map: Arc::new(DashMap::new()),
+        }
+    }
+}
+
+/// Ratelimit Bucket
+#[derive(Clone, Copy, Debug)]
+pub struct Entry {
+    used: u32,
+    reset: u128,
+}
+
+/// 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(map: &DashMap<u64, Entry>, 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) {
+        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;
+        }
+    }
+
+    /// Save information
+    pub fn save(self, map: &DashMap<u64, Entry>, key: u64) {
+        map.insert(key, self);
+    }
+
+    /// Get remaining units in the bucket
+    pub fn get_remaining(&self, limit: u32) -> u32 {
+        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();
+        self.reset.saturating_sub(current_time)
+    }
+}
+
+/// Ratelimit Guard
+#[derive(Serialize, Clone, Copy, Debug)]
+#[allow(dead_code)]
+pub struct Ratelimiter {
+    pub key: u64,
+    pub limit: u32,
+    pub remaining: u32,
+    pub reset: u128,
+}
+
+impl Ratelimiter {
+    /// Generate guard from identifier and target bucket
+    pub fn from(
+        map: &DashMap<u64, Entry>,
+        identifier: &str,
+        limit: u32,
+        (bucket, resource): (&str, Option<&str>),
+    ) -> Result<Ratelimiter, Ratelimiter> {
+        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 mut entry = Entry::from(map, key);
+
+        let remaining = entry.get_remaining(limit);
+        let reset = entry.left_until_reset();
+        let mut ratelimiter = Ratelimiter {
+            key,
+            limit,
+            remaining,
+            reset,
+        };
+        if remaining == 0 {
+            return Err(ratelimiter);
+        }
+
+        entry.deduct();
+        entry.save(map, key);
+        ratelimiter.remaining -= 1;
+        ratelimiter.reset = entry.left_until_reset();
+
+        Ok(ratelimiter)
+    }
+}
+
+#[derive(Serialize)]
+#[serde(untagged)]
+pub enum RatelimitInformation {
+    Success(Ratelimiter),
+    Failure { retry_after: u128 },
+}

crates/core/ratelimits/src/rocket.rs

@@ -0,0 +1,162 @@
+use async_trait::async_trait;
+use log::info;
+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, State};
+
+use revolt_rocket_okapi::r#gen::OpenApiGenerator;
+use revolt_rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
+
+use authifier::models::Session;
+
+use crate::ratelimiter::RequestKind;
+use crate::ratelimiter::{RatelimitInformation, Ratelimiter};
+
+#[derive(Clone, Copy)]
+pub struct RocketRequestKind;
+
+impl RequestKind for RocketRequestKind {
+    type R<'a> = Request<'a>;
+}
+
+pub type RatelimitStorage = crate::ratelimiter::RatelimitStorage<RocketRequestKind>;
+
+/// Find the remote IP of the client
+fn to_ip(request: &'_ rocket::Request<'_>) -> String {
+    request
+        .remote()
+        .map(|x| x.ip().to_string())
+        .unwrap_or_default()
+}
+
+/// Find the actual IP of the client
+fn to_real_ip(request: &'_ rocket::Request<'_>) -> String {
+    if let Ok(true) = std::env::var("TRUST_CLOUDFLARE").map(|x| x == "1") {
+        request
+            .headers()
+            .get_one("CF-Connecting-IP")
+            .map(|x| x.to_string())
+            .unwrap_or_else(|| to_ip(request))
+    } else {
+        to_ip(request)
+    }
+}
+
+#[async_trait]
+impl<'r> FromRequest<'r> for Ratelimiter {
+    type Error = Ratelimiter;
+
+    async fn from_request<'a>(request: &'r rocket::Request<'a>) -> Outcome<Self, Self::Error> {
+        let ratelimiter = request
+            .local_cache_async(async {
+                use rocket::outcome::Outcome;
+
+                let storage = request.guard::<&State<RatelimitStorage>>().await.unwrap();
+
+                let identifier = if let Outcome::Success(session) = request.guard::<Session>().await
+                {
+                    session.id
+                } else {
+                    to_real_ip(request)
+                };
+
+                let (bucket, resource) = storage.resolver.resolve_bucket(request);
+                let limit = storage.resolver.resolve_bucket_limit(bucket);
+
+                Ratelimiter::from(&storage.map, &identifier, limit, (bucket, resource))
+            })
+            .await;
+
+        match ratelimiter {
+            Ok(ratelimiter) => Outcome::Success(*ratelimiter),
+            Err(ratelimiter) => Outcome::Error((Status::TooManyRequests, *ratelimiter)),
+        }
+    }
+}
+
+impl OpenApiFromRequest<'_> for Ratelimiter {
+    fn from_request_input(
+        _gen: &mut OpenApiGenerator,
+        _name: String,
+        _required: bool,
+    ) -> revolt_rocket_okapi::Result<RequestHeaderInput> {
+        Ok(RequestHeaderInput::None)
+    }
+}
+
+/// Attach ratelimiter to the Rocket application
+pub struct RatelimitFairing;
+
+#[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::Error(_) = request.guard::<Ratelimiter>().await {
+            info!(
+                "User rate-limited on route {}! (IP = {:?})",
+                request.uri(),
+                to_real_ip(request)
+            );
+
+            request.set_method(Method::Get);
+            request.set_uri(Origin::parse("/ratelimit").unwrap())
+        }
+    }
+
+    async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
+        let guard = request.guard::<Ratelimiter>().await;
+        let (Outcome::Success(ratelimiter) | Outcome::Error((_, ratelimiter))) = guard else {
+            unreachable!()
+        };
+        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());
+
+        if guard.is_error() {
+            response.set_status(Status::TooManyRequests);
+        }
+    }
+}
+
+#[async_trait]
+impl<'r> FromRequest<'r> for RatelimitInformation {
+    type Error = u128;
+
+    async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
+        let info = match request.guard::<Ratelimiter>().await {
+            Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter),
+            Outcome::Error((_, ratelimiter)) => RatelimitInformation::Failure {
+                retry_after: ratelimiter.reset,
+            },
+            _ => unreachable!(),
+        };
+        Outcome::Success(info)
+    }
+}
+
+#[rocket::get("/ratelimit")]
+fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
+    Json(info)
+}
+
+pub fn routes() -> Vec<rocket::Route> {
+    rocket::routes![ratelimit_info]
+}

crates/delta/Cargo.toml

@@ -81,6 +81,7 @@ revolt-models = { path = "../core/models", features = [
 revolt-presence = { path = "../core/presence" }
 revolt-result = { path = "../core/result", features = ["rocket", "okapi"] }
 revolt-permissions = { path = "../core/permissions", features = ["schemas"] }
+revolt-ratelimits = { path = "../core/ratelimits", features = ["rocket"] }
 
 [build-dependencies]
 vergen = "7.5.0"

crates/delta/src/main.rs

@@ -11,6 +11,7 @@ pub mod util;
 use revolt_config::config;
 use revolt_database::events::client::EventV1;
 use revolt_database::AMQP;
+use revolt_ratelimits::rocket as ratelimiter;
 use rocket::{Build, Rocket};
 use rocket_cors::{AllowedOrigins, CorsOptions};
 use rocket_prometheus::PrometheusMetrics;
@@ -122,18 +123,22 @@ pub async fn web() -> Rocket<Build> {
     let rocket = rocket::build();
     let prometheus = PrometheusMetrics::new();
 
+    // Ratelimits
+    let ratelimits = ratelimiter::RatelimitStorage::new(util::ratelimits::DeltaRatelimits);
+
     routes::mount(config, rocket)
         .attach(prometheus.clone())
         .mount("/metrics", prometheus)
         .mount("/", rocket_cors::catch_all_options_routes())
-        .mount("/", util::ratelimiter::routes())
+        .mount("/", ratelimiter::routes())
         .mount("/swagger/", swagger)
         .mount("/0.8/swagger/", swagger_0_8)
         .manage(authifier)
         .manage(db)
         .manage(amqp)
         .manage(cors.clone())
-        .attach(util::ratelimiter::RatelimitFairing)
+        .manage(ratelimits)
+        .attach(ratelimiter::RatelimitFairing)
         .attach(cors)
         .configure(rocket::Config {
             limits: rocket::data::Limits::default().limit("string", 5.megabytes()),

crates/delta/src/util/mod.rs

@@ -1,2 +1,2 @@
-pub mod ratelimiter;
+pub mod ratelimits;
 pub mod test;

crates/delta/src/util/ratelimiter.rs

@@ -1,347 +0,0 @@
-use std::collections::hash_map::DefaultHasher;
-use std::hash::Hasher;
-use std::ops::Add;
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
-
-use authifier::models::Session;
-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 revolt_rocket_okapi::gen::OpenApiGenerator;
-use revolt_rocket_okapi::request::{OpenApiFromRequest, RequestHeaderInput};
-
-use serde::Serialize;
-
-use dashmap::DashMap;
-use once_cell::sync::Lazy;
-
-/// Ratelimit Bucket
-#[derive(Clone, Copy, Debug)]
-struct Entry {
-    used: u8,
-    reset: u128,
-}
-
-static MAP: Lazy<DashMap<u64, Entry>> = Lazy::new(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) {
-        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;
-        }
-    }
-
-    /// Save information
-    pub fn save(self, key: u64) {
-        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>) {
-    let (segment, resource, extra) = if matches!(request.routed_segment(0), Some("0.8")) {
-        (
-            request.routed_segment(1),
-            request.routed_segment(2),
-            request.routed_segment(3),
-        )
-    } else {
-        (
-            request.routed_segment(0),
-            request.routed_segment(1),
-            request.routed_segment(2),
-        )
-    };
-
-    if let Some(segment) = segment {
-        #[allow(clippy::redundant_locals)]
-        let resource = resource;
-
-        let method = request.method();
-        match (segment, resource, method) {
-            ("users", target, Method::Patch) => ("user_edit", target),
-            ("users", _, _) => {
-                if let Some("default_avatar") = extra {
-                    return ("default_avatar", None);
-                }
-
-                ("users", None)
-            }
-            ("bots", _, _) => ("bots", None),
-            ("channels", Some(id), _) => {
-                if request.method() == Method::Post {
-                    if let Some("messages") = extra {
-                        return ("messaging", Some(id));
-                    }
-                }
-
-                ("channels", Some(id))
-            }
-            ("servers", Some(id), _) => ("servers", Some(id)),
-            ("auth", _, _) => {
-                if request.method() == Method::Delete {
-                    ("auth_delete", None)
-                } else {
-                    ("auth", None)
-                }
-            }
-            ("swagger", _, _) => ("swagger", None),
-            ("safety", Some("report"), _) => ("safety_report", Some("report")),
-            ("safety", _, _) => ("safety", None),
-            _ => ("any", None),
-        }
-    } else {
-        ("any", None)
-    }
-}
-
-/// Resolve per-bucket limits
-fn resolve_bucket_limit(bucket: &str) -> u8 {
-    match bucket {
-        "user_edit" => 2,
-        "users" => 20,
-        "bots" => 10,
-        "messaging" => 10,
-        "channels" => 15,
-        "servers" => 5,
-        "auth" => 15,
-        "auth_delete" => 255,
-        "default_avatar" => 255,
-        "swagger" => 100,
-        "safety" => 15,
-        "safety_report" => 3,
-        _ => 20,
-    }
-}
-
-/// Find the remote IP of the client
-fn to_ip(request: &'_ rocket::Request<'_>) -> String {
-    request
-        .remote()
-        .map(|x| x.ip().to_string())
-        .unwrap_or_default()
-}
-
-/// Find the actual IP of the client
-fn to_real_ip(request: &'_ rocket::Request<'_>) -> String {
-    if let Ok(true) = std::env::var("TRUST_CLOUDFLARE").map(|x| x == "1") {
-        request
-            .headers()
-            .get_one("CF-Connecting-IP")
-            .map(|x| x.to_string())
-            .unwrap_or_else(|| to_ip(request))
-    } else {
-        to_ip(request)
-    }
-}
-
-impl Ratelimiter {
-    /// Generate guard from identifier and target bucket
-    pub fn from(
-        identifier: &str,
-        (bucket, resource): (&str, Option<&str>),
-    ) -> Result<Ratelimiter, Ratelimiter> {
-        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 mut entry = Entry::from(key);
-
-        let remaining = entry.get_remaining(limit);
-        let reset = entry.left_until_reset();
-        let mut ratelimiter = Ratelimiter {
-            key,
-            limit,
-            remaining,
-            reset,
-        };
-        if remaining == 0 {
-            return Err(ratelimiter);
-        }
-
-        entry.deduct();
-        entry.save(key);
-        ratelimiter.remaining -= 1;
-        ratelimiter.reset = entry.left_until_reset();
-
-        Ok(ratelimiter)
-    }
-}
-
-#[async_trait]
-impl<'r> FromRequest<'r> for Ratelimiter {
-    type Error = Ratelimiter;
-
-    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::<Session>().await
-                {
-                    session.id
-                } else {
-                    to_real_ip(request)
-                };
-
-                Ratelimiter::from(&identifier, resolve_bucket(request))
-            })
-            .await;
-
-        match ratelimiter {
-            Ok(ratelimiter) => Outcome::Success(*ratelimiter),
-            Err(ratelimiter) => Outcome::Error((Status::TooManyRequests, *ratelimiter)),
-        }
-    }
-}
-
-impl OpenApiFromRequest<'_> for Ratelimiter {
-    fn from_request_input(
-        _gen: &mut OpenApiGenerator,
-        _name: String,
-        _required: bool,
-    ) -> revolt_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::Error(_) = request.guard::<Ratelimiter>().await {
-            info!(
-                "User rate-limited on route {}! (IP = {:?})",
-                request.uri(),
-                to_real_ip(request)
-            );
-
-            request.set_method(Method::Get);
-            request.set_uri(Origin::parse("/ratelimit").unwrap())
-        }
-    }
-
-    async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
-        let guard = request.guard::<Ratelimiter>().await;
-        let (Outcome::Success(ratelimiter) | Outcome::Error((_, ratelimiter))) = guard else {
-            unreachable!()
-        };
-        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());
-
-        if guard.is_error() {
-            response.set_status(Status::TooManyRequests);
-        }
-    }
-}
-
-#[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> {
-        let info = match request.guard::<Ratelimiter>().await {
-            Outcome::Success(ratelimiter) => RatelimitInformation::Success(ratelimiter),
-            Outcome::Error((_, ratelimiter)) => RatelimitInformation::Failure {
-                retry_after: ratelimiter.reset,
-            },
-            _ => unreachable!(),
-        };
-        Outcome::Success(info)
-    }
-}
-
-#[rocket::get("/ratelimit")]
-fn ratelimit_info(info: RatelimitInformation) -> Json<RatelimitInformation> {
-    Json(info)
-}
-
-pub fn routes() -> Vec<rocket::Route> {
-    rocket::routes![ratelimit_info]
-}

crates/delta/src/util/ratelimits.rs

@@ -0,0 +1,81 @@
+use revolt_ratelimits::ratelimiter::RatelimitResolver;
+use rocket::{http::Method, Request};
+
+pub struct DeltaRatelimits;
+
+impl<'a> RatelimitResolver<Request<'a>> for DeltaRatelimits {
+    fn resolve_bucket<'r>(&self, request: &'r Request<'_>) -> (&'r str, Option<&'r str>) {
+        let (segment, resource, extra) = if request.routed_segment(0) == Some("0.8") {
+            (
+                request.routed_segment(1),
+                request.routed_segment(2),
+                request.routed_segment(3),
+            )
+        } else {
+            (
+                request.routed_segment(0),
+                request.routed_segment(1),
+                request.routed_segment(2),
+            )
+        };
+
+        if let Some(segment) = segment {
+            #[allow(clippy::redundant_locals)]
+            let resource = resource;
+
+            let method = request.method();
+            match (segment, resource, method) {
+                ("users", target, Method::Patch) => ("user_edit", target),
+                ("users", _, _) => {
+                    if let Some("default_avatar") = extra {
+                        return ("default_avatar", None);
+                    }
+
+                    ("users", None)
+                }
+                ("bots", _, _) => ("bots", None),
+                ("channels", Some(id), _) => {
+                    if request.method() == Method::Post {
+                        if let Some("messages") = extra {
+                            return ("messaging", Some(id));
+                        }
+                    }
+
+                    ("channels", Some(id))
+                }
+                ("servers", Some(id), _) => ("servers", Some(id)),
+                ("auth", _, _) => {
+                    if request.method() == Method::Delete {
+                        ("auth_delete", None)
+                    } else {
+                        ("auth", None)
+                    }
+                }
+                ("swagger", _, _) => ("swagger", None),
+                ("safety", Some("report"), _) => ("safety_report", Some("report")),
+                ("safety", _, _) => ("safety", None),
+                _ => ("any", None),
+            }
+        } else {
+            ("any", None)
+        }
+    }
+
+    fn resolve_bucket_limit(&self, bucket: &str) -> u32 {
+        match bucket {
+            "user_edit" => 2,
+            "users" => 20,
+            "bots" => 10,
+            "messaging" => 10,
+            "channels" => 15,
+            "servers" => 5,
+            "auth" => 15,
+            "auth_delete" => 255,
+            "default_avatar" => 255,
+            "swagger" => 100,
+            "safety" => 15,
+            "safety_report" => 3,
+            _ => 20,
+        }
+    }
+}

crates/services/autumn/Cargo.toml

@@ -52,6 +52,7 @@ revolt-result = { version = "0.8.8", path = "../../core/result", features = [
     "utoipa",
     "axum",
 ] }
+revolt-ratelimits = { version = "0.8.8", path = "../../core/ratelimits", features = ["axum"] }
 
 # Axum / web server
 tempfile = "3.12.0"

crates/services/autumn/src/api.rs

@@ -25,10 +25,10 @@ use tokio::time::Instant;
 use tower_http::cors::{AllowHeaders, Any, CorsLayer};
 use utoipa::ToSchema;
 
-use crate::{exif::strip_metadata, metadata::generate_metadata, mime_type::determine_mime_type};
+use crate::{exif::strip_metadata, metadata::generate_metadata, mime_type::determine_mime_type, AppState};
 
 /// Build the API router
-pub async fn router() -> Router<Database> {
+pub async fn router() -> Router<AppState> {
     let config = config().await;
 
     let cors = CorsLayer::new()

crates/services/autumn/src/main.rs

@@ -1,8 +1,10 @@
 use std::net::{Ipv4Addr, SocketAddr};
 
-use axum::Router;
+use axum::{middleware::from_fn_with_state, Router};
 
-use revolt_database::DatabaseInfo;
+use axum_macros::FromRef;
+use revolt_database::{Database, DatabaseInfo};
+use revolt_ratelimits::axum as ratelimiter;
 use tokio::net::TcpListener;
 use utoipa::{
     openapi::security::{ApiKey, ApiKeyValue, SecurityScheme},
@@ -15,6 +17,13 @@ pub mod clamav;
 pub mod exif;
 pub mod metadata;
 pub mod mime_type;
+mod ratelimits;
+
+#[derive(FromRef, Clone)]
+struct AppState {
+    database: Database,
+    ratelimit_storage: ratelimiter::RatelimitStorage,
+}
 
 #[tokio::main]
 async fn main() -> Result<(), std::io::Error> {
@@ -69,12 +78,23 @@ async fn main() -> Result<(), std::io::Error> {
 
     // Connect to the database
     let db = DatabaseInfo::Auto.connect().await.unwrap();
+    let ratelimits = ratelimiter::RatelimitStorage::new(ratelimits::AutumnRatelimits);
+
+    let state = AppState {
+        database: db,
+        ratelimit_storage: ratelimits,
+    };
 
     // Configure Axum and router
     let app = Router::new()
         .merge(Scalar::with_url("/scalar", ApiDoc::openapi()))
         .nest("/", api::router().await)
-        .with_state(db);
+        .nest("/", ratelimiter::routes())
+        .layer(from_fn_with_state(
+            state.clone(),
+            ratelimiter::ratelimit_middleware,
+        ))
+        .with_state(state);
 
     // Configure TCP listener and bind
     let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, 14704));

crates/services/autumn/src/ratelimits.rs

@@ -0,0 +1,28 @@
+use axum::http::{request::Parts, Method};
+use revolt_ratelimits::ratelimiter::RatelimitResolver;
+
+pub struct AutumnRatelimits;
+
+impl RatelimitResolver<Parts> for AutumnRatelimits {
+    fn resolve_bucket<'a>(&self, parts: &'a Parts) -> (&'a str, Option<&'a str>) {
+        let path = parts
+            .uri
+            .path()
+            .trim_matches('/')
+            .split_terminator("/")
+            .collect::<Vec<&str>>();
+
+        match (&parts.method, path.as_slice()) {
+            (&Method::POST, &[tag]) => ("upload", Some(tag)),
+            _ => ("any", None),
+        }
+    }
+
+    fn resolve_bucket_limit(&self, bucket: &str) -> u32 {
+        match bucket {
+            "upload" => 10,
+            "any" => u32::MAX,
+            _ => unreachable!("Bucket defined but no limit set"),
+        }
+    }
+}