From cd87c9f4abf5776d355c06982563551546ba017c Mon Sep 17 00:00:00 2001 From: Trezy Date: Thu, 16 Apr 2026 09:47:19 -0500 Subject: [PATCH] feat: add support for multiple domains --- ...20260414000000_drop_rate_limits_tables.sql | 2 + .../20260415000000_create_domains.sql | 7 + ...20260414000000_drop_rate_limits_tables.sql | 2 + .../sqlite/20260415000000_create_domains.sql | 7 + src/admin/api_clients.rs | 16 +- src/admin/domains.rs | 306 ++++++++++++ src/admin/mod.rs | 10 +- src/admin/permissions.rs | 16 - src/admin/rate_limits.rs | 158 ------- src/admin/types.rs | 46 +- src/auth/client_registry.rs | 63 ++- src/auth/routes.rs | 21 +- src/domain.rs | 198 ++++++++ src/domain_middleware.rs | 38 ++ src/external_auth/routes.rs | 7 +- src/lib.rs | 3 + src/lua/atproto_api.rs | 12 +- src/lua/db_api.rs | 12 +- src/lua/execute.rs | 12 +- src/lua/http_api.rs | 12 +- src/lua/xrpc_api.rs | 12 +- src/main.rs | 203 +++++++- src/rate_limit.rs | 444 +++++------------- src/repo/session.rs | 2 +- src/repo/upload_blob.rs | 21 +- src/server.rs | 71 ++- src/xrpc/mod.rs | 12 +- tests/common/app.rs | 12 +- tests/common/db.rs | 3 +- tests/e2e_api_clients.rs | 8 +- tests/e2e_domains.rs | 352 ++++++++++++++ tests/lua_atproto_api.rs | 12 +- tests/lua_db_api.rs | 12 +- 33 files changed, 1450 insertions(+), 662 deletions(-) create mode 100644 migrations/postgres/20260414000000_drop_rate_limits_tables.sql create mode 100644 migrations/postgres/20260415000000_create_domains.sql create mode 100644 migrations/sqlite/20260414000000_drop_rate_limits_tables.sql create mode 100644 migrations/sqlite/20260415000000_create_domains.sql create mode 100644 src/admin/domains.rs delete mode 100644 src/admin/rate_limits.rs create mode 100644 src/domain.rs create mode 100644 src/domain_middleware.rs create mode 100644 tests/e2e_domains.rs diff --git a/migrations/postgres/20260414000000_drop_rate_limits_tables.sql b/migrations/postgres/20260414000000_drop_rate_limits_tables.sql new file mode 100644 index 0000000..e679247 --- /dev/null +++ b/migrations/postgres/20260414000000_drop_rate_limits_tables.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS rate_limits; +DROP TABLE IF EXISTS rate_limit_settings; diff --git a/migrations/postgres/20260415000000_create_domains.sql b/migrations/postgres/20260415000000_create_domains.sql new file mode 100644 index 0000000..edfddaa --- /dev/null +++ b/migrations/postgres/20260415000000_create_domains.sql @@ -0,0 +1,7 @@ +CREATE TABLE IF NOT EXISTS domains ( + id TEXT PRIMARY KEY, + url TEXT NOT NULL UNIQUE, + is_primary INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL +); diff --git a/migrations/sqlite/20260414000000_drop_rate_limits_tables.sql b/migrations/sqlite/20260414000000_drop_rate_limits_tables.sql new file mode 100644 index 0000000..e679247 --- /dev/null +++ b/migrations/sqlite/20260414000000_drop_rate_limits_tables.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS rate_limits; +DROP TABLE IF EXISTS rate_limit_settings; diff --git a/migrations/sqlite/20260415000000_create_domains.sql b/migrations/sqlite/20260415000000_create_domains.sql new file mode 100644 index 0000000..edfddaa --- /dev/null +++ b/migrations/sqlite/20260415000000_create_domains.sql @@ -0,0 +1,7 @@ +CREATE TABLE IF NOT EXISTS domains ( + id TEXT PRIMARY KEY, + url TEXT NOT NULL UNIQUE, + is_primary INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL +); diff --git a/src/admin/api_clients.rs b/src/admin/api_clients.rs index 4597158..2960cba 100644 --- a/src/admin/api_clients.rs +++ b/src/admin/api_clients.rs @@ -94,15 +94,15 @@ pub(super) async fn create_api_client( if let (Some(capacity), Some(refill_rate)) = (body.rate_limit_capacity, body.rate_limit_refill_rate) { - let global = state.rate_limiter.global_config(); + let defaults = state.rate_limiter.defaults(); state.rate_limiter.register_client_config( client_key.clone(), crate::rate_limit::RateLimitConfig { capacity: capacity as u32, refill_rate, - default_query_cost: global.default_query_cost, - default_procedure_cost: global.default_procedure_cost, - default_proxy_cost: global.default_proxy_cost, + default_query_cost: defaults.query_cost, + default_procedure_cost: defaults.procedure_cost, + default_proxy_cost: defaults.proxy_cost, }, ); } @@ -397,15 +397,15 @@ pub(super) async fn update_api_client( }, ); if let (Some(cap), Some(refill)) = (capacity, refill_rate) { - let global = state.rate_limiter.global_config(); + let defaults = state.rate_limiter.defaults(); state.rate_limiter.register_client_config( client_key, crate::rate_limit::RateLimitConfig { capacity: cap as u32, refill_rate: refill, - default_query_cost: global.default_query_cost, - default_procedure_cost: global.default_procedure_cost, - default_proxy_cost: global.default_proxy_cost, + default_query_cost: defaults.query_cost, + default_procedure_cost: defaults.procedure_cost, + default_proxy_cost: defaults.proxy_cost, }, ); } else { diff --git a/src/admin/domains.rs b/src/admin/domains.rs new file mode 100644 index 0000000..e91b1a5 --- /dev/null +++ b/src/admin/domains.rs @@ -0,0 +1,306 @@ +use axum::Json; +use axum::extract::{Path, State}; +use axum::http::StatusCode; + +use crate::AppState; +use crate::db::{adapt_sql, now_rfc3339}; +use crate::domain::Domain; +use crate::error::AppError; +use crate::event_log::{EventLog, Severity, log_event}; + +use super::auth::UserAuth; +use super::permissions::Permission; +use super::types::{CreateDomainBody, DomainResponse}; + +fn domain_to_response(d: &Domain) -> DomainResponse { + DomainResponse { + id: d.id.clone(), + url: d.url.clone(), + is_primary: d.is_primary, + created_at: d.created_at.clone(), + updated_at: d.updated_at.clone(), + } +} + +/// GET /admin/domains +pub(super) async fn list( + State(state): State, + auth: UserAuth, +) -> Result>, AppError> { + auth.require(Permission::SettingsManage).await?; + + let sql = adapt_sql( + "SELECT id, url, is_primary, created_at, updated_at FROM domains ORDER BY created_at", + state.db_backend, + ); + let rows: Vec<(String, String, i32, String, String)> = sqlx::query_as(&sql) + .fetch_all(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to list domains: {e}")))?; + + let domains: Vec = rows + .into_iter() + .map( + |(id, url, is_primary, created_at, updated_at)| DomainResponse { + id, + url, + is_primary: is_primary != 0, + created_at, + updated_at, + }, + ) + .collect(); + + Ok(Json(domains)) +} + +/// POST /admin/domains +pub(super) async fn create( + State(state): State, + auth: UserAuth, + Json(body): Json, +) -> Result<(StatusCode, Json), AppError> { + auth.require(Permission::SettingsManage).await?; + + let url = body.url.trim_end_matches('/').to_string(); + + let parsed = + reqwest::Url::parse(&url).map_err(|_| AppError::BadRequest("invalid URL".into()))?; + + if parsed.path() != "/" && !parsed.path().is_empty() { + return Err(AppError::BadRequest("URL must not contain a path".into())); + } + + if parsed.host_str().is_none() { + return Err(AppError::BadRequest("URL must contain a host".into())); + } + + let is_loopback = state.config.public_url.contains("127.0.0.1") + || state.config.public_url.contains("[::1]") + || state.config.public_url.contains("localhost"); + + if parsed.scheme() != "https" && !is_loopback { + return Err(AppError::BadRequest("URL scheme must be https".into())); + } + + // Check for duplicates + let existing: Option<(String,)> = sqlx::query_as(&adapt_sql( + "SELECT id FROM domains WHERE url = ?", + state.db_backend, + )) + .bind(&url) + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to check domain: {e}")))?; + + if existing.is_some() { + return Err(AppError::BadRequest(format!( + "domain '{url}' already exists" + ))); + } + + let id = uuid::Uuid::new_v4().to_string(); + let now = now_rfc3339(); + + let sql = adapt_sql( + "INSERT INTO domains (id, url, is_primary, created_at, updated_at) VALUES (?, ?, 0, ?, ?)", + state.db_backend, + ); + sqlx::query(&sql) + .bind(&id) + .bind(&url) + .bind(&now) + .bind(&now) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to create domain: {e}")))?; + + let domain = Domain { + id: id.clone(), + url: url.clone(), + is_primary: false, + created_at: now.clone(), + updated_at: now, + }; + + // Register the OAuth client for this domain + state + .oauth + .register_domain_client(url.clone(), state.oauth.primary_client()); + + // Build a proper OAuth client if not loopback + let domain_is_loopback = + url.contains("127.0.0.1") || url.contains("[::1]") || url.contains("localhost"); + if !domain_is_loopback { + let client_id_url = format!("{}/oauth-client-metadata.json", url.trim_end_matches('/')); + let callback = format!("{}/auth/callback", url.trim_end_matches('/')); + if let Err(e) = state.oauth.register_api_client( + &client_id_url, + &url, + vec![callback], + "atproto", + &crate::auth::client_registry::ApiClientOAuthParams { + plc_url: state.config.plc_url.clone(), + state_store: state.oauth_state_store.clone(), + session_store_pool: state.db.clone(), + db_backend: state.db_backend, + }, + ) { + tracing::error!(domain = %url, error = %e, "Failed to create OAuth client for domain"); + } else { + // Move from `clients` (where register_api_client puts it) to domain_clients + clients + if let Some(client) = state.oauth.get(&client_id_url) { + state.oauth.remove(&client_id_url); + state.oauth.register_domain_client(url.clone(), client); + } + } + } + + // Update in-memory cache + state.domain_cache.insert(domain.clone()).await; + + log_event( + &state.db, + EventLog { + event_type: "domain.created".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(url), + detail: serde_json::json!({ "id": id }), + }, + state.db_backend, + ) + .await; + + let response = domain_to_response(&domain); + Ok((StatusCode::CREATED, Json(response))) +} + +/// DELETE /admin/domains/{id} +pub(super) async fn delete( + State(state): State, + auth: UserAuth, + Path(id): Path, +) -> Result { + auth.require(Permission::SettingsManage).await?; + + let sql = adapt_sql( + "SELECT id, url, is_primary, created_at, updated_at FROM domains WHERE id = ?", + state.db_backend, + ); + let row: Option<(String, String, i32, String, String)> = sqlx::query_as(&sql) + .bind(&id) + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to find domain: {e}")))?; + + let (_, url, is_primary, _, _) = + row.ok_or_else(|| AppError::NotFound("domain not found".into()))?; + + if is_primary != 0 { + return Err(AppError::BadRequest( + "cannot delete the primary domain — set a different domain as primary first".into(), + )); + } + + let delete_sql = adapt_sql("DELETE FROM domains WHERE id = ?", state.db_backend); + sqlx::query(&delete_sql) + .bind(&id) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to delete domain: {e}")))?; + + // Remove OAuth client and cache entry + state.oauth.remove_domain_client(&url); + let host = url + .strip_prefix("https://") + .or_else(|| url.strip_prefix("http://")) + .unwrap_or(&url); + state.domain_cache.remove(host).await; + + log_event( + &state.db, + EventLog { + event_type: "domain.deleted".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(url), + detail: serde_json::json!({ "id": id }), + }, + state.db_backend, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} + +/// POST /admin/domains/{id}/primary +pub(super) async fn set_primary( + State(state): State, + auth: UserAuth, + Path(id): Path, +) -> Result { + auth.require(Permission::SettingsManage).await?; + + let sql = adapt_sql( + "SELECT id, url, is_primary, created_at, updated_at FROM domains WHERE id = ?", + state.db_backend, + ); + let row: Option<(String, String, i32, String, String)> = sqlx::query_as(&sql) + .bind(&id) + .fetch_optional(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to find domain: {e}")))?; + + let (_, url, _, _, _) = row.ok_or_else(|| AppError::NotFound("domain not found".into()))?; + + let now = now_rfc3339(); + + let unset_sql = adapt_sql( + "UPDATE domains SET is_primary = 0, updated_at = ? WHERE is_primary = 1", + state.db_backend, + ); + sqlx::query(&unset_sql) + .bind(&now) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to unset primary: {e}")))?; + + let set_sql = adapt_sql( + "UPDATE domains SET is_primary = 1, updated_at = ? WHERE id = ?", + state.db_backend, + ); + sqlx::query(&set_sql) + .bind(&now) + .bind(&id) + .execute(&state.db) + .await + .map_err(|e| AppError::Internal(format!("failed to set primary: {e}")))?; + + // Update cache + let host = url + .strip_prefix("https://") + .or_else(|| url.strip_prefix("http://")) + .unwrap_or(&url); + state.domain_cache.set_primary(host).await; + + // Update OAuth primary client + if let Some(client) = state.oauth.get_domain_client(&url) { + state.oauth.set_primary_client(client); + } + + log_event( + &state.db, + EventLog { + event_type: "domain.primary_changed".to_string(), + severity: Severity::Info, + actor_did: Some(auth.did.clone()), + subject: Some(url), + detail: serde_json::json!({ "id": id }), + }, + state.db_backend, + ) + .await; + + Ok(StatusCode::NO_CONTENT) +} diff --git a/src/admin/mod.rs b/src/admin/mod.rs index 355e885..4df9974 100644 --- a/src/admin/mod.rs +++ b/src/admin/mod.rs @@ -2,13 +2,13 @@ mod api_clients; mod api_keys; pub(crate) mod auth; mod backfill; +mod domains; mod events; mod labelers; mod lexicons; mod network_lexicons; pub(crate) mod permissions; mod plugins; -mod rate_limits; mod records; mod script_variables; pub mod settings; @@ -70,11 +70,6 @@ pub fn admin_routes(_state: AppState) -> Router { "/labelers/{did}", patch(labelers::update).delete(labelers::delete), ) - .route( - "/rate-limits", - post(rate_limits::upsert).get(rate_limits::list), - ) - .route("/rate-limits/enabled", put(rate_limits::set_enabled)) .route("/settings", get(settings::list)) .route( "/settings/logo", @@ -104,4 +99,7 @@ pub fn admin_routes(_state: AppState) -> Router { .put(api_clients::update_api_client) .delete(api_clients::delete_api_client), ) + .route("/domains", post(domains::create).get(domains::list)) + .route("/domains/{id}", delete(domains::delete)) + .route("/domains/{id}/primary", post(domains::set_primary)) } diff --git a/src/admin/permissions.rs b/src/admin/permissions.rs index 6bd2525..129ae62 100644 --- a/src/admin/permissions.rs +++ b/src/admin/permissions.rs @@ -60,13 +60,6 @@ pub enum Permission { #[serde(rename = "labelers:delete")] LabelersDelete, - #[serde(rename = "rate-limits:read")] - RateLimitsRead, - #[serde(rename = "rate-limits:create")] - RateLimitsCreate, - #[serde(rename = "rate-limits:delete")] - RateLimitsDelete, - #[serde(rename = "settings:manage")] SettingsManage, @@ -114,9 +107,6 @@ impl Permission { 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", Self::SettingsManage => "settings:manage", Self::PluginsRead => "plugins:read", Self::PluginsCreate => "plugins:create", @@ -154,9 +144,6 @@ impl Permission { Self::LabelersCreate, Self::LabelersRead, Self::LabelersDelete, - Self::RateLimitsRead, - Self::RateLimitsCreate, - Self::RateLimitsDelete, Self::SettingsManage, Self::PluginsRead, Self::PluginsCreate, @@ -209,9 +196,6 @@ impl Template { 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.insert(Permission::SettingsManage); perms.insert(Permission::PluginsRead); perms.insert(Permission::PluginsCreate); diff --git a/src/admin/rate_limits.rs b/src/admin/rate_limits.rs deleted file mode 100644 index e8a79dc..0000000 --- a/src/admin/rate_limits.rs +++ /dev/null @@ -1,158 +0,0 @@ -use axum::Json; -use axum::extract::State; -use axum::http::StatusCode; - -use crate::AppState; -use crate::db::{adapt_sql, now_rfc3339}; -use crate::error::AppError; -use crate::event_log::{EventLog, Severity, log_event}; - -use super::auth::UserAuth; -use super::permissions::Permission; -use super::types::{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 backend = state.db_backend; - - let enabled_sql = adapt_sql( - "SELECT value FROM rate_limit_settings WHERE key = 'enabled'", - backend, - ); - let enabled: String = sqlx::query_scalar(&enabled_sql) - .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_sql = adapt_sql( - "SELECT capacity, refill_rate, default_query_cost, default_procedure_cost, default_proxy_cost FROM rate_limits WHERE method IS NULL", - backend, - ); - let row: Option<(i32, f64, i32, i32, i32)> = sqlx::query_as(&limits_sql) - .fetch_optional(&state.db) - .await - .map_err(|e| AppError::Internal(format!("failed to read rate limits: {e}")))?; - - let (capacity, refill_rate, default_query_cost, default_procedure_cost, default_proxy_cost) = - row.unwrap_or((100, 2.0, 1, 1, 1)); - - Ok(Json(RateLimitsResponse { - enabled: enabled == "true", - capacity, - refill_rate, - default_query_cost, - default_procedure_cost, - default_proxy_cost, - })) -} - -/// POST /admin/rate-limits — upsert the global rate limit config. -pub(super) async fn upsert( - State(state): State, - auth: UserAuth, - Json(body): Json, -) -> Result { - auth.require(Permission::RateLimitsCreate).await?; - - let backend = state.db_backend; - let now = now_rfc3339(); - let sql = adapt_sql( - r#" - INSERT INTO rate_limits (method, capacity, refill_rate, default_query_cost, default_procedure_cost, default_proxy_cost, created_at) - VALUES (NULL, ?, ?, ?, ?, ?, ?) - ON CONFLICT (method) DO UPDATE SET - capacity = EXCLUDED.capacity, - refill_rate = EXCLUDED.refill_rate, - default_query_cost = EXCLUDED.default_query_cost, - default_procedure_cost = EXCLUDED.default_procedure_cost, - default_proxy_cost = EXCLUDED.default_proxy_cost, - updated_at = ? - "#, - backend, - ); - sqlx::query(&sql) - .bind(body.capacity as i32) - .bind(body.refill_rate) - .bind(body.default_query_cost as i32) - .bind(body.default_procedure_cost as i32) - .bind(body.default_proxy_cost as i32) - .bind(&now) - .bind(&now) - .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: None, - detail: serde_json::json!({ - "capacity": body.capacity, - "refill_rate": body.refill_rate, - "default_query_cost": body.default_query_cost, - "default_procedure_cost": body.default_procedure_cost, - "default_proxy_cost": body.default_proxy_cost, - }), - }, - state.db_backend, - ) - .await; - - Ok(StatusCode::CREATED) -} - -/// 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" }; - - let backend = state.db_backend; - let now = now_rfc3339(); - let sql = adapt_sql( - r#" - INSERT INTO rate_limit_settings (key, value) - VALUES ('enabled', ?) - ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = ? - "#, - backend, - ); - sqlx::query(&sql) - .bind(value) - .bind(&now) - .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 }), - }, - state.db_backend, - ) - .await; - - Ok(StatusCode::NO_CONTENT) -} diff --git a/src/admin/types.rs b/src/admin/types.rs index 6d7988c..98bfd55 100644 --- a/src/admin/types.rs +++ b/src/admin/types.rs @@ -324,6 +324,24 @@ pub(super) struct UpdatePluginSecretsBody { pub(super) secrets: std::collections::HashMap, } +// --------------------------------------------------------------------------- +// Domain types +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +pub(super) struct DomainResponse { + pub(super) id: String, + pub(super) url: String, + pub(super) is_primary: bool, + pub(super) created_at: String, + pub(super) updated_at: String, +} + +#[derive(Deserialize)] +pub(super) struct CreateDomainBody { + pub(super) url: String, +} + // --------------------------------------------------------------------------- // API client types // --------------------------------------------------------------------------- @@ -380,31 +398,3 @@ pub(super) struct CreateApiClientResponse { pub(super) name: String, pub(super) client_id_url: String, } - -// --------------------------------------------------------------------------- -// Rate limit types -// --------------------------------------------------------------------------- - -#[derive(Deserialize)] -pub(super) struct UpsertRateLimitBody { - pub(super) capacity: u32, - pub(super) refill_rate: f64, - pub(super) default_query_cost: u32, - pub(super) default_procedure_cost: u32, - pub(super) default_proxy_cost: u32, -} - -#[derive(Deserialize)] -pub(super) struct SetEnabledBody { - pub(super) enabled: bool, -} - -#[derive(Serialize)] -pub(super) struct RateLimitsResponse { - pub(super) enabled: bool, - pub(super) capacity: i32, - pub(super) refill_rate: f64, - pub(super) default_query_cost: i32, - pub(super) default_procedure_cost: i32, - pub(super) default_proxy_cost: i32, -} diff --git a/src/auth/client_registry.rs b/src/auth/client_registry.rs index e08f6e4..c46a9d0 100644 --- a/src/auth/client_registry.rs +++ b/src/auth/client_registry.rs @@ -1,3 +1,4 @@ +use arc_swap::ArcSwap; use dashmap::DashMap; use std::sync::Arc; @@ -27,14 +28,16 @@ pub struct ApiClientOAuthParams { /// shows the correct domain. The default client is HappyView's own identity, /// used for dashboard auth. pub struct OAuthClientRegistry { - default_client: Arc, + primary_client: ArcSwap, + domain_clients: DashMap>, clients: DashMap>, } impl OAuthClientRegistry { - pub fn new(default_client: Arc) -> Self { + pub fn new(primary_client: Arc) -> Self { Self { - default_client, + primary_client: ArcSwap::new(primary_client), + domain_clients: DashMap::new(), clients: DashMap::new(), } } @@ -54,21 +57,63 @@ impl OAuthClientRegistry { self.clients.get(client_id_url).map(|r| r.value().clone()) } - /// Look up a client by `client_id_url`, falling back to the default. + /// Look up a client by `client_id_url`, falling back to the primary client. pub fn get_or_default(&self, client_id_url: Option<&str>) -> Arc { if let Some(url) = client_id_url { self.clients .get(url) .map(|r| r.value().clone()) - .unwrap_or_else(|| self.default_client.clone()) + .unwrap_or_else(|| self.primary_client.load_full()) } else { - self.default_client.clone() + self.primary_client.load_full() } } - /// Get the default (HappyView dashboard) client. - pub fn default_client(&self) -> &Arc { - &self.default_client + /// Get the primary (HappyView dashboard) client. + pub fn primary_client(&self) -> Arc { + self.primary_client.load_full() + } + + /// Register a domain-specific OAuth client. + /// Inserts into both `domain_clients` (keyed by domain URL, for `get_for_domain`) + /// and `clients` (keyed by client_id_url, for `get_or_default`). + pub fn register_domain_client(&self, domain_url: String, client: Arc) { + let client_id_url = format!( + "{}/oauth-client-metadata.json", + domain_url.trim_end_matches('/') + ); + self.domain_clients.insert(domain_url, Arc::clone(&client)); + self.clients.insert(client_id_url, client); + } + + /// Remove a domain-specific OAuth client from both maps. + pub fn remove_domain_client(&self, domain_url: &str) { + self.domain_clients.remove(domain_url); + let client_id_url = format!( + "{}/oauth-client-metadata.json", + domain_url.trim_end_matches('/') + ); + self.clients.remove(&client_id_url); + } + + /// Look up a domain-specific OAuth client. + pub fn get_domain_client(&self, domain_url: &str) -> Option> { + self.domain_clients + .get(domain_url) + .map(|r| r.value().clone()) + } + + /// Get the OAuth client for a domain, falling back to the primary client. + pub fn get_for_domain(&self, domain_url: &str) -> Arc { + self.domain_clients + .get(domain_url) + .map(|r| r.value().clone()) + .unwrap_or_else(|| self.primary_client.load_full()) + } + + /// Replace the primary OAuth client (e.g. when admin changes the primary domain). + pub fn set_primary_client(&self, client: Arc) { + self.primary_client.store(client); } /// Build and register a single OAuth client from API client metadata. diff --git a/src/auth/routes.rs b/src/auth/routes.rs index f6c0281..3bafaa6 100644 --- a/src/auth/routes.rs +++ b/src/auth/routes.rs @@ -57,6 +57,7 @@ pub fn routes() -> Router { async fn login( State(state): State, jar: SignedCookieJar, + domain: Option>>, Query(query): Query, ) -> Result<(SignedCookieJar, Json), AppError> { tracing::debug!(handle = %query.handle, redirect_uri = ?query.redirect_uri, scope = ?query.scope, "login request"); @@ -77,8 +78,18 @@ async fn login( tracing::debug!(scopes = ?scopes, client_id = ?query.client_id, "resolved oauth scopes"); + // For dashboard logins (no explicit client_id), use the domain's OAuth client + let domain_url = domain.map(|d| d.0.url.clone()); + let effective_client_id = if query.client_id.is_some() { + query.client_id.clone() + } else { + domain_url + .as_ref() + .map(|du| format!("{}/oauth-client-metadata.json", du.trim_end_matches('/'))) + }; + // Select the appropriate OAuth client based on client_id - let oauth_client = state.oauth.get_or_default(query.client_id.as_deref()); + let oauth_client = state.oauth.get_or_default(effective_client_id.as_deref()); // Hold the authorize lock so that authorize() + take_last_state_key() are atomic. // This prevents concurrent logins from swapping each other's state keys. @@ -106,9 +117,9 @@ async fn login( // Store the redirect URI in the database, keyed by the OAuth state parameter. // This avoids third-party cookie issues when Pentaract (cross-origin) calls this endpoint. // Store redirect URI and client_id for the callback to use - if query.redirect_uri.is_some() || query.client_id.is_some() { + if query.redirect_uri.is_some() || effective_client_id.is_some() { let redirect_uri = query.redirect_uri.as_deref().unwrap_or(""); - tracing::debug!(oauth_state = ?oauth_state, redirect_uri = %redirect_uri, client_id = ?query.client_id, "storing redirect for state"); + tracing::debug!(oauth_state = ?oauth_state, redirect_uri = %redirect_uri, client_id = ?effective_client_id, "storing redirect for state"); if let Some(oauth_state) = oauth_state { let now = now_rfc3339(); @@ -120,7 +131,7 @@ async fn login( let _ = sqlx::query(&sql) .bind(&oauth_state) .bind(redirect_uri) - .bind(query.client_id.as_deref()) + .bind(effective_client_id.as_deref()) .bind(&now) .bind(&expires_at) .execute(&state.db) @@ -257,7 +268,7 @@ async fn logout( let raw = cookie.value().to_string(); let did_str = raw.split('\n').next().unwrap_or(&raw).to_string(); if let Ok(did) = atrium_api::types::string::Did::new(did_str) { - let _ = state.oauth.default_client().revoke(&did).await; + let _ = state.oauth.primary_client().revoke(&did).await; } } diff --git a/src/domain.rs b/src/domain.rs new file mode 100644 index 0000000..4bfe539 --- /dev/null +++ b/src/domain.rs @@ -0,0 +1,198 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::RwLock; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Domain { + pub id: String, + pub url: String, + pub is_primary: bool, + pub created_at: String, + pub updated_at: String, +} + +impl Domain { + pub fn host(&self) -> &str { + let after_scheme = self + .url + .strip_prefix("https://") + .or_else(|| self.url.strip_prefix("http://")) + .unwrap_or(&self.url); + after_scheme.split(':').next().unwrap_or(after_scheme) + } +} + +#[derive(Clone)] +pub struct DomainCache { + by_host: Arc>>>, + primary: Arc>>>, +} + +impl DomainCache { + pub fn new() -> Self { + Self { + by_host: Arc::new(RwLock::new(HashMap::new())), + primary: Arc::new(RwLock::new(None)), + } + } + + pub async fn load(&self, domains: Vec) { + let mut by_host = self.by_host.write().await; + let mut primary = self.primary.write().await; + + by_host.clear(); + *primary = None; + + for domain in domains { + let arc = Arc::new(domain); + if arc.is_primary { + *primary = Some(arc.clone()); + } + by_host.insert(arc.host().to_string(), arc); + } + } + + pub async fn get(&self, host: &str) -> Option> { + let by_host = self.by_host.read().await; + by_host.get(host).cloned() + } + + pub async fn primary(&self) -> Option> { + let primary = self.primary.read().await; + primary.clone() + } + + pub async fn insert(&self, domain: Domain) { + let arc = Arc::new(domain); + let mut by_host = self.by_host.write().await; + let mut primary = self.primary.write().await; + + if arc.is_primary { + *primary = Some(arc.clone()); + } + by_host.insert(arc.host().to_string(), arc); + } + + pub async fn remove(&self, host: &str) { + let mut by_host = self.by_host.write().await; + let removed = by_host.remove(host); + + if let Some(domain) = removed + && domain.is_primary + { + let mut primary = self.primary.write().await; + *primary = None; + } + } + + pub async fn set_primary(&self, host: &str) { + let by_host = self.by_host.read().await; + if let Some(domain) = by_host.get(host).cloned() { + drop(by_host); + let mut primary = self.primary.write().await; + *primary = Some(domain); + } + } + + pub async fn all(&self) -> Vec> { + let by_host = self.by_host.read().await; + by_host.values().cloned().collect() + } +} + +impl Default for DomainCache { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use uuid::Uuid; + + fn make_domain(url: &str, is_primary: bool) -> Domain { + Domain { + id: Uuid::new_v4().to_string(), + url: url.to_string(), + is_primary, + created_at: "2024-01-01T00:00:00Z".to_string(), + updated_at: "2024-01-01T00:00:00Z".to_string(), + } + } + + #[test] + fn host_strips_https() { + let domain = make_domain("https://example.com", false); + assert_eq!(domain.host(), "example.com"); + } + + #[test] + fn host_strips_http() { + let domain = make_domain("http://localhost:3000", false); + assert_eq!(domain.host(), "localhost"); + } + + #[tokio::test] + async fn load_and_get() { + let cache = DomainCache::new(); + let domains = vec![ + make_domain("https://example.com", true), + make_domain("https://other.com", false), + ]; + cache.load(domains).await; + + let found = cache.get("example.com").await; + assert!(found.is_some()); + assert_eq!(found.unwrap().url, "https://example.com"); + + assert!(cache.get("other.com").await.is_some()); + + let missing = cache.get("unknown.com").await; + assert!(missing.is_none()); + } + + #[tokio::test] + async fn primary_returns_primary_domain() { + let cache = DomainCache::new(); + let domains = vec![ + make_domain("https://example.com", false), + make_domain("https://primary.com", true), + ]; + cache.load(domains).await; + + let primary = cache.primary().await; + assert!(primary.is_some()); + assert_eq!(primary.unwrap().url, "https://primary.com"); + } + + #[tokio::test] + async fn insert_and_remove() { + let cache = DomainCache::new(); + let domain = make_domain("https://example.com", false); + cache.insert(domain).await; + + assert!(cache.get("example.com").await.is_some()); + + cache.remove("example.com").await; + assert!(cache.get("example.com").await.is_none()); + } + + #[tokio::test] + async fn set_primary_updates() { + let cache = DomainCache::new(); + let domains = vec![ + make_domain("https://example.com", true), + make_domain("https://other.com", false), + ]; + cache.load(domains).await; + + // Initially example.com is primary + assert_eq!(cache.primary().await.unwrap().url, "https://example.com"); + + // Change primary to other.com + cache.set_primary("other.com").await; + assert_eq!(cache.primary().await.unwrap().url, "https://other.com"); + } +} diff --git a/src/domain_middleware.rs b/src/domain_middleware.rs new file mode 100644 index 0000000..ad4422b --- /dev/null +++ b/src/domain_middleware.rs @@ -0,0 +1,38 @@ +use axum::{ + extract::{Request, State}, + http::StatusCode, + middleware::Next, + response::Response, +}; +use std::sync::Arc; + +use crate::AppState; +use crate::domain::Domain; + +pub async fn resolve_domain( + State(state): State, + mut req: Request, + next: Next, +) -> Result { + let host = req + .headers() + .get("x-forwarded-host") + .or_else(|| req.headers().get("host")) + .and_then(|v| v.to_str().ok()) + .map(|h| h.split(':').next().unwrap_or(h)) + .unwrap_or(""); + + let domain = state.domain_cache.get(host).await; + + match domain { + Some(domain) => { + req.extensions_mut().insert(domain); + Ok(next.run(req).await) + } + None => Err((StatusCode::MISDIRECTED_REQUEST, "Unknown host")), + } +} + +pub fn extract_domain(req: &Request) -> Option> { + req.extensions().get::>().cloned() +} diff --git a/src/external_auth/routes.rs b/src/external_auth/routes.rs index 833c79f..d5d9fcd 100644 --- a/src/external_auth/routes.rs +++ b/src/external_auth/routes.rs @@ -75,6 +75,7 @@ async fn authorize( State(app_state): State, Path(plugin_id): Path, Query(query): Query, + domain: Option>>, claims: Claims, ) -> Result, AppError> { let _plugin = app_state @@ -128,9 +129,13 @@ async fn authorize( // Build the backend callback URL for OpenID/OAuth return_to // This ensures the auth provider redirects back to the backend, not the frontend + let domain_url = domain + .map(|d| d.0.url.clone()) + .unwrap_or_else(|| app_state.config.public_url.clone()); + let callback_url = format!( "{}/external-auth/{}/callback", - app_state.config.public_url.trim_end_matches('/'), + domain_url.trim_end_matches('/'), plugin_id ); diff --git a/src/lib.rs b/src/lib.rs index 5b4baa5..9cd6c8d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,8 @@ pub mod auth; pub mod config; pub mod db; pub mod dns; +pub mod domain; +pub mod domain_middleware; pub mod error; pub mod event_log; pub mod external_auth; @@ -54,6 +56,7 @@ pub struct AppState { pub http: reqwest::Client, pub db: sqlx::AnyPool, pub db_backend: DatabaseBackend, + pub domain_cache: domain::DomainCache, pub lexicons: LexiconRegistry, pub collections_tx: watch::Sender>, pub labeler_subscriptions_tx: watch::Sender<()>, diff --git a/src/lua/atproto_api.rs b/src/lua/atproto_api.rs index e0f2d99..c536db1 100644 --- a/src/lua/atproto_api.rs +++ b/src/lua/atproto_api.rs @@ -336,17 +336,15 @@ mod tests { http: reqwest::Client::new(), db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, + domain_cache: crate::domain::DomainCache::new(), 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + crate::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(crate::auth::OAuthClientRegistry::new(std::sync::Arc::new( diff --git a/src/lua/db_api.rs b/src/lua/db_api.rs index e2b407b..30c01e8 100644 --- a/src/lua/db_api.rs +++ b/src/lua/db_api.rs @@ -687,17 +687,15 @@ mod tests { http: reqwest::Client::new(), db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, + domain_cache: crate::domain::DomainCache::new(), 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + crate::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(crate::auth::OAuthClientRegistry::new(std::sync::Arc::new( diff --git a/src/lua/execute.rs b/src/lua/execute.rs index 0f8c01c..973e17d 100644 --- a/src/lua/execute.rs +++ b/src/lua/execute.rs @@ -1068,17 +1068,15 @@ mod tests { http: reqwest::Client::new(), db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, + domain_cache: crate::domain::DomainCache::new(), 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + crate::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(crate::auth::OAuthClientRegistry::new(std::sync::Arc::new( diff --git a/src/lua/http_api.rs b/src/lua/http_api.rs index 653ab85..d6356d1 100644 --- a/src/lua/http_api.rs +++ b/src/lua/http_api.rs @@ -153,17 +153,15 @@ mod tests { http: reqwest::Client::new(), db: test_db.clone(), db_backend: crate::db::DatabaseBackend::Sqlite, + domain_cache: crate::domain::DomainCache::new(), 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + crate::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(crate::auth::OAuthClientRegistry::new(std::sync::Arc::new( diff --git a/src/lua/xrpc_api.rs b/src/lua/xrpc_api.rs index b94a515..64958c8 100644 --- a/src/lua/xrpc_api.rs +++ b/src/lua/xrpc_api.rs @@ -257,17 +257,15 @@ mod tests { http: reqwest::Client::new(), db: test_db.clone(), db_backend: DatabaseBackend::Sqlite, + domain_cache: crate::domain::DomainCache::new(), 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + crate::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(crate::auth::OAuthClientRegistry::new(std::sync::Arc::new( diff --git a/src/main.rs b/src/main.rs index e39160d..092453f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,9 +5,10 @@ use happyview::config::Config; use happyview::db; use happyview::dns::NativeDnsResolver; use happyview::lexicon::{LexiconRegistry, ParsedLexicon, ProcedureAction}; -use happyview::rate_limit::{RateLimitConfig, RateLimiter}; +use happyview::rate_limit::{RateLimitDefaults, RateLimiter}; use happyview::resolve::{fetch_lexicon_from_pds, resolve_nsid_authority}; use happyview::{AppState, jetstream, labeler, server}; +use sqlx::Row; use tokio::sync::watch; use tracing::{info, warn}; @@ -264,9 +265,9 @@ async fn main() { } } - // Initialize rate limiter from DB. - let rl_state = RateLimiter::load_from_db(&db_pool).await; - let rate_limiter = RateLimiter::new(rl_state.enabled, rl_state.global); + // Seed and load per-instance default token costs from instance_settings. + let defaults = seed_and_load_rate_limit_defaults(&db_pool, db_backend).await; + let rate_limiter = RateLimiter::new(defaults); tokio::spawn(rate_limiter.clone().spawn_cleanup()); // Load per-client rate limit configs and identities from api_clients table. @@ -279,7 +280,6 @@ async fn main() { .await .unwrap_or_default(); - let global = rate_limiter.global_config(); for (client_key, secret_hash, client_uri, capacity, refill_rate) in client_rows { rate_limiter.register_client_identity( client_key.clone(), @@ -291,18 +291,74 @@ async fn main() { if let (Some(cap), Some(refill)) = (capacity, refill_rate) { rate_limiter.register_client_config( client_key, - RateLimitConfig { + happyview::rate_limit::RateLimitConfig { capacity: cap as u32, refill_rate: refill, - default_query_cost: global.default_query_cost, - default_procedure_cost: global.default_procedure_cost, - default_proxy_cost: global.default_proxy_cost, + default_query_cost: defaults.query_cost, + default_procedure_cost: defaults.procedure_cost, + default_proxy_cost: defaults.proxy_cost, }, ); } } } + // Seed and load domain cache + let domain_cache = happyview::domain::DomainCache::new(); + { + let count_sql = happyview::db::adapt_sql("SELECT COUNT(*) FROM domains", db_backend); + let row = sqlx::query(&count_sql) + .fetch_one(&db_pool) + .await + .expect("Failed to count domains"); + let count: i64 = row.try_get(0).unwrap_or(0); + + if count == 0 { + let id = uuid::Uuid::new_v4().to_string(); + let now = happyview::db::now_rfc3339(); + let insert_sql = happyview::db::adapt_sql( + "INSERT INTO domains (id, url, is_primary, created_at, updated_at) VALUES (?, ?, 1, ?, ?)", + db_backend, + ); + sqlx::query(&insert_sql) + .bind(&id) + .bind(&config.public_url) + .bind(&now) + .bind(&now) + .execute(&db_pool) + .await + .expect("Failed to insert primary domain"); + info!("Seeded primary domain: {}", config.public_url); + } + + let select_sql = happyview::db::adapt_sql( + "SELECT id, url, is_primary, created_at, updated_at FROM domains", + db_backend, + ); + let rows = sqlx::query(&select_sql) + .fetch_all(&db_pool) + .await + .expect("Failed to load domains"); + + let domains: Vec = rows + .into_iter() + .map(|row| { + let is_primary_int: i32 = row.try_get("is_primary").unwrap_or(0); + happyview::domain::Domain { + id: row.try_get("id").unwrap_or_default(), + url: row.try_get("url").unwrap_or_default(), + is_primary: is_primary_int != 0, + created_at: row.try_get("created_at").unwrap_or_default(), + updated_at: row.try_get("updated_at").unwrap_or_default(), + } + }) + .collect(); + + let domain_count = domains.len(); + domain_cache.load(domains).await; + info!("Loaded {} domain(s) into cache", domain_count); + } + // Build atrium-oauth client let dns = NativeDnsResolver::new(); let callback_url = format!("{}/auth/callback", config.public_url.trim_end_matches('/')); @@ -394,8 +450,9 @@ async fn main() { let (labeler_subscriptions_tx, labeler_subscriptions_rx) = watch::channel(()); // Build the OAuth client registry and load API clients from DB - let oauth_registry = Arc::new(happyview::auth::OAuthClientRegistry::new(Arc::new( - oauth_client, + let oauth_client_arc = Arc::new(oauth_client); + let oauth_registry = Arc::new(happyview::auth::OAuthClientRegistry::new(Arc::clone( + &oauth_client_arc, ))); oauth_registry .load_from_db( @@ -407,6 +464,64 @@ async fn main() { ) .await; + // Register the primary domain's OAuth client in domain_clients + if let Some(ref pd) = domain_cache.primary().await { + oauth_registry.register_domain_client(pd.url.clone(), Arc::clone(&oauth_client_arc)); + } + + // Build OAuth clients for all non-primary domains + let all_domains = domain_cache.all().await; + for domain in &all_domains { + if domain.is_primary { + continue; // Already registered above + } + + let domain_callback_url = format!("{}/auth/callback", domain.url.trim_end_matches('/')); + let domain_client_id = format!( + "{}/oauth-client-metadata.json", + domain.url.trim_end_matches('/') + ); + + let domain_http = Arc::new(DefaultHttpClient::default()); + let domain_resolver = OAuthResolverConfig { + did_resolver: CommonDidResolver::new(CommonDidResolverConfig { + plc_directory_url: config.plc_url.clone(), + http_client: Arc::clone(&domain_http), + }), + handle_resolver: AtprotoHandleResolver::new(AtprotoHandleResolverConfig { + dns_txt_resolver: NativeDnsResolver::new(), + http_client: Arc::clone(&domain_http), + }), + authorization_server_metadata: Default::default(), + protected_resource_metadata: Default::default(), + }; + + match atrium_oauth::OAuthClient::new(OAuthClientConfig { + client_metadata: AtprotoClientMetadata { + client_id: domain_client_id, + client_uri: Some(domain.url.clone()), + redirect_uris: vec![domain_callback_url], + token_endpoint_auth_method: AuthMethod::None, + grant_types: vec![GrantType::AuthorizationCode, GrantType::RefreshToken], + scopes: vec![Scope::Known(KnownScope::Atproto)], + jwks_uri: None, + token_endpoint_auth_signing_alg: None, + }, + keys: None, + state_store: oauth_state_store.clone(), + session_store: DbSessionStore::new(db_pool.clone(), db_backend), + resolver: domain_resolver, + }) { + Ok(client) => { + info!(domain = %domain.url, "Registered domain OAuth client"); + oauth_registry.register_domain_client(domain.url.clone(), Arc::new(client)); + } + Err(e) => { + tracing::error!(domain = %domain.url, error = %e, "Failed to create domain OAuth client"); + } + } + } + let official_registry: happyview::plugin::official_registry::SharedRegistry = std::sync::Arc::new(tokio::sync::RwLock::new( happyview::plugin::official_registry::OfficialRegistryState::default(), @@ -424,6 +539,7 @@ async fn main() { http, db: db_pool, db_backend, + domain_cache: domain_cache.clone(), lexicons, collections_tx, labeler_subscriptions_tx, @@ -460,3 +576,68 @@ async fn main() { axum::serve(listener, app).await.expect("server error"); } + +async fn seed_and_load_rate_limit_defaults( + pool: &sqlx::AnyPool, + backend: happyview::db::DatabaseBackend, +) -> RateLimitDefaults { + use happyview::rate_limit::{ + SEED_DEFAULT_PROCEDURE_COST, SEED_DEFAULT_PROXY_COST, SEED_DEFAULT_QUERY_COST, + SETTING_DEFAULT_PROCEDURE_COST, SETTING_DEFAULT_PROXY_COST, SETTING_DEFAULT_QUERY_COST, + }; + + async fn seed_and_read( + pool: &sqlx::AnyPool, + backend: happyview::db::DatabaseBackend, + key: &str, + seed: u32, + ) -> u32 { + if happyview::admin::settings::get_setting(pool, key, backend) + .await + .is_none() + { + let now = happyview::db::now_rfc3339(); + let sql = happyview::db::adapt_sql( + "INSERT INTO instance_settings (key, value, updated_at) VALUES (?, ?, ?) ON CONFLICT (key) DO NOTHING", + backend, + ); + if let Err(e) = sqlx::query(&sql) + .bind(key) + .bind(seed.to_string()) + .bind(&now) + .execute(pool) + .await + { + warn!(error = %e, key = key, "failed to seed rate-limit default"); + } + } + happyview::admin::settings::get_setting(pool, key, backend) + .await + .and_then(|s| s.parse::().ok()) + .unwrap_or(seed) + } + + RateLimitDefaults { + query_cost: seed_and_read( + pool, + backend, + SETTING_DEFAULT_QUERY_COST, + SEED_DEFAULT_QUERY_COST, + ) + .await, + procedure_cost: seed_and_read( + pool, + backend, + SETTING_DEFAULT_PROCEDURE_COST, + SEED_DEFAULT_PROCEDURE_COST, + ) + .await, + proxy_cost: seed_and_read( + pool, + backend, + SETTING_DEFAULT_PROXY_COST, + SEED_DEFAULT_PROXY_COST, + ) + .await, + } +} diff --git a/src/rate_limit.rs b/src/rate_limit.rs index eb3453b..f89df59 100644 --- a/src/rate_limit.rs +++ b/src/rate_limit.rs @@ -1,10 +1,29 @@ -use arc_swap::ArcSwap; use dashmap::DashMap; -use sqlx::AnyPool; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Instant, SystemTime, UNIX_EPOCH}; +/// Hardcoded seed values for the per-instance default token costs. These are +/// only used by the startup seeding step in `main.rs` to populate fresh +/// `instance_settings` rows; at runtime the values are read from the DB into +/// `RateLimitDefaults`. +pub const SEED_DEFAULT_QUERY_COST: u32 = 1; +pub const SEED_DEFAULT_PROCEDURE_COST: u32 = 1; +pub const SEED_DEFAULT_PROXY_COST: u32 = 1; + +/// `instance_settings` keys for the seeded defaults. +pub const SETTING_DEFAULT_QUERY_COST: &str = "rate_limit.default_query_cost"; +pub const SETTING_DEFAULT_PROCEDURE_COST: &str = "rate_limit.default_procedure_cost"; +pub const SETTING_DEFAULT_PROXY_COST: &str = "rate_limit.default_proxy_cost"; + +/// Default token costs per XRPC request type, loaded from `instance_settings` +/// at startup. Owned by the `RateLimiter`. +#[derive(Clone, Copy)] +pub struct RateLimitDefaults { + pub query_cost: u32, + pub procedure_cost: u32, + pub proxy_cost: u32, +} + pub struct RateLimitConfig { pub capacity: u32, pub refill_rate: f64, @@ -44,20 +63,16 @@ pub struct ClientIdentity { } pub struct RateLimiter { - enabled: AtomicBool, + defaults: RateLimitDefaults, buckets: DashMap, - global_config: ArcSwap, - /// Per-client config overrides, keyed by client_key (e.g. "hvc_...") + /// Per-client config, keyed by client_key (e.g. "hvc_..."). Presence in + /// this map is the *only* thing that enables rate limiting for a key — + /// unregistered keys are always allowed. client_configs: DashMap, /// Registered client identities, keyed by client_key client_identities: DashMap, } -pub struct RateLimiterState { - pub enabled: bool, - pub global: RateLimitConfig, -} - fn now_unix() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -66,27 +81,23 @@ fn now_unix() -> u64 { } impl RateLimiter { - pub fn new(enabled: bool, global: RateLimitConfig) -> Arc { + pub fn new(defaults: RateLimitDefaults) -> Arc { Arc::new(Self { - enabled: AtomicBool::new(enabled), + defaults, buckets: DashMap::new(), - global_config: ArcSwap::new(Arc::new(global)), client_configs: DashMap::new(), client_identities: DashMap::new(), }) } - pub fn check(&self, key: &str, cost: u32) -> CheckResult { - if !self.enabled.load(Ordering::Relaxed) { - return CheckResult::Disabled; - } + pub fn defaults(&self) -> RateLimitDefaults { + self.defaults + } - // Use per-client config if available, otherwise fall back to global - let (capacity, refill_rate) = if let Some(client_cfg) = self.client_configs.get(key) { - (client_cfg.capacity, client_cfg.refill_rate) - } else { - let global = self.global_config.load(); - (global.capacity, global.refill_rate) + pub fn check(&self, key: &str, cost: u32) -> CheckResult { + let (capacity, refill_rate) = match self.client_configs.get(key) { + Some(cfg) => (cfg.capacity, cfg.refill_rate), + None => return CheckResult::Disabled, }; let cost_f64 = cost as f64; @@ -103,11 +114,9 @@ impl RateLimiter { 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; @@ -137,56 +146,42 @@ impl RateLimiter { } } - /// Get the default cost for a given request type. - pub fn default_cost_for_type(&self, request_type: &str) -> u32 { - let config = self.global_config.load(); + /// Get the default cost for a request type. Looks up the per-client + /// override if one is registered, otherwise falls back to the seeded + /// instance defaults. + pub fn default_cost_for_type(&self, client_key: &str, request_type: &str) -> u32 { + if let Some(cfg) = self.client_configs.get(client_key) { + return match request_type { + "query" => cfg.default_query_cost, + "procedure" => cfg.default_procedure_cost, + "proxy" => cfg.default_proxy_cost, + _ => 1, + }; + } match request_type { - "query" => config.default_query_cost, - "procedure" => config.default_procedure_cost, - "proxy" => config.default_proxy_cost, + "query" => self.defaults.query_cost, + "procedure" => self.defaults.procedure_cost, + "proxy" => self.defaults.proxy_cost, _ => 1, } } - /// Get a snapshot of the current global config. - pub fn global_config(&self) -> Arc { - self.global_config.load_full() - } - - 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) { - self.global_config.store(Arc::new(global)); - } - - /// Register a per-client rate limit config override. pub fn register_client_config(&self, client_key: String, config: RateLimitConfig) { self.client_configs.insert(client_key, config); } - /// Remove a per-client rate limit config override. pub fn remove_client_config(&self, client_key: &str) { self.client_configs.remove(client_key); } - /// Register a client identity (key, secret hash, and client URI). pub fn register_client_identity(&self, client_key: String, identity: ClientIdentity) { self.client_identities.insert(client_key, identity); } - /// Remove a client identity. pub fn remove_client_identity(&self, client_key: &str) { self.client_identities.remove(client_key); } - /// Validate a client key + secret combination. Returns true if the secret - /// hash matches the stored hash for this client key. pub fn validate_client_secret(&self, client_key: &str, secret: &str) -> bool { use sha2::{Digest, Sha256}; if let Some(identity) = self.client_identities.get(client_key) { @@ -197,11 +192,8 @@ impl RateLimiter { } } - /// Validate a client key + origin combination. Returns true if the origin - /// matches the registered client_uri for this client key. pub fn validate_client_origin(&self, client_key: &str, origin: &str) -> bool { if let Some(identity) = self.client_identities.get(client_key) { - // Compare origins: strip trailing slash for consistency let registered = identity.client_uri.trim_end_matches('/'); let provided = origin.trim_end_matches('/'); registered == provided @@ -210,14 +202,13 @@ impl RateLimiter { } } - /// Check whether a client key is registered. pub fn is_valid_client_key(&self, client_key: &str) -> bool { self.client_identities.contains_key(client_key) } 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 + let stale_threshold = std::time::Duration::from_secs(300); loop { tokio::time::sleep(interval).await; let now = Instant::now(); @@ -225,245 +216,84 @@ impl RateLimiter { .retain(|_, bucket| now.duration_since(bucket.last_access) < stale_threshold); } } - - pub async fn load_from_db(db: &AnyPool) -> 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 global rate limit config (method IS NULL row) - let row: Option<(i32, f64, i32, i32, i32)> = sqlx::query_as( - "SELECT capacity, refill_rate, default_query_cost, default_procedure_cost, default_proxy_cost FROM rate_limits WHERE method IS NULL", - ) - .fetch_optional(db) - .await - .unwrap_or(None); - - let global = match row { - Some((capacity, refill_rate, query_cost, procedure_cost, proxy_cost)) => { - RateLimitConfig { - capacity: capacity as u32, - refill_rate, - default_query_cost: query_cost as u32, - default_procedure_cost: procedure_cost as u32, - default_proxy_cost: proxy_cost as u32, - } - } - None => RateLimitConfig { - capacity: 100, - refill_rate: 2.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - }; - - RateLimiterState { enabled, global } - } - - /// Reload all config from DB and apply to the live limiter. - pub async fn reload_from_db(&self, db: &AnyPool) { - let state = Self::load_from_db(db).await; - self.set_enabled(state.enabled); - self.update_config(state.global); - } } #[cfg(test)] mod tests { use super::*; - #[test] - fn basic_allow_and_exhaust() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 3, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - // Should allow 3 requests (bucket starts full, cost=1 each) - for _ in 0..3 { - assert!(matches!(rl.check("k", 1), CheckResult::Allowed { .. })); + fn defaults() -> RateLimitDefaults { + RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, } - // 4th should be limited - assert!(matches!(rl.check("k", 1), CheckResult::Limited { .. })); } - #[test] - fn cost_deducts_multiple_tokens() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 10, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - // Cost of 5 should allow 2 requests (10 tokens total) - assert!(matches!( - rl.check("k", 5), - CheckResult::Allowed { remaining: 5, .. } - )); - assert!(matches!( - rl.check("k", 5), - CheckResult::Allowed { remaining: 0, .. } - )); - // 3rd should be limited - assert!(matches!(rl.check("k", 5), CheckResult::Limited { .. })); + fn cfg(capacity: u32, refill_rate: f64) -> RateLimitConfig { + RateLimitConfig { + capacity, + refill_rate, + default_query_cost: 1, + default_procedure_cost: 1, + default_proxy_cost: 1, + } } #[test] - fn disabled_returns_disabled() { - let rl = RateLimiter::new( - false, - RateLimitConfig { - capacity: 1, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - assert!(matches!(rl.check("k", 1), CheckResult::Disabled)); + fn unregistered_key_is_not_rate_limited() { + let rl = RateLimiter::new(defaults()); + for _ in 0..1000 { + assert!(matches!(rl.check("anything", 1), CheckResult::Disabled)); + } } #[test] - fn default_cost_for_type() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 100, - refill_rate: 10.0, - default_query_cost: 2, - default_procedure_cost: 5, - default_proxy_cost: 3, - }, - ); + fn registered_client_is_rate_limited() { + let rl = RateLimiter::new(defaults()); + rl.register_client_config("hvc_a".to_string(), cfg(3, 0.001)); - assert_eq!(rl.default_cost_for_type("query"), 2); - assert_eq!(rl.default_cost_for_type("procedure"), 5); - assert_eq!(rl.default_cost_for_type("proxy"), 3); - assert_eq!(rl.default_cost_for_type("unknown"), 1); + for _ in 0..3 { + assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); + } + assert!(matches!(rl.check("hvc_a", 1), CheckResult::Limited { .. })); } #[test] - fn per_client_config_override() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 10, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - // Register a client with lower capacity - rl.register_client_config( - "hvc_client1".to_string(), - RateLimitConfig { - capacity: 2, - refill_rate: 0.001, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); + fn cost_deducts_multiple_tokens() { + let rl = RateLimiter::new(defaults()); + rl.register_client_config("hvc_a".to_string(), cfg(10, 0.001)); - // Client key should use client config (capacity=2) - assert!(matches!( - rl.check("hvc_client1", 1), - CheckResult::Allowed { .. } - )); assert!(matches!( - rl.check("hvc_client1", 1), - CheckResult::Allowed { .. } - )); - assert!(matches!( - rl.check("hvc_client1", 1), - CheckResult::Limited { .. } + rl.check("hvc_a", 5), + CheckResult::Allowed { remaining: 5, .. } )); - - // Other key should use global config (capacity=10) - for _ in 0..10 { - assert!(matches!( - rl.check("other_key", 1), - CheckResult::Allowed { .. } - )); - } assert!(matches!( - rl.check("other_key", 1), - CheckResult::Limited { .. } + rl.check("hvc_a", 5), + CheckResult::Allowed { remaining: 0, .. } )); + assert!(matches!(rl.check("hvc_a", 5), CheckResult::Limited { .. })); } #[test] - fn per_client_config_fallback_to_global() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 3, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); + fn different_clients_get_separate_buckets() { + let rl = RateLimiter::new(defaults()); + rl.register_client_config("hvc_a".to_string(), cfg(2, 0.001)); + rl.register_client_config("hvc_b".to_string(), cfg(2, 0.001)); - // No client config registered — should use global (capacity=3) - for _ in 0..3 { - assert!(matches!( - rl.check("hvc_unregistered", 1), - CheckResult::Allowed { .. } - )); - } - assert!(matches!( - rl.check("hvc_unregistered", 1), - CheckResult::Limited { .. } - )); + assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); + assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); + assert!(matches!(rl.check("hvc_a", 1), CheckResult::Limited { .. })); + + assert!(matches!(rl.check("hvc_b", 1), CheckResult::Allowed { .. })); + assert!(matches!(rl.check("hvc_b", 1), CheckResult::Allowed { .. })); + assert!(matches!(rl.check("hvc_b", 1), CheckResult::Limited { .. })); } #[test] - fn register_and_remove_client_config() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 10, - refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - rl.register_client_config( - "hvc_temp".to_string(), - RateLimitConfig { - capacity: 1, - refill_rate: 0.001, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - // Should be limited after 1 request (client config capacity=1) + fn remove_client_config_disables_limiting() { + let rl = RateLimiter::new(defaults()); + rl.register_client_config("hvc_temp".to_string(), cfg(1, 0.001)); assert!(matches!( rl.check("hvc_temp", 1), CheckResult::Allowed { .. } @@ -473,69 +303,41 @@ mod tests { CheckResult::Limited { .. } )); - // Remove client config — new bucket should use global (capacity=10) rl.remove_client_config("hvc_temp"); - // Note: the old bucket still exists and is exhausted, but capacity was - // updated to global. A new bucket would get global capacity. + assert!(matches!(rl.check("hvc_temp", 1), CheckResult::Disabled)); } #[test] - fn different_clients_get_separate_buckets() { - let rl = RateLimiter::new( - true, - RateLimitConfig { - capacity: 2, - refill_rate: 0.001, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, - }, - ); - - // Exhaust client A - assert!(matches!( - rl.check("clientA", 1), - CheckResult::Allowed { .. } - )); - assert!(matches!( - rl.check("clientA", 1), - CheckResult::Allowed { .. } - )); - assert!(matches!( - rl.check("clientA", 1), - CheckResult::Limited { .. } - )); - - // Client B should still have tokens - assert!(matches!( - rl.check("clientB", 1), - CheckResult::Allowed { .. } - )); - assert!(matches!( - rl.check("clientB", 1), - CheckResult::Allowed { .. } - )); - assert!(matches!( - rl.check("clientB", 1), - CheckResult::Limited { .. } - )); + fn default_cost_for_type_uses_seeded_defaults_when_no_client_override() { + let rl = RateLimiter::new(RateLimitDefaults { + query_cost: 2, + procedure_cost: 5, + proxy_cost: 3, + }); + assert_eq!(rl.default_cost_for_type("nope", "query"), 2); + assert_eq!(rl.default_cost_for_type("nope", "procedure"), 5); + assert_eq!(rl.default_cost_for_type("nope", "proxy"), 3); } #[test] - fn toggle_enabled() { - let rl = RateLimiter::new( - true, + fn default_cost_for_type_uses_per_client_override() { + let rl = RateLimiter::new(RateLimitDefaults { + query_cost: 2, + procedure_cost: 5, + proxy_cost: 3, + }); + rl.register_client_config( + "hvc_a".to_string(), RateLimitConfig { - capacity: 1, + capacity: 100, refill_rate: 1.0, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + default_query_cost: 7, + default_procedure_cost: 8, + default_proxy_cost: 9, }, ); - assert!(rl.is_enabled()); - rl.set_enabled(false); - assert!(!rl.is_enabled()); - assert!(matches!(rl.check("k", 1), CheckResult::Disabled)); + assert_eq!(rl.default_cost_for_type("hvc_a", "query"), 7); + assert_eq!(rl.default_cost_for_type("hvc_a", "procedure"), 8); + assert_eq!(rl.default_cost_for_type("hvc_a", "proxy"), 9); } } diff --git a/src/repo/session.rs b/src/repo/session.rs index 87335c4..be38a5e 100644 --- a/src/repo/session.rs +++ b/src/repo/session.rs @@ -14,7 +14,7 @@ pub(crate) async fn get_oauth_session( Did::new(did.to_string()).map_err(|_| AppError::Auth(format!("invalid DID: {did}")))?; state .oauth - .default_client() + .primary_client() .restore(&did) .await .map_err(|e| AppError::Auth(format!("no OAuth session for {}: {e}", did.as_ref()))) diff --git a/src/repo/upload_blob.rs b/src/repo/upload_blob.rs index 4364d35..03c5afd 100644 --- a/src/repo/upload_blob.rs +++ b/src/repo/upload_blob.rs @@ -17,17 +17,20 @@ pub async fn upload_blob( headers: HeaderMap, body: Bytes, ) -> Result { - let rate_key = claims.did().to_string(); - let check = state.rate_limiter.check( - &rate_key, - state.rate_limiter.default_cost_for_type("procedure"), - ); + let check = if let Some(client_key) = claims.client_key() { + let cost = state + .rate_limiter + .default_cost_for_type(client_key, "procedure"); + Some(state.rate_limiter.check(client_key, cost)) + } else { + None + }; - if let CheckResult::Limited { + if let Some(CheckResult::Limited { retry_after, limit, reset, - } = check + }) = check { return Err(AppError::RateLimited { retry_after, @@ -45,11 +48,11 @@ pub async fn upload_blob( let mut response = pds_post_blob(&state, &session, content_type, body).await?; - if let CheckResult::Allowed { + if let Some(CheckResult::Allowed { remaining, limit, reset, - } = check + }) = check { let h = response.headers_mut(); h.insert("RateLimit-Limit", limit.into()); diff --git a/src/server.rs b/src/server.rs index 0757475..0b03c8d 100644 --- a/src/server.rs +++ b/src/server.rs @@ -13,6 +13,7 @@ use tower_http::trace::TraceLayer; use crate::AppState; use crate::admin; use crate::auth::Claims; +use crate::domain_middleware::resolve_domain; use crate::error::AppError; use crate::profile; use crate::rate_limit::CheckResult; @@ -60,10 +61,7 @@ pub fn router(state: AppState) -> Router { let serve_dir = ServeDir::new(&static_dir).not_found_service(spa_fallback); - Router::new() - .route("/health", get(health)) - .route("/settings/logo", get(crate::admin::settings::serve_logo)) - .nest("/admin", admin::admin_routes(state.clone())) + let domain_routes = Router::new() .nest("/auth", crate::auth::routes::routes()) .nest("/external-auth", crate::external_auth::routes()) // https://atproto.com/specs/oauth#types-of-clients @@ -76,6 +74,16 @@ pub fn router(state: AppState) -> Router { // Catch-all for dynamically registered lexicons .route("/xrpc/{method}", get(xrpc::xrpc_get).post(xrpc::xrpc_post)) .route("/config", get(config_endpoint)) + .route("/settings/logo", get(crate::admin::settings::serve_logo)) + .layer(axum::middleware::from_fn_with_state( + state.clone(), + resolve_domain, + )); + + Router::new() + .route("/health", get(health)) + .nest("/admin", admin::admin_routes(state.clone())) + .merge(domain_routes) .fallback_service(serve_dir) .layer(TraceLayer::new_for_http()) .layer( @@ -98,7 +106,14 @@ async fn health() -> &'static str { "ok" } -async fn config_endpoint(State(state): State) -> Json { +async fn config_endpoint( + State(state): State, + req: axum::extract::Request, +) -> Json { + let domain_url = crate::domain_middleware::extract_domain(&req) + .map(|d| d.url.clone()) + .unwrap_or_else(|| state.config.public_url.clone()); + let pool = &state.db; let backend = state.db_backend; @@ -112,7 +127,7 @@ async fn config_endpoint(State(state): State) -> Json) -> Json) -> Json) -> Json { - let mut metadata = - serde_json::to_value(&state.oauth.default_client().client_metadata).unwrap_or_default(); +async fn client_metadata( + State(state): State, + req: axum::extract::Request, +) -> Json { + let domain_url = crate::domain_middleware::extract_domain(&req) + .map(|d| d.url.clone()) + .unwrap_or_else(|| state.config.public_url.clone()); + + let oauth_client = state.oauth.get_for_domain(&domain_url); + let mut metadata = serde_json::to_value(&oauth_client.client_metadata).unwrap_or_default(); // The `client_id` field in the response must exactly match the URL the // authorization server fetched. let client_id = format!( "{}/oauth-client-metadata.json", - state.config.public_url.trim_end_matches('/') + domain_url.trim_end_matches('/') ); metadata["client_id"] = serde_json::Value::String(client_id); @@ -169,7 +191,7 @@ async fn client_metadata(State(state): State) -> Json) -> Json, claims: Claims) -> Result { - let rate_key = claims - .client_key() - .map(|k| k.to_string()) - .unwrap_or_else(|| claims.did().to_string()); - let check = state - .rate_limiter - .check(&rate_key, state.rate_limiter.default_cost_for_type("query")); - - if let CheckResult::Limited { + let check = if let Some(client_key) = claims.client_key() { + let cost = state + .rate_limiter + .default_cost_for_type(client_key, "query"); + Some(state.rate_limiter.check(client_key, cost)) + } else { + None + }; + + if let Some(CheckResult::Limited { retry_after, limit, reset, - } = check + }) = check { return Err(AppError::RateLimited { retry_after, @@ -212,11 +235,11 @@ async fn get_profile(State(state): State, claims: Claims) -> Result { sqlx::query( - "TRUNCATE records, lexicons, backfill_jobs, users, user_permissions, api_keys, event_logs, script_variables, dead_letter_hooks, record_refs, labeler_subscriptions, labels, instance_settings RESTART IDENTITY CASCADE", + "TRUNCATE records, lexicons, backfill_jobs, users, user_permissions, api_keys, event_logs, script_variables, dead_letter_hooks, record_refs, labeler_subscriptions, labels, instance_settings, domains RESTART IDENTITY CASCADE", ) .execute(pool) .await @@ -41,6 +41,7 @@ pub async fn truncate_all(pool: &AnyPool) { "labeler_subscriptions", "labels", "instance_settings", + "domains", ]; for table in tables { sqlx::query(&format!("DELETE FROM {table}")) diff --git a/tests/e2e_api_clients.rs b/tests/e2e_api_clients.rs index b5db164..cf0089b 100644 --- a/tests/e2e_api_clients.rs +++ b/tests/e2e_api_clients.rs @@ -580,10 +580,10 @@ async fn oauth_registry_get_or_default_returns_default_for_unknown() { .state .oauth .get_or_default(Some("https://unknown.example.com/metadata.json")); - let default = app.state.oauth.default_client(); + let default = app.state.oauth.primary_client(); // Should be the same Arc (default client) - assert!(std::sync::Arc::ptr_eq(&client, default)); + assert!(std::sync::Arc::ptr_eq(&client, &default)); } #[tokio::test] @@ -593,9 +593,9 @@ async fn oauth_registry_get_or_default_returns_default_for_none() { let app = TestApp::new().await; let client = app.state.oauth.get_or_default(None); - let default = app.state.oauth.default_client(); + let default = app.state.oauth.primary_client(); - assert!(std::sync::Arc::ptr_eq(&client, default)); + assert!(std::sync::Arc::ptr_eq(&client, &default)); } // --------------------------------------------------------------------------- diff --git a/tests/e2e_domains.rs b/tests/e2e_domains.rs new file mode 100644 index 0000000..6ec08e3 --- /dev/null +++ b/tests/e2e_domains.rs @@ -0,0 +1,352 @@ +mod common; + +use axum::body::Body; +use axum::http::{Method, Request, StatusCode}; +use http_body_util::BodyExt; +use serde_json::{Value, json}; +use serial_test::serial; +use tower::ServiceExt; + +use common::app::TestApp; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +async fn json_body(resp: axum::response::Response) -> Value { + let body = resp.into_body().collect().await.unwrap().to_bytes(); + serde_json::from_slice(&body).unwrap() +} + +fn admin_get( + uri: &str, + cookie: (axum::http::HeaderName, axum::http::HeaderValue), +) -> Request { + Request::builder() + .uri(uri) + .header(cookie.0, cookie.1) + .body(Body::empty()) + .unwrap() +} + +fn admin_post( + uri: &str, + cookie: (axum::http::HeaderName, axum::http::HeaderValue), + body: &Value, +) -> Request { + Request::builder() + .method(Method::POST) + .uri(uri) + .header(cookie.0, cookie.1) + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(body).unwrap())) + .unwrap() +} + +fn admin_delete( + uri: &str, + cookie: (axum::http::HeaderName, axum::http::HeaderValue), +) -> Request { + Request::builder() + .method(Method::DELETE) + .uri(uri) + .header(cookie.0, cookie.1) + .body(Body::empty()) + .unwrap() +} + +fn get_with_host(uri: &str, host: &str) -> Request { + Request::builder() + .uri(uri) + .header("host", host) + .body(Body::empty()) + .unwrap() +} + +async fn seed_domain(app: &TestApp, id: &str, url: &str, is_primary: bool) { + let now = happyview::db::now_rfc3339(); + let sql = happyview::db::adapt_sql( + "INSERT INTO domains (id, url, is_primary, created_at, updated_at) VALUES (?, ?, ?, ?, ?)", + app.state.db_backend, + ); + sqlx::query(&sql) + .bind(id) + .bind(url) + .bind(if is_primary { 1i32 } else { 0i32 }) + .bind(&now) + .bind(&now) + .execute(&app.state.db) + .await + .unwrap(); + app.state + .domain_cache + .insert(happyview::domain::Domain { + id: id.into(), + url: url.into(), + is_primary, + created_at: now.clone(), + updated_at: now, + }) + .await; +} + +// --------------------------------------------------------------------------- +// Domains tests +// --------------------------------------------------------------------------- + +#[tokio::test] +#[serial] +#[ignore] +async fn domains_list_returns_seeded_domain() { + let app = TestApp::new().await; + + seed_domain(&app, "primary-id", "http://127.0.0.1:0", true).await; + + let resp = app + .router + .clone() + .oneshot(admin_get("/admin/domains", app.admin_cookie())) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + let json = json_body(resp).await; + let domains = json.as_array().expect("expected array"); + assert_eq!(domains.len(), 1, "expected 1 domain, got {}", domains.len()); + assert_eq!(domains[0]["url"], "http://127.0.0.1:0"); + assert_eq!(domains[0]["is_primary"], true); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn domains_create_and_delete() { + let app = TestApp::new().await; + + seed_domain(&app, "primary-id", "http://127.0.0.1:0", true).await; + + // Create a new domain + let resp = app + .router + .clone() + .oneshot(admin_post( + "/admin/domains", + app.admin_cookie(), + &json!({ "url": "http://127.0.0.1:9999" }), + )) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::CREATED, + "expected 201 on create, got {}", + resp.status() + ); + let json = json_body(resp).await; + assert!( + json["id"].is_string(), + "expected id in response, got {:?}", + json + ); + assert_eq!(json["url"], "http://127.0.0.1:9999"); + assert_eq!(json["is_primary"], false); + + let new_id = json["id"].as_str().unwrap().to_string(); + + // Delete the newly created domain + let resp = app + .router + .clone() + .oneshot(admin_delete( + &format!("/admin/domains/{new_id}"), + app.admin_cookie(), + )) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::NO_CONTENT, + "expected 204 on delete, got {}", + resp.status() + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn domains_duplicate_url_returns_400() { + let app = TestApp::new().await; + + seed_domain(&app, "primary-id", "http://127.0.0.1:0", true).await; + + // Attempt to create a domain with the same URL + let resp = app + .router + .clone() + .oneshot(admin_post( + "/admin/domains", + app.admin_cookie(), + &json!({ "url": "http://127.0.0.1:0" }), + )) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::BAD_REQUEST, + "expected 400 on duplicate URL, got {}", + resp.status() + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn domains_cannot_delete_primary() { + let app = TestApp::new().await; + + seed_domain(&app, "primary-id", "http://127.0.0.1:0", true).await; + + let resp = app + .router + .clone() + .oneshot(admin_delete( + "/admin/domains/primary-id", + app.admin_cookie(), + )) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::BAD_REQUEST, + "expected 400 when deleting primary domain, got {}", + resp.status() + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn domains_set_primary() { + let app = TestApp::new().await; + + seed_domain(&app, "id-a", "http://127.0.0.1:0", true).await; + seed_domain(&app, "id-b", "http://127.0.0.1:9999", false).await; + + // Set domain b as primary + let resp = app + .router + .clone() + .oneshot(admin_post( + "/admin/domains/id-b/primary", + app.admin_cookie(), + &json!({}), + )) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::NO_CONTENT, + "expected 204 on set primary, got {}", + resp.status() + ); + + // Verify domain b is now primary + let resp = app + .router + .clone() + .oneshot(admin_get("/admin/domains", app.admin_cookie())) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + let json = json_body(resp).await; + let domains = json.as_array().expect("expected array"); + let domain_b = domains + .iter() + .find(|d| d["id"] == "id-b") + .expect("domain b not found"); + assert_eq!( + domain_b["is_primary"], true, + "expected domain b to be primary" + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn unknown_host_returns_421_on_domain_scoped_routes() { + let app = TestApp::new().await; + + // No domains seeded — cache is empty + let resp = app + .router + .clone() + .oneshot(get_with_host("/config", "unknown.example.com")) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::MISDIRECTED_REQUEST, + "expected 421 for unknown host, got {}", + resp.status() + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn health_check_bypasses_domain_resolution() { + let app = TestApp::new().await; + + // No domains seeded — cache is empty + let resp = app + .router + .clone() + .oneshot(get_with_host("/health", "unknown.example.com")) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::OK, + "expected 200 on /health regardless of host, got {}", + resp.status() + ); +} + +#[tokio::test] +#[serial] +#[ignore] +async fn domain_scoped_route_works_with_known_host() { + let app = TestApp::new().await; + + // Domain.host() for "http://localhost:3000" is "localhost:3000" + seed_domain(&app, "local-id", "http://localhost:3000", true).await; + + let resp = app + .router + .clone() + .oneshot(get_with_host("/config", "localhost:3000")) + .await + .unwrap(); + + assert_eq!( + resp.status(), + StatusCode::OK, + "expected 200 on /config with known host, got {}", + resp.status() + ); + + let json = json_body(resp).await; + assert_eq!( + json["public_url"], "http://localhost:3000", + "expected public_url to match domain URL, got {:?}", + json["public_url"] + ); +} diff --git a/tests/lua_atproto_api.rs b/tests/lua_atproto_api.rs index c3f6e39..52123f8 100644 --- a/tests/lua_atproto_api.rs +++ b/tests/lua_atproto_api.rs @@ -73,13 +73,10 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + happyview::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(happyview::auth::OAuthClientRegistry::new( @@ -97,6 +94,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> )), official_registry_config: happyview::plugin::official_registry::RegistryConfig::production( ), + domain_cache: happyview::domain::DomainCache::new(), } } diff --git a/tests/lua_db_api.rs b/tests/lua_db_api.rs index 6c97858..abf9754 100644 --- a/tests/lua_db_api.rs +++ b/tests/lua_db_api.rs @@ -76,13 +76,10 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> 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, - default_query_cost: 1, - default_procedure_cost: 1, - default_proxy_cost: 1, + happyview::rate_limit::RateLimitDefaults { + query_cost: 1, + procedure_cost: 1, + proxy_cost: 1, }, ), oauth: std::sync::Arc::new(happyview::auth::OAuthClientRegistry::new( @@ -100,6 +97,7 @@ async fn test_state_with_pool(pool: sqlx::AnyPool, backend: DatabaseBackend) -> )), official_registry_config: happyview::plugin::official_registry::RegistryConfig::production( ), + domain_cache: happyview::domain::DomainCache::new(), } } -- 2.51.2