From ca65897d6c5d73fa5033554da83d1860c22623ae Mon Sep 17 00:00:00 2001 From: Trezy Date: Tue, 17 Mar 2026 01:31:12 +0000 Subject: [PATCH] feat: add configurable rate limiting --- Cargo.lock | 32 ++++++++++++++++++++++++++++++++ Cargo.toml | 3 +++ migrations/20260316000000_create_rate_limits.sql | 28 ++++++++++++++++++++++++++++ src/admin/mod.rs | 14 +++++++++++++- src/admin/permissions.rs | 18 +++++++++++++++++- src/admin/rate_limits.rs | 274 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/admin/types.rs | 47 +++++++++++++++++++++++++++++++++++++++++++++++ src/aip.rs | 9 +++++++++ src/error.rs | 27 +++++++++++++++++++++++++++ src/lib.rs | 4 ++++ src/lua/atproto_api.rs | 9 +++++++++ src/lua/db_api.rs | 9 +++++++++ src/lua/execute.rs | 9 +++++++++ src/lua/http_api.rs | 9 +++++++++ src/main.rs | 12 ++++++++++++ src/rate_limit.rs | 364 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/repo/upload_blob.rs | 42 +++++++++++++++++++++++++++++++++++++++++- src/server.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++-- src/xrpc/mod.rs | 138 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----- tests/common/app.rs | 9 +++++++++ tests/lua_atproto_api.rs | 9 +++++++++ tests/lua_db_api.rs | 9 +++++++++ 22 file(s) changed, 1114 insertion(s)(+), 10 deletion(s)(-) diff --git a/Cargo.lock b/Cargo.lock --- a/Cargo.lock +++ b/Cargo.lock @@ -33,6 +33,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" [[package]] +name = "arc-swap" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9f3647c145568cec02c42054e07bdf9a5a698e15b466fb2341bfc393cd24aa5" +dependencies = [ + "rustversion", +] + +[[package]] name = "assert-json-diff" version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -334,6 +343,20 @@ checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", "typenum", +] + +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", ] [[package]] @@ -780,15 +803,18 @@ [[package]] name = "happyview" version = "0.1.0" dependencies = [ + "arc-swap", "axum", "base64", "bytes", "chrono", + "dashmap", "dotenvy", "futures-util", "hex", "hickory-resolver", "http-body-util", + "ipnet", "jsonwebtoken", "mlua", "p256", @@ -809,6 +835,12 @@ "urlencoding", "uuid", "wiremock", ] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" [[package]] name = "hashbrown" diff --git a/Cargo.toml b/Cargo.toml --- a/Cargo.toml +++ b/Cargo.toml @@ -4,8 +4,10 @@ version = "0.1.0" edition = "2024" [dependencies] +arc-swap = "1" axum = "0.8" base64 = "0.22" +dashmap = "6" dotenvy = "0.15" hex = "0.4" futures-util = "0.3" @@ -26,6 +28,7 @@ tower = { version = "0.5", features = ["util"] } tower-http = { version = "0.6", features = ["cors", "fs", "trace"] } http-body-util = "0.1" hickory-resolver = "0.25" +ipnet = "2" mlua = { version = "0.11", features = ["lua54", "async", "serialize", "vendored", "send"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } diff --git a/migrations/20260316000000_create_rate_limits.sql b/migrations/20260316000000_create_rate_limits.sql new file mode 100644 --- /dev/null +++ b/migrations/20260316000000_create_rate_limits.sql @@ -0,0 +1,28 @@ +CREATE TABLE rate_limits ( + id SERIAL PRIMARY KEY, + method TEXT UNIQUE, -- NULL = global default, otherwise XRPC method NSID + capacity INTEGER NOT NULL, -- max tokens in bucket + refill_rate REAL NOT NULL, -- tokens added per second + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Seed global default: 100 token capacity, refills at 2/sec +INSERT INTO rate_limits (method, capacity, refill_rate) VALUES (NULL, 100, 2.0); + +-- Global enabled flag +CREATE TABLE rate_limit_settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +INSERT INTO rate_limit_settings (key, value) VALUES ('enabled', 'true'); + +-- IP/CIDR allowlist: exempt IPs from rate limiting +CREATE TABLE rate_limit_allowlist ( + id SERIAL PRIMARY KEY, + cidr TEXT NOT NULL UNIQUE, -- IP or CIDR, e.g. '10.0.0.0/8' or '203.0.113.5/32' + note TEXT, -- human-readable reason for the exemption + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); diff --git a/src/admin/mod.rs b/src/admin/mod.rs --- a/src/admin/mod.rs +++ b/src/admin/mod.rs @@ -6,6 +6,7 @@ mod labelers; mod lexicons; mod network_lexicons; pub(crate) mod permissions; +mod rate_limits; mod records; mod script_variables; mod stats; @@ -14,7 +15,7 @@ mod types; mod users; use axum::Router; -use axum::routing::{delete, get, patch, post}; +use axum::routing::{delete, get, patch, post, put}; use crate::AppState; @@ -67,5 +68,16 @@ .route("/labelers", post(labelers::add).get(labelers::list)) .route( "/labelers/{did}", patch(labelers::update).delete(labelers::delete), + ) + .route( + "/rate-limits", + post(rate_limits::upsert).get(rate_limits::list), + ) + .route("/rate-limits/{id}", delete(rate_limits::delete)) + .route("/rate-limits/enabled", put(rate_limits::set_enabled)) + .route("/rate-limits/allowlist", post(rate_limits::add_allowlist)) + .route( + "/rate-limits/allowlist/{id}", + delete(rate_limits::remove_allowlist), ) } diff --git a/src/admin/permissions.rs b/src/admin/permissions.rs --- a/src/admin/permissions.rs +++ b/src/admin/permissions.rs @@ -59,6 +59,13 @@ #[serde(rename = "labelers:read")] LabelersRead, #[serde(rename = "labelers:delete")] LabelersDelete, + + #[serde(rename = "rate-limits:read")] + RateLimitsRead, + #[serde(rename = "rate-limits:create")] + RateLimitsCreate, + #[serde(rename = "rate-limits:delete")] + RateLimitsDelete, } impl Permission { @@ -88,10 +95,13 @@ Self::EventsRead => "events:read", Self::LabelersCreate => "labelers:create", Self::LabelersRead => "labelers:read", Self::LabelersDelete => "labelers:delete", + Self::RateLimitsRead => "rate-limits:read", + Self::RateLimitsCreate => "rate-limits:create", + Self::RateLimitsDelete => "rate-limits:delete", } } - /// All 23 permissions. + /// All 26 permissions. pub fn all() -> HashSet { HashSet::from([ Self::LexiconsCreate, @@ -117,6 +127,9 @@ Self::EventsRead, Self::LabelersCreate, Self::LabelersRead, Self::LabelersDelete, + Self::RateLimitsRead, + Self::RateLimitsCreate, + Self::RateLimitsDelete, ]) } } @@ -161,6 +174,9 @@ perms.insert(Permission::RecordsDelete); perms.insert(Permission::LabelersCreate); perms.insert(Permission::LabelersRead); perms.insert(Permission::LabelersDelete); + perms.insert(Permission::RateLimitsRead); + perms.insert(Permission::RateLimitsCreate); + perms.insert(Permission::RateLimitsDelete); perms } Self::FullAccess => Permission::all(), diff --git a/src/admin/rate_limits.rs b/src/admin/rate_limits.rs new file mode 100644 --- /dev/null +++ b/src/admin/rate_limits.rs @@ -0,0 +1,274 @@ +use axum::Json; +use axum::extract::{Path, State}; +use axum::http::StatusCode; + +use crate::AppState; +use crate::error::AppError; +use crate::event_log::{EventLog, Severity, log_event}; + +use super::auth::UserAuth; +use super::permissions::Permission; +use super::types::{ + AddAllowlistBody, AllowlistEntry, RateLimitSummary, RateLimitsResponse, SetEnabledBody, + UpsertRateLimitBody, +}; + +/// GET /admin/rate-limits — list rate limit config. +pub(super) async fn list( + State(state): State, + auth: UserAuth, +) -> Result, AppError> { + auth.require(Permission::RateLimitsRead).await?; + + let enabled: String = + sqlx::query_scalar("SELECT value FROM rate_limit_settings WHERE key = 'enabled'") + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to read rate limit settings: {e}")))? + .unwrap_or_else(|| "true".to_string()); + + let limits: Vec = sqlx::query_as( + "SELECT id, method, capacity, refill_rate, created_at, updated_at FROM rate_limits ORDER BY id", + ) + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to list rate limits: {e}")))?; + + let allowlist: Vec = + sqlx::query_as("SELECT id, cidr, note, created_at FROM rate_limit_allowlist ORDER BY id") + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to list allowlist: {e}")))?; + + Ok(Json(RateLimitsResponse { + enabled: enabled == "true", + limits, + allowlist, + })) +} + +/// POST /admin/rate-limits — upsert a rate limit rule. +pub(super) async fn upsert( + State(state): State, + auth: UserAuth, + Json(body): Json, +) -> Result { + auth.require(Permission::RateLimitsCreate).await?; + + sqlx::query( + r#" + INSERT INTO rate_limits (method, capacity, refill_rate) + VALUES ($1, $2, $3) + ON CONFLICT (method) DO UPDATE SET + capacity = EXCLUDED.capacity, + refill_rate = EXCLUDED.refill_rate, + updated_at = NOW() + "#, + ) + .bind(&body.method) + .bind(body.capacity as i32) + .bind(body.refill_rate as f32) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to upsert rate limit: {e}")))?; + + state.rate_limiter.reload_from_db(&state.db).await; + + log_event( + &state.db, + EventLog { + event_type: "rate_limit.upserted".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: body.method.clone(), + detail: serde_json::json!({ + "capacity": body.capacity, + "refill_rate": body.refill_rate, + }), + }, + ) + .await; + + Ok(StatusCode::CREATED) +} + +/// DELETE /admin/rate-limits/{id} — delete a rate limit rule. +pub(super) async fn delete( + State(state): State, + auth: UserAuth, + Path(id): Path, +) -> Result { + auth.require(Permission::RateLimitsDelete).await?; + + // Prevent deleting the global default (method IS NULL) + let is_global: Option<(bool,)> = + sqlx::query_as("SELECT (method IS NULL) FROM rate_limits WHERE id = $1") + .bind(id) + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to check rate limit: {e}")))?; + + match is_global { + None => { + return Err(AppError::NotFound(format!( + "rate limit rule {id} not found" + ))); + } + Some((true,)) => { + return Err(AppError::BadRequest( + "cannot delete the global default rate limit".to_string(), + )); + } + Some((false,)) => {} + } + + sqlx::query("DELETE FROM rate_limits WHERE id = $1") + .bind(id) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to delete rate limit: {e}")))?; + + state.rate_limiter.reload_from_db(&state.db).await; + + log_event( + &state.db, + EventLog { + event_type: "rate_limit.deleted".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(id.to_string()), + detail: serde_json::json!({}), + }, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} + +/// PUT /admin/rate-limits/enabled — toggle rate limiting. +pub(super) async fn set_enabled( + State(state): State, + auth: UserAuth, + Json(body): Json, +) -> Result { + auth.require(Permission::RateLimitsCreate).await?; + + let value = if body.enabled { "true" } else { "false" }; + + sqlx::query( + r#" + INSERT INTO rate_limit_settings (key, value) + VALUES ('enabled', $1) + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = NOW() + "#, + ) + .bind(value) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to update rate limit settings: {e}")))?; + + state.rate_limiter.set_enabled(body.enabled); + + log_event( + &state.db, + EventLog { + event_type: "rate_limit.toggled".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: None, + detail: serde_json::json!({ "enabled": body.enabled }), + }, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} + +/// POST /admin/rate-limits/allowlist — add an IP/CIDR to the allowlist. +pub(super) async fn add_allowlist( + State(state): State, + auth: UserAuth, + Json(body): Json, +) -> Result { + auth.require(Permission::RateLimitsCreate).await?; + + // Validate CIDR syntax; if it's a bare IP, append /32 or /128 + let cidr_str = if body.cidr.contains('/') { + body.cidr.clone() + } else if let Ok(ip) = body.cidr.parse::() { + match ip { + std::net::IpAddr::V4(_) => format!("{}/32", body.cidr), + std::net::IpAddr::V6(_) => format!("{}/128", body.cidr), + } + } else { + return Err(AppError::BadRequest(format!( + "invalid IP or CIDR: {}", + body.cidr + ))); + }; + + // Validate it parses as IpNet + if cidr_str.parse::().is_err() { + return Err(AppError::BadRequest(format!("invalid CIDR: {}", cidr_str))); + } + + sqlx::query("INSERT INTO rate_limit_allowlist (cidr, note) VALUES ($1, $2)") + .bind(&cidr_str) + .bind(&body.note) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to add allowlist entry: {e}")))?; + + state.rate_limiter.reload_from_db(&state.db).await; + + log_event( + &state.db, + EventLog { + event_type: "rate_limit.allowlist_added".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(cidr_str), + detail: serde_json::json!({ "note": body.note }), + }, + ) + .await; + + Ok(StatusCode::CREATED) +} + +/// DELETE /admin/rate-limits/allowlist/{id} — remove an allowlist entry. +pub(super) async fn remove_allowlist( + State(state): State, + auth: UserAuth, + Path(id): Path, +) -> Result { + auth.require(Permission::RateLimitsDelete).await?; + + let result = sqlx::query("DELETE FROM rate_limit_allowlist WHERE id = $1") + .bind(id) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to delete allowlist entry: {e}")))?; + + if result.rows_affected() == 0 { + return Err(AppError::NotFound(format!( + "allowlist entry {id} not found" + ))); + } + + state.rate_limiter.reload_from_db(&state.db).await; + + log_event( + &state.db, + EventLog { + event_type: "rate_limit.allowlist_removed".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(id.to_string()), + detail: serde_json::json!({}), + }, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} diff --git a/src/admin/types.rs b/src/admin/types.rs --- a/src/admin/types.rs +++ b/src/admin/types.rs @@ -204,3 +204,50 @@ #[derive(Deserialize)] pub(super) struct UpdateLabelerBody { pub(super) status: String, } + +// --------------------------------------------------------------------------- +// Rate limit types +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +pub(super) struct UpsertRateLimitBody { + pub(super) method: Option, + pub(super) capacity: u32, + pub(super) refill_rate: f64, +} + +#[derive(Deserialize)] +pub(super) struct SetEnabledBody { + pub(super) enabled: bool, +} + +#[derive(Deserialize)] +pub(super) struct AddAllowlistBody { + pub(super) cidr: String, + pub(super) note: Option, +} + +#[derive(Serialize)] +pub(super) struct RateLimitsResponse { + pub(super) enabled: bool, + pub(super) limits: Vec, + pub(super) allowlist: Vec, +} + +#[derive(Serialize, sqlx::FromRow)] +pub(super) struct RateLimitSummary { + pub(super) id: i32, + pub(super) method: Option, + pub(super) capacity: i32, + pub(super) refill_rate: f32, + pub(super) created_at: chrono::DateTime, + pub(super) updated_at: chrono::DateTime, +} + +#[derive(Serialize, sqlx::FromRow)] +pub(super) struct AllowlistEntry { + pub(super) id: i32, + pub(super) cidr: String, + pub(super) note: Option, + pub(super) created_at: chrono::DateTime, +} diff --git a/src/aip.rs b/src/aip.rs --- a/src/aip.rs +++ b/src/aip.rs @@ -119,6 +119,15 @@ db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: crate::lexicon::LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: crate::rate_limit::RateLimiter::new( + false, + crate::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/src/error.rs b/src/error.rs --- a/src/error.rs +++ b/src/error.rs @@ -56,6 +56,11 @@ InsufficientPermissions(String), Internal(String), NotFound(String), PdsError(StatusCode, Bytes), + RateLimited { + retry_after: u64, + limit: u32, + reset: u64, + }, ScriptError { error_type: ScriptErrorType, message: String, @@ -76,6 +81,9 @@ AppError::InsufficientPermissions(perm) => write!(f, "Missing permission: {perm}"), AppError::Internal(msg) => write!(f, "internal error: {msg}"), AppError::NotFound(msg) => write!(f, "not found: {msg}"), AppError::PdsError(status, _) => write!(f, "PDS error: {status}"), + AppError::RateLimited { retry_after, .. } => { + write!(f, "rate limited: retry after {retry_after}s") + } AppError::ScriptError { error_type, message, @@ -139,6 +147,24 @@ "message": format!("Missing permission: {perm}"), }); (StatusCode::FORBIDDEN, axum::Json(body)).into_response() } + AppError::RateLimited { + retry_after, + limit, + reset, + } => { + let body = serde_json::json!({ + "error": "RateLimited", + "message": "Too many requests", + }); + let mut response = + (StatusCode::TOO_MANY_REQUESTS, axum::Json(body)).into_response(); + let headers = response.headers_mut(); + headers.insert("RateLimit-Limit", limit.into()); + headers.insert("RateLimit-Remaining", 0u32.into()); + headers.insert("RateLimit-Reset", reset.into()); + headers.insert("Retry-After", retry_after.into()); + response + } other => { let (status, message) = match &other { AppError::Auth(msg) => (StatusCode::UNAUTHORIZED, msg.clone()), @@ -154,6 +180,7 @@ AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()), AppError::PdsError(..) | AppError::AuthDpopNonce(..) | AppError::InsufficientPermissions(..) + | AppError::RateLimited { .. } | AppError::ScriptError { .. } => unreachable!(), }; diff --git a/src/lib.rs b/src/lib.rs --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,7 @@ pub mod labeler; pub mod lexicon; pub mod lua; pub mod profile; +pub mod rate_limit; pub mod record_refs; pub mod repo; pub mod resolve; @@ -17,6 +18,8 @@ pub mod xrpc; use config::Config; use lexicon::LexiconRegistry; +use rate_limit::RateLimiter; +use std::sync::Arc; use tokio::sync::watch; #[derive(Clone)] @@ -27,4 +30,5 @@ pub db: sqlx::PgPool, pub lexicons: LexiconRegistry, pub collections_tx: watch::Sender>, pub labeler_subscriptions_tx: watch::Sender<()>, + pub rate_limiter: Arc, } diff --git a/src/lua/atproto_api.rs b/src/lua/atproto_api.rs --- a/src/lua/atproto_api.rs +++ b/src/lua/atproto_api.rs @@ -199,6 +199,15 @@ db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: crate::rate_limit::RateLimiter::new( + false, + crate::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/src/lua/db_api.rs b/src/lua/db_api.rs --- a/src/lua/db_api.rs +++ b/src/lua/db_api.rs @@ -592,6 +592,15 @@ db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: crate::rate_limit::RateLimiter::new( + false, + crate::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/src/lua/execute.rs b/src/lua/execute.rs --- a/src/lua/execute.rs +++ b/src/lua/execute.rs @@ -927,6 +927,15 @@ db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: crate::rate_limit::RateLimiter::new( + false, + crate::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/src/lua/http_api.rs b/src/lua/http_api.rs --- a/src/lua/http_api.rs +++ b/src/lua/http_api.rs @@ -109,6 +109,15 @@ db: sqlx::PgPool::connect_lazy("postgres://localhost/fake").unwrap(), lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: crate::rate_limit::RateLimiter::new( + false, + crate::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/src/main.rs b/src/main.rs --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ use happyview::config::Config; use happyview::lexicon::{LexiconRegistry, ParsedLexicon, ProcedureAction}; +use happyview::rate_limit::RateLimiter; use happyview::resolve::{fetch_lexicon_from_pds, resolve_nsid_authority}; use happyview::{AppState, labeler, server, tap}; use tokio::sync::watch; @@ -154,6 +155,16 @@ "processed network lexicons on startup" ); } + // Initialize rate limiter from DB. + let rl_state = RateLimiter::load_from_db(&db).await; + let rate_limiter = RateLimiter::new( + rl_state.enabled, + rl_state.global, + rl_state.overrides, + rl_state.allowlist, + ); + tokio::spawn(rate_limiter.clone().spawn_cleanup()); + let initial_collections = lexicons.get_record_collections().await; let initial_collections_for_sync = initial_collections.clone(); let (collections_tx, collections_rx) = watch::channel(initial_collections); @@ -166,6 +177,7 @@ db, lexicons, collections_tx, labeler_subscriptions_tx, + rate_limiter, }; // Sync initial collections to Tap on startup. diff --git a/src/rate_limit.rs b/src/rate_limit.rs new file mode 100644 --- /dev/null +++ b/src/rate_limit.rs @@ -0,0 +1,364 @@ +use arc_swap::ArcSwap; +use dashmap::DashMap; +use ipnet::IpNet; +use sqlx::PgPool; +use std::collections::HashMap; +use std::net::IpAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +pub struct RateLimitConfig { + pub capacity: u32, + pub refill_rate: f64, +} + +pub enum CheckResult { + Allowed { + remaining: u32, + limit: u32, + reset: u64, + }, + Limited { + retry_after: u64, + limit: u32, + reset: u64, + }, + Disabled, +} + +struct TokenBucket { + tokens: f64, + capacity: u32, + refill_rate: f64, + last_refill: Instant, + last_access: Instant, +} + +pub struct RateLimiter { + enabled: AtomicBool, + buckets: DashMap, + global_config: ArcSwap, + overrides: ArcSwap>, + allowlist: ArcSwap>, +} + +pub struct RateLimiterState { + pub enabled: bool, + pub global: RateLimitConfig, + pub overrides: HashMap, + pub allowlist: Vec, +} + +fn now_unix() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +impl RateLimiter { + pub fn new( + enabled: bool, + global: RateLimitConfig, + overrides: HashMap, + allowlist: Vec, + ) -> Arc { + Arc::new(Self { + enabled: AtomicBool::new(enabled), + buckets: DashMap::new(), + global_config: ArcSwap::new(Arc::new(global)), + overrides: ArcSwap::new(Arc::new(overrides)), + allowlist: ArcSwap::new(Arc::new(allowlist)), + }) + } + + pub fn check(&self, key: &str, method: Option<&str>, client_ip: Option) -> CheckResult { + if !self.enabled.load(Ordering::Relaxed) { + return CheckResult::Disabled; + } + + if let Some(ip) = client_ip { + let list = self.allowlist.load(); + for net in list.iter() { + if net.contains(&ip) { + return CheckResult::Disabled; + } + } + } + + let overrides = self.overrides.load(); + let global = self.global_config.load(); + + let (capacity, refill_rate) = if let Some(method) = method { + if let Some(cfg) = overrides.get(method) { + (cfg.capacity, cfg.refill_rate) + } else { + (global.capacity, global.refill_rate) + } + } else { + (global.capacity, global.refill_rate) + }; + + let now = Instant::now(); + + let mut bucket = self + .buckets + .entry(key.to_string()) + .or_insert_with(|| TokenBucket { + tokens: capacity as f64, + capacity, + refill_rate, + last_refill: now, + last_access: now, + }); + + // Hot-reload config changes + bucket.capacity = capacity; + bucket.refill_rate = refill_rate; + + // Refill tokens + let elapsed = now.duration_since(bucket.last_refill).as_secs_f64(); + bucket.tokens = (bucket.tokens + elapsed * refill_rate).min(capacity as f64); + bucket.last_refill = now; + bucket.last_access = now; + + let reset_secs = if bucket.tokens < capacity as f64 { + ((capacity as f64 - bucket.tokens) / refill_rate).ceil() as u64 + } else { + 0 + }; + let reset = now_unix() + reset_secs; + + if bucket.tokens >= 1.0 { + bucket.tokens -= 1.0; + CheckResult::Allowed { + remaining: bucket.tokens.floor() as u32, + limit: capacity, + reset, + } + } else { + let retry_after = ((1.0 - bucket.tokens) / refill_rate).ceil() as u64; + CheckResult::Limited { + retry_after, + limit: capacity, + reset: now_unix() + ((capacity as f64) / refill_rate).ceil() as u64, + } + } + } + + pub fn set_enabled(&self, enabled: bool) { + self.enabled.store(enabled, Ordering::Relaxed); + } + + pub fn is_enabled(&self) -> bool { + self.enabled.load(Ordering::Relaxed) + } + + pub fn update_config( + &self, + global: RateLimitConfig, + overrides: HashMap, + ) { + self.global_config.store(Arc::new(global)); + self.overrides.store(Arc::new(overrides)); + } + + pub fn update_allowlist(&self, entries: Vec) { + self.allowlist.store(Arc::new(entries)); + } + + pub async fn spawn_cleanup(self: Arc) { + let interval = tokio::time::Duration::from_secs(60); + let stale_threshold = std::time::Duration::from_secs(300); // 5 minutes + loop { + tokio::time::sleep(interval).await; + let now = Instant::now(); + self.buckets + .retain(|_, bucket| now.duration_since(bucket.last_access) < stale_threshold); + } + } + + pub async fn load_from_db(db: &PgPool) -> RateLimiterState { + // Load enabled flag + let enabled: bool = sqlx::query_scalar::<_, String>( + "SELECT value FROM rate_limit_settings WHERE key = 'enabled'", + ) + .fetch_optional(db) + .await + .ok() + .flatten() + .map(|v| v == "true") + .unwrap_or(true); + + // Load rate limit configs + let rows: Vec<(Option, i32, f32)> = + sqlx::query_as("SELECT method, capacity, refill_rate FROM rate_limits") + .fetch_all(db) + .await + .unwrap_or_default(); + + let mut global = RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }; + let mut overrides = HashMap::new(); + + for (method, capacity, refill_rate) in rows { + let config = RateLimitConfig { + capacity: capacity as u32, + refill_rate: refill_rate as f64, + }; + match method { + None => global = config, + Some(m) => { + overrides.insert(m, config); + } + } + } + + // Load allowlist + let cidr_rows: Vec<(String,)> = sqlx::query_as("SELECT cidr FROM rate_limit_allowlist") + .fetch_all(db) + .await + .unwrap_or_default(); + + let allowlist: Vec = cidr_rows + .into_iter() + .filter_map(|(cidr,)| cidr.parse().ok()) + .collect(); + + RateLimiterState { + enabled, + global, + overrides, + allowlist, + } + } + + /// Reload all config from DB and apply to the live limiter. + pub async fn reload_from_db(&self, db: &PgPool) { + let state = Self::load_from_db(db).await; + self.set_enabled(state.enabled); + self.update_config(state.global, state.overrides); + self.update_allowlist(state.allowlist); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn basic_allow_and_exhaust() { + let rl = RateLimiter::new( + true, + RateLimitConfig { + capacity: 3, + refill_rate: 1.0, + }, + HashMap::new(), + vec![], + ); + + // Should allow 3 requests (bucket starts full) + for _ in 0..3 { + assert!(matches!( + rl.check("k", None, None), + CheckResult::Allowed { .. } + )); + } + // 4th should be limited + assert!(matches!( + rl.check("k", None, None), + CheckResult::Limited { .. } + )); + } + + #[test] + fn disabled_returns_disabled() { + let rl = RateLimiter::new( + false, + RateLimitConfig { + capacity: 1, + refill_rate: 1.0, + }, + HashMap::new(), + vec![], + ); + assert!(matches!(rl.check("k", None, None), CheckResult::Disabled)); + } + + #[test] + fn allowlisted_ip_bypasses() { + let rl = RateLimiter::new( + true, + RateLimitConfig { + capacity: 1, + refill_rate: 0.001, + }, + HashMap::new(), + vec!["10.0.0.0/8".parse().unwrap()], + ); + + let ip: IpAddr = "10.0.0.5".parse().unwrap(); + // Even after exhausting, allowlisted IP gets Disabled + assert!(matches!( + rl.check("k", None, Some(ip)), + CheckResult::Disabled + )); + } + + #[test] + fn method_override_applies() { + let mut overrides = HashMap::new(); + overrides.insert( + "com.atproto.repo.uploadBlob".to_string(), + RateLimitConfig { + capacity: 2, + refill_rate: 0.001, + }, + ); + + let rl = RateLimiter::new( + true, + RateLimitConfig { + capacity: 100, + refill_rate: 100.0, + }, + overrides, + vec![], + ); + + // Override has capacity 2 + assert!(matches!( + rl.check("k", Some("com.atproto.repo.uploadBlob"), None), + CheckResult::Allowed { limit: 2, .. } + )); + assert!(matches!( + rl.check("k", Some("com.atproto.repo.uploadBlob"), None), + CheckResult::Allowed { limit: 2, .. } + )); + assert!(matches!( + rl.check("k", Some("com.atproto.repo.uploadBlob"), None), + CheckResult::Limited { limit: 2, .. } + )); + } + + #[test] + fn toggle_enabled() { + let rl = RateLimiter::new( + true, + RateLimitConfig { + capacity: 1, + refill_rate: 1.0, + }, + HashMap::new(), + vec![], + ); + assert!(rl.is_enabled()); + rl.set_enabled(false); + assert!(!rl.is_enabled()); + assert!(matches!(rl.check("k", None, None), CheckResult::Disabled)); + } +} diff --git a/src/repo/upload_blob.rs b/src/repo/upload_blob.rs --- a/src/repo/upload_blob.rs +++ b/src/repo/upload_blob.rs @@ -2,10 +2,12 @@ use axum::body::Bytes; use axum::extract::State; use axum::http::HeaderMap; use axum::response::Response; +use std::net::IpAddr; use crate::AppState; use crate::auth::Claims; use crate::error::AppError; +use crate::rate_limit::CheckResult; use super::pds::pds_post_blob; use super::session::get_atp_session; @@ -16,6 +18,30 @@ claims: Claims, headers: HeaderMap, body: Bytes, ) -> Result { + let client_ip: Option = headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.split(',').next()) + .and_then(|s| s.trim().parse().ok()); + + let rate_key = claims.did().to_string(); + let check = state + .rate_limiter + .check(&rate_key, Some("com.atproto.repo.uploadBlob"), client_ip); + + if let CheckResult::Limited { + retry_after, + limit, + reset, + } = check + { + return Err(AppError::RateLimited { + retry_after, + limit, + reset, + }); + } + let session = get_atp_session(&state, claims.token()).await?; let content_type = headers @@ -23,5 +49,19 @@ .get("content-type") .and_then(|v| v.to_str().ok()) .unwrap_or("application/octet-stream"); - pds_post_blob(&state, &session, content_type, body).await + let mut response = pds_post_blob(&state, &session, content_type, body).await?; + + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + let h = response.headers_mut(); + h.insert("RateLimit-Limit", limit.into()); + h.insert("RateLimit-Remaining", remaining.into()); + h.insert("RateLimit-Reset", reset.into()); + } + + Ok(response) } diff --git a/src/server.rs b/src/server.rs --- a/src/server.rs +++ b/src/server.rs @@ -1,9 +1,12 @@ use axum::extract::{DefaultBodyLimit, State}; +use axum::http::HeaderMap; +use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; use bytes::Bytes; use http_body_util::Full; use std::convert::Infallible; +use std::net::IpAddr; use tower_http::cors::CorsLayer; use tower_http::services::ServeDir; use tower_http::trace::TraceLayer; @@ -14,6 +17,7 @@ use crate::aip; use crate::auth::Claims; use crate::error::AppError; use crate::profile; +use crate::rate_limit::CheckResult; use crate::repo; use crate::xrpc; @@ -84,11 +88,52 @@ async fn config_endpoint(State(state): State) -> Json { Json(serde_json::json!({ "aip_url": state.config.aip_public_url })) } +fn ip_from_forwarded_for(value: Option<&str>) -> Option { + let forwarded = value?; + let first = forwarded.split(',').next()?; + first.trim().parse::().ok() +} + async fn get_profile( State(state): State, claims: Claims, -) -> Result, AppError> { + headers: HeaderMap, +) -> Result { + let client_ip = + ip_from_forwarded_for(headers.get("x-forwarded-for").and_then(|v| v.to_str().ok())); + let rate_key = claims.did().to_string(); + let check = state + .rate_limiter + .check(&rate_key, Some("app.bsky.actor.getProfile"), client_ip); + + if let CheckResult::Limited { + retry_after, + limit, + reset, + } = check + { + return Err(AppError::RateLimited { + retry_after, + limit, + reset, + }); + } + let profile = profile::resolve_profile(&state.http, &state.config.plc_url, claims.did()).await?; - Ok(Json(profile)) + let mut response = Json(profile).into_response(); + + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + let h = response.headers_mut(); + h.insert("RateLimit-Limit", limit.into()); + h.insert("RateLimit-Remaining", remaining.into()); + h.insert("RateLimit-Reset", reset.into()); + } + + Ok(response) } diff --git a/src/xrpc/mod.rs b/src/xrpc/mod.rs --- a/src/xrpc/mod.rs +++ b/src/xrpc/mod.rs @@ -3,17 +3,19 @@ mod query; use axum::Json; use axum::body::Body; -use axum::extract::{FromRequestParts, Path, RawQuery, State}; +use axum::extract::{ConnectInfo, FromRequestParts, Path, RawQuery, State}; use axum::http::StatusCode; use axum::http::request::Parts; use axum::response::Response; use serde_json::Value; use std::collections::HashMap; +use std::net::{IpAddr, SocketAddr}; use crate::AppState; use crate::auth::Claims; use crate::error::AppError; use crate::lexicon::LexiconType; +use crate::rate_limit::CheckResult; use crate::resolve::resolve_nsid_authority; /// Parse a raw query string into a map where repeated keys become JSON arrays. @@ -105,6 +107,31 @@ .body(Body::from(bytes)) .unwrap()) } +/// Extract client IP from X-Forwarded-For header or ConnectInfo. +fn extract_client_ip(parts: &Parts) -> Option { + if let Some(forwarded) = parts + .headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + && let Some(first) = forwarded.split(',').next() + && let Ok(ip) = first.trim().parse::() + { + return Some(ip); + } + parts + .extensions + .get::>() + .map(|ci| ci.0.ip()) +} + +/// Apply rate limit headers to a response. +fn apply_rate_limit_headers(response: &mut Response, remaining: u32, limit: u32, reset: u64) { + let headers = response.headers_mut(); + headers.insert("RateLimit-Limit", limit.into()); + headers.insert("RateLimit-Remaining", remaining.into()); + headers.insert("RateLimit-Reset", reset.into()); +} + /// Catch-all GET handler for XRPC queries. pub async fn xrpc_get( State(state): State, @@ -114,12 +141,51 @@ mut parts: Parts, ) -> Result { let raw_query = raw_query.unwrap_or_default(); let params = parse_query_params(&raw_query); + let client_ip = extract_client_ip(&parts); let claims = Claims::from_request_parts(&mut parts, &state).await.ok(); + // Rate limit check + let rate_key = claims + .as_ref() + .map(|c| c.did().to_string()) + .unwrap_or_else(|| { + client_ip + .map(|ip| ip.to_string()) + .unwrap_or_else(|| "unknown".to_string()) + }); + + let check = state + .rate_limiter + .check(&rate_key, Some(&method), client_ip); + + match check { + CheckResult::Limited { + retry_after, + limit, + reset, + } => { + return Err(AppError::RateLimited { + retry_after, + limit, + reset, + }); + } + CheckResult::Allowed { .. } | CheckResult::Disabled => {} + } + let lexicon = match state.lexicons.get(&method).await { Some(l) => l, None => { - return proxy_to_authority(&state, &method, &raw_query, None).await; + let mut response = proxy_to_authority(&state, &method, &raw_query, None).await?; + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + apply_rate_limit_headers(&mut response, remaining, limit, reset); + } + return Ok(response); } }; @@ -129,7 +195,24 @@ "{method} is not a query endpoint" ))); } - query::handle_query(&state, &method, ¶ms, &lexicon, claims.as_ref()).await + let mut response = + query::handle_query(&state, &method, ¶ms, &lexicon, claims.as_ref()).await?; + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + apply_rate_limit_headers(&mut response, remaining, limit, reset); + } + Ok(response) +} + +/// Extract client IP from X-Forwarded-For header value. +fn ip_from_forwarded_for(value: Option<&str>) -> Option { + let forwarded = value?; + let first = forwarded.split(',').next()?; + first.trim().parse::().ok() } /// Catch-all POST handler for XRPC procedures. @@ -137,11 +220,46 @@ pub async fn xrpc_post( State(state): State, Path(method): Path, claims: Claims, + headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { + let client_ip = + ip_from_forwarded_for(headers.get("x-forwarded-for").and_then(|v| v.to_str().ok())); + let rate_key = claims.did().to_string(); + + let check = state + .rate_limiter + .check(&rate_key, Some(&method), client_ip); + + match check { + CheckResult::Limited { + retry_after, + limit, + reset, + } => { + return Err(AppError::RateLimited { + retry_after, + limit, + reset, + }); + } + CheckResult::Allowed { .. } | CheckResult::Disabled => {} + } + let lexicon = match state.lexicons.get(&method).await { Some(l) => l, - None => return proxy_to_authority(&state, &method, "", Some(&body)).await, + None => { + let mut response = proxy_to_authority(&state, &method, "", Some(&body)).await?; + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + apply_rate_limit_headers(&mut response, remaining, limit, reset); + } + return Ok(response); + } }; if lexicon.lexicon_type != LexiconType::Procedure { @@ -150,5 +268,15 @@ "{method} is not a procedure endpoint" ))); } - procedure::handle_procedure(&state, &method, &claims, &body, &lexicon).await + let mut response = + procedure::handle_procedure(&state, &method, &claims, &body, &lexicon).await?; + if let CheckResult::Allowed { + remaining, + limit, + reset, + } = check + { + apply_rate_limit_headers(&mut response, remaining, limit, reset); + } + Ok(response) } diff --git a/tests/common/app.rs b/tests/common/app.rs --- a/tests/common/app.rs +++ b/tests/common/app.rs @@ -68,6 +68,15 @@ db: pool, lexicons, collections_tx, labeler_subscriptions_tx, + rate_limiter: happyview::rate_limit::RateLimiter::new( + false, + happyview::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), }; let router = server::router(state.clone()); diff --git a/tests/lua_atproto_api.rs b/tests/lua_atproto_api.rs --- a/tests/lua_atproto_api.rs +++ b/tests/lua_atproto_api.rs @@ -32,6 +32,15 @@ db: pool, lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: happyview::rate_limit::RateLimiter::new( + false, + happyview::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } diff --git a/tests/lua_db_api.rs b/tests/lua_db_api.rs --- a/tests/lua_db_api.rs +++ b/tests/lua_db_api.rs @@ -35,6 +35,15 @@ db: pool, lexicons: LexiconRegistry::new(), collections_tx: tx, labeler_subscriptions_tx: labeler_tx, + rate_limiter: happyview::rate_limit::RateLimiter::new( + false, + happyview::rate_limit::RateLimitConfig { + capacity: 100, + refill_rate: 2.0, + }, + std::collections::HashMap::new(), + vec![], + ), } } -- tangled.sh