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

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