feat: add ratelimits to gifbox

Cargo.lock

@@ -6519,6 +6519,7 @@ dependencies = [
  "revolt-config",
  "revolt-database",
  "revolt-models",
+ "revolt-ratelimits",
  "revolt-result",
  "serde",
  "serde_json",

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

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<AppState> for Database {
+struct SecurityAddon;
-    fn from_ref(state: &AppState) -> Self {
-        state.database.clone()
-    }
-}
 
-impl FromRef<AppState> for Tenor {
+impl Modify for SecurityAddon {
-    fn from_ref(state: &AppState) -> Self {
-        state.tenor.clone()
-    }
-}
-
-struct TokenAddon;
-
-impl Modify for TokenAddon {
     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
+        database,
-            .connect()
+        tenor,
-            .await
+        ratelimit_storage,
-            .expect("Unable to connect to database"),
-        tenor: tenor::Tenor::new(&config.api.security.tenor_key),
     };
 
     // 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

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<Parts> 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"),
+        }
+    }
+}