From d651ccfee07e12f53c3323f338f51d21889cf583 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:37:47 +0800 Subject: [PATCH] refactor(server): move the access handlers' SQL into repositories The API key, auth, SSO, setup and OIDC/settings handlers no longer carry SQL. Their 48 statements move verbatim into services::api_key_repository, auth_repository, setup_repository and settings_repository; the key-rotation and account-deletion transactions move with them. Handlers keep permission checks, validation, audit, and session / cache work. admin_access.rs adds integration tests for what no test reached: key create / read / PATCH semantics and validation, list paging and the archived view, revoke / force-revoke, expiring keys, cost centers, the default-expiry setting, registration with a default role, /me roles and teams, TOTP status, and SSO sign-in against a mock identity provider (activation, provisioning, re-login, refused accounts). Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + crates/server/src/handlers/admin/oidc.rs | 9 +- crates/server/src/handlers/admin/settings.rs | 16 +- crates/server/src/handlers/api_keys.rs | 298 ++------ crates/server/src/handlers/auth.rs | 163 +--- crates/server/src/handlers/setup.rs | 71 +- crates/server/src/handlers/sso.rs | 35 +- .../server/src/services/api_key_repository.rs | 345 +++++++++ crates/server/src/services/auth_repository.rs | 240 ++++++ crates/server/src/services/mod.rs | 4 + .../src/services/settings_repository.rs | 24 + .../server/src/services/setup_repository.rs | 98 +++ crates/test-support/Cargo.toml | 1 + crates/test-support/tests/admin_access.rs | 722 ++++++++++++++++++ 14 files changed, 1587 insertions(+), 440 deletions(-) create mode 100644 crates/server/src/services/api_key_repository.rs create mode 100644 crates/server/src/services/auth_repository.rs create mode 100644 crates/server/src/services/settings_repository.rs create mode 100644 crates/server/src/services/setup_repository.rs create mode 100644 crates/test-support/tests/admin_access.rs diff --git a/Cargo.lock b/Cargo.lock index dd577288..843be17d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4237,6 +4237,7 @@ dependencies = [ "futures", "hex", "hmac 0.13.0", + "jsonwebtoken", "once_cell", "p256", "rand 0.10.0", diff --git a/crates/server/src/handlers/admin/oidc.rs b/crates/server/src/handlers/admin/oidc.rs index 9d423dc5..fc541903 100644 --- a/crates/server/src/handlers/admin/oidc.rs +++ b/crates/server/src/handlers/admin/oidc.rs @@ -29,6 +29,7 @@ use think_watch_common::errors::AppError; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::settings_repository; const OIDC_TEST_RESULT_KEY: &str = "oidc:test:result"; @@ -306,9 +307,7 @@ pub async fn delete_oidc_draft( auth_user .require_global_permission(&state.db, "system:configure_oidc") .await?; - sqlx::query("DELETE FROM system_settings WHERE key = 'oidc.draft'") - .execute(&state.db) - .await?; + settings_repository::delete_oidc_draft(&state.db).await?; state .dynamic_config .reload() @@ -564,9 +563,7 @@ pub async fn activate_oidc_draft( .map_err(AppError::Internal)?; } - sqlx::query("DELETE FROM system_settings WHERE key = 'oidc.draft'") - .execute(&state.db) - .await?; + settings_repository::delete_oidc_draft(&state.db).await?; dc.reload().await.map_err(AppError::Internal)?; let _: Result<(), _> = fred::interfaces::KeysInterface::del::<(), _>(&state.redis, OIDC_TEST_RESULT_KEY).await; diff --git a/crates/server/src/handlers/admin/settings.rs b/crates/server/src/handlers/admin/settings.rs index f0f42f0a..e27afb11 100644 --- a/crates/server/src/handlers/admin/settings.rs +++ b/crates/server/src/handlers/admin/settings.rs @@ -15,6 +15,7 @@ use think_watch_common::errors::AppError; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::settings_repository; use super::retention::{MAX_RETENTION_DAYS, apply_blob_lifecycle, apply_clickhouse_ttls}; @@ -270,17 +271,10 @@ pub async fn update_settings( // DB-level validation for settings that reference other entities if let Some(role_val) = req.settings.get("auth.default_role") { let role_name = role_val.as_str().unwrap_or(""); - if !role_name.is_empty() { - let exists: Option<(String,)> = - sqlx::query_as("SELECT name FROM rbac_roles WHERE name = $1") - .bind(role_name) - .fetch_optional(&state.db) - .await?; - if exists.is_none() { - return Err(AppError::BadRequest(format!( - "Role '{role_name}' does not exist" - ))); - } + if !role_name.is_empty() && !settings_repository::role_exists(&state.db, role_name).await? { + return Err(AppError::BadRequest(format!( + "Role '{role_name}' does not exist" + ))); } } diff --git a/crates/server/src/handlers/api_keys.rs b/crates/server/src/handlers/api_keys.rs index 2bab6006..128f5685 100644 --- a/crates/server/src/handlers/api_keys.rs +++ b/crates/server/src/handlers/api_keys.rs @@ -13,6 +13,7 @@ use think_watch_common::models::ApiKey; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::api_key_repository::{self as repo, ApiKeyPatch, NewApiKey}; /// Resolve whether the caller sees only their own keys or the whole /// table. API keys are user-owned, so scope collapses to two cases: @@ -51,12 +52,9 @@ async fn assert_owner_or_admin( // and the downstream UPDATE no-op'd because it carried its own // `AND deleted_at IS NULL` guard, but an audit entry still fired // claiming the operation happened. - let owner: Option = - sqlx::query_scalar("SELECT user_id FROM api_keys WHERE id = $1 AND deleted_at IS NULL") - .bind(key_id) - .fetch_optional(pool) - .await?; - let owner = owner.ok_or_else(|| AppError::NotFound("API key not found".into()))?; + let owner = repo::owner_of_live(pool, key_id) + .await? + .ok_or_else(|| AppError::NotFound("API key not found".into()))?; if auth_user.claims.sub == owner { return Ok(()); } @@ -205,16 +203,7 @@ async fn validate_mcp_account_overrides( let label = label_val.as_str().ok_or_else(|| { AppError::BadRequest("mcp_account_overrides values must be strings".into()) })?; - let exists: Option = sqlx::query_scalar( - "SELECT 1 FROM mcp_user_credentials - WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3", - ) - .bind(server_id) - .bind(user_id) - .bind(label) - .fetch_optional(pool) - .await?; - if exists.is_none() { + if !repo::mcp_credential_exists(pool, server_id, user_id, label).await? { return Err(AppError::BadRequest(format!( "mcp_account_overrides points at '{label}' for server {server_id_str}, \ but you have no credential with that label" @@ -291,50 +280,22 @@ pub async fn list_keys( // every key in the system. let global = caller_is_admin_tier(&auth_user); - // Two view modes: - // live: deleted_at IS NULL (default) - // archived: deleted_at IS NOT NULL AND revoke variant - // The archived predicate excludes user_deleted / account_deleted - // soft-deletes — those are cascades from a user wipe, not - // intentional key revocations, and shouldn't appear in a - // "revoked keys" tab. - let visibility_clause = if params.archived { - "deleted_at IS NOT NULL \ - AND (disabled_reason = 'revoked' OR disabled_reason LIKE 'force_revoked:%')" - } else { - "deleted_at IS NULL" - }; - + // Two view modes: live keys (default), or `archived` — revoked + // keys only, not the soft-deletes cascaded from a user wipe. + let archived = params.archived; let (total, keys): (i64, Vec) = if global { - let total: i64 = sqlx::query_scalar(&format!( - "SELECT COUNT(*) FROM api_keys WHERE {visibility_clause}" - )) - .fetch_one(&state.db) - .await?; - let keys = sqlx::query_as::<_, ApiKey>(&format!( - "SELECT * FROM api_keys WHERE {visibility_clause} \ - ORDER BY created_at DESC LIMIT $1 OFFSET $2" - )) - .bind(per_page as i64) - .bind(offset as i64) - .fetch_all(&state.db) - .await?; + let total = repo::count_all(&state.db, archived).await?; + let keys = repo::list_all_page(&state.db, archived, per_page as i64, offset as i64).await?; (total, keys) } else { - let total: i64 = sqlx::query_scalar(&format!( - "SELECT COUNT(*) FROM api_keys WHERE {visibility_clause} AND user_id = $1" - )) - .bind(caller_id) - .fetch_one(&state.db) - .await?; - let keys = sqlx::query_as::<_, ApiKey>(&format!( - "SELECT * FROM api_keys WHERE {visibility_clause} AND user_id = $1 \ - ORDER BY created_at DESC LIMIT $2 OFFSET $3" - )) - .bind(caller_id) - .bind(per_page as i64) - .bind(offset as i64) - .fetch_all(&state.db) + let total = repo::count_for_user(&state.db, archived, caller_id).await?; + let keys = repo::list_for_user_page( + &state.db, + archived, + caller_id, + per_page as i64, + offset as i64, + ) .await?; (total, keys) }; @@ -475,25 +436,23 @@ pub async fn create_key( // key. Subsequent rotations carry over the same lineage_id, // so descendants will have id != lineage_id. let id = uuid::Uuid::new_v4(); - let row = sqlx::query_as::<_, ApiKey>( - r#"INSERT INTO api_keys (id, lineage_id, key_prefix, key_hash, name, user_id, surfaces, - allowed_models, allowed_mcp_tools, mcp_account_overrides, expires_at, - cost_center, rotation_period_days) - VALUES ($1, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING *"#, + let row = repo::insert( + &state.db, + &NewApiKey { + id, + key_prefix: &generated.prefix, + key_hash: &generated.hash, + name: &req.name, + user_id: auth_user.claims.sub, + surfaces: &surfaces, + allowed_models: &req.allowed_models, + allowed_mcp_tools: &req.allowed_mcp_tools, + mcp_account_overrides: &mcp_account_overrides, + expires_at, + cost_center: cost_center.as_deref(), + rotation_period_days, + }, ) - .bind(id) - .bind(&generated.prefix) - .bind(&generated.hash) - .bind(&req.name) - .bind(auth_user.claims.sub) - .bind(&surfaces) - .bind(&req.allowed_models) - .bind(&req.allowed_mcp_tools) - .bind(&mcp_account_overrides) - .bind(expires_at) - .bind(cost_center.as_deref()) - .bind(rotation_period_days) - .fetch_one(&state.db) .await?; Ok(Json(CreateApiKeyResponse { @@ -526,12 +485,9 @@ pub async fn get_key( ) -> Result, AppError> { auth_user.require_permission("api_keys:read")?; assert_owner_or_admin(&auth_user, &state.db, id).await?; - let key = - sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys WHERE id = $1 AND deleted_at IS NULL") - .bind(id) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("API key not found".into()))?; + let key = repo::find_live(&state.db, id) + .await? + .ok_or(AppError::NotFound("API key not found".into()))?; Ok(Json(key)) } @@ -569,16 +525,7 @@ pub async fn revoke_key( // disappears from the default list view and is hard-deleted by // the retention sweep ~30 days later. The `?archived=true` view // surfaces it in the meantime for audit / oh-shit lookups. - let result = sqlx::query( - "UPDATE api_keys SET is_active = false, grace_period_ends_at = NULL, \ - disabled_reason = 'revoked', deleted_at = now() \ - WHERE id = $1 AND deleted_at IS NULL", - ) - .bind(id) - .execute(&state.db) - .await?; - - if result.rows_affected() == 0 { + if repo::revoke(&state.db, id).await? == 0 { return Err(AppError::NotFound("API key not found".into())); } @@ -650,17 +597,7 @@ pub async fn force_revoke_key( "force_revoked:{}", reason.chars().take(64).collect::() ); - let result = sqlx::query( - "UPDATE api_keys SET is_active = false, grace_period_ends_at = NULL, \ - disabled_reason = $1, deleted_at = now() \ - WHERE id = $2 AND deleted_at IS NULL", - ) - .bind(&disabled_reason) - .bind(id) - .execute(&state.db) - .await?; - - if result.rows_affected() == 0 { + if repo::force_revoke(&state.db, id, &disabled_reason).await? == 0 { return Err(AppError::NotFound("API key not found".into())); } @@ -762,12 +699,9 @@ pub async fn update_key( return Err(AppError::BadRequest(format!("{name} must be >= 0"))); } } - let key = - sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys WHERE id = $1 AND deleted_at IS NULL") - .bind(id) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("API key not found".into()))?; + let key = repo::find_live(&state.db, id) + .await? + .ok_or(AppError::NotFound("API key not found".into()))?; // Subset check against the *key owner's* roles, not the caller's — // a super-admin editing someone else's key still can't grant tools @@ -851,35 +785,25 @@ pub async fn update_key( } }; - let updated = sqlx::query_as::<_, ApiKey>( - r#"UPDATE api_keys SET - allowed_models = CASE WHEN $11 THEN $1 ELSE allowed_models END, - allowed_mcp_tools = CASE WHEN $12 THEN $10 ELSE allowed_mcp_tools END, - surfaces = COALESCE($2, surfaces), - expires_at = $3, - rotation_period_days = COALESCE($4, rotation_period_days), - inactivity_timeout_days = COALESCE($5, inactivity_timeout_days), - cost_center = CASE WHEN $7 THEN $6 ELSE cost_center END, - mcp_account_overrides = CASE WHEN $13 THEN $14 ELSE mcp_account_overrides END, - last_expiry_warning_days = CASE WHEN $9 THEN NULL - ELSE last_expiry_warning_days END - WHERE id = $8 RETURNING *"#, + let updated = repo::update( + &state.db, + id, + &ApiKeyPatch { + allowed_models_set: models_set, + allowed_models: models_value, + allowed_mcp_tools_set: mcp_tools_set, + allowed_mcp_tools: mcp_tools_value, + surfaces: normalized_surfaces.as_ref(), + expires_at, + rotation_period_days: req.rotation_period_days, + inactivity_timeout_days: req.inactivity_timeout_days, + cost_center_set, + cost_center: cost_center_value.as_deref(), + mcp_account_overrides_set: overrides_set, + mcp_account_overrides: &overrides_value, + expiry_extended, + }, ) - .bind(models_value) - .bind(normalized_surfaces.as_ref()) - .bind(expires_at) - .bind(req.rotation_period_days) - .bind(req.inactivity_timeout_days) - .bind(cost_center_value.as_deref()) - .bind(cost_center_set) - .bind(id) - .bind(expiry_extended) - .bind(mcp_tools_value) - .bind(models_set) - .bind(mcp_tools_set) - .bind(overrides_set) - .bind(&overrides_value) - .fetch_one(&state.db) .await?; // Record what actually changed in the audit detail. Surfaces / @@ -978,12 +902,9 @@ pub async fn rotate_key( ) -> Result, AppError> { auth_user.require_permission("api_keys:rotate")?; assert_owner_or_admin(&auth_user, &state.db, id).await?; - let old_key = - sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys WHERE id = $1 AND deleted_at IS NULL") - .bind(id) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("API key not found".into()))?; + let old_key = repo::find_live(&state.db, id) + .await? + .ok_or(AppError::NotFound("API key not found".into()))?; if !old_key.is_active { return Err(AppError::BadRequest("Cannot rotate an inactive key".into())); @@ -1013,56 +934,20 @@ pub async fn rotate_key( // when rotation happens. let generated = api_key::generate_api_key(); - // INSERT new key + UPDATE old key's grace period must be atomic. - // Without the transaction, an error between the two leaves the - // old key with no grace_period_ends_at — meaning it never enters - // the rotation grace window and both keys remain valid forever. - let mut tx = state.db.begin().await?; - - let new_key = sqlx::query_as::<_, ApiKey>( - r#"INSERT INTO api_keys (key_prefix, key_hash, name, user_id, surfaces, allowed_models, - allowed_mcp_tools, expires_at, rotation_period_days, inactivity_timeout_days, - cost_center, rotated_from_id, last_rotation_at, lineage_id) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, now(), $13) - RETURNING *"#, - ) - .bind(&generated.prefix) - .bind(&generated.hash) - // Carry the original name verbatim. The provenance / generation - // chain is already captured by `rotated_from_id` + `last_rotation_at`, - // and the row's status badge ("已轮换" / "活跃") tells the operator - // which generation is which. An earlier version stamped a literal - // " (rotated)" suffix into `name`, which (a) outlived the old key - // (the suffix has no removal logic), and (b) stacked on every - // subsequent rotation — "Foo (rotated) (rotated) (rotated)…". - .bind(&old_key.name) - .bind(old_key.user_id) - .bind(&old_key.surfaces) - .bind(&old_key.allowed_models) - .bind(&old_key.allowed_mcp_tools) - .bind(old_key.expires_at) - .bind(old_key.rotation_period_days) - .bind(old_key.inactivity_timeout_days) - .bind(old_key.cost_center.as_deref()) - .bind(id) - // Inherit the parent's lineage_id so every generation in the - // rotation chain shares one stable identity. Per-key analytics - // can then group on `api_key_lineage_id` instead of recursing - // on `rotated_from_id`. - .bind(old_key.lineage_id) - .fetch_one(&mut *tx) - .await?; - - sqlx::query( - "UPDATE api_keys SET grace_period_ends_at = $1, disabled_reason = 'rotated' WHERE id = $2", + // The new key carries the old one's name verbatim — the chain is + // in `rotated_from_id` + `last_rotation_at`, so a " (rotated)" + // suffix would only stack up on every rotation — and its lineage + // id, so per-key analytics group every generation together. The + // insert and the old key's grace window are one transaction. + let new_key = repo::rotate( + &state.db, + &old_key, + &generated.prefix, + &generated.hash, + grace_period_ends_at, ) - .bind(grace_period_ends_at) - .bind(id) - .execute(&mut *tx) .await?; - tx.commit().await?; - state.audit.log( auth_user .audit("api_key.rotate") @@ -1119,31 +1004,9 @@ pub async fn list_expiring_keys( let global = caller_is_admin_tier(&auth_user); let keys = if global { - sqlx::query_as::<_, ApiKey>( - r#"SELECT * FROM api_keys - WHERE is_active = true - AND deleted_at IS NULL - AND expires_at IS NOT NULL - AND expires_at <= $1 - ORDER BY expires_at ASC"#, - ) - .bind(threshold) - .fetch_all(&state.db) - .await? + repo::list_expiring_all(&state.db, threshold).await? } else { - sqlx::query_as::<_, ApiKey>( - r#"SELECT * FROM api_keys - WHERE is_active = true - AND deleted_at IS NULL - AND expires_at IS NOT NULL - AND expires_at <= $1 - AND user_id = $2 - ORDER BY expires_at ASC"#, - ) - .bind(threshold) - .bind(caller_id) - .fetch_all(&state.db) - .await? + repo::list_expiring_for_user(&state.db, threshold, caller_id).await? }; Ok(Json(keys)) @@ -1186,14 +1049,7 @@ pub async fn list_cost_centers( State(state): State, ) -> Result>, AppError> { auth_user.require_permission("api_keys:read")?; - let rows: Vec<(String,)> = sqlx::query_as( - "SELECT DISTINCT cost_center FROM api_keys \ - WHERE cost_center IS NOT NULL AND deleted_at IS NULL \ - ORDER BY cost_center ASC", - ) - .fetch_all(&state.db) - .await?; - Ok(Json(rows.into_iter().map(|(s,)| s).collect())) + Ok(Json(repo::cost_centers(&state.db).await?)) } #[derive(Debug, serde::Serialize, utoipa::ToSchema)] diff --git a/crates/server/src/handlers/auth.rs b/crates/server/src/handlers/auth.rs index ada72eaf..119551a7 100644 --- a/crates/server/src/handlers/auth.rs +++ b/crates/server/src/handlers/auth.rs @@ -12,13 +12,13 @@ use think_watch_common::dto::{ RefreshRequest, UserResponse, }; use think_watch_common::errors::AppError; -use think_watch_common::models::User; use think_watch_common::validation::{normalize_email, validate_email, validate_password}; use crate::middleware::verify_signature; use crate::app::AppState; use crate::middleware::auth_guard::AuthUser; +use crate::services::auth_repository as repo; /// Parse a JSON request body from a raw `axum::extract::Request`, /// enforcing a maximum byte limit. Shared by login, register, and @@ -495,12 +495,7 @@ pub async fn login( // lookup in this file already filters `deleted_at IS NULL`; the // login path was the lone exception, leaving a 30-day window after // soft-delete where the credential still worked. - let maybe_user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE email = $1 AND is_active = true AND deleted_at IS NULL", - ) - .bind(&email) - .fetch_optional(&state.db) - .await?; + let maybe_user = repo::find_active_by_email(&state.db, &email).await?; let (user, password_hash) = match maybe_user { Some(u) => { @@ -639,16 +634,13 @@ pub async fn login( // plaintext: two concurrent requests reading the // same `codes_blob` and racing to update will see // exactly one rows_affected==1. - let rows = sqlx::query( - "UPDATE users SET totp_recovery_codes = $1 \ - WHERE id = $2 AND totp_recovery_codes = $3", + let rows = repo::swap_recovery_codes( + &state.db, + user.id, + codes_blob, + &updated_blob, ) - .bind(&updated_blob) - .bind(user.id) - .bind(codes_blob) - .execute(&state.db) - .await? - .rows_affected(); + .await?; if rows == 1 { recovery_used = true; // Actor is identified (credentials passed) @@ -1057,17 +1049,8 @@ pub async fn register( let mut tx = state.db.begin().await?; // Use INSERT ... ON CONFLICT to avoid leaking whether email exists (user enumeration) - let user = sqlx::query_as::<_, User>( - r#"INSERT INTO users (email, display_name, password_hash) - VALUES ($1, $2, $3) - ON CONFLICT (email) DO NOTHING - RETURNING *"#, - ) - .bind(&email) - .bind(&req.display_name) - .bind(&password_hash) - .fetch_optional(&mut *tx) - .await?; + let user = + repo::insert_user_unless_taken(&mut tx, &email, &req.display_name, &password_hash).await?; let user = match user { Some(u) => u, @@ -1089,14 +1072,7 @@ pub async fn register( // Assign default role (configurable via settings; empty = no role) if let Some(role_name) = state.dynamic_config.default_role().await { - sqlx::query( - r#"INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) - SELECT $1, id, 'global', $1 FROM rbac_roles WHERE name = $2"#, - ) - .bind(user.id) - .bind(&role_name) - .execute(&mut *tx) - .await?; + repo::assign_default_role(&mut *tx, user.id, &role_name).await?; } tx.commit().await?; @@ -1227,12 +1203,9 @@ pub async fn refresh( // check covers the cold-start / flushed-Redis case. Without it, a // disabled-then-cache-cleared user can mint fresh access tokens // for up to refresh_ttl_days (default 7). - let user_active: Option = - sqlx::query_scalar("SELECT is_active FROM users WHERE id = $1 AND deleted_at IS NULL") - .bind(claims.sub) - .fetch_optional(&state.db) - .await - .map_err(|_| AppError::Unauthorized)?; + let user_active = repo::is_active(&state.db, claims.sub) + .await + .map_err(|_| AppError::Unauthorized)?; if !matches!(user_active, Some(true)) { return Err(AppError::Unauthorized); } @@ -1351,13 +1324,9 @@ pub async fn me( auth_user: AuthUser, State(state): State, ) -> Result, AppError> { - let user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE id = $1 AND is_active = true AND deleted_at IS NULL", - ) - .bind(auth_user.claims.sub) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("User not found".into()))?; + let user = repo::find_active(&state.db, auth_user.claims.sub) + .await? + .ok_or(AppError::NotFound("User not found".into()))?; let role_assignments = fetch_user_role_assignments(&state, user.id).await; @@ -1379,17 +1348,7 @@ pub async fn me( // Team memberships — used by the frontend permission cache // and the team-context badge in the header. - type TeamRow = (uuid::Uuid, String); - let team_rows: Vec = sqlx::query_as( - "SELECT t.id, t.name FROM team_members tm \ - JOIN teams t ON t.id = tm.team_id \ - WHERE tm.user_id = $1 \ - ORDER BY t.name ASC", - ) - .bind(user.id) - .fetch_all(&state.db) - .await - .unwrap_or_default(); + let team_rows = repo::teams_of(&state.db, user.id).await.unwrap_or_default(); let teams: Vec = team_rows .into_iter() .map(|(id, name)| think_watch_common::dto::UserTeamSummary { id, name }) @@ -1417,18 +1376,9 @@ async fn fetch_user_role_assignments( state: &AppState, user_id: uuid::Uuid, ) -> Vec { - type Row = (uuid::Uuid, String, bool, String, Option); - let rows: Vec = sqlx::query_as( - "SELECT r.id, r.name, r.is_system, ra.scope_kind, ra.scope_id \ - FROM rbac_role_assignments ra \ - JOIN rbac_roles r ON r.id = ra.role_id \ - WHERE ra.user_id = $1 \ - ORDER BY r.is_system DESC, r.name ASC", - ) - .bind(user_id) - .fetch_all(&state.db) - .await - .unwrap_or_default(); + let rows = repo::role_assignments_of(&state.db, user_id) + .await + .unwrap_or_default(); rows.into_iter() .map(|(role_id, name, is_system, scope_kind, scope_id)| { let scope = match (scope_kind.as_str(), scope_id) { @@ -1470,13 +1420,9 @@ pub async fn change_password( // here without a matching row is a deleted/disabled account on // a still-valid token. 404 leaks existence info AND contradicted // the OpenAPI contract (only 200/400/401 were documented). - let user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE id = $1 AND is_active = true AND deleted_at IS NULL", - ) - .bind(auth_user.claims.sub) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::Unauthorized)?; + let user = repo::find_active(&state.db, auth_user.claims.sub) + .await? + .ok_or(AppError::Unauthorized)?; let current_hash = user .password_hash @@ -1509,11 +1455,7 @@ pub async fn change_password( .await; let new_hash = password::hash_password(&req.new_password)?; - sqlx::query("UPDATE users SET password_hash = $1, password_change_required = false, updated_at = now() WHERE id = $2") - .bind(&new_hash) - .bind(user.id) - .execute(&state.db) - .await?; + repo::set_own_password(&state.db, user.id, &new_hash).await?; // Revoke all signing public keys for this user (invalidates sessions) let pubkey_key = format!("signing_pubkey:{}", user.id); @@ -1583,16 +1525,7 @@ pub async fn delete_account( let user_id = auth_user.claims.sub; // Soft-delete in a transaction: mark keys + user as deleted atomically - let mut tx = state.db.begin().await?; - sqlx::query("UPDATE api_keys SET is_active = false, deleted_at = now(), disabled_reason = 'account_deleted' WHERE user_id = $1") - .bind(user_id) - .execute(&mut *tx) - .await?; - sqlx::query("UPDATE users SET is_active = false, deleted_at = now() WHERE id = $1") - .bind(user_id) - .execute(&mut *tx) - .await?; - tx.commit().await?; + repo::soft_delete_account(&state.db, user_id).await?; // Revoke all sessions. Proceed even if Redis is unreachable — the // account's DB flags (is_active = false) already invalidate future @@ -1702,13 +1635,9 @@ pub async fn totp_setup( ) -> Result, AppError> { use think_watch_auth::totp; - let user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE id = $1 AND is_active = true AND deleted_at IS NULL", - ) - .bind(auth_user.claims.sub) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("User not found".into()))?; + let user = repo::find_active(&state.db, auth_user.claims.sub) + .await? + .ok_or(AppError::NotFound("User not found".into()))?; if user.totp_enabled { return Err(AppError::BadRequest("TOTP is already enabled".into())); @@ -1816,13 +1745,12 @@ pub async fn totp_verify_setup( let encrypted_recovery_codes = crate::services::totp_service::encrypt_recovery_codes(&state, &pending.recovery_codes)?; - sqlx::query( - "UPDATE users SET totp_secret = $1, totp_enabled = true, totp_recovery_codes = $2, updated_at = now() WHERE id = $3", + repo::enable_totp( + &state.db, + user_id, + &encrypted_secret, + &encrypted_recovery_codes, ) - .bind(&encrypted_secret) - .bind(&encrypted_recovery_codes) - .bind(user_id) - .execute(&state.db) .await?; // Clean up pending @@ -1854,13 +1782,9 @@ pub async fn totp_disable( State(state): State, Json(req): Json, ) -> Result, AppError> { - let user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE id = $1 AND is_active = true AND deleted_at IS NULL", - ) - .bind(auth_user.claims.sub) - .fetch_optional(&state.db) - .await? - .ok_or(AppError::NotFound("User not found".into()))?; + let user = repo::find_active(&state.db, auth_user.claims.sub) + .await? + .ok_or(AppError::NotFound("User not found".into()))?; if !user.totp_enabled { return Err(AppError::BadRequest("TOTP is not enabled".into())); @@ -1874,12 +1798,7 @@ pub async fn totp_disable( return Err(AppError::Unauthorized); } - sqlx::query( - "UPDATE users SET totp_secret = NULL, totp_enabled = false, totp_recovery_codes = NULL, updated_at = now() WHERE id = $1", - ) - .bind(user.id) - .execute(&state.db) - .await?; + repo::disable_totp(&state.db, user.id).await?; state .audit @@ -1902,11 +1821,7 @@ pub async fn totp_status( auth_user: AuthUser, State(state): State, ) -> Result, AppError> { - let enabled: bool = - sqlx::query_scalar("SELECT totp_enabled FROM users WHERE id = $1 AND deleted_at IS NULL") - .bind(auth_user.claims.sub) - .fetch_one(&state.db) - .await?; + let enabled = repo::totp_enabled(&state.db, auth_user.claims.sub).await?; // Check if platform requires TOTP let required: bool = state diff --git a/crates/server/src/handlers/setup.rs b/crates/server/src/handlers/setup.rs index 893609d6..0084f427 100644 --- a/crates/server/src/handlers/setup.rs +++ b/crates/server/src/handlers/setup.rs @@ -9,6 +9,7 @@ use think_watch_common::validation::{normalize_email, validate_email, validate_p use utoipa::ToSchema; use crate::app::AppState; +use crate::services::setup_repository::{self as repo, FirstAdmin}; #[derive(Debug, Serialize, ToSchema)] pub struct SetupStatusResponse { @@ -123,14 +124,9 @@ pub async fn setup_initialize( let mut tx = state.db.begin().await?; // Acquire an advisory lock (key = 1 for setup). This blocks concurrent setup attempts. - sqlx::query("SELECT pg_advisory_xact_lock(1)") - .execute(&mut *tx) - .await?; + repo::lock_setup(&mut tx).await?; - let db_initialized: Option = - sqlx::query_scalar("SELECT value FROM system_settings WHERE key = 'setup.initialized'") - .fetch_optional(&mut *tx) - .await?; + let db_initialized = repo::initialized_flag(&mut tx).await?; if db_initialized .as_ref() @@ -147,57 +143,24 @@ pub async fn setup_initialize( let admin_email = normalize_email(&req.admin.email); validate_email(&admin_email)?; - // 1. Create super_admin user + // Create the super_admin user with the first API key, and mark + // setup done. let password_hash = password::hash_password(&req.admin.password)?; - let admin_user = sqlx::query_as::<_, (uuid::Uuid, String)>( - r#"INSERT INTO users (email, display_name, password_hash) - VALUES ($1, $2, $3) RETURNING id, email"#, - ) - .bind(&admin_email) - .bind(&req.admin.display_name) - .bind(&password_hash) - .fetch_one(&mut *tx) - .await?; - // sqlx unique-violations now map to `AppError::Conflict` globally - // via `From`. No per-site string-sniffing needed. - - // Assign super_admin role. - sqlx::query( - r#"INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) - SELECT $1, id, 'global', $1 FROM rbac_roles WHERE name = 'super_admin'"#, - ) - .bind(admin_user.0) - .execute(&mut *tx) - .await?; - - // 2. Generate first API key for admin user let generated = api_key::generate_api_key(); - sqlx::query( - r#"INSERT INTO api_keys (key_prefix, key_hash, name, user_id, surfaces) - VALUES ($1, $2, $3, $4, $5)"#, - ) - .bind(&generated.prefix) - .bind(&generated.hash) - .bind("Default Admin Key") - .bind(admin_user.0) - .bind(super::api_keys::ALLOWED_SURFACES) - .execute(&mut *tx) - .await?; - - // 3. Mark as initialized let site_name = req.site_name.as_deref().unwrap_or("ThinkWatch"); - sqlx::query( - "UPDATE system_settings SET value = $1, updated_at = now() WHERE key = 'setup.initialized'", - ) - .bind(serde_json::json!(true)) - .execute(&mut *tx) - .await?; - - sqlx::query( - "UPDATE system_settings SET value = $1, updated_at = now() WHERE key = 'setup.site_name'", + let admin_user = repo::create_first_admin( + &mut tx, + &FirstAdmin { + email: &admin_email, + display_name: &req.admin.display_name, + password_hash: &password_hash, + key_prefix: &generated.prefix, + key_hash: &generated.hash, + key_name: "Default Admin Key", + key_surfaces: super::api_keys::ALLOWED_SURFACES, + site_name, + }, ) - .bind(serde_json::json!(site_name)) - .execute(&mut *tx) .await?; tx.commit().await?; diff --git a/crates/server/src/handlers/sso.rs b/crates/server/src/handlers/sso.rs index 18cc3928..9ac20bc9 100644 --- a/crates/server/src/handlers/sso.rs +++ b/crates/server/src/handlers/sso.rs @@ -9,9 +9,9 @@ use think_watch_common::audit::AuditActor; use think_watch_common::config::AppConfig; use think_watch_common::crypto::parse_encryption_key; use think_watch_common::errors::AppError; -use think_watch_common::models::User; use crate::app::AppState; +use crate::services::auth_repository; const OIDC_STATE_KEY_PREFIX: &str = "oidc:state:"; const OIDC_STATE_TTL_SECS: i64 = 600; @@ -322,13 +322,9 @@ async fn handle_live_callback( .await .map_err(|e| AppError::BadRequest(format!("SSO authentication failed: {e}")))?; - let user = sqlx::query_as::<_, User>( - "SELECT * FROM users WHERE oidc_subject = $1 AND oidc_issuer = $2", - ) - .bind(&user_info.subject) - .bind(&user_info.issuer) - .fetch_optional(&state.db) - .await?; + let user = + auth_repository::find_by_oidc_identity(&state.db, &user_info.subject, &user_info.issuer) + .await?; let user = match user { Some(u) if u.deleted_at.is_some() => { @@ -395,26 +391,17 @@ async fn handle_live_callback( .as_deref() .unwrap_or(user_info.email.as_deref().unwrap_or(&email)); - let u = sqlx::query_as::<_, User>( - r#"INSERT INTO users (email, display_name, oidc_subject, oidc_issuer) - VALUES ($1, $2, $3, $4) RETURNING *"#, + let u = auth_repository::insert_oidc_user( + &state.db, + &email, + display_name, + &user_info.subject, + &user_info.issuer, ) - .bind(&email) - .bind(display_name) - .bind(&user_info.subject) - .bind(&user_info.issuer) - .fetch_one(&state.db) .await?; if let Some(role_name) = state.dynamic_config.default_role().await { - sqlx::query( - r#"INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) - SELECT $1, id, 'global', $1 FROM rbac_roles WHERE name = $2"#, - ) - .bind(u.id) - .bind(&role_name) - .execute(&state.db) - .await?; + auth_repository::assign_default_role(&state.db, u.id, &role_name).await?; } u diff --git a/crates/server/src/services/api_key_repository.rs b/crates/server/src/services/api_key_repository.rs new file mode 100644 index 00000000..0c8390d5 --- /dev/null +++ b/crates/server/src/services/api_key_repository.rs @@ -0,0 +1,345 @@ +//! API key repository — the `api_keys` table. Keys are soft-deleted +//! (`deleted_at`); "live" means not deleted. Revoked keys stay in the +//! table, archived, until the retention sweep hard-deletes them. +//! +//! Key generation, permission checks and allow-list validation stay in +//! `handlers::api_keys`. + +use chrono::{DateTime, Utc}; +use sqlx::PgPool; +use think_watch_common::errors::AppError; +use think_watch_common::models::ApiKey; +use uuid::Uuid; + +/// The owner of a live key. +pub async fn owner_of_live(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok( + sqlx::query_scalar("SELECT user_id FROM api_keys WHERE id = $1 AND deleted_at IS NULL") + .bind(id) + .fetch_optional(pool) + .await?, + ) +} + +/// Whether `user_id` holds an MCP credential labelled `account_label` +/// for the server — what a key's `mcp_account_overrides` may point at. +pub async fn mcp_credential_exists( + pool: &PgPool, + server_id: Uuid, + user_id: Uuid, + account_label: &str, +) -> Result { + let exists: Option = sqlx::query_scalar( + "SELECT 1 FROM mcp_user_credentials + WHERE mcp_server_id = $1 AND user_id = $2 AND account_label = $3", + ) + .bind(server_id) + .bind(user_id) + .bind(account_label) + .fetch_optional(pool) + .await?; + Ok(exists.is_some()) +} + +/// The list's row filter: live keys, or the revoke-archived ones. The +/// archived view leaves out keys soft-deleted along with their user — +/// those were not revoked. +fn visibility_clause(archived: bool) -> &'static str { + if archived { + "deleted_at IS NOT NULL \ + AND (disabled_reason = 'revoked' OR disabled_reason LIKE 'force_revoked:%')" + } else { + "deleted_at IS NULL" + } +} + +/// How many keys the list shows, across every user. +pub async fn count_all(pool: &PgPool, archived: bool) -> Result { + let visibility_clause = visibility_clause(archived); + Ok(sqlx::query_scalar(&format!( + "SELECT COUNT(*) FROM api_keys WHERE {visibility_clause}" + )) + .fetch_one(pool) + .await?) +} + +/// One page of the list across every user, newest first. +pub async fn list_all_page( + pool: &PgPool, + archived: bool, + limit: i64, + offset: i64, +) -> Result, AppError> { + let visibility_clause = visibility_clause(archived); + Ok(sqlx::query_as::<_, ApiKey>(&format!( + "SELECT * FROM api_keys WHERE {visibility_clause} \ + ORDER BY created_at DESC LIMIT $1 OFFSET $2" + )) + .bind(limit) + .bind(offset) + .fetch_all(pool) + .await?) +} + +/// How many of one user's keys the list shows. +pub async fn count_for_user(pool: &PgPool, archived: bool, user_id: Uuid) -> Result { + let visibility_clause = visibility_clause(archived); + Ok(sqlx::query_scalar(&format!( + "SELECT COUNT(*) FROM api_keys WHERE {visibility_clause} AND user_id = $1" + )) + .bind(user_id) + .fetch_one(pool) + .await?) +} + +/// One page of one user's keys, newest first. +pub async fn list_for_user_page( + pool: &PgPool, + archived: bool, + user_id: Uuid, + limit: i64, + offset: i64, +) -> Result, AppError> { + let visibility_clause = visibility_clause(archived); + Ok(sqlx::query_as::<_, ApiKey>(&format!( + "SELECT * FROM api_keys WHERE {visibility_clause} AND user_id = $1 \ + ORDER BY created_at DESC LIMIT $2 OFFSET $3" + )) + .bind(user_id) + .bind(limit) + .bind(offset) + .fetch_all(pool) + .await?) +} + +/// A key about to be created. `id` doubles as the lineage id: a new key +/// is the root of its own rotation chain. +pub struct NewApiKey<'a> { + pub id: Uuid, + pub key_prefix: &'a str, + pub key_hash: &'a str, + pub name: &'a str, + pub user_id: Uuid, + pub surfaces: &'a [String], + pub allowed_models: &'a Option>, + pub allowed_mcp_tools: &'a Option>, + pub mcp_account_overrides: &'a serde_json::Value, + pub expires_at: Option>, + pub cost_center: Option<&'a str>, + pub rotation_period_days: Option, +} + +pub async fn insert(pool: &PgPool, key: &NewApiKey<'_>) -> Result { + Ok(sqlx::query_as::<_, ApiKey>( + r#"INSERT INTO api_keys (id, lineage_id, key_prefix, key_hash, name, user_id, surfaces, + allowed_models, allowed_mcp_tools, mcp_account_overrides, expires_at, + cost_center, rotation_period_days) + VALUES ($1, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING *"#, + ) + .bind(key.id) + .bind(key.key_prefix) + .bind(key.key_hash) + .bind(key.name) + .bind(key.user_id) + .bind(key.surfaces) + .bind(key.allowed_models) + .bind(key.allowed_mcp_tools) + .bind(key.mcp_account_overrides) + .bind(key.expires_at) + .bind(key.cost_center) + .bind(key.rotation_period_days) + .fetch_one(pool) + .await?) +} + +pub async fn find_live(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok( + sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys WHERE id = $1 AND deleted_at IS NULL") + .bind(id) + .fetch_optional(pool) + .await?, + ) +} + +/// Revoke and archive a live key, ending any rotation grace window. +/// Returns how many rows changed (0 when the key is gone). +pub async fn revoke(pool: &PgPool, id: Uuid) -> Result { + Ok(sqlx::query( + "UPDATE api_keys SET is_active = false, grace_period_ends_at = NULL, \ + disabled_reason = 'revoked', deleted_at = now() \ + WHERE id = $1 AND deleted_at IS NULL", + ) + .bind(id) + .execute(pool) + .await? + .rows_affected()) +} + +/// Like [`revoke`], recording `disabled_reason` (a `force_revoked:` tag). +pub async fn force_revoke(pool: &PgPool, id: Uuid, disabled_reason: &str) -> Result { + Ok(sqlx::query( + "UPDATE api_keys SET is_active = false, grace_period_ends_at = NULL, \ + disabled_reason = $1, deleted_at = now() \ + WHERE id = $2 AND deleted_at IS NULL", + ) + .bind(disabled_reason) + .bind(id) + .execute(pool) + .await? + .rows_affected()) +} + +/// A PATCH to a key's settings. Each `*_set` flag says whether its value +/// replaces the column (a `None` value then clears it); `surfaces`, +/// `rotation_period_days` and `inactivity_timeout_days` keep the column +/// when `None`. `expires_at` is always written. +pub struct ApiKeyPatch<'a> { + pub allowed_models_set: bool, + pub allowed_models: Option<&'a [String]>, + pub allowed_mcp_tools_set: bool, + pub allowed_mcp_tools: Option<&'a [String]>, + pub surfaces: Option<&'a Vec>, + pub expires_at: Option>, + pub rotation_period_days: Option, + pub inactivity_timeout_days: Option, + pub cost_center_set: bool, + pub cost_center: Option<&'a str>, + pub mcp_account_overrides_set: bool, + pub mcp_account_overrides: &'a serde_json::Value, + /// Reset the expiry-warning dedupe, for an expiry pushed later. + pub expiry_extended: bool, +} + +pub async fn update(pool: &PgPool, id: Uuid, patch: &ApiKeyPatch<'_>) -> Result { + Ok(sqlx::query_as::<_, ApiKey>( + r#"UPDATE api_keys SET + allowed_models = CASE WHEN $11 THEN $1 ELSE allowed_models END, + allowed_mcp_tools = CASE WHEN $12 THEN $10 ELSE allowed_mcp_tools END, + surfaces = COALESCE($2, surfaces), + expires_at = $3, + rotation_period_days = COALESCE($4, rotation_period_days), + inactivity_timeout_days = COALESCE($5, inactivity_timeout_days), + cost_center = CASE WHEN $7 THEN $6 ELSE cost_center END, + mcp_account_overrides = CASE WHEN $13 THEN $14 ELSE mcp_account_overrides END, + last_expiry_warning_days = CASE WHEN $9 THEN NULL + ELSE last_expiry_warning_days END + WHERE id = $8 RETURNING *"#, + ) + .bind(patch.allowed_models) + .bind(patch.surfaces) + .bind(patch.expires_at) + .bind(patch.rotation_period_days) + .bind(patch.inactivity_timeout_days) + .bind(patch.cost_center) + .bind(patch.cost_center_set) + .bind(id) + .bind(patch.expiry_extended) + .bind(patch.allowed_mcp_tools) + .bind(patch.allowed_models_set) + .bind(patch.allowed_mcp_tools_set) + .bind(patch.mcp_account_overrides_set) + .bind(patch.mcp_account_overrides) + .fetch_one(pool) + .await?) +} + +/// Rotate `old_key`: insert its successor (same name, owner, scope, +/// expiry and lineage; `rotated_from_id` pointing back) and put the old +/// key into its grace window, in one transaction — otherwise a failure +/// between the two would leave both keys valid with no grace end. +/// Returns the new key. +pub async fn rotate( + pool: &PgPool, + old_key: &ApiKey, + key_prefix: &str, + key_hash: &str, + grace_period_ends_at: DateTime, +) -> Result { + let mut tx = pool.begin().await?; + + let new_key = sqlx::query_as::<_, ApiKey>( + r#"INSERT INTO api_keys (key_prefix, key_hash, name, user_id, surfaces, allowed_models, + allowed_mcp_tools, expires_at, rotation_period_days, inactivity_timeout_days, + cost_center, rotated_from_id, last_rotation_at, lineage_id) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, now(), $13) + RETURNING *"#, + ) + .bind(key_prefix) + .bind(key_hash) + .bind(&old_key.name) + .bind(old_key.user_id) + .bind(&old_key.surfaces) + .bind(&old_key.allowed_models) + .bind(&old_key.allowed_mcp_tools) + .bind(old_key.expires_at) + .bind(old_key.rotation_period_days) + .bind(old_key.inactivity_timeout_days) + .bind(old_key.cost_center.as_deref()) + .bind(old_key.id) + .bind(old_key.lineage_id) + .fetch_one(&mut *tx) + .await?; + + sqlx::query( + "UPDATE api_keys SET grace_period_ends_at = $1, disabled_reason = 'rotated' WHERE id = $2", + ) + .bind(grace_period_ends_at) + .bind(old_key.id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(new_key) +} + +/// Active live keys across every user that expire by `threshold`, +/// soonest first. +pub async fn list_expiring_all( + pool: &PgPool, + threshold: DateTime, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, ApiKey>( + r#"SELECT * FROM api_keys + WHERE is_active = true + AND deleted_at IS NULL + AND expires_at IS NOT NULL + AND expires_at <= $1 + ORDER BY expires_at ASC"#, + ) + .bind(threshold) + .fetch_all(pool) + .await?) +} + +/// [`list_expiring_all`], for one user's keys. +pub async fn list_expiring_for_user( + pool: &PgPool, + threshold: DateTime, + user_id: Uuid, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, ApiKey>( + r#"SELECT * FROM api_keys + WHERE is_active = true + AND deleted_at IS NULL + AND expires_at IS NOT NULL + AND expires_at <= $1 + AND user_id = $2 + ORDER BY expires_at ASC"#, + ) + .bind(threshold) + .bind(user_id) + .fetch_all(pool) + .await?) +} + +/// Distinct cost-center tags on live keys, alphabetical. +pub async fn cost_centers(pool: &PgPool) -> Result, AppError> { + let rows: Vec<(String,)> = sqlx::query_as( + "SELECT DISTINCT cost_center FROM api_keys \ + WHERE cost_center IS NOT NULL AND deleted_at IS NULL \ + ORDER BY cost_center ASC", + ) + .fetch_all(pool) + .await?; + Ok(rows.into_iter().map(|(s,)| s).collect()) +} diff --git a/crates/server/src/services/auth_repository.rs b/crates/server/src/services/auth_repository.rs new file mode 100644 index 00000000..800cf97a --- /dev/null +++ b/crates/server/src/services/auth_repository.rs @@ -0,0 +1,240 @@ +//! Auth repository — what login, registration, SSO and the account +//! endpoints (`/api/auth/*`) read and write: the caller's `users` row, +//! their role assignments and team memberships, and self-service +//! account deletion. +//! +//! Password hashing, TOTP crypto, lockouts and sessions stay in +//! `handlers::auth` / `handlers::sso` and their services. + +use sqlx::{PgConnection, PgExecutor, PgPool}; +use think_watch_common::errors::AppError; +use think_watch_common::models::User; +use uuid::Uuid; + +/// A role assignment as `/me` lists it: role id, name, whether it is a +/// system role, scope kind and scope id. +pub type RoleAssignmentRow = (Uuid, String, bool, String, Option); + +/// The active, not deleted user with this (normalized) email. +pub async fn find_active_by_email(pool: &PgPool, email: &str) -> Result, AppError> { + Ok(sqlx::query_as::<_, User>( + "SELECT * FROM users WHERE email = $1 AND is_active = true AND deleted_at IS NULL", + ) + .bind(email) + .fetch_optional(pool) + .await?) +} + +/// The active, not deleted user with this id. +pub async fn find_active(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok(sqlx::query_as::<_, User>( + "SELECT * FROM users WHERE id = $1 AND is_active = true AND deleted_at IS NULL", + ) + .bind(id) + .fetch_optional(pool) + .await?) +} + +/// The user an SSO identity maps to — deleted or deactivated included, +/// so the caller can refuse those explicitly. +pub async fn find_by_oidc_identity( + pool: &PgPool, + subject: &str, + issuer: &str, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, User>( + "SELECT * FROM users WHERE oidc_subject = $1 AND oidc_issuer = $2", + ) + .bind(subject) + .bind(issuer) + .fetch_optional(pool) + .await?) +} + +/// `is_active` of a not deleted user. +pub async fn is_active(pool: &PgPool, id: Uuid) -> Result, AppError> { + Ok( + sqlx::query_scalar("SELECT is_active FROM users WHERE id = $1 AND deleted_at IS NULL") + .bind(id) + .fetch_optional(pool) + .await?, + ) +} + +/// `totp_enabled` of a not deleted user. +pub async fn totp_enabled(pool: &PgPool, id: Uuid) -> Result { + Ok( + sqlx::query_scalar("SELECT totp_enabled FROM users WHERE id = $1 AND deleted_at IS NULL") + .bind(id) + .fetch_one(pool) + .await?, + ) +} + +/// Replace the (encrypted) recovery codes if they are still `expected` +/// — a compare-and-swap, so of two requests spending the same code only +/// one wins. Returns how many rows changed (1 = this request won). +pub async fn swap_recovery_codes( + pool: &PgPool, + id: Uuid, + expected: &str, + updated: &str, +) -> Result { + Ok(sqlx::query( + "UPDATE users SET totp_recovery_codes = $1 \ + WHERE id = $2 AND totp_recovery_codes = $3", + ) + .bind(updated) + .bind(id) + .bind(expected) + .execute(pool) + .await? + .rows_affected()) +} + +/// Insert a self-registered user, or nothing when the email is taken. +pub async fn insert_user_unless_taken( + conn: &mut PgConnection, + email: &str, + display_name: &str, + password_hash: &str, +) -> Result, AppError> { + Ok(sqlx::query_as::<_, User>( + r#"INSERT INTO users (email, display_name, password_hash) + VALUES ($1, $2, $3) + ON CONFLICT (email) DO NOTHING + RETURNING *"#, + ) + .bind(email) + .bind(display_name) + .bind(password_hash) + .fetch_optional(conn) + .await?) +} + +/// Insert a user provisioned by SSO. +pub async fn insert_oidc_user( + pool: &PgPool, + email: &str, + display_name: &str, + subject: &str, + issuer: &str, +) -> Result { + Ok(sqlx::query_as::<_, User>( + r#"INSERT INTO users (email, display_name, oidc_subject, oidc_issuer) + VALUES ($1, $2, $3, $4) RETURNING *"#, + ) + .bind(email) + .bind(display_name) + .bind(subject) + .bind(issuer) + .fetch_one(pool) + .await?) +} + +/// Give a new user the named role at global scope, self-assigned. A +/// role name that does not exist assigns nothing. +pub async fn assign_default_role<'e>( + executor: impl PgExecutor<'e>, + user_id: Uuid, + role_name: &str, +) -> Result<(), AppError> { + sqlx::query( + r#"INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) + SELECT $1, id, 'global', $1 FROM rbac_roles WHERE name = $2"#, + ) + .bind(user_id) + .bind(role_name) + .execute(executor) + .await?; + Ok(()) +} + +/// The user's teams, by name. +pub async fn teams_of(pool: &PgPool, user_id: Uuid) -> Result, AppError> { + Ok(sqlx::query_as( + "SELECT t.id, t.name FROM team_members tm \ + JOIN teams t ON t.id = tm.team_id \ + WHERE tm.user_id = $1 \ + ORDER BY t.name ASC", + ) + .bind(user_id) + .fetch_all(pool) + .await?) +} + +/// The user's role assignments, system roles first, then by name. +pub async fn role_assignments_of( + pool: &PgPool, + user_id: Uuid, +) -> Result, AppError> { + Ok(sqlx::query_as( + "SELECT r.id, r.name, r.is_system, ra.scope_kind, ra.scope_id \ + FROM rbac_role_assignments ra \ + JOIN rbac_roles r ON r.id = ra.role_id \ + WHERE ra.user_id = $1 \ + ORDER BY r.is_system DESC, r.name ASC", + ) + .bind(user_id) + .fetch_all(pool) + .await?) +} + +/// Set a password the user chose themselves (no forced change next +/// login). +pub async fn set_own_password( + pool: &PgPool, + id: Uuid, + password_hash: &str, +) -> Result<(), AppError> { + sqlx::query("UPDATE users SET password_hash = $1, password_change_required = false, updated_at = now() WHERE id = $2") + .bind(password_hash) + .bind(id) + .execute(pool) + .await?; + Ok(()) +} + +/// Soft-delete the user and every key they own, in one transaction. +pub async fn soft_delete_account(pool: &PgPool, user_id: Uuid) -> Result<(), AppError> { + let mut tx = pool.begin().await?; + sqlx::query("UPDATE api_keys SET is_active = false, deleted_at = now(), disabled_reason = 'account_deleted' WHERE user_id = $1") + .bind(user_id) + .execute(&mut *tx) + .await?; + sqlx::query("UPDATE users SET is_active = false, deleted_at = now() WHERE id = $1") + .bind(user_id) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(()) +} + +/// Turn TOTP on with the (encrypted) secret and recovery codes. +pub async fn enable_totp( + pool: &PgPool, + id: Uuid, + encrypted_secret: &str, + encrypted_recovery_codes: &str, +) -> Result<(), AppError> { + sqlx::query( + "UPDATE users SET totp_secret = $1, totp_enabled = true, totp_recovery_codes = $2, updated_at = now() WHERE id = $3", + ) + .bind(encrypted_secret) + .bind(encrypted_recovery_codes) + .bind(id) + .execute(pool) + .await?; + Ok(()) +} + +/// Turn TOTP off, dropping the secret and recovery codes. +pub async fn disable_totp(pool: &PgPool, id: Uuid) -> Result<(), AppError> { + sqlx::query( + "UPDATE users SET totp_secret = NULL, totp_enabled = false, totp_recovery_codes = NULL, updated_at = now() WHERE id = $1", + ) + .bind(id) + .execute(pool) + .await?; + Ok(()) +} diff --git a/crates/server/src/services/mod.rs b/crates/server/src/services/mod.rs index 8ae5afba..d84cc2df 100644 --- a/crates/server/src/services/mod.rs +++ b/crates/server/src/services/mod.rs @@ -26,7 +26,9 @@ //! a service yet; the migration is iterative. pub mod analytics_repository; +pub mod api_key_repository; pub mod auth_lockout; +pub mod auth_repository; pub mod limits_repository; pub mod log_forwarder_repository; pub mod mcp_credential_repository; @@ -41,6 +43,8 @@ pub mod rbac_service; pub mod refresh_blacklist; pub mod role_repository; pub mod session_service; +pub mod settings_repository; +pub mod setup_repository; pub mod team_repository; pub mod totp_service; pub mod user_repository; diff --git a/crates/server/src/services/settings_repository.rs b/crates/server/src/services/settings_repository.rs new file mode 100644 index 00000000..ce0a958c --- /dev/null +++ b/crates/server/src/services/settings_repository.rs @@ -0,0 +1,24 @@ +//! Settings repository — the admin settings handlers' direct reads and +//! writes. Ordinary settings go through `DynamicConfig`; these are the +//! few statements that bypass it. + +use sqlx::PgPool; +use think_watch_common::errors::AppError; + +/// Drop the OIDC wizard's draft (`oidc.draft`), if any. +pub async fn delete_oidc_draft(pool: &PgPool) -> Result<(), AppError> { + sqlx::query("DELETE FROM system_settings WHERE key = 'oidc.draft'") + .execute(pool) + .await?; + Ok(()) +} + +/// Whether a role with this name exists — what `auth.default_role` may +/// name. +pub async fn role_exists(pool: &PgPool, name: &str) -> Result { + let exists: Option<(String,)> = sqlx::query_as("SELECT name FROM rbac_roles WHERE name = $1") + .bind(name) + .fetch_optional(pool) + .await?; + Ok(exists.is_some()) +} diff --git a/crates/server/src/services/setup_repository.rs b/crates/server/src/services/setup_repository.rs new file mode 100644 index 00000000..000cf97e --- /dev/null +++ b/crates/server/src/services/setup_repository.rs @@ -0,0 +1,98 @@ +//! Setup repository — the first-boot wizard's writes: the first super +//! admin, their first API key, and the `setup.*` settings. +//! +//! Everything here runs on the caller's transaction, which holds the +//! setup advisory lock for its whole length so two concurrent setups +//! cannot both pass the "not initialized yet" check. + +use sqlx::PgConnection; +use think_watch_common::errors::AppError; +use uuid::Uuid; + +/// Take the setup advisory lock (key 1) until the transaction ends. +pub async fn lock_setup(conn: &mut PgConnection) -> Result<(), AppError> { + sqlx::query("SELECT pg_advisory_xact_lock(1)") + .execute(conn) + .await?; + Ok(()) +} + +/// The stored `setup.initialized` value, read from the database rather +/// than the settings cache. +pub async fn initialized_flag( + conn: &mut PgConnection, +) -> Result, AppError> { + Ok( + sqlx::query_scalar("SELECT value FROM system_settings WHERE key = 'setup.initialized'") + .fetch_optional(conn) + .await?, + ) +} + +/// The first super admin and their first key. +pub struct FirstAdmin<'a> { + pub email: &'a str, + pub display_name: &'a str, + pub password_hash: &'a str, + pub key_prefix: &'a str, + pub key_hash: &'a str, + pub key_name: &'a str, + pub key_surfaces: &'a [&'a str], + pub site_name: &'a str, +} + +/// Create the super admin (global `super_admin` role) and their key, +/// then mark setup done and store the site name. Returns the admin's id +/// and email. +pub async fn create_first_admin( + conn: &mut PgConnection, + admin: &FirstAdmin<'_>, +) -> Result<(Uuid, String), AppError> { + let admin_user = sqlx::query_as::<_, (uuid::Uuid, String)>( + r#"INSERT INTO users (email, display_name, password_hash) + VALUES ($1, $2, $3) RETURNING id, email"#, + ) + .bind(admin.email) + .bind(admin.display_name) + .bind(admin.password_hash) + .fetch_one(&mut *conn) + .await?; + // A taken email surfaces as `AppError::Conflict` via + // `From`. + + sqlx::query( + r#"INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) + SELECT $1, id, 'global', $1 FROM rbac_roles WHERE name = 'super_admin'"#, + ) + .bind(admin_user.0) + .execute(&mut *conn) + .await?; + + sqlx::query( + r#"INSERT INTO api_keys (key_prefix, key_hash, name, user_id, surfaces) + VALUES ($1, $2, $3, $4, $5)"#, + ) + .bind(admin.key_prefix) + .bind(admin.key_hash) + .bind(admin.key_name) + .bind(admin_user.0) + .bind(admin.key_surfaces) + .execute(&mut *conn) + .await?; + + sqlx::query( + "UPDATE system_settings SET value = $1, updated_at = now() WHERE key = 'setup.initialized'", + ) + .bind(serde_json::json!(true)) + .execute(&mut *conn) + .await?; + + sqlx::query( + "UPDATE system_settings SET value = $1, updated_at = now() WHERE key = 'setup.site_name'", + ) + .bind(serde_json::json!(admin.site_name)) + .execute(&mut *conn) + .await?; + + Ok(admin_user) +} diff --git a/crates/test-support/Cargo.toml b/crates/test-support/Cargo.toml index 238b3802..d491dce3 100644 --- a/crates/test-support/Cargo.toml +++ b/crates/test-support/Cargo.toml @@ -54,4 +54,5 @@ bytes = { workspace = true } rust_decimal = { workspace = true } [dev-dependencies] +jsonwebtoken = { workspace = true } tokio = { workspace = true, features = ["test-util"] } diff --git a/crates/test-support/tests/admin_access.rs b/crates/test-support/tests/admin_access.rs new file mode 100644 index 00000000..cdefe90b --- /dev/null +++ b/crates/test-support/tests/admin_access.rs @@ -0,0 +1,722 @@ +//! The access endpoints end to end: API keys, the account endpoints +//! under `/api/auth`, the default-role setting, and SSO sign-in against +//! a mock identity provider. +//! +//! Login, refresh, password change, TOTP set-up / disable / recovery, +//! account deletion, first-boot setup and key rotation have their own +//! files; this one covers what they don't reach, so moving the access +//! handlers' SQL around (into `services::*_repository`) is checked +//! rather than assumed. + +use serde_json::Value; +use think_watch_test_support::prelude::*; +use wiremock::matchers::{body_string_contains, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +async fn login(app: &TestApp, user: &fixtures::SeededUser) -> TestClient { + let con = app.console_client(); + con.post( + "/api/auth/login", + json!({"email": user.user.email, "password": user.plaintext_password}), + ) + .await + .unwrap() + .assert_ok(); + con +} + +async fn create_key(con: &TestClient, body: Value) -> Value { + let resp = con.post("/api/keys", body).await.unwrap(); + resp.assert_ok(); + resp.json().unwrap() +} + +async fn get(con: &TestClient, path: &str) -> Value { + let resp = con.get(path).await.unwrap(); + resp.assert_ok(); + resp.json().unwrap() +} + +async fn patch(con: &TestClient, path: &str, body: Value) -> Value { + let resp = con.patch(path, body).await.unwrap(); + resp.assert_ok(); + resp.json().unwrap() +} + +fn ids(list: &Value) -> Vec { + list.as_array() + .expect("array") + .iter() + .map(|k| k["id"].as_str().unwrap().to_string()) + .collect() +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_key_is_created_read_and_patched() { + let app = TestApp::spawn().await; + let (con, admin) = admin_session_with_user(&app).await; + + let name = unique_name("access-key"); + let created = create_key( + &con, + json!({ + "name": name, + "surfaces": ["mcp_gateway", "ai_gateway", "ai_gateway"], + "allowed_models": ["gpt-4o"], + "expires_in_days": 30, + "cost_center": " team-a ", + }), + ) + .await; + let id = created["id"].as_str().unwrap().to_string(); + assert_eq!(created["name"], name.as_str()); + let plaintext = created["key"].as_str().unwrap(); + assert!( + plaintext.starts_with(created["key_prefix"].as_str().unwrap()), + "{created}" + ); + + let key = get(&con, &format!("/api/keys/{id}")).await; + assert_eq!(key["user_id"], admin.user.id.to_string()); + assert_eq!(key["surfaces"], json!(["ai_gateway", "mcp_gateway"])); + assert_eq!(key["allowed_models"], json!(["gpt-4o"])); + assert!(key["allowed_mcp_tools"].is_null(), "{key}"); + assert_eq!(key["cost_center"], "team-a"); + assert_eq!(key["mcp_account_overrides"], json!({})); + assert!(key["expires_at"].is_string(), "{key}"); + assert!(key["is_active"].as_bool().unwrap()); + assert!(key.get("key_hash").is_none(), "{key}"); + assert!(key.get("lineage_id").is_none(), "{key}"); + // A new key is the root of its own lineage. + let (lineage_id,): (Uuid,) = sqlx::query_as("SELECT lineage_id FROM api_keys WHERE id = $1") + .bind(Uuid::parse_str(&id).unwrap()) + .fetch_one(&app.db) + .await + .unwrap(); + assert_eq!(lineage_id.to_string(), id); + + // Set and clear in one PATCH: null clears a list, "" clears the + // cost center, 0 clears the expiry. + let patched = patch( + &con, + &format!("/api/keys/{id}"), + json!({ + "allowed_models": null, + "allowed_mcp_tools": ["github__list_issues"], + "surfaces": ["console"], + "expires_in_days": 0, + "rotation_period_days": 45, + "inactivity_timeout_days": 10, + "cost_center": "", + }), + ) + .await; + assert!(patched["allowed_models"].is_null(), "{patched}"); + assert_eq!(patched["allowed_mcp_tools"], json!(["github__list_issues"])); + assert_eq!(patched["surfaces"], json!(["console"])); + assert!(patched["expires_at"].is_null(), "{patched}"); + assert_eq!(patched["rotation_period_days"], 45); + assert_eq!(patched["inactivity_timeout_days"], 10); + assert!(patched["cost_center"].is_null(), "{patched}"); + + // Absent fields are left alone. + let untouched = patch(&con, &format!("/api/keys/{id}"), json!({})).await; + assert_eq!( + untouched["allowed_mcp_tools"], + json!(["github__list_issues"]) + ); + assert_eq!(untouched["surfaces"], json!(["console"])); + assert_eq!(untouched["rotation_period_days"], 45); + assert_eq!(untouched["inactivity_timeout_days"], 10); + assert!(untouched["expires_at"].is_null(), "{untouched}"); + + let later = patch( + &con, + &format!("/api/keys/{id}"), + json!({"expires_in_days": 5, "cost_center": "team-b", "mcp_account_overrides": {}}), + ) + .await; + assert!(later["expires_at"].is_string(), "{later}"); + assert_eq!(later["cost_center"], "team-b"); + assert_eq!(later["mcp_account_overrides"], json!({})); + + // Refused input. + for body in [ + json!({"expires_in_days": -1}), + json!({"rotation_period_days": -1}), + json!({"surfaces": []}), + json!({"surfaces": ["nope"]}), + json!({"cost_center": "x".repeat(65)}), + json!({"mcp_account_overrides": ["not", "an", "object"]}), + json!({"mcp_account_overrides": {"not-a-uuid": "work"}}), + json!({"mcp_account_overrides": {Uuid::new_v4().to_string(): "work"}}), + json!({"mcp_account_overrides": {Uuid::new_v4().to_string(): 7}}), + ] { + con.patch(&format!("/api/keys/{id}"), body.clone()) + .await + .unwrap() + .assert_status(400); + } + for body in [ + json!({"name": "k", "surfaces": []}), + json!({"name": "k", "surfaces": ["nope"]}), + json!({"name": "k", "surfaces": ["ai_gateway"], "expires_in_days": -1}), + json!({"name": "k", "surfaces": ["ai_gateway"], "mcp_account_overrides": {Uuid::new_v4().to_string(): "work"}}), + ] { + con.post("/api/keys", body) + .await + .unwrap() + .assert_status(400); + } + + let missing = format!("/api/keys/{}", Uuid::new_v4()); + con.get(&missing).await.unwrap().assert_status(404); + con.patch(&missing, json!({})) + .await + .unwrap() + .assert_status(404); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_key_is_created_with_the_default_expiry_setting() { + let app = TestApp::spawn().await; + let con = admin_session(&app).await; + + app.set_setting("api_keys.default_expiry_days", json!(0)) + .await; + let never = create_key(&con, json!({"name": "never", "surfaces": ["ai_gateway"]})).await; + let never = get( + &con, + &format!("/api/keys/{}", never["id"].as_str().unwrap()), + ) + .await; + assert!(never["expires_at"].is_null(), "{never}"); + + app.set_setting("api_keys.default_expiry_days", json!(10)) + .await; + app.set_setting("api_keys.rotation_period_days", json!(20)) + .await; + let expiring = create_key( + &con, + json!({"name": "expiring", "surfaces": ["ai_gateway"]}), + ) + .await; + let expiring = get( + &con, + &format!("/api/keys/{}", expiring["id"].as_str().unwrap()), + ) + .await; + assert!(expiring["expires_at"].is_string(), "{expiring}"); + assert_eq!(expiring["rotation_period_days"], 20); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_key_list_pages_revokes_and_archives() { + let app = TestApp::spawn().await; + let (admin, _) = admin_session_with_user(&app).await; + let dev_user = fixtures::create_random_user(&app.db).await.unwrap(); + let dev = login(&app, &dev_user).await; + + let dev_key = create_key(&dev, json!({"name": "dev", "surfaces": ["ai_gateway"]})).await; + let dev_key = dev_key["id"].as_str().unwrap().to_string(); + let admin_key = create_key(&admin, json!({"name": "admin", "surfaces": ["ai_gateway"]})).await; + let admin_key = admin_key["id"].as_str().unwrap().to_string(); + + // A developer sees their own keys; the admin tier sees everyone's, + // newest first, a page at a time. + let mine = get(&dev, "/api/keys").await; + assert_eq!(mine["total"], 1); + assert_eq!(ids(&mine["data"]), vec![dev_key.clone()]); + let page1 = get(&admin, "/api/keys?per_page=1").await; + assert_eq!(page1["total"], 2); + assert_eq!(page1["page"], 1); + assert_eq!(page1["per_page"], 1); + assert_eq!(ids(&page1["data"]), vec![admin_key.clone()]); + let page2 = get(&admin, "/api/keys?per_page=1&page=2").await; + assert_eq!(ids(&page2["data"]), vec![dev_key.clone()]); + + // Somebody else's key is out of a developer's reach. + dev.get(&format!("/api/keys/{admin_key}")) + .await + .unwrap() + .assert_status(403); + + // Revoke (developers lack `api_keys:delete`; the admin tier may + // revoke anyone's key): once, then it's gone. + dev.delete(&format!("/api/keys/{dev_key}")) + .await + .unwrap() + .assert_status(403); + let resp: Value = admin + .delete(&format!("/api/keys/{dev_key}")) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(resp["status"], "revoked"); + admin + .delete(&format!("/api/keys/{dev_key}")) + .await + .unwrap() + .assert_status(404); + + // Force-revoke needs a reason, and records it. + admin + .post( + &format!("/api/admin/keys/{admin_key}/force-revoke"), + json!({"reason": " "}), + ) + .await + .unwrap() + .assert_status(400); + let resp: Value = admin + .post( + &format!("/api/admin/keys/{admin_key}/force-revoke"), + json!({"reason": "suspected leak"}), + ) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(resp["status"], "force_revoked"); + assert_eq!(resp["reason"], "suspected leak"); + admin + .post( + &format!("/api/admin/keys/{admin_key}/force-revoke"), + json!({"reason": "again"}), + ) + .await + .unwrap() + .assert_status(404); + + // A key that went with its deleted account is not "revoked". + let leaver = fixtures::create_random_user(&app.db).await.unwrap(); + fixtures::create_api_key( + &app.db, + leaver.user.id, + "leaver", + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + let leaver_con = login(&app, &leaver).await; + leaver_con + .delete("/api/auth/account") + .await + .unwrap() + .assert_ok(); + + let live = get(&admin, "/api/keys").await; + assert_eq!(live["total"], 0, "{live}"); + let archived = get(&admin, "/api/keys?archived=true").await; + assert_eq!(archived["total"], 2, "{archived}"); + let mut reasons: Vec = archived["data"] + .as_array() + .unwrap() + .iter() + .map(|k| k["disabled_reason"].as_str().unwrap().to_string()) + .collect(); + reasons.sort(); + assert_eq!(reasons, vec!["force_revoked:suspected leak", "revoked"]); + let dev_archived = get(&dev, "/api/keys?archived=true").await; + assert_eq!(dev_archived["total"], 1); + assert_eq!(ids(&dev_archived["data"]), vec![dev_key]); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn expiring_keys_cost_centers_and_policy_scope() { + let app = TestApp::spawn().await; + let admin = admin_session(&app).await; + let dev_user = fixtures::create_random_user(&app.db).await.unwrap(); + let dev = login(&app, &dev_user).await; + + let key = |name: &str, days: i32, cost_center: &str| { + json!({ + "name": name, + "surfaces": ["ai_gateway"], + "expires_in_days": days, + "cost_center": cost_center, + }) + }; + let in3 = create_key(&admin, key("in3", 3, "zeta")).await; + let in20 = create_key(&admin, key("in20", 20, "alpha")).await; + let dev_in2 = create_key(&dev, key("dev-in2", 2, "alpha")).await; + let gone = create_key(&admin, key("gone", 1, "beta")).await; + admin + .delete(&format!("/api/keys/{}", gone["id"].as_str().unwrap())) + .await + .unwrap() + .assert_ok(); + let id = |v: &Value| v["id"].as_str().unwrap().to_string(); + + // Soonest first; revoked keys never show. + let week = get(&admin, "/api/keys/expiring").await; + assert_eq!(ids(&week), vec![id(&dev_in2), id(&in3)]); + let month = get(&admin, "/api/keys/expiring?days=30").await; + assert_eq!(ids(&month), vec![id(&dev_in2), id(&in3), id(&in20)]); + let none = get(&admin, "/api/keys/expiring?days=-5").await; + assert_eq!(ids(&none), Vec::::new()); + let dev_month = get(&dev, "/api/keys/expiring?days=30").await; + assert_eq!(ids(&dev_month), vec![id(&dev_in2)]); + + let centers = get(&admin, "/api/keys/cost-centers").await; + assert_eq!(centers, json!(["alpha", "zeta"])); + + let scope = get(&dev, "/api/keys/policy-scope").await; + assert!(scope.get("allowed_models").is_some(), "{scope}"); + assert!(scope.get("allowed_mcp_tools").is_some(), "{scope}"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn registration_assigns_the_default_role_and_me_lists_roles_and_teams() { + let app = TestApp::spawn().await; + let admin = admin_session(&app).await; + + // The setting has no seeded row; the admin API only changes rows + // that exist. + app.set_setting("auth.default_role", json!("")).await; + assert_eq!(app.state.dynamic_config.default_role().await, None); + // The default role must name a role that exists. + admin + .patch( + "/api/admin/settings", + json!({"settings": {"auth.default_role": "no-such-role"}}), + ) + .await + .unwrap() + .assert_status(400); + admin + .patch( + "/api/admin/settings", + json!({"settings": {"auth.default_role": "viewer"}}), + ) + .await + .unwrap() + .assert_ok(); + assert_eq!( + app.state.dynamic_config.default_role().await.as_deref(), + Some("viewer") + ); + app.set_setting("auth.allow_registration", json!(true)) + .await; + + let email = unique_email(); + let con = app.console_client(); + con.post( + "/api/auth/register", + json!({"email": email, "display_name": "Newcomer", "password": "Test_password_12345!"}), + ) + .await + .unwrap() + .assert_ok(); + assert!(con.cookie("__Host-access_token").is_some()); + + // Registering the same address again answers the same way, without + // a session. + let again = app.console_client(); + let resp: Value = again + .post( + "/api/auth/register", + json!({"email": email, "display_name": "Twin", "password": "Test_password_12345!"}), + ) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(resp["expires_in"], 0); + assert!(again.cookie("__Host-access_token").is_none()); + + let (user_id,): (Uuid,) = sqlx::query_as("SELECT id FROM users WHERE email = $1") + .bind(&email) + .fetch_one(&app.db) + .await + .unwrap(); + fixtures::assign_role_global(&app.db, user_id, "developer") + .await + .unwrap(); + for team in ["b-team", "a-team"] { + let team_name = unique_name(team); + sqlx::query( + "WITH t AS (INSERT INTO teams (name) VALUES ($1) RETURNING id) \ + INSERT INTO team_members (user_id, team_id) SELECT $2, id FROM t", + ) + .bind(&team_name) + .bind(user_id) + .execute(&app.db) + .await + .unwrap(); + } + + let me = get(&con, "/api/auth/me").await; + assert_eq!(me["id"], user_id.to_string()); + assert_eq!(me["email"], email.as_str()); + assert_eq!(me["display_name"], "Newcomer"); + let roles: Vec<(String, String)> = me["role_assignments"] + .as_array() + .unwrap() + .iter() + .map(|r| { + assert!(r["is_system"].as_bool().unwrap(), "{r}"); + ( + r["name"].as_str().unwrap().to_string(), + r["scope"].as_str().unwrap().to_string(), + ) + }) + .collect(); + assert_eq!( + roles, + vec![ + ("developer".to_string(), "global".to_string()), + ("viewer".to_string(), "global".to_string()), + ] + ); + let teams: Vec<&str> = me["teams"] + .as_array() + .unwrap() + .iter() + .map(|t| t["name"].as_str().unwrap()) + .collect(); + assert_eq!(teams.len(), 2, "{me}"); + assert!(teams[0].starts_with("a-team-") && teams[1].starts_with("b-team-")); + assert!(!me["permissions"].as_array().unwrap().is_empty(), "{me}"); + + let status = get(&con, "/api/auth/totp/status").await; + assert_eq!(status, json!({"enabled": false, "required": false})); +} + +// --- SSO against a mock identity provider --- + +/// Test-only RSA key the mock identity provider signs ID tokens with. +const IDP_KEY_PEM: &str = "\ +-----BEGIN RSA PRIVATE KEY-----\n\ +MIIEpAIBAAKCAQEAvDyl7CDoXwG8SqSFodeEVK3aGjpg7cmsweIZIblHOLA/Ftd5\n\ +D1XEOaAt7AuXjYQINz4zy3Nvcd3DCx42mCw4tdeGobZOSpGI7z5dq3rJFV+pjVGh\n\ +J5o7nLO3hipNbaiZKMrCFwybh/pSF9jtf6aP1nzMKEH9kTabxbnHyiEZuNJW07oH\n\ +jQhi7JdrM9+l+dfxWpnHSdyZ+6DBIG9jV9NB5fT8yJ+oGUuakUm38+7TwE1rR19L\n\ +x2HW7r41s09t2Qkzjo0E7McSF1nwJM+Ek0VS4eD3zqdoz1aHLUdWavzQZJ6lyCkN\n\ +A/FAdmQ912oMdlIqvfOp/2RiuVQ+QH/UJY8S5QIDAQABAoIBACUkZGr0zVUdzwz9\n\ +bJ7UGzjoOv5s4X5aCnwRRHs6h1qgsDouFyWW+0qRmC4Y1XUnhcV8wRSWePmDU/6I\n\ +HialpyT+W4LiKY2eLOJkMHBrIG1WvGp1nnJlhPi1H3PaOf/2wg3iAC0zICdTFcq9\n\ +05MaBwzAADq7VrDGETORJmJ0aJJmm+APvLRgu/3CnAT5ja/RPpcCcRgA9Wt1oy0Q\n\ +SaF5Kb15cHfWjtix+FmFCpLKLQPOI2dBf64PFRsglUMstL0pqJ7noihzVLpdS6Pb\n\ ++a2cndgL4bH2FpgJTaO6z3whqmVjamZfksRpAU9s8f3/puCISidYz9qCgNC0BvyS\n\ +XJjWBJECgYEA5ijbrld40JIIzgZ7zW0W4tv1MKY9EsC23IGl7ZesxY5Rsgmsa0Ka\n\ +lhNMRugd58NqG/XnAIBWZRLQM40qYBHyARW+cJxh7+3onvOHYwNkcrdQ/kcBbvme\n\ +aSFRWeGn0QLo/4pGGpCssWoO+OvYKHUbH4qbPHdEFKOMwjsmKvHQX30CgYEA0V7d\n\ ++hqea/fb6gvBDhosMqDe4y8Ry7zqu1IOVBxoG3cAhnoJn7zPn0rCLbYXgCp5/Nj9\n\ +VqjQpUHCgevTE+W/ac/lq73e3oXaUX8M7baD/aPkNF5Y09AeabVw1ITLrYL9zWdU\n\ +zfclwsIhWq2Dt/5DE4/bD9FtD1r7P7Vigba/LYkCgYBr1/E3e50MfaDKiJcx5k+2\n\ +9MGqjfpH8yy7nbQV49/8oXb+KTI0//xXHau7/b8lfZcWit42ievxaCNORHL6mO4A\n\ +PCQDuALb3WoGMK3bYxeJ+QNmYfb1/NiRAh+QMf/kG6z5L90xTWDdsIhbcobSTizr\n\ +VpLufiPUV934lKaJsMymMQKBgQDByxKh7kOW4jwO/cQ6/mTMk/TayfWp5HpM2p3i\n\ +osyGJ3c4AfuofEadRcBIOVS1UBvLuzl7HhTJ8f1M7nBY6X5sPX9zoPKKe9DhQD1C\n\ +Rn8TpcCT7IRBwlB0Pfpq62PvfeDYX/2yC0JLbA8ddKAIDXQexjfZA1r0LJ2EkarV\n\ +L8bzKQKBgQCbH3O8Dd1XQJ13SqDiWt7PJ37PalawTerYlFnw7jWOSooAzjuWw/La\n\ +/jOy0BPlmkjxjAXAP1jY5Kq/UdkeXnlvQVne4F8oRKL2PD8iqYSNWAAyO4v+zHdX\n\ +WtuDNx5B6NeOBk3E4n28oYStdHw20B+mHCxdr/wQ2iosqQfqEugKmg==\n\ +-----END RSA PRIVATE KEY-----\n"; +/// Its modulus, base64url, for the JWKS. +const IDP_KEY_N: &str = "vDyl7CDoXwG8SqSFodeEVK3aGjpg7cmsweIZIblHOLA_Ftd5D1XEOaAt7AuXjYQINz4zy3Nvcd3DCx42mCw4tdeGobZOSpGI7z5dq3rJFV-pjVGhJ5o7nLO3hipNbaiZKMrCFwybh_pSF9jtf6aP1nzMKEH9kTabxbnHyiEZuNJW07oHjQhi7JdrM9-l-dfxWpnHSdyZ-6DBIG9jV9NB5fT8yJ-oGUuakUm38-7TwE1rR19Lx2HW7r41s09t2Qkzjo0E7McSF1nwJM-Ek0VS4eD3zqdoz1aHLUdWavzQZJ6lyCkNA_FAdmQ912oMdlIqvfOp_2RiuVQ-QH_UJY8S5Q"; +const IDP_KID: &str = "access-test"; +const CLIENT_ID: &str = "tw-access-client"; + +/// Serve discovery and the JWKS; `/token` is mounted per login. +async fn mock_idp() -> MockServer { + let idp = MockServer::start().await; + let issuer = idp.uri(); + Mock::given(method("GET")) + .and(path("/.well-known/openid-configuration")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "issuer": issuer, + "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint": format!("{issuer}/token"), + "jwks_uri": format!("{issuer}/jwks"), + "response_types_supported": ["code"], + "subject_types_supported": ["public"], + "id_token_signing_alg_values_supported": ["RS256"], + }))) + .mount(&idp) + .await; + Mock::given(method("GET")) + .and(path("/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "keys": [{ + "kty": "RSA", "use": "sig", "alg": "RS256", "kid": IDP_KID, + "n": IDP_KEY_N, "e": "AQAB", + }] + }))) + .mount(&idp) + .await; + idp +} + +/// Draft the mock provider in the wizard and activate it. +async fn activate_sso(app: &TestApp, admin: &TestClient, idp: &MockServer) { + admin + .patch( + "/api/admin/settings/oidc/draft", + json!({ + "issuer_url": idp.uri(), + "client_id": CLIENT_ID, + "client_secret": "access-client-secret-123", + "redirect_url": "http://localhost:3001/api/auth/sso/callback", + }), + ) + .await + .unwrap() + .assert_ok(); + let passed = json!({"passed": true, "at": chrono::Utc::now().timestamp()}); + let _: () = fred::interfaces::KeysInterface::set( + &app.state.redis, + "oidc:test:result", + passed.to_string(), + Some(fred::types::Expiration::EX(1800)), + None, + false, + ) + .await + .unwrap(); + admin + .post("/api/admin/settings/oidc/activate", json!({})) + .await + .unwrap() + .assert_ok(); +} + +fn query_param(url: &str, name: &str) -> String { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find(|(k, _)| k == name) + .unwrap_or_else(|| panic!("{name} in {url}")) + .1 + .into_owned() +} + +/// Sign in through `/api/auth/sso/authorize` → the provider → the +/// callback, as the identity `claims` describes. Returns the client and +/// the callback's response. +async fn sso_login(app: &TestApp, idp: &MockServer, claims: Value) -> (TestClient, u16) { + let con = app.console_client(); + let resp = con.get("/api/auth/sso/authorize").await.unwrap(); + resp.assert_status(307); + let location = resp.headers["location"].to_str().unwrap().to_string(); + let nonce = query_param(&location, "nonce"); + let state = query_param(&location, "state"); + + let now = chrono::Utc::now().timestamp(); + let mut claims = claims; + let obj = claims.as_object_mut().unwrap(); + obj.insert("iss".into(), json!(idp.uri())); + obj.insert("aud".into(), json!(CLIENT_ID)); + obj.insert("iat".into(), json!(now)); + obj.insert("exp".into(), json!(now + 300)); + obj.insert("nonce".into(), json!(nonce)); + let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some(IDP_KID.into()); + let id_token = jsonwebtoken::encode( + &header, + &claims, + &jsonwebtoken::EncodingKey::from_rsa_pem(IDP_KEY_PEM.as_bytes()).unwrap(), + ) + .unwrap(); + + let code = Uuid::new_v4().simple().to_string(); + Mock::given(method("POST")) + .and(path("/token")) + .and(body_string_contains(code.as_str())) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "idp-access-token", + "token_type": "Bearer", + "expires_in": 300, + "id_token": id_token, + }))) + .mount(idp) + .await; + + let resp = con + .get(&format!("/api/auth/sso/callback?code={code}&state={state}")) + .await + .unwrap(); + (con, resp.status.as_u16()) +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn sso_provisions_a_user_then_signs_them_in_again() { + let app = TestApp::spawn_reaching_loopback().await; + let admin = admin_session(&app).await; + let idp = mock_idp().await; + activate_sso(&app, &admin, &idp).await; + + // Activation promotes the draft and drops it. + let oidc = get(&admin, "/api/admin/settings/oidc").await; + assert!(oidc["draft"].is_null(), "{oidc}"); + assert_eq!(oidc["active"]["enabled"], true); + assert_eq!(oidc["active"]["configured"], true); + + app.set_setting("auth.default_role", json!("viewer")).await; + + let subject = unique_name("sub"); + let identity = json!({"sub": subject, "email": "Person.One@Example.com", "name": "Person One"}); + let (con, status) = sso_login(&app, &idp, identity.clone()).await; + assert_eq!(status, 307); + assert!(con.cookie("__Host-access_token").is_some()); + let me = get(&con, "/api/auth/me").await; + assert_eq!(me["email"], "person.one@example.com"); + assert_eq!(me["display_name"], "Person One"); + assert_eq!(me["oidc_subject"], subject.as_str()); + let roles: Vec<&str> = me["role_assignments"] + .as_array() + .unwrap() + .iter() + .map(|r| r["name"].as_str().unwrap()) + .collect(); + assert_eq!(roles, vec!["viewer"]); + + // The same identity signs in to the same row. + let (_, status) = sso_login(&app, &idp, identity.clone()).await; + assert_eq!(status, 307); + let rows: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM users WHERE oidc_subject = $1") + .bind(&subject) + .fetch_one(&app.db) + .await + .unwrap(); + assert_eq!(rows, 1); + + // An identity with no email gets a stable placeholder address. + let anonymous = unique_name("sub"); + let (con, status) = sso_login(&app, &idp, json!({"sub": anonymous})).await; + assert_eq!(status, 307); + let me = get(&con, "/api/auth/me").await; + let email = me["email"].as_str().unwrap(); + assert!( + email.starts_with("sso-") && email.ends_with("@oidc.invalid"), + "{email}" + ); + assert_eq!(me["display_name"], email); + + // Deactivated, then deleted: refused. + sqlx::query("UPDATE users SET is_active = false WHERE oidc_subject = $1") + .bind(&subject) + .execute(&app.db) + .await + .unwrap(); + let (con, status) = sso_login(&app, &idp, identity.clone()).await; + assert_eq!(status, 403); + assert!(con.cookie("__Host-access_token").is_none()); + sqlx::query("UPDATE users SET deleted_at = now() WHERE oidc_subject = $1") + .bind(&subject) + .execute(&app.db) + .await + .unwrap(); + let (_, status) = sso_login(&app, &idp, identity).await; + assert_eq!(status, 403); +}