diff --git a/Cargo.lock b/Cargo.lock index 4d310e03..a4fafabf 100644 --- a/Cargo.lock +++ b/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", @@ -6414,6 +6416,7 @@ dependencies = [ "revolt-models", "revolt-permissions", "revolt-presence", + "revolt-ratelimits", "revolt-result", "revolt_rocket_okapi", "rocket", @@ -6562,6 +6565,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" diff --git a/crates/core/database/src/models/users/axum.rs b/crates/core/database/src/models/users/axum.rs index da83c5b1..a6d8b5a3 100644 --- a/crates/core/database/src/models/users/axum.rs +++ b/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 for User { +impl FromRequestParts for User +where + Database: FromRef, +{ type Rejection = Error; - async fn from_request_parts(parts: &mut Parts, db: &Database) -> Result { + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + 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 diff --git a/crates/core/ratelimits/Cargo.toml b/crates/core/ratelimits/Cargo.toml new file mode 100644 index 00000000..49110e1e --- /dev/null +++ b/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" diff --git a/crates/core/ratelimits/src/axum.rs b/crates/core/ratelimits/src/axum.rs new file mode 100644 index 00000000..39b0c3bb --- /dev/null +++ b/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; + +fn to_ip(parts: &Parts) -> String { + parts + .extensions + .get::>() + .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 FromRequestParts for Ratelimiter +where + Database: FromRef, + RatelimitStorage: FromRef, +{ + type Rejection = Json; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + if parts + .extensions + .get::>>() + .is_none() + { + let storage = RatelimitStorage::from_ref(state); + + let identifier = if let Ok(user) = parts.extract_with_state::(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::>>() + .unwrap() + } +} + +#[async_trait] +impl FromRequestParts for RatelimitInformation +where + Database: FromRef, + RatelimitStorage: FromRef, +{ + type Rejection = Json; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + if parts + .extensions + .get::>>() + .is_none() + { + let ratelimiter = parts.extract_with_state::(state).await; + + parts.extensions.insert(ratelimiter); + }; + + let ratelimiter = *parts + .extensions + .get::>>() + .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, + State(ratelimit_storage): State, + request: Request, + 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 { + Json(info) +} + +pub fn routes() -> Router +where + Database: FromRef, + RatelimitStorage: FromRef, +{ + Router::new().route("/ratelimit", get(ratelimit_info)) +} diff --git a/crates/core/ratelimits/src/lib.rs b/crates/core/ratelimits/src/lib.rs new file mode 100644 index 00000000..5eaa8108 --- /dev/null +++ b/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; diff --git a/crates/core/ratelimits/src/ratelimiter.rs b/crates/core/ratelimits/src/ratelimiter.rs new file mode 100644 index 00000000..a232feb3 --- /dev/null +++ b/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: 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 { + pub resolver: Arc RatelimitResolver>>, + pub map: Arc>, +} + +impl RatelimitStorage { + pub fn new RatelimitResolver> + '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, 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, 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, + identifier: &str, + limit: u32, + (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 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 }, +} diff --git a/crates/core/ratelimits/src/rocket.rs b/crates/core/ratelimits/src/rocket.rs new file mode 100644 index 00000000..6e7a0c9d --- /dev/null +++ b/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; + +/// 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 { + let ratelimiter = request + .local_cache_async(async { + use rocket::outcome::Outcome; + + let storage = request.guard::<&State>().await.unwrap(); + + let identifier = if let Outcome::Success(session) = request.guard::().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 { + 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::().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::().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 { + let info = match request.guard::().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 { + Json(info) +} + +pub fn routes() -> Vec { + rocket::routes![ratelimit_info] +} diff --git a/crates/delta/Cargo.toml b/crates/delta/Cargo.toml index d07b44bb..889b1167 100644 --- a/crates/delta/Cargo.toml +++ b/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" diff --git a/crates/delta/src/main.rs b/crates/delta/src/main.rs index 9cce39ce..6d0eaa60 100644 --- a/crates/delta/src/main.rs +++ b/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 { 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()), diff --git a/crates/delta/src/util/mod.rs b/crates/delta/src/util/mod.rs index 6ee3a7f6..61e6d0da 100644 --- a/crates/delta/src/util/mod.rs +++ b/crates/delta/src/util/mod.rs @@ -1,2 +1,2 @@ -pub mod ratelimiter; +pub mod ratelimits; pub mod test; diff --git a/crates/delta/src/util/ratelimiter.rs b/crates/delta/src/util/ratelimiter.rs deleted file mode 100644 index 7d00e96d..00000000 --- a/crates/delta/src/util/ratelimiter.rs +++ /dev/null @@ -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> = 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 { - 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 { - let ratelimiter = request - .local_cache_async(async { - use rocket::outcome::Outcome; - let identifier = if let Outcome::Success(session) = request.guard::().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 { - 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::().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::().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 { - let info = match request.guard::().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 { - Json(info) -} - -pub fn routes() -> Vec { - rocket::routes![ratelimit_info] -} diff --git a/crates/delta/src/util/ratelimits.rs b/crates/delta/src/util/ratelimits.rs new file mode 100644 index 00000000..178d15d2 --- /dev/null +++ b/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> 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, + } + } +} diff --git a/crates/services/autumn/Cargo.toml b/crates/services/autumn/Cargo.toml index 464da056..3514f501 100644 --- a/crates/services/autumn/Cargo.toml +++ b/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" diff --git a/crates/services/autumn/src/api.rs b/crates/services/autumn/src/api.rs index d068c2dc..5c567146 100644 --- a/crates/services/autumn/src/api.rs +++ b/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 { +pub async fn router() -> Router { let config = config().await; let cors = CorsLayer::new() diff --git a/crates/services/autumn/src/main.rs b/crates/services/autumn/src/main.rs index ff29b2e3..cb936e0c 100644 --- a/crates/services/autumn/src/main.rs +++ b/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)); diff --git a/crates/services/autumn/src/ratelimits.rs b/crates/services/autumn/src/ratelimits.rs new file mode 100644 index 00000000..609d8e1e --- /dev/null +++ b/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 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::>(); + + 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"), + } + } +}