diff --git a/Cargo.lock b/Cargo.lock index 1274b75a..3270fbf3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6519,6 +6519,7 @@ dependencies = [ "revolt-config", "revolt-database", "revolt-models", + "revolt-ratelimits", "revolt-result", "serde", "serde_json", diff --git a/crates/services/gifbox/Cargo.toml b/crates/services/gifbox/Cargo.toml index f67a04e8..b640766e 100644 --- a/crates/services/gifbox/Cargo.toml +++ b/crates/services/gifbox/Cargo.toml @@ -28,6 +28,9 @@ revolt-coalesced = { version = "0.8.8", path = "../../core/coalesced", features revolt-database = { version = "0.8.8", path = "../../core/database", features = [ "axum-impl", ] } +revolt-ratelimits = { version = "0.8.8", path = "../../core/ratelimits", features = [ + "axum", +] } # Axum / web server axum = { version = "0.7.5" } diff --git a/crates/services/gifbox/src/main.rs b/crates/services/gifbox/src/main.rs index 86588518..059cb22f 100644 --- a/crates/services/gifbox/src/main.rs +++ b/crates/services/gifbox/src/main.rs @@ -1,9 +1,10 @@ use std::net::{Ipv4Addr, SocketAddr}; -use axum::{extract::FromRef, Router}; +use axum::{extract::FromRef, middleware::from_fn_with_state, Router}; use revolt_config::config; use revolt_database::{Database, DatabaseInfo}; +use revolt_ratelimits::axum as ratelimiter; use tokio::net::TcpListener; use utoipa::{ openapi::security::{ApiKey, ApiKeyValue, SecurityScheme}, @@ -13,31 +14,21 @@ use utoipa_scalar::{Scalar, Servable as ScalarServable}; use crate::tenor::Tenor; +mod ratelimits; mod routes; mod tenor; mod types; -#[derive(Clone)] +#[derive(Clone, FromRef)] struct AppState { pub database: Database, pub tenor: Tenor, + pub ratelimit_storage: ratelimiter::RatelimitStorage, } -impl FromRef for Database { - fn from_ref(state: &AppState) -> Self { - state.database.clone() - } -} +struct SecurityAddon; -impl FromRef for Tenor { - fn from_ref(state: &AppState) -> Self { - state.tenor.clone() - } -} - -struct TokenAddon; - -impl Modify for TokenAddon { +impl Modify for SecurityAddon { fn modify(&self, openapi: &mut utoipa::openapi::OpenApi) { let components = openapi.components.get_or_insert_default(); @@ -63,7 +54,7 @@ async fn main() -> Result<(), std::io::Error> { // Configure API schema #[derive(OpenApi)] #[openapi( - modifiers(&TokenAddon), + modifiers(&SecurityAddon), paths( routes::categories::categories, routes::root::root, @@ -86,18 +77,30 @@ async fn main() -> Result<(), std::io::Error> { struct ApiDoc; let config = config().await; + + let database = DatabaseInfo::Auto + .connect() + .await + .expect("Unable to connect to database"); + + let tenor = tenor::Tenor::new(&config.api.security.tenor_key); + + let ratelimit_storage = ratelimiter::RatelimitStorage::new(ratelimits::GifboxRatelimits); + let state = AppState { - database: DatabaseInfo::Auto - .connect() - .await - .expect("Unable to connect to database"), - tenor: tenor::Tenor::new(&config.api.security.tenor_key), + database, + tenor, + ratelimit_storage, }; // Configure Axum and router let app = Router::new() .merge(Scalar::with_url("/scalar", ApiDoc::openapi())) .nest("/", routes::router()) + .layer(from_fn_with_state( + state.clone(), + ratelimiter::ratelimit_middleware, + )) .with_state(state); // Configure TCP listener and bind diff --git a/crates/services/gifbox/src/ratelimits.rs b/crates/services/gifbox/src/ratelimits.rs new file mode 100644 index 00000000..10c904d4 --- /dev/null +++ b/crates/services/gifbox/src/ratelimits.rs @@ -0,0 +1,32 @@ +use axum::http::request::Parts; +use revolt_ratelimits::ratelimiter::RatelimitResolver; + +pub struct GifboxRatelimits; + +impl RatelimitResolver for GifboxRatelimits { + fn resolve_bucket<'a>(&self, parts: &'a Parts) -> (&'a str, Option<&'a str>) { + let path = parts + .uri + .path() + .trim_matches('/') + .split_terminator("/") + .next(); + + match path { + Some("categories") => ("categories", None), + Some("trending") => ("trending", None), + Some("search") => ("search", None), + _ => ("any", None), + } + } + + fn resolve_bucket_limit(&self, bucket: &str) -> u32 { + match bucket { + "categories" => 2, + "trending" => 5, + "search" => 10, + "any" => u32::MAX, + _ => unreachable!("Bucket defined but no limit set"), + } + } +}