From c1f02d548f47ef9df0309fc8080f7c32b5569772 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 12:25:08 +0200 Subject: [PATCH 01/10] fix(auth): audit permission denials safely Mark genuine authorization failures explicitly, aggregate them through a bounded non-blocking recorder, and expose only normalized route and principal metadata. Add backend and console regression coverage. --- crates/temps-auth/src/audit.rs | 65 +++ crates/temps-auth/src/lib.rs | 1 + .../src/permission_denial_recorder.rs | 402 ++++++++++++++++++ crates/temps-auth/src/permission_guard.rs | 110 ++++- crates/temps-auth/src/plugin.rs | 10 + crates/temps-auth/src/temps_middleware.rs | 248 ++++++++++- crates/temps-core/src/error_builder.rs | 20 + crates/temps-core/src/problemdetails/mod.rs | 124 +++++- web/src/components/audit/AuditLogItem.test.ts | 25 ++ web/src/components/audit/AuditLogItem.tsx | 6 +- web/src/lib/audit-operation-filters.ts | 4 + web/src/lib/permission-denial-display.ts | 25 ++ web/src/pages/AuditLogs.test.ts | 12 + web/src/pages/AuditLogs.tsx | 2 + 14 files changed, 1040 insertions(+), 14 deletions(-) create mode 100644 crates/temps-auth/src/permission_denial_recorder.rs create mode 100644 web/src/components/audit/AuditLogItem.test.ts create mode 100644 web/src/lib/audit-operation-filters.ts create mode 100644 web/src/lib/permission-denial-display.ts create mode 100644 web/src/pages/AuditLogs.test.ts diff --git a/crates/temps-auth/src/audit.rs b/crates/temps-auth/src/audit.rs index 8545e4b1b..4daf7a4e9 100644 --- a/crates/temps-auth/src/audit.rs +++ b/crates/temps-auth/src/audit.rs @@ -680,6 +680,48 @@ impl_oidc_audit_op!(OidcProviderDeletedAudit, "OIDC_PROVIDER_DELETED"); impl_oidc_audit_op!(OidcRoleMappingCreatedAudit, "OIDC_ROLE_MAPPING_CREATED"); impl_oidc_audit_op!(OidcRoleMappingDeletedAudit, "OIDC_ROLE_MAPPING_DELETED"); +/// Aggregated authorization-guard denials. Only stable, server-derived +/// principal metadata is recorded: credential names and secrets are never +/// accepted by this type. +#[derive(Debug, Clone, Serialize)] +pub struct PermissionDeniedAudit { + pub user_id: Option, + pub auth_source: String, + pub credential_id: Option, + pub method: String, + /// Axum route template (for example `/projects/{project_id}`), or the + /// fixed value `unmatched`. Never a raw request URI. + pub route: String, + pub denial_kind: String, + pub required_permission: Option, + pub attempt_count: u64, + pub ip_address: Option, + pub user_agent: String, +} + +impl AuditOperation for PermissionDeniedAudit { + fn operation_type(&self) -> String { + "PERMISSION_DENIED".to_string() + } + + fn user_id(&self) -> Option { + self.user_id + } + + fn ip_address(&self) -> Option { + self.ip_address.clone() + } + + fn user_agent(&self) -> &str { + &self.user_agent + } + + fn serialize(&self) -> Result { + serde_json::to_string(self) + .map_err(|e| anyhow::anyhow!("Failed to serialize audit operation {}", e)) + } +} + /// Recorded when a login attempt is rejected before a session exists. /// /// Unlike most audit events the actor may be unknown (an attempt against an @@ -837,4 +879,27 @@ mod failure_audit_tests { assert!(!json.contains("123456")); assert!(json.contains("invalid_code")); } + + #[test] + fn permission_denied_audit_supports_machine_principals_and_aggregation() { + let audit = PermissionDeniedAudit { + user_id: None, + auth_source: "deployment_token".to_string(), + credential_id: Some(17), + method: "POST".to_string(), + route: "/projects/{project_id}/deployments".to_string(), + denial_kind: "cross_project_scope".to_string(), + required_permission: None, + attempt_count: 4, + ip_address: Some("203.0.113.8".to_string()), + user_agent: "test-agent".to_string(), + }; + + assert_eq!(audit.operation_type(), "PERMISSION_DENIED"); + assert_eq!(audit.user_id(), None); + let json = AuditOperation::serialize(&audit).expect("permission denial audit serializes"); + assert!(json.contains("\"attempt_count\":4")); + assert!(json.contains("\"credential_id\":17")); + assert!(!json.contains("token_name")); + } } diff --git a/crates/temps-auth/src/lib.rs b/crates/temps-auth/src/lib.rs index 81ab4b5b1..53b004492 100644 --- a/crates/temps-auth/src/lib.rs +++ b/crates/temps-auth/src/lib.rs @@ -22,6 +22,7 @@ mod oidc_service; mod oidc_types; mod permission_attribute; mod permission_decorator; +mod permission_denial_recorder; mod permission_guard; pub mod permissions; mod plugin; diff --git a/crates/temps-auth/src/permission_denial_recorder.rs b/crates/temps-auth/src/permission_denial_recorder.rs new file mode 100644 index 000000000..22fc3b1f5 --- /dev/null +++ b/crates/temps-auth/src/permission_denial_recorder.rs @@ -0,0 +1,402 @@ +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::mpsc; + +use crate::audit::PermissionDeniedAudit; + +const DEFAULT_QUEUE_CAPACITY: usize = 1_024; +const DEFAULT_MAX_AGGREGATIONS: usize = 1_024; +const DEFAULT_WINDOW: Duration = Duration::from_secs(5); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum AuthSourceKind { + Anonymous, + Session, + CliToken, + ApiKey, + DeploymentToken, +} + +impl AuthSourceKind { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Anonymous => "anonymous", + Self::Session => "session", + Self::CliToken => "cli_token", + Self::ApiKey => "api_key", + Self::DeploymentToken => "deployment_token", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SafePrincipal { + pub(crate) user_id: Option, + pub(crate) source: AuthSourceKind, + pub(crate) credential_id: Option, +} + +impl SafePrincipal { + pub(crate) const fn anonymous() -> Self { + Self { + user_id: None, + source: AuthSourceKind::Anonymous, + credential_id: None, + } + } +} + +#[derive(Debug, Clone)] +pub(crate) struct PermissionDenialEvent { + pub(crate) principal: SafePrincipal, + pub(crate) method: String, + pub(crate) route: String, + pub(crate) denial_kind: String, + pub(crate) required_permission: Option, + pub(crate) ip_address: Option, + pub(crate) user_agent: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct AggregationKey { + user_id: Option, + source: AuthSourceKind, + credential_id: Option, + method: String, + route: String, + denial_kind: String, + required_permission: Option, +} + +impl From<&PermissionDenialEvent> for AggregationKey { + fn from(event: &PermissionDenialEvent) -> Self { + Self { + user_id: event.principal.user_id, + source: event.principal.source, + credential_id: event.principal.credential_id, + method: event.method.clone(), + route: event.route.clone(), + denial_kind: event.denial_kind.clone(), + required_permission: event.required_permission.clone(), + } + } +} + +impl PermissionDenialEvent { + fn into_audit(self) -> PermissionDeniedAudit { + PermissionDeniedAudit { + user_id: self.principal.user_id, + auth_source: self.principal.source.as_str().to_string(), + credential_id: self.principal.credential_id, + method: self.method, + route: self.route, + denial_kind: self.denial_kind, + required_permission: self.required_permission, + attempt_count: 1, + ip_address: self.ip_address, + user_agent: self.user_agent, + } + } +} + +struct RecorderCounters { + queue_drops: AtomicU64, + aggregation_overflows: AtomicU64, + write_failures: AtomicU64, +} + +impl RecorderCounters { + fn new() -> Self { + Self { + queue_drops: AtomicU64::new(0), + aggregation_overflows: AtomicU64::new(0), + write_failures: AtomicU64::new(0), + } + } +} + +#[derive(Debug, Clone, Copy)] +struct RecorderConfig { + queue_capacity: usize, + max_aggregations: usize, + window: Duration, +} + +/// Non-blocking, explicitly bounded entry point for permission-denial audit +/// events. Request handling performs only `try_send`; aggregation and audit +/// logger I/O happen in the spawned worker. +pub struct PermissionDenialRecorder { + sender: mpsc::Sender, + counters: Arc, +} + +impl PermissionDenialRecorder { + pub fn new(logger: Arc) -> Arc { + Self::with_config( + logger, + RecorderConfig { + queue_capacity: DEFAULT_QUEUE_CAPACITY, + max_aggregations: DEFAULT_MAX_AGGREGATIONS, + window: DEFAULT_WINDOW, + }, + ) + } + + fn with_config(logger: Arc, config: RecorderConfig) -> Arc { + let (sender, receiver) = mpsc::channel(config.queue_capacity); + let counters = Arc::new(RecorderCounters::new()); + let worker = PermissionDenialWorker { + receiver, + logger, + counters: counters.clone(), + max_aggregations: config.max_aggregations, + window: config.window, + pending: HashMap::with_capacity(config.max_aggregations), + }; + tokio::spawn(worker.run()); + + Arc::new(Self { sender, counters }) + } + + pub(crate) fn record(&self, event: PermissionDenialEvent) { + if let Err(error) = self.sender.try_send(event) { + let event = error.into_inner(); + let total = self.counters.queue_drops.fetch_add(1, Ordering::Relaxed) + 1; + if should_log_counter(total) { + tracing::warn!( + total_queue_drops = total, + user_id = event.principal.user_id, + auth_source = event.principal.source.as_str(), + credential_id = event.principal.credential_id, + method = event.method, + route = event.route, + denial_kind = event.denial_kind, + required_permission = event.required_permission, + "permission-denial audit queue saturated or closed; event dropped" + ); + } + } + } +} + +fn should_log_counter(total: u64) -> bool { + total.is_power_of_two() +} + +struct PermissionDenialWorker { + receiver: mpsc::Receiver, + logger: Arc, + counters: Arc, + max_aggregations: usize, + window: Duration, + pending: HashMap, +} + +impl PermissionDenialWorker { + async fn run(mut self) { + let start = tokio::time::Instant::now() + self.window; + let mut ticker = tokio::time::interval_at(start, self.window); + + loop { + tokio::select! { + event = self.receiver.recv() => { + match event { + Some(event) => self.aggregate(event), + None => { + self.flush().await; + break; + } + } + } + _ = ticker.tick() => self.flush().await, + } + } + } + + fn aggregate(&mut self, event: PermissionDenialEvent) { + let key = AggregationKey::from(&event); + if let Some(audit) = self.pending.get_mut(&key) { + audit.attempt_count = audit.attempt_count.saturating_add(1); + return; + } + + if self.pending.len() >= self.max_aggregations { + let total = self + .counters + .aggregation_overflows + .fetch_add(1, Ordering::Relaxed) + + 1; + if should_log_counter(total) { + tracing::warn!( + total_aggregation_overflows = total, + max_aggregations = self.max_aggregations, + user_id = event.principal.user_id, + auth_source = event.principal.source.as_str(), + credential_id = event.principal.credential_id, + method = event.method, + route = event.route, + denial_kind = event.denial_kind, + required_permission = event.required_permission, + "permission-denial aggregation map full; new key dropped" + ); + } + return; + } + + self.pending.insert(key, event.into_audit()); + } + + async fn flush(&mut self) { + // At most `max_aggregations` entries are moved into this temporary + // vector. The queue and map remain bounded while logger I/O is in flight. + let audits: Vec<_> = self.pending.drain().map(|(_, audit)| audit).collect(); + for audit in audits { + if let Err(error) = self.logger.create_audit_log(&audit).await { + let total = self.counters.write_failures.fetch_add(1, Ordering::Relaxed) + 1; + if should_log_counter(total) { + tracing::warn!( + total_write_failures = total, + user_id = audit.user_id, + auth_source = audit.auth_source, + credential_id = audit.credential_id, + method = audit.method, + route = audit.route, + denial_kind = audit.denial_kind, + required_permission = audit.required_permission, + error = %error, + "failed to persist permission-denial audit; denial response was unaffected" + ); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use anyhow::Result; + use temps_core::{AuditLogger, AuditOperation}; + + use super::*; + + #[derive(Default)] + struct RecordingLogger { + records: Mutex>, + } + + #[async_trait::async_trait] + impl AuditLogger for RecordingLogger { + async fn create_audit_log(&self, operation: &dyn AuditOperation) -> Result<()> { + let serialized = operation.serialize()?; + let value = serde_json::from_str(&serialized)?; + self.records + .lock() + .map_err(|_| anyhow::anyhow!("test logger lock poisoned"))? + .push(value); + Ok(()) + } + } + + struct FailingLogger; + + #[async_trait::async_trait] + impl AuditLogger for FailingLogger { + async fn create_audit_log(&self, _operation: &dyn AuditOperation) -> Result<()> { + Err(anyhow::anyhow!("intentional logger failure")) + } + } + + fn event(route: &str) -> PermissionDenialEvent { + PermissionDenialEvent { + principal: SafePrincipal { + user_id: Some(42), + source: AuthSourceKind::ApiKey, + credential_id: Some(9), + }, + method: "PATCH".to_string(), + route: route.to_string(), + denial_kind: "insufficient_permission".to_string(), + required_permission: Some("projects:write".to_string()), + ip_address: Some("203.0.113.7".to_string()), + user_agent: "test-agent".to_string(), + } + } + + fn test_config(queue_capacity: usize, max_aggregations: usize) -> RecorderConfig { + RecorderConfig { + queue_capacity, + max_aggregations, + window: Duration::from_millis(20), + } + } + + #[tokio::test] + async fn repeated_safe_key_is_aggregated_and_ip_is_not_part_of_key() { + let logger = Arc::new(RecordingLogger::default()); + let recorder = PermissionDenialRecorder::with_config(logger.clone(), test_config(8, 8)); + recorder.record(event("/projects/{project_id}")); + let mut second = event("/projects/{project_id}"); + second.ip_address = Some("198.51.100.99".to_string()); + recorder.record(second); + + tokio::time::sleep(Duration::from_millis(60)).await; + let records = logger.records.lock().expect("test logger lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0]["attempt_count"], 2); + assert_eq!(records[0]["route"], "/projects/{project_id}"); + assert_eq!(records[0]["auth_source"], "api_key"); + assert_eq!(records[0]["credential_id"], 9); + assert!(records[0].get("key_name").is_none()); + assert!(records[0].get("token_name").is_none()); + } + + #[tokio::test] + async fn aggregation_cardinality_is_bounded() { + let logger = Arc::new(RecordingLogger::default()); + let recorder = PermissionDenialRecorder::with_config(logger.clone(), test_config(8, 1)); + recorder.record(event("/first")); + recorder.record(event("/second")); + + tokio::time::sleep(Duration::from_millis(60)).await; + assert_eq!( + recorder + .counters + .aggregation_overflows + .load(Ordering::Relaxed), + 1 + ); + assert_eq!(logger.records.lock().expect("test logger lock").len(), 1); + } + + #[tokio::test] + async fn queue_overflow_is_counted_without_blocking() { + let (sender, _receiver) = mpsc::channel(1); + let counters = Arc::new(RecorderCounters::new()); + let recorder = PermissionDenialRecorder { + sender, + counters: counters.clone(), + }; + + recorder.record(event("/first")); + recorder.record(event("/second")); + assert_eq!(counters.queue_drops.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn logger_failure_is_counted_and_worker_continues() { + let recorder = + PermissionDenialRecorder::with_config(Arc::new(FailingLogger), test_config(8, 8)); + recorder.record(event("/first")); + tokio::time::sleep(Duration::from_millis(60)).await; + assert_eq!(recorder.counters.write_failures.load(Ordering::Relaxed), 1); + + recorder.record(event("/second")); + tokio::time::sleep(Duration::from_millis(60)).await; + assert_eq!(recorder.counters.write_failures.load(Ordering::Relaxed), 2); + } +} diff --git a/crates/temps-auth/src/permission_guard.rs b/crates/temps-auth/src/permission_guard.rs index 6d78ff48a..b779d3332 100644 --- a/crates/temps-auth/src/permission_guard.rs +++ b/crates/temps-auth/src/permission_guard.rs @@ -30,6 +30,10 @@ macro_rules! permission_guard { $crate::permissions::Permission::$permission.to_string(), ) .value("user_role", $auth.effective_role.to_string()) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::InsufficientPermission, + Some($crate::permissions::Permission::$permission.to_string()), + ) .build()); } }; @@ -74,6 +78,10 @@ macro_rules! project_scope_guard { "This deployment token is scoped to a different project and \ cannot access this resource", ) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::CrossProjectScope, + None, + ) .build()); } }; @@ -113,6 +121,10 @@ macro_rules! deny_deployment_token { "This endpoint requires user or API-key authentication; \ deployment tokens are not permitted", ) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::DeploymentTokenNotAllowed, + None, + ) .build()); } }; @@ -184,6 +196,10 @@ macro_rules! project_access_guard { "Your team membership does not include access to \ this project", ) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::ProjectAccess, + None, + ) .build()); } Err(__e) => { @@ -291,6 +307,10 @@ macro_rules! project_permission_guard { $crate::permissions::Permission::$permission.to_string(), ) .value("user_role", $auth.effective_role.to_string()) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::InsufficientPermission, + Some($crate::permissions::Permission::$permission.to_string()), + ) .build()); } @@ -324,6 +344,10 @@ macro_rules! project_permission_guard { __required )) .value("required_permission", __required) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::ProjectPermission, + Some($crate::permissions::Permission::$permission.to_string()), + ) .build()); } } @@ -362,6 +386,10 @@ macro_rules! project_permission_guard { .detail( "Could not resolve caller identity for project permission check", ) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::MissingPrincipal, + Some($crate::permissions::Permission::$permission.to_string()), + ) .build()); } } @@ -401,6 +429,10 @@ macro_rules! permission_check { )) .value("required_permission", $permission.to_string()) .value("user_role", $auth.effective_role.to_string()) + .permission_denial( + temps_core::problemdetails::PermissionDenialKind::InsufficientPermission, + Some($permission.to_string()), + ) .build()); } }; @@ -411,13 +443,14 @@ mod tests { use std::sync::Arc; use async_trait::async_trait; + use axum::response::IntoResponse; use chrono::Utc; - use temps_core::problemdetails::Problem; + use temps_core::problemdetails::{PermissionDenialKind, PermissionDenialMarker, Problem}; use temps_core::ProjectAccessChecker; use temps_entities::users; use crate::context::AuthContext; - use crate::permissions::Role; + use crate::permissions::{Permission, Role}; // --------------------------------------------------------------------------- // Test helpers @@ -462,6 +495,65 @@ mod tests { ) } + fn denial_marker(problem: Problem) -> Option { + problem + .into_response() + .extensions() + .get::() + .cloned() + } + + fn run_instance_permission_guard(auth: &AuthContext) -> Result<(), Problem> { + permission_guard!(auth, UsersWrite); + Ok(()) + } + + fn run_permission_check(auth: &AuthContext) -> Result<(), Problem> { + permission_check!(auth, Permission::UsersWrite); + Ok(()) + } + + fn run_project_scope_guard(auth: &AuthContext, project_id: i32) -> Result<(), Problem> { + project_scope_guard!(auth, project_id); + Ok(()) + } + + fn run_deny_deployment_token(auth: &AuthContext) -> Result<(), Problem> { + deny_deployment_token!(auth); + Ok(()) + } + + #[test] + fn synchronous_guard_denials_carry_stable_markers() { + let auth = user_auth(Role::User); + for problem in [ + run_instance_permission_guard(&auth).expect_err("user lacks users:write"), + run_permission_check(&auth).expect_err("user lacks users:write"), + ] { + let marker = denial_marker(problem).expect("permission denial marker"); + assert_eq!(marker.kind(), PermissionDenialKind::InsufficientPermission); + assert_eq!(marker.required_permission(), Some("users:write")); + } + + let token = deployment_token_auth(); + let marker = denial_marker( + run_project_scope_guard(&token, 8).expect_err("token is bound to project 7"), + ) + .expect("project scope denial marker"); + assert_eq!(marker.kind(), PermissionDenialKind::CrossProjectScope); + assert_eq!(marker.required_permission(), None); + + let marker = denial_marker( + run_deny_deployment_token(&token).expect_err("deployment token is denied"), + ) + .expect("deployment token denial marker"); + assert_eq!( + marker.kind(), + PermissionDenialKind::DeploymentTokenNotAllowed + ); + assert_eq!(marker.required_permission(), None); + } + /// A mock [`ProjectAccessChecker`] that returns a fixed outcome. struct MockChecker { result: fn() -> Result>, @@ -546,6 +638,9 @@ mod tests { axum::http::StatusCode::FORBIDDEN, "denial should be HTTP 403" ); + let marker = denial_marker(err).expect("project access denial marker"); + assert_eq!(marker.kind(), PermissionDenialKind::ProjectAccess); + assert_eq!(marker.required_permission(), None); } // --------------------------------------------------------------------------- @@ -563,6 +658,10 @@ mod tests { axum::http::StatusCode::INTERNAL_SERVER_ERROR, "infrastructure failure should be HTTP 500" ); + assert!( + denial_marker(err).is_none(), + "infrastructure failures must not be marked as auth denials" + ); } // --------------------------------------------------------------------------- @@ -760,6 +859,9 @@ mod tests { err.body.get("type").and_then(|v| v.as_str()), Some("https://temps.sh/probs/project-permission-denied") ); + let marker = denial_marker(err).expect("project permission denial marker"); + assert_eq!(marker.kind(), PermissionDenialKind::ProjectPermission); + assert_eq!(marker.required_permission(), Some("deployments:create")); } #[tokio::test] @@ -807,6 +909,10 @@ mod tests { err.status_code, axum::http::StatusCode::INTERNAL_SERVER_ERROR ); + assert!( + denial_marker(err).is_none(), + "resolver failures must not be marked as auth denials" + ); } #[tokio::test] diff --git a/crates/temps-auth/src/plugin.rs b/crates/temps-auth/src/plugin.rs index f017ec64b..1be15711f 100644 --- a/crates/temps-auth/src/plugin.rs +++ b/crates/temps-auth/src/plugin.rs @@ -50,6 +50,12 @@ impl TempsPlugin for AuthPlugin { let encryption_service = context.require_service::(); let cookie_crypto = context.require_service::(); + let permission_denial_recorder = + crate::permission_denial_recorder::PermissionDenialRecorder::new( + audit_service.clone(), + ); + context.register_service(permission_denial_recorder); + // Require notification service let notification_service = context.require_service::(); @@ -221,6 +227,9 @@ impl TempsPlugin for AuthPlugin { let cookie_crypto = context.require_service::(); let api_key_service = context.require_service::(); let db = context.require_service::(); + let permission_denial_recorder = + context + .require_service::(); // Request-metadata middleware (runs on BOTH admin and public routers // because public ingest endpoints — session-replay init, analytics @@ -237,6 +246,7 @@ impl TempsPlugin for AuthPlugin { user_service, cookie_crypto, db, + permission_denial_recorder, ); middleware_collection.add_temps_middleware(Arc::new(auth_middleware)); diff --git a/crates/temps-auth/src/temps_middleware.rs b/crates/temps-auth/src/temps_middleware.rs index 6afc24b4c..bd6a4d3cd 100644 --- a/crates/temps-auth/src/temps_middleware.rs +++ b/crates/temps-auth/src/temps_middleware.rs @@ -8,11 +8,19 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use axum::{extract::Request, http::StatusCode, middleware::Next, response::Response}; +use axum::{ + extract::{MatchedPath, Request}, + http::StatusCode, + middleware::Next, + response::{IntoResponse, Response}, +}; use temps_core::plugin::{ MiddlewareCondition, MiddlewarePriority, PluginContext, PluginError, TempsMiddleware, }; +use crate::permission_denial_recorder::{ + AuthSourceKind, PermissionDenialEvent, PermissionDenialRecorder, SafePrincipal, +}; use crate::{ apikey_service::ApiKeyService, auth_service::AuthService, deployment_token_service::DeploymentTokenValidationService, user_service::UserService, @@ -26,6 +34,7 @@ pub struct AuthMiddleware { user_service: Arc, cookie_crypto: Arc, deployment_token_service: DeploymentTokenValidationService, + permission_denial_recorder: Arc, } impl AuthMiddleware { @@ -35,6 +44,7 @@ impl AuthMiddleware { user_service: Arc, cookie_crypto: Arc, db: Arc, + permission_denial_recorder: Arc, ) -> Self { let deployment_token_service = DeploymentTokenValidationService::new(db); Self { @@ -43,6 +53,7 @@ impl AuthMiddleware { user_service, cookie_crypto, deployment_token_service, + permission_denial_recorder, } } } @@ -82,10 +93,7 @@ impl TempsMiddleware for AuthMiddleware { // Call the simplified auth middleware match self.execute_auth_middleware_logic(req, next).await { Ok(response) => Ok(response), - Err(status) => Ok(Response::builder() - .status(status) - .body(axum::body::Body::empty()) - .unwrap()), + Err(status) => Ok(status.into_response()), } }) } @@ -195,12 +203,35 @@ impl AuthMiddleware { if let Some(user) = user { req.extensions_mut().insert(user); } + let principal = auth_context + .as_ref() + .map(safe_principal) + .unwrap_or_else(SafePrincipal::anonymous); if let Some(auth_ctx) = auth_context { req.extensions_mut().insert(auth_ctx); } + let method = req.method().as_str().to_string(); + let route = matched_route_template(&req); + let (ip_address, user_agent) = req + .extensions() + .get::() + .map(|metadata| { + ( + Some(metadata.ip_address.clone()), + metadata.user_agent.clone(), + ) + }) + .unwrap_or_else(|| (None, "unknown".to_string())); + // Run the next middleware/handler - Ok(next.run(req).await) + let response = next.run(req).await; + if let Some(event) = + permission_denial_event(&response, principal, method, route, ip_address, user_agent) + { + self.permission_denial_recorder.record(event); + } + Ok(response) } /// Extract user session from "session" cookie only (for authentication) @@ -228,3 +259,208 @@ impl AuthMiddleware { None } } + +fn safe_principal(auth: &crate::context::AuthContext) -> SafePrincipal { + let (source, credential_id) = match &auth.source { + crate::context::AuthSource::Session { .. } => (AuthSourceKind::Session, None), + crate::context::AuthSource::CliToken { .. } => (AuthSourceKind::CliToken, None), + crate::context::AuthSource::ApiKey { key_id, .. } => { + (AuthSourceKind::ApiKey, Some(*key_id)) + } + crate::context::AuthSource::DeploymentToken { token_id, .. } => { + (AuthSourceKind::DeploymentToken, Some(*token_id)) + } + }; + + SafePrincipal { + user_id: auth.user_id_opt(), + source, + credential_id, + } +} + +fn matched_route_template(req: &Request) -> String { + req.extensions() + .get::() + .map(|matched| matched.as_str().to_string()) + .unwrap_or_else(|| "unmatched".to_string()) +} + +fn permission_denial_event( + response: &Response, + principal: SafePrincipal, + method: String, + route: String, + ip_address: Option, + user_agent: String, +) -> Option { + if response.status() != StatusCode::FORBIDDEN { + return None; + } + + let marker = response + .extensions() + .get::()?; + Some(PermissionDenialEvent { + principal, + method, + route, + denial_kind: marker.kind().as_str().to_string(), + required_permission: marker.required_permission().map(str::to_string), + ip_address, + user_agent, + }) +} + +#[cfg(test)] +mod tests { + use axum::http::StatusCode; + use axum::middleware::from_fn; + use axum::routing::get; + use axum::Router; + use chrono::Utc; + use temps_core::problemdetails::PermissionDenialKind; + use temps_entities::deployment_tokens::DeploymentTokenPermission; + use temps_entities::users; + use tower::ServiceExt; + + use super::*; + use crate::permissions::{Permission, Role}; + + fn test_user() -> users::Model { + let now = Utc::now(); + users::Model { + id: 42, + name: "User-chosen name must not leak".to_string(), + email: "private@example.com".to_string(), + password_hash: None, + email_verified: true, + email_verification_token: None, + email_verification_expires: None, + password_reset_token: None, + password_reset_expires: None, + must_change_password: false, + deleted_at: None, + mfa_secret: None, + mfa_enabled: false, + mfa_recovery_codes: None, + oidc_subject: None, + oidc_provider_id: None, + created_at: now, + updated_at: now, + } + } + + fn event_from(response: &Response) -> Option { + permission_denial_event( + response, + SafePrincipal::anonymous(), + "GET".to_string(), + "/projects/{project_id}".to_string(), + Some("203.0.113.5".to_string()), + "test-agent".to_string(), + ) + } + + #[test] + fn marked_guard_403_becomes_an_audit_event() { + let response = temps_core::error_builder::ErrorBuilder::new(StatusCode::FORBIDDEN) + .permission_denial( + PermissionDenialKind::InsufficientPermission, + Some("projects:write".to_string()), + ) + .build() + .into_response(); + + let event = event_from(&response).expect("marked guard denial should be audited"); + assert_eq!(event.denial_kind, "insufficient_permission"); + assert_eq!(event.required_permission.as_deref(), Some("projects:write")); + assert_eq!(event.route, "/projects/{project_id}"); + } + + #[test] + fn generic_business_403_is_ignored() { + let response = StatusCode::FORBIDDEN.into_response(); + assert!(event_from(&response).is_none()); + } + + #[test] + fn marked_non_403_is_ignored() { + let response = temps_core::problemdetails::new(StatusCode::BAD_REQUEST) + .with_permission_denial(PermissionDenialKind::ProjectAccess, None) + .into_response(); + assert!(event_from(&response).is_none()); + } + + #[test] + fn safe_principal_keeps_only_source_kind_and_opaque_ids() { + let api_key = crate::context::AuthContext::new_api_key( + test_user(), + Some(Role::User), + Some(vec![Permission::ProjectsRead]), + "user-chosen key name".to_string(), + 99, + ); + let principal = safe_principal(&api_key); + assert_eq!(principal.user_id, Some(42)); + assert_eq!(principal.source, AuthSourceKind::ApiKey); + assert_eq!(principal.credential_id, Some(99)); + + let token = crate::context::AuthContext::new_deployment_token( + 7, + None, + None, + 123, + "user-chosen deployment token name".to_string(), + vec![DeploymentTokenPermission::AnalyticsRead], + ); + let principal = safe_principal(&token); + assert_eq!(principal.user_id, None); + assert_eq!(principal.source, AuthSourceKind::DeploymentToken); + assert_eq!(principal.credential_id, Some(123)); + } + + #[test] + fn missing_matched_path_never_falls_back_to_raw_uri() { + let request = Request::builder() + .uri("/projects/secret-project?token=secret") + .body(axum::body::Body::empty()) + .expect("test request should build"); + assert_eq!(matched_route_template(&request), "unmatched"); + } + + #[tokio::test] + async fn matched_route_uses_axum_template_instead_of_concrete_uri() { + #[derive(Clone)] + struct CapturedRoute(String); + + async fn capture_route(req: Request, next: Next) -> Response { + let route = matched_route_template(&req); + let mut response = next.run(req).await; + response.extensions_mut().insert(CapturedRoute(route)); + response + } + + let app = Router::new() + .route("/projects/{project_id}", get(|| async { "ok" })) + .layer(from_fn(capture_route)); + let response = app + .oneshot( + Request::builder() + .uri("/projects/secret-project?token=secret") + .body(axum::body::Body::empty()) + .expect("test request should build"), + ) + .await + .expect("test router should respond"); + + assert_eq!( + response + .extensions() + .get::() + .expect("capture middleware should attach route") + .0, + "/projects/{project_id}" + ); + } +} diff --git a/crates/temps-core/src/error_builder.rs b/crates/temps-core/src/error_builder.rs index b7e6ff97a..2fd307d6b 100644 --- a/crates/temps-core/src/error_builder.rs +++ b/crates/temps-core/src/error_builder.rs @@ -3,6 +3,8 @@ use axum::http::StatusCode; use serde::Serialize; use std::collections::HashMap; +use crate::problemdetails::PermissionDenialKind; + pub struct ErrorBuilder { status: StatusCode, type_: String, @@ -10,6 +12,7 @@ pub struct ErrorBuilder { detail: String, instance: String, values: HashMap, + permission_denial: Option<(PermissionDenialKind, Option)>, } impl ErrorBuilder { @@ -21,6 +24,7 @@ impl ErrorBuilder { detail: String::new(), instance: String::new(), values: HashMap::new(), + permission_denial: None, } } @@ -51,6 +55,18 @@ impl ErrorBuilder { self } + /// Attach internal authorization-denial metadata for response middleware. + /// This metadata is never serialized into the problem body. + #[doc(hidden)] + pub fn permission_denial( + mut self, + kind: PermissionDenialKind, + required_permission: Option, + ) -> Self { + self.permission_denial = Some((kind, required_permission)); + self + } + pub fn build(self) -> problemdetails::Problem { let mut problem = problemdetails::new(self.status) .with_type(self.type_) @@ -63,6 +79,10 @@ impl ErrorBuilder { problem = problem.with_value(&key, value); } + if let Some((kind, required_permission)) = self.permission_denial { + problem = problem.with_permission_denial(kind, required_permission); + } + problem } } diff --git a/crates/temps-core/src/problemdetails/mod.rs b/crates/temps-core/src/problemdetails/mod.rs index bd8139114..9f440c319 100644 --- a/crates/temps-core/src/problemdetails/mod.rs +++ b/crates/temps-core/src/problemdetails/mod.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use serde; use serde_json::Value; -use axum::http::StatusCode; +use axum::http::{HeaderValue, StatusCode}; use axum::{http::header::CONTENT_TYPE, response::IntoResponse, Json}; use serde::Serialize; @@ -48,6 +48,53 @@ pub struct Problem { pub status_code: StatusCode, /// The actual body of the problem. pub body: BTreeMap, + /// Internal-only metadata carried to response middleware. This is never + /// serialized into the RFC 7807 response body. + permission_denial: Option, +} + +/// Stable authorization guard denial categories used by security auditing. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PermissionDenialKind { + InsufficientPermission, + CrossProjectScope, + DeploymentTokenNotAllowed, + ProjectAccess, + ProjectPermission, + MissingPrincipal, +} + +impl PermissionDenialKind { + pub const fn as_str(self) -> &'static str { + match self { + Self::InsufficientPermission => "insufficient_permission", + Self::CrossProjectScope => "cross_project_scope", + Self::DeploymentTokenNotAllowed => "deployment_token_not_allowed", + Self::ProjectAccess => "project_access", + Self::ProjectPermission => "project_permission", + Self::MissingPrincipal => "missing_principal", + } + } +} + +/// A server-generated response extension proving that a 403 came from an +/// authorization guard. Its fields are intentionally private so callers can +/// inspect, but cannot forge or mutate, the marker. +#[doc(hidden)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PermissionDenialMarker { + kind: PermissionDenialKind, + required_permission: Option, +} + +impl PermissionDenialMarker { + pub fn kind(&self) -> PermissionDenialKind { + self.kind + } + + pub fn required_permission(&self) -> Option<&str> { + self.required_permission.as_deref() + } } /// Create a new `Problem` response to send to the client. @@ -58,10 +105,26 @@ where Problem { status_code: status_code.into(), body: BTreeMap::new(), + permission_denial: None, } } impl Problem { + /// Mark this problem as a genuine authorization-guard denial. The marker + /// is attached only to the resulting HTTP response extensions. + #[doc(hidden)] + pub fn with_permission_denial( + mut self, + kind: PermissionDenialKind, + required_permission: Option, + ) -> Self { + self.permission_denial = Some(PermissionDenialMarker { + kind, + required_permission, + }); + self + } + /// Specify the "type" to use for the problem. pub fn with_type(self, value: S) -> Self where @@ -122,16 +185,67 @@ pub type Result = std::result::Result; impl IntoResponse for Problem { fn into_response(self) -> axum::response::Response { - if self.body.is_empty() { + let mut response = if self.body.is_empty() { self.status_code.into_response() } else { let body = Json(self.body); let mut response = (self.status_code, body).into_response(); + response.headers_mut().insert( + CONTENT_TYPE, + HeaderValue::from_static("application/problem+json"), + ); response - .headers_mut() - .insert(CONTENT_TYPE, "application/problem+json".parse().unwrap()); - response + }; + + if let Some(marker) = self.permission_denial { + response.extensions_mut().insert(marker); } + + response + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::to_bytes; + + #[tokio::test] + async fn permission_denial_marker_is_internal_and_preserves_body() { + let problem = new(StatusCode::FORBIDDEN) + .with_title("Forbidden") + .with_value("required_permission", "users:write") + .with_permission_denial( + PermissionDenialKind::InsufficientPermission, + Some("users:write".to_string()), + ); + + let response = problem.into_response(); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + let marker = response + .extensions() + .get::() + .expect("guard marker should propagate through IntoResponse"); + assert_eq!(marker.kind(), PermissionDenialKind::InsufficientPermission); + assert_eq!(marker.required_permission(), Some("users:write")); + + let body = to_bytes(response.into_body(), 1024) + .await + .expect("problem body should be readable"); + let json: serde_json::Value = + serde_json::from_slice(&body).expect("problem body should remain JSON"); + assert_eq!(json["title"], "Forbidden"); + assert_eq!(json["required_permission"], "users:write"); + assert!(json.get("permission_denial").is_none()); + } + + #[test] + fn ordinary_problem_has_no_permission_denial_marker() { + let response = new(StatusCode::FORBIDDEN).into_response(); + assert!(response + .extensions() + .get::() + .is_none()); } } diff --git a/web/src/components/audit/AuditLogItem.test.ts b/web/src/components/audit/AuditLogItem.test.ts new file mode 100644 index 000000000..d8648d6e1 --- /dev/null +++ b/web/src/components/audit/AuditLogItem.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, test } from 'bun:test' + +import { describePermissionDenial } from '@/lib/permission-denial-display' +import { categorize } from './AuditLogItem' + +describe('permission-denial audit presentation', () => { + test('categorizes permission denials as authentication events', () => { + expect(categorize('PERMISSION_DENIED')).toBe('auth') + }) + + test('renders only normalized, redacted denial metadata', () => { + expect( + describePermissionDenial({ + method: 'DELETE', + route: '/projects/{project_id}', + auth_source: 'api_key', + attempt_count: 3, + }) + ).toBe('Denied DELETE /projects/{project_id} for api key (3 attempts)') + }) + + test('handles missing optional denial metadata', () => { + expect(describePermissionDenial()).toBe('Denied a request') + }) +}) diff --git a/web/src/components/audit/AuditLogItem.tsx b/web/src/components/audit/AuditLogItem.tsx index e41e11d59..10faba8a1 100644 --- a/web/src/components/audit/AuditLogItem.tsx +++ b/web/src/components/audit/AuditLogItem.tsx @@ -2,6 +2,7 @@ import { AuditLogIpInfo, AuditLogUserInfo } from '@/api/client' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { TableCell, TableRow } from '@/components/ui/table' +import { describePermissionDenial } from '@/lib/permission-denial-display' import { cn } from '@/lib/utils' import { format } from 'date-fns' import { @@ -84,7 +85,8 @@ function categorize(op: string): Category { op.startsWith('AUTH_') || op === 'USER_LOGOUT' || op === 'PASSWORD_RESET' || - op === 'EMAIL_VERIFIED' + op === 'EMAIL_VERIFIED' || + op === 'PERMISSION_DENIED' ) return 'auth' if (op.startsWith('USER_') || op.startsWith('ROLE_')) return 'user' @@ -295,6 +297,8 @@ function describe( return 'Logged in successfully' case 'LOGIN_FAILURE': return `Failed login attempt${displayedAttemptedEmail ? ` for ${displayedAttemptedEmail}` : ''}${failureReason ? ` (${humanize(failureReason).toLowerCase()})` : ''}` + case 'PERMISSION_DENIED': + return describePermissionDenial(data) case 'USER_LOGOUT': return 'Logged out' case 'AUTH_INITIATED': diff --git a/web/src/lib/audit-operation-filters.ts b/web/src/lib/audit-operation-filters.ts new file mode 100644 index 000000000..b9eeb2189 --- /dev/null +++ b/web/src/lib/audit-operation-filters.ts @@ -0,0 +1,4 @@ +export const PERMISSION_DENIED_FILTER = { + value: 'PERMISSION_DENIED', + label: 'Permission Denied', +} as const diff --git a/web/src/lib/permission-denial-display.ts b/web/src/lib/permission-denial-display.ts new file mode 100644 index 000000000..628f6f00f --- /dev/null +++ b/web/src/lib/permission-denial-display.ts @@ -0,0 +1,25 @@ +function humanize(value: string): string { + return value.toLowerCase().replace(/_/g, ' ') +} + +function get( + data: Record | undefined, + key: string +): T | undefined { + return data?.[key] as T | undefined +} + +export function describePermissionDenial( + data?: Record +): string { + const method = get(data, 'method') + const route = get(data, 'route') + const authSource = get(data, 'auth_source') + const attemptCount = get(data, 'attempt_count') + const target = [method, route].filter(Boolean).join(' ') + const source = authSource ? ` for ${humanize(authSource)}` : '' + const attempts = + attemptCount && attemptCount > 1 ? ` (${attemptCount} attempts)` : '' + + return `Denied ${target || 'a request'}${source}${attempts}` +} diff --git a/web/src/pages/AuditLogs.test.ts b/web/src/pages/AuditLogs.test.ts new file mode 100644 index 000000000..ba9c4aaae --- /dev/null +++ b/web/src/pages/AuditLogs.test.ts @@ -0,0 +1,12 @@ +import { describe, expect, test } from 'bun:test' + +import { PERMISSION_DENIED_FILTER } from '@/lib/audit-operation-filters' + +describe('audit operation filters', () => { + test('offers permission denials in the authentication group', () => { + expect(PERMISSION_DENIED_FILTER).toEqual({ + value: 'PERMISSION_DENIED', + label: 'Permission Denied', + }) + }) +}) diff --git a/web/src/pages/AuditLogs.tsx b/web/src/pages/AuditLogs.tsx index 255e9069f..d93455aef 100644 --- a/web/src/pages/AuditLogs.tsx +++ b/web/src/pages/AuditLogs.tsx @@ -23,6 +23,7 @@ import { import { useBreadcrumbs } from '@/contexts/BreadcrumbContext' import { useCanViewAuditLogs } from '@/hooks/useAuditAccess' import { usePageTitle } from '@/hooks/usePageTitle' +import { PERMISSION_DENIED_FILTER } from '@/lib/audit-operation-filters' import { useQuery } from '@tanstack/react-query' import { ScrollText, X } from 'lucide-react' import { useEffect, useMemo, useState } from 'react' @@ -45,6 +46,7 @@ const OPERATION_GROUPS: OperationGroup[] = [ { value: 'USER_LOGOUT', label: 'User Logout' }, { value: 'PASSWORD_RESET', label: 'Password Reset' }, { value: 'EMAIL_VERIFIED', label: 'Email Verified' }, + PERMISSION_DENIED_FILTER, ], }, { From 2f684aa8b8200f9aea3db65cfcaabcd0545eb748 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 12:43:49 +0200 Subject: [PATCH 02/10] fix(auth): bound permission denial audit storage --- Cargo.lock | 1 + crates/temps-audit/Cargo.toml | 1 + crates/temps-audit/src/plugin.rs | 41 ++++ .../temps-audit/src/services/audit_service.rs | 178 +++++++++++++++- crates/temps-auth/src/audit.rs | 12 ++ .../src/permission_denial_recorder.rs | 198 +++++++++++++++++- ...00001_index_permission_denied_retention.rs | 52 +++++ crates/temps-migrations/src/migration/mod.rs | 2 + 8 files changed, 478 insertions(+), 7 deletions(-) create mode 100644 crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs diff --git a/Cargo.lock b/Cargo.lock index 3e6c59ebc..89ae38be6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10423,6 +10423,7 @@ dependencies = [ "temps-database", "temps-entities", "temps-geo", + "thiserror 2.0.19", "tokio", "tracing", "utoipa", diff --git a/crates/temps-audit/Cargo.toml b/crates/temps-audit/Cargo.toml index b040bfac6..41cb48487 100644 --- a/crates/temps-audit/Cargo.toml +++ b/crates/temps-audit/Cargo.toml @@ -21,6 +21,7 @@ sea-orm = { workspace = true } log = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } +thiserror = { workspace = true } axum = { workspace = true } serde_json = { workspace = true } diff --git a/crates/temps-audit/src/plugin.rs b/crates/temps-audit/src/plugin.rs index b588106c8..4074566b7 100644 --- a/crates/temps-audit/src/plugin.rs +++ b/crates/temps-audit/src/plugin.rs @@ -42,6 +42,47 @@ impl TempsPlugin for AuditPlugin { // Create AuditService let audit_service = Arc::new(AuditService::new(db.clone(), ip_address_service.clone())); context.register_service(audit_service.clone()); + + // At the recorder's hard ceiling of 17 rows/minute, at most 1,020 + // permission-denial rows are created per hour. The bounded 2,048-row + // hourly prune therefore cannot be outrun by allowed production + // writes, without ever issuing an unbounded delete transaction. + let retention_service = audit_service.clone(); + tokio::spawn(async move { + let mut ticker = tokio::time::interval( + crate::services::audit_service::PERMISSION_DENIED_PRUNE_INTERVAL, + ); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; + loop { + ticker.tick().await; + let cutoff = chrono::Utc::now() + - chrono::Duration::days( + crate::services::audit_service::PERMISSION_DENIED_RETENTION_DAYS, + ); + match retention_service + .prune_permission_denied_before( + cutoff, + crate::services::audit_service::PERMISSION_DENIED_PRUNE_BATCH_SIZE, + ) + .await + { + Ok(rows_deleted) => tracing::debug!( + rows_deleted, + retention_days = + crate::services::audit_service::PERMISSION_DENIED_RETENTION_DAYS, + "pruned expired permission-denial audit rows" + ), + Err(error) => tracing::warn!( + error = %error, + retention_days = crate::services::audit_service::PERMISSION_DENIED_RETENTION_DAYS, + batch_size = crate::services::audit_service::PERMISSION_DENIED_PRUNE_BATCH_SIZE, + "permission-denial audit retention pass failed; will retry" + ), + } + } + }); + let initial_logger: Arc = audit_service.clone(); let audit_slot = Arc::new(AuditLoggerSlot::new(initial_logger)); context.register_service(audit_slot.clone()); diff --git a/crates/temps-audit/src/services/audit_service.rs b/crates/temps-audit/src/services/audit_service.rs index a637c71d4..3503673ab 100644 --- a/crates/temps-audit/src/services/audit_service.rs +++ b/crates/temps-audit/src/services/audit_service.rs @@ -1,5 +1,8 @@ use chrono::Utc; -use sea_orm::{prelude::*, ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect, Set}; +use sea_orm::{ + prelude::*, ColumnTrait, ConnectionTrait, DatabaseBackend, EntityTrait, QueryFilter, + QueryOrder, QuerySelect, Set, Statement, +}; use serde::Serialize; use std::sync::Arc; use temps_core::{AuditLogger, AuditOperation, UtcDateTime}; @@ -8,6 +11,31 @@ use temps_entities::{audit_logs, ip_geolocations, users}; use temps_geo::IpAddressService; use tracing::warn; +pub const PERMISSION_DENIED_RETENTION_DAYS: i64 = 90; +pub const PERMISSION_DENIED_PRUNE_BATCH_SIZE: u64 = 2_048; +pub const PERMISSION_DENIED_PRUNE_INTERVAL: std::time::Duration = + std::time::Duration::from_secs(60 * 60); +const PERMISSION_DENIED_OPERATION: &str = "PERMISSION_DENIED"; + +#[derive(Debug, thiserror::Error)] +pub enum AuditMaintenanceError { + #[error("permission-denied retention batch size {batch_size} is outside 1..={max_batch_size}")] + InvalidBatchSize { + batch_size: u64, + max_batch_size: u64, + }, + #[error( + "failed to prune permission-denied audit rows older than {cutoff} with batch size \ + {batch_size}: {source}" + )] + PrunePermissionDenied { + cutoff: DateTimeUtc, + batch_size: u64, + #[source] + source: DbErr, + }, +} + /// Audit log with enriched user and IP geolocation data #[derive(Debug, Clone, Serialize)] pub struct AuditLogWithDetails { @@ -65,6 +93,54 @@ impl AuditService { Ok(result) } + /// Delete one bounded batch of expired permission-denial rows. + /// + /// Production runs this hourly with a 2,048-row batch and 90-day cutoff. + /// The recorder can persist at most 17 rows/minute (1,020/hour), so one + /// successful pass removes more than twice the maximum rows created in a + /// cadence interval while keeping every delete transaction bounded. + pub async fn prune_permission_denied_before( + &self, + cutoff: DateTimeUtc, + batch_size: u64, + ) -> Result { + if batch_size == 0 || batch_size > PERMISSION_DENIED_PRUNE_BATCH_SIZE { + return Err(AuditMaintenanceError::InvalidBatchSize { + batch_size, + max_batch_size: PERMISSION_DENIED_PRUNE_BATCH_SIZE, + }); + } + + let statement = Statement::from_sql_and_values( + DatabaseBackend::Postgres, + r#" +WITH expired AS ( + SELECT id + FROM audit_logs + WHERE operation_type = $1 AND audit_date < $2 + ORDER BY audit_date ASC, id ASC + LIMIT $3 +) +DELETE FROM audit_logs +WHERE id IN (SELECT id FROM expired) +"#, + vec![ + PERMISSION_DENIED_OPERATION.into(), + cutoff.into(), + (batch_size as i64).into(), + ], + ); + self.db + .execute(statement) + .await + .map(|result| result.rows_affected()) + .map_err(|source| AuditMaintenanceError::PrunePermissionDenied { + cutoff, + batch_size, + source, + }) + } + pub async fn get_user_audit_logs( &self, user_id_param: i32, @@ -210,7 +286,7 @@ impl AuditLogger for AuditService { #[cfg(test)] mod tests { use super::*; - use sea_orm::{DatabaseBackend, MockDatabase, Value}; + use sea_orm::{DatabaseBackend, MockDatabase, MockExecResult, Value}; use temps_geo::geoip_service::{GeoIpService, MockGeoIpService}; fn service_with(db: sea_orm::DatabaseConnection) -> AuditService { @@ -356,4 +432,102 @@ mod tests { assert_eq!(details.log.user_id, Some(7)); assert!(details.user.is_none()); } + + #[tokio::test] + async fn permission_denied_retention_is_bounded_and_parameterized() { + let db = Arc::new( + MockDatabase::new(DatabaseBackend::Postgres) + .append_exec_results([MockExecResult { + last_insert_id: 0, + rows_affected: 37, + }]) + .into_connection(), + ); + let geoip = Arc::new(GeoIpService::Mock(MockGeoIpService::new())); + let ip_service = Arc::new(IpAddressService::new(db.clone(), geoip)); + let service = AuditService::new(db.clone(), ip_service); + let cutoff = Utc::now() - chrono::Duration::days(PERMISSION_DENIED_RETENTION_DAYS); + + let deleted = service + .prune_permission_denied_before(cutoff, 128) + .await + .expect("bounded retention delete should succeed"); + assert_eq!(deleted, 37); + + drop(service); + let transactions = Arc::try_unwrap(db) + .expect("audit service should release the database connection") + .into_transaction_log(); + let statement = &transactions + .first() + .expect("retention should execute one statement") + .statements()[0]; + assert!(statement.sql.contains("operation_type = $1")); + assert!(statement.sql.contains("audit_date < $2")); + assert!(statement.sql.contains("LIMIT $3")); + assert!(statement.sql.contains("ORDER BY audit_date ASC, id ASC")); + assert_eq!( + statement + .values + .as_ref() + .expect("retention statement binds values") + .0 + .len(), + 3 + ); + assert!(!statement.sql.contains(PERMISSION_DENIED_OPERATION)); + } + + #[tokio::test] + async fn permission_denied_retention_rejects_unbounded_batches() { + let service = service_with(MockDatabase::new(DatabaseBackend::Postgres).into_connection()); + + for batch_size in [0, PERMISSION_DENIED_PRUNE_BATCH_SIZE + 1] { + let error = service + .prune_permission_denied_before(Utc::now(), batch_size) + .await + .expect_err("invalid batch must fail before querying"); + assert!(matches!( + error, + AuditMaintenanceError::InvalidBatchSize { + batch_size: actual, + max_batch_size: PERMISSION_DENIED_PRUNE_BATCH_SIZE, + } if actual == batch_size + )); + } + } + + #[tokio::test] + async fn permission_denied_retention_database_error_has_context() { + let service = service_with( + MockDatabase::new(DatabaseBackend::Postgres) + .append_exec_errors([DbErr::Custom("retention unavailable".to_string())]) + .into_connection(), + ); + let cutoff = Utc::now() - chrono::Duration::days(PERMISSION_DENIED_RETENTION_DAYS); + + let error = service + .prune_permission_denied_before(cutoff, 256) + .await + .expect_err("database failure should be typed"); + assert!(matches!( + error, + AuditMaintenanceError::PrunePermissionDenied { + cutoff: actual_cutoff, + batch_size: 256, + .. + } if actual_cutoff == cutoff + )); + } + + #[test] + fn retention_capacity_exceeds_maximum_permission_denied_creation_rate() { + const MAX_ROWS_PER_MINUTE: u64 = 17; + let intervals_per_hour = PERMISSION_DENIED_PRUNE_INTERVAL.as_secs() / 60; + assert_eq!(intervals_per_hour, 60); + assert!( + PERMISSION_DENIED_PRUNE_BATCH_SIZE > MAX_ROWS_PER_MINUTE * intervals_per_hour, + "cleanup must delete faster than the recorder can create rows" + ); + } } diff --git a/crates/temps-auth/src/audit.rs b/crates/temps-auth/src/audit.rs index 4daf7a4e9..d522aa9be 100644 --- a/crates/temps-auth/src/audit.rs +++ b/crates/temps-auth/src/audit.rs @@ -688,6 +688,9 @@ pub struct PermissionDeniedAudit { pub user_id: Option, pub auth_source: String, pub credential_id: Option, + /// True when attempts in this aggregate used more than one credential. + /// In that case `credential_id` is cleared to avoid false attribution. + pub multiple_credentials: bool, pub method: String, /// Axum route template (for example `/projects/{project_id}`), or the /// fixed value `unmatched`. Never a raw request URI. @@ -695,6 +698,12 @@ pub struct PermissionDeniedAudit { pub denial_kind: String, pub required_permission: Option, pub attempt_count: u64, + /// True when the aggregate spans multiple IPs or user agents. Singular + /// origin fields are neutralized when this is set. + pub multiple_origins: bool, + /// Identifies the single reserved row summarizing attempts omitted by the + /// per-window persistence budgets. + pub suppressed_by_budget: bool, pub ip_address: Option, pub user_agent: String, } @@ -886,11 +895,14 @@ mod failure_audit_tests { user_id: None, auth_source: "deployment_token".to_string(), credential_id: Some(17), + multiple_credentials: false, method: "POST".to_string(), route: "/projects/{project_id}/deployments".to_string(), denial_kind: "cross_project_scope".to_string(), required_permission: None, attempt_count: 4, + multiple_origins: false, + suppressed_by_budget: false, ip_address: Some("203.0.113.8".to_string()), user_agent: "test-agent".to_string(), }; diff --git a/crates/temps-auth/src/permission_denial_recorder.rs b/crates/temps-auth/src/permission_denial_recorder.rs index 22fc3b1f5..086aa080f 100644 --- a/crates/temps-auth/src/permission_denial_recorder.rs +++ b/crates/temps-auth/src/permission_denial_recorder.rs @@ -9,7 +9,12 @@ use crate::audit::PermissionDeniedAudit; const DEFAULT_QUEUE_CAPACITY: usize = 1_024; const DEFAULT_MAX_AGGREGATIONS: usize = 1_024; -const DEFAULT_WINDOW: Duration = Duration::from_secs(5); +/// A one-minute window limits normal persistence to at most 16 detailed rows +/// plus one reserved suppression-summary row (17 writes/minute globally). +const DEFAULT_WINDOW: Duration = Duration::from_secs(60); +const DEFAULT_MAX_DETAIL_ROWS: usize = 16; +const DEFAULT_MAX_DETAIL_ROWS_PER_ACTOR: usize = 4; +const MIXED_VALUE: &str = "multiple"; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub(crate) enum AuthSourceKind { @@ -64,7 +69,6 @@ pub(crate) struct PermissionDenialEvent { struct AggregationKey { user_id: Option, source: AuthSourceKind, - credential_id: Option, method: String, route: String, denial_kind: String, @@ -76,7 +80,6 @@ impl From<&PermissionDenialEvent> for AggregationKey { Self { user_id: event.principal.user_id, source: event.principal.source, - credential_id: event.principal.credential_id, method: event.method.clone(), route: event.route.clone(), denial_kind: event.denial_kind.clone(), @@ -91,11 +94,14 @@ impl PermissionDenialEvent { user_id: self.principal.user_id, auth_source: self.principal.source.as_str().to_string(), credential_id: self.principal.credential_id, + multiple_credentials: false, method: self.method, route: self.route, denial_kind: self.denial_kind, required_permission: self.required_permission, attempt_count: 1, + multiple_origins: false, + suppressed_by_budget: false, ip_address: self.ip_address, user_agent: self.user_agent, } @@ -106,6 +112,7 @@ struct RecorderCounters { queue_drops: AtomicU64, aggregation_overflows: AtomicU64, write_failures: AtomicU64, + budget_suppressed_attempts: AtomicU64, } impl RecorderCounters { @@ -114,6 +121,7 @@ impl RecorderCounters { queue_drops: AtomicU64::new(0), aggregation_overflows: AtomicU64::new(0), write_failures: AtomicU64::new(0), + budget_suppressed_attempts: AtomicU64::new(0), } } } @@ -123,6 +131,8 @@ struct RecorderConfig { queue_capacity: usize, max_aggregations: usize, window: Duration, + max_detail_rows: usize, + max_detail_rows_per_actor: usize, } /// Non-blocking, explicitly bounded entry point for permission-denial audit @@ -141,6 +151,8 @@ impl PermissionDenialRecorder { queue_capacity: DEFAULT_QUEUE_CAPACITY, max_aggregations: DEFAULT_MAX_AGGREGATIONS, window: DEFAULT_WINDOW, + max_detail_rows: DEFAULT_MAX_DETAIL_ROWS, + max_detail_rows_per_actor: DEFAULT_MAX_DETAIL_ROWS_PER_ACTOR, }, ) } @@ -154,6 +166,8 @@ impl PermissionDenialRecorder { counters: counters.clone(), max_aggregations: config.max_aggregations, window: config.window, + max_detail_rows: config.max_detail_rows, + max_detail_rows_per_actor: config.max_detail_rows_per_actor, pending: HashMap::with_capacity(config.max_aggregations), }; tokio::spawn(worker.run()); @@ -192,6 +206,8 @@ struct PermissionDenialWorker { counters: Arc, max_aggregations: usize, window: Duration, + max_detail_rows: usize, + max_detail_rows_per_actor: usize, pending: HashMap, } @@ -220,6 +236,7 @@ impl PermissionDenialWorker { let key = AggregationKey::from(&event); if let Some(audit) = self.pending.get_mut(&key) { audit.attempt_count = audit.attempt_count.saturating_add(1); + merge_attribution(audit, &event); return; } @@ -252,7 +269,54 @@ impl PermissionDenialWorker { async fn flush(&mut self) { // At most `max_aggregations` entries are moved into this temporary // vector. The queue and map remain bounded while logger I/O is in flight. - let audits: Vec<_> = self.pending.drain().map(|(_, audit)| audit).collect(); + let mut audits: Vec<_> = self.pending.drain().map(|(_, audit)| audit).collect(); + audits.sort_by(|left, right| { + right + .attempt_count + .cmp(&left.attempt_count) + .then_with(|| left.user_id.cmp(&right.user_id)) + .then_with(|| left.auth_source.cmp(&right.auth_source)) + .then_with(|| left.method.cmp(&right.method)) + .then_with(|| left.route.cmp(&right.route)) + .then_with(|| left.denial_kind.cmp(&right.denial_kind)) + .then_with(|| left.required_permission.cmp(&right.required_permission)) + }); + + let mut actor_rows: HashMap<(Option, String), usize> = HashMap::new(); + let mut selected = Vec::with_capacity(self.max_detail_rows); + let mut suppressed_attempts = 0_u64; + for audit in audits { + let actor = (audit.user_id, audit.auth_source.clone()); + let actor_count = actor_rows.entry(actor).or_default(); + if selected.len() < self.max_detail_rows + && *actor_count < self.max_detail_rows_per_actor + { + *actor_count += 1; + selected.push(audit); + } else { + suppressed_attempts = suppressed_attempts.saturating_add(audit.attempt_count); + } + } + + if suppressed_attempts > 0 { + let total = self + .counters + .budget_suppressed_attempts + .fetch_add(suppressed_attempts, Ordering::Relaxed) + .saturating_add(suppressed_attempts); + if should_log_counter(total) || total == suppressed_attempts { + tracing::warn!( + suppressed_attempts, + total_budget_suppressed_attempts = total, + max_detail_rows = self.max_detail_rows, + max_detail_rows_per_actor = self.max_detail_rows_per_actor, + "permission-denial detail persistence budget reached" + ); + } + selected.push(suppression_summary(suppressed_attempts)); + } + + let audits = selected; for audit in audits { if let Err(error) = self.logger.create_audit_log(&audit).await { let total = self.counters.write_failures.fetch_add(1, Ordering::Relaxed) + 1; @@ -275,6 +339,39 @@ impl PermissionDenialWorker { } } +fn merge_attribution(audit: &mut PermissionDeniedAudit, event: &PermissionDenialEvent) { + if !audit.multiple_credentials && audit.credential_id != event.principal.credential_id { + audit.credential_id = None; + audit.multiple_credentials = true; + } + + let mixed_ip = audit.ip_address != event.ip_address; + let mixed_user_agent = audit.user_agent != event.user_agent; + if mixed_ip || mixed_user_agent { + audit.ip_address = None; + audit.user_agent = MIXED_VALUE.to_string(); + audit.multiple_origins = true; + } +} + +fn suppression_summary(attempt_count: u64) -> PermissionDeniedAudit { + PermissionDeniedAudit { + user_id: None, + auth_source: MIXED_VALUE.to_string(), + credential_id: None, + multiple_credentials: true, + method: MIXED_VALUE.to_string(), + route: MIXED_VALUE.to_string(), + denial_kind: "persistence_budget_suppressed".to_string(), + required_permission: None, + attempt_count, + multiple_origins: true, + suppressed_by_budget: true, + ip_address: None, + user_agent: MIXED_VALUE.to_string(), + } +} + #[cfg(test)] mod tests { use std::sync::Mutex; @@ -332,11 +429,13 @@ mod tests { queue_capacity, max_aggregations, window: Duration::from_millis(20), + max_detail_rows: 8, + max_detail_rows_per_actor: 8, } } #[tokio::test] - async fn repeated_safe_key_is_aggregated_and_ip_is_not_part_of_key() { + async fn mixed_ip_clears_singular_origin_attribution() { let logger = Arc::new(RecordingLogger::default()); let recorder = PermissionDenialRecorder::with_config(logger.clone(), test_config(8, 8)); recorder.record(event("/projects/{project_id}")); @@ -351,10 +450,99 @@ mod tests { assert_eq!(records[0]["route"], "/projects/{project_id}"); assert_eq!(records[0]["auth_source"], "api_key"); assert_eq!(records[0]["credential_id"], 9); + assert_eq!(records[0]["ip_address"], serde_json::Value::Null); + assert_eq!(records[0]["user_agent"], MIXED_VALUE); + assert_eq!(records[0]["multiple_origins"], true); assert!(records[0].get("key_name").is_none()); assert!(records[0].get("token_name").is_none()); } + #[tokio::test] + async fn mixed_user_agent_clears_singular_origin_attribution() { + let logger = Arc::new(RecordingLogger::default()); + let recorder = PermissionDenialRecorder::with_config(logger.clone(), test_config(8, 8)); + recorder.record(event("/projects/{project_id}")); + let mut second = event("/projects/{project_id}"); + second.user_agent = "different-agent".to_string(); + recorder.record(second); + + tokio::time::sleep(Duration::from_millis(60)).await; + let records = logger.records.lock().expect("test logger lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0]["ip_address"], serde_json::Value::Null); + assert_eq!(records[0]["user_agent"], MIXED_VALUE); + assert_eq!(records[0]["multiple_origins"], true); + } + + #[tokio::test] + async fn mixed_credential_ids_cannot_multiply_aggregation_keys() { + let logger = Arc::new(RecordingLogger::default()); + let recorder = PermissionDenialRecorder::with_config(logger.clone(), test_config(8, 8)); + recorder.record(event("/projects/{project_id}")); + let mut second = event("/projects/{project_id}"); + second.principal.credential_id = Some(10); + recorder.record(second); + + tokio::time::sleep(Duration::from_millis(60)).await; + let records = logger.records.lock().expect("test logger lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0]["attempt_count"], 2); + assert_eq!(records[0]["credential_id"], serde_json::Value::Null); + assert_eq!(records[0]["multiple_credentials"], true); + } + + #[tokio::test] + async fn adversarial_cardinality_obeys_actor_and_global_write_budgets() { + let logger = Arc::new(RecordingLogger::default()); + let recorder = PermissionDenialRecorder::with_config( + logger.clone(), + RecorderConfig { + queue_capacity: 256, + max_aggregations: 256, + window: Duration::from_millis(30), + max_detail_rows: 16, + max_detail_rows_per_actor: 4, + }, + ); + + for actor_offset in 0..5 { + for route_offset in 0..10 { + let mut attempt = event(&format!("/route/{actor_offset}/{route_offset}")); + attempt.principal.user_id = Some(40 + actor_offset); + attempt.principal.credential_id = Some(1_000 + actor_offset * 10 + route_offset); + recorder.record(attempt); + } + } + + tokio::time::sleep(Duration::from_millis(100)).await; + let records = logger.records.lock().expect("test logger lock"); + let details: Vec<_> = records + .iter() + .filter(|record| record["suppressed_by_budget"] == false) + .collect(); + let summaries: Vec<_> = records + .iter() + .filter(|record| record["suppressed_by_budget"] == true) + .collect(); + + assert_eq!(details.len(), 16, "global detail ceiling must hold"); + let mut rows_by_actor = HashMap::::new(); + for detail in details { + let actor = detail["user_id"].as_i64().expect("detail has actor"); + *rows_by_actor.entry(actor).or_default() += 1; + } + assert!(rows_by_actor.values().all(|count| *count <= 4)); + assert_eq!(summaries.len(), 1, "summary uses one reserved row"); + assert_eq!(summaries[0]["attempt_count"], 34); + assert_eq!( + recorder + .counters + .budget_suppressed_attempts + .load(Ordering::Relaxed), + 34 + ); + } + #[tokio::test] async fn aggregation_cardinality_is_bounded() { let logger = Arc::new(RecordingLogger::default()); diff --git a/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs b/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs new file mode 100644 index 000000000..be081563d --- /dev/null +++ b/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs @@ -0,0 +1,52 @@ +//! Adds the partial index used by the bounded permission-denial retention +//! worker. Without it, each small delete batch could scan the entire audit +//! history as the table grows, defeating the bounded-maintenance contract. + +use sea_orm_migration::prelude::*; + +const INDEX_NAME: &str = "idx_audit_logs_permission_denied_retention"; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + // Migrations run transactionally, so PostgreSQL does not permit + // CONCURRENTLY. Bound both lock acquisition and build time rather than + // allowing startup to hang indefinitely on a busy audit table. + manager + .get_connection() + .execute_unprepared("SET LOCAL lock_timeout = '5s'") + .await?; + manager + .get_connection() + .execute_unprepared("SET LOCAL statement_timeout = '30s'") + .await?; + manager + .get_connection() + .execute_unprepared(&format!( + "CREATE INDEX IF NOT EXISTS {INDEX_NAME} \ + ON audit_logs (audit_date ASC, id ASC) \ + WHERE operation_type = 'PERMISSION_DENIED'" + )) + .await?; + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .get_connection() + .execute_unprepared("SET LOCAL lock_timeout = '5s'") + .await?; + manager + .get_connection() + .execute_unprepared("SET LOCAL statement_timeout = '30s'") + .await?; + manager + .get_connection() + .execute_unprepared(&format!("DROP INDEX IF EXISTS {INDEX_NAME}")) + .await?; + Ok(()) + } +} diff --git a/crates/temps-migrations/src/migration/mod.rs b/crates/temps-migrations/src/migration/mod.rs index 9ad65cdd3..a139d9a74 100644 --- a/crates/temps-migrations/src/migration/mod.rs +++ b/crates/temps-migrations/src/migration/mod.rs @@ -176,6 +176,7 @@ mod m20260804_000001_add_ai_data_access_to_external_services; mod m20260804_000001_add_must_change_password_to_users; pub mod m20260805_000001_index_normalized_managed_domains; pub mod m20260806_000001_sandbox_workspace_lifecycle; +mod m20260806_000001_index_permission_denied_retention; pub struct Migrator; @@ -367,6 +368,7 @@ impl MigratorTrait for Migrator { ), Box::new(m20260805_000001_index_normalized_managed_domains::Migration), Box::new(m20260806_000001_sandbox_workspace_lifecycle::Migration), + Box::new(m20260806_000001_index_permission_denied_retention::Migration), ] } } From 6368dc79a0bc67a5d79edeeb8c1c9b045efef8c7 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 13:22:00 +0200 Subject: [PATCH 03/10] fix(security): bound query memory and audit data --- Cargo.lock | 1 + .../temps-audit/src/services/audit_service.rs | 120 +++++++- crates/temps-providers/src/handlers/audit.rs | 34 ++- .../src/handlers/query_handlers.rs | 266 +++++++++--------- crates/temps-providers/src/mariadb_query.rs | 234 ++++++++------- crates/temps-query-mongodb/src/lib.rs | 69 +++-- crates/temps-query-postgres/Cargo.toml | 1 + crates/temps-query-postgres/src/lib.rs | 229 +++++++-------- crates/temps-query-redis/src/lib.rs | 200 ++++++++++--- crates/temps-query/src/budget.rs | 253 +++++++++++++++++ crates/temps-query/src/error.rs | 21 ++ crates/temps-query/src/lib.rs | 4 +- crates/temps-query/src/types.rs | 32 +++ 13 files changed, 1041 insertions(+), 423 deletions(-) create mode 100644 crates/temps-query/src/budget.rs diff --git a/Cargo.lock b/Cargo.lock index 89ae38be6..416d2af4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12306,6 +12306,7 @@ version = "0.1.0-beta.55" dependencies = [ "async-trait", "chrono", + "futures-util", "rustls", "serde_json", "temps-query", diff --git a/crates/temps-audit/src/services/audit_service.rs b/crates/temps-audit/src/services/audit_service.rs index 3503673ab..92fb2ae46 100644 --- a/crates/temps-audit/src/services/audit_service.rs +++ b/crates/temps-audit/src/services/audit_service.rs @@ -59,16 +59,22 @@ impl AuditService { operation: &T, ) -> anyhow::Result { let now = Utc::now(); + let operation_type = operation.operation_type(); let ip_address = operation.ip_address(); - let ip_address_id_val = match ip_address { - Some(ip_address) => match self.ip_service.get_or_create_ip(&ip_address).await { + // Permission denials are attacker-amplifiable and already retain their + // safe origin string inside the bounded audit JSON. Creating a durable + // geolocation row for each rotating denial IP would outlive the 90-day + // audit row and turn the security signal into an unbounded side table. + let ip_address_id_val = match (operation_type.as_str(), ip_address) { + (PERMISSION_DENIED_OPERATION, _) => None, + (_, Some(ip_address)) => match self.ip_service.get_or_create_ip(&ip_address).await { Ok(ip_address) => Some(ip_address.id), Err(err) => { warn!("Error getting ip address {:?}: {}", ip_address, err); None } }, - None => None, + (_, None) => None, }; // Serialize the operation to JSON @@ -76,7 +82,7 @@ impl AuditService { let new_audit_log = audit_logs::ActiveModel { user_id: Set(operation.user_id()), - operation_type: Set(operation.operation_type()), + operation_type: Set(operation_type), user_agent: Set(operation.user_agent().to_string()), ip_address_id: Set(ip_address_id_val), audit_date: Set(now), @@ -313,11 +319,13 @@ mod tests { #[derive(Serialize)] struct TestAuditOperation { user_id: Option, + operation_type: &'static str, + ip_address: Option, } impl AuditOperation for TestAuditOperation { fn operation_type(&self) -> String { - "TEST_OPERATION".to_string() + self.operation_type.to_string() } fn user_id(&self) -> Option { @@ -325,7 +333,7 @@ mod tests { } fn ip_address(&self) -> Option { - None + self.ip_address.clone() } fn user_agent(&self) -> &str { @@ -376,7 +384,11 @@ mod tests { #[tokio::test] async fn test_create_audit_log_persists_null_actor() { - let operation = TestAuditOperation { user_id: None }; + let operation = TestAuditOperation { + user_id: None, + operation_type: "TEST_OPERATION", + ip_address: None, + }; assert_eq!( persisted_actor_value(&operation, operation.user_id()).await, Value::Int(None) @@ -385,13 +397,105 @@ mod tests { #[tokio::test] async fn test_create_audit_log_persists_known_actor() { - let operation = TestAuditOperation { user_id: Some(42) }; + let operation = TestAuditOperation { + user_id: Some(42), + operation_type: "TEST_OPERATION", + ip_address: None, + }; assert_eq!( persisted_actor_value(&operation, operation.user_id()).await, Value::Int(Some(42)) ); } + #[tokio::test] + async fn permission_denial_keeps_ip_in_json_without_geolocation_row() { + let db = Arc::new( + MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results([vec![log_row(Some(42))]]) + .into_connection(), + ); + let geoip = Arc::new(GeoIpService::Mock(MockGeoIpService::new())); + let ip_service = Arc::new(IpAddressService::new(db.clone(), geoip)); + let service = AuditService::new(db.clone(), ip_service); + let operation = TestAuditOperation { + user_id: Some(42), + operation_type: PERMISSION_DENIED_OPERATION, + ip_address: Some("203.0.113.91".to_string()), + }; + + service + .create_audit_log_typed(&operation) + .await + .expect("permission denial audit should persist"); + drop(service); + let transactions = Arc::try_unwrap(db) + .expect("audit service should release database") + .into_transaction_log(); + let statements: Vec<_> = transactions + .iter() + .flat_map(|transaction| transaction.statements()) + .collect(); + assert_eq!(statements.len(), 1, "denial must only insert the audit row"); + assert!(statements[0].sql.contains("INSERT INTO \"audit_logs\"")); + assert!(!statements[0].sql.contains("ip_geolocations")); + assert!(AuditOperation::serialize(&operation) + .expect("test audit should serialize") + .contains("203.0.113.91")); + } + + #[tokio::test] + async fn normal_audit_still_enriches_existing_ip() { + let now = Utc::now(); + let ip = temps_entities::ip_geolocations::Model { + id: 77, + ip_address: "203.0.113.92".to_string(), + latitude: None, + longitude: None, + region: None, + city: None, + country: "".to_string(), + country_code: None, + timezone: None, + is_eu: false, + asn_org: None, + is_hosting_provider: None, + created_at: now, + updated_at: now, + }; + let db = Arc::new( + MockDatabase::new(DatabaseBackend::Postgres) + .append_query_results([vec![ip]]) + .append_query_results([vec![log_row(Some(42))]]) + .into_connection(), + ); + let geoip = Arc::new(GeoIpService::Mock(MockGeoIpService::new())); + let ip_service = Arc::new(IpAddressService::new(db.clone(), geoip)); + let service = AuditService::new(db.clone(), ip_service); + let operation = TestAuditOperation { + user_id: Some(42), + operation_type: "NORMAL_AUDIT", + ip_address: Some("203.0.113.92".to_string()), + }; + + service + .create_audit_log_typed(&operation) + .await + .expect("normal audit should enrich IP"); + drop(service); + let transactions = Arc::try_unwrap(db) + .expect("audit service should release database") + .into_transaction_log(); + let sql = transactions + .iter() + .flat_map(|transaction| transaction.statements()) + .map(|statement| statement.sql.as_str()) + .collect::>() + .join("\n"); + assert!(sql.contains("FROM \"ip_geolocations\"")); + assert!(sql.contains("INSERT INTO \"audit_logs\"")); + } + #[tokio::test] async fn test_get_log_by_id_without_user_skips_user_lookup() { // Only the audit row itself is prepared. If the service issued a diff --git a/crates/temps-providers/src/handlers/audit.rs b/crates/temps-providers/src/handlers/audit.rs index dd84bbe43..352c036d0 100644 --- a/crates/temps-providers/src/handlers/audit.rs +++ b/crates/temps-providers/src/handlers/audit.rs @@ -111,9 +111,9 @@ pub struct AiRowsReadAudit { pub entity: String, pub returned_rows: usize, pub truncated: bool, - /// The filter the agent supplied, if any. Recorded because it is the part - /// of the request a prompt injection would be steering. - pub filter: Option, + /// Bounded structural category only; literal filter values are never + /// persisted because they commonly contain emails, tokens, and other PII. + pub filter_shape: Option, } impl AuditOperation for AiRowsReadAudit { @@ -464,3 +464,31 @@ impl AuditOperation for ExternalServiceClusterMemberPromotedAudit { .map_err(|e| anyhow::anyhow!("Failed to serialize audit operation {}", e)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ai_row_read_audit_never_serializes_filter_literals() { + let audit = AiRowsReadAudit { + context: AuditContext { + user_id: 42, + ip_address: None, + user_agent: "test-agent".to_string(), + }, + service_id: 7, + service_name: "postgres".to_string(), + container_path: "app/public".to_string(), + entity: "users".to_string(), + returned_rows: 1, + truncated: false, + filter_shape: Some("sql_where".to_string()), + }; + + let serialized = AuditOperation::serialize(&audit).expect("audit serializes"); + assert!(serialized.contains("sql_where")); + assert!(!serialized.contains("alice@example.com")); + assert!(!serialized.contains("tok_live_secret")); + } +} diff --git a/crates/temps-providers/src/handlers/query_handlers.rs b/crates/temps-providers/src/handlers/query_handlers.rs index 4f9e1423c..feaa50993 100644 --- a/crates/temps-providers/src/handlers/query_handlers.rs +++ b/crates/temps-providers/src/handlers/query_handlers.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use std::sync::Arc; use temps_auth::{permission_guard, RequireAuth}; use temps_core::problemdetails::Problem; -use temps_query::{ContainerInfo, ContainerPath, EntityInfo, QueryOptions}; +use temps_query::{ContainerInfo, ContainerPath, EntityInfo, QueryBudget, QueryOptions}; use utoipa::ToSchema; use super::audit::{AiDataAccessChangedAudit, AiRowsReadAudit}; @@ -294,40 +294,95 @@ const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; /// that there is more. const AI_MAX_RESPONSE_BYTES: usize = 256 * 1024; -/// Serialise rows, stopping before the response exceeds `budget` bytes. -/// -/// Returns the rows that fit and whether anything was dropped. Truncating here -/// rather than letting a downstream layer cut the body in half keeps every -/// response valid JSON, and the caller reports the truncation explicitly so a -/// partial page is never mistaken for a complete one — by a human reading the -/// console, by a script reading the CLI, or by a model reading a tool result. -fn take_rows_within_budget( - rows: Vec, - budget: usize, -) -> (Vec, bool) { - let mut out = Vec::with_capacity(rows.len()); - let mut used = 0usize; - let mut truncated = false; - - for row in rows { - let value = serde_json::to_value(row).unwrap_or_default(); - // Cheap proxy for the serialised size. Exact to within the separators - // the array itself adds, which is far inside the slack in the budget. - let size = serde_json::to_string(&value).map(|s| s.len()).unwrap_or(0); - - // Always emit at least one row. A single row larger than the whole - // budget would otherwise produce an empty page that looks like "no - // results" and cannot be paged past. - if !out.is_empty() && used + size > budget { - truncated = true; - break; +fn query_budget(is_ai_call: bool) -> QueryBudget { + if is_ai_call { + QueryBudget { + max_bytes: AI_MAX_RESPONSE_BYTES, + max_cells: 4_096, + max_cells_per_row: 128, + max_cell_bytes: 64 * 1024, + max_value_elements_per_row: 1_000, + max_nesting_depth: 32, + } + } else { + QueryBudget { + max_bytes: MAX_RESPONSE_BYTES, + ..QueryBudget::default() } - - used += size; - out.push(value); } +} + +fn data_rows_to_json(rows: Vec) -> Vec { + rows.into_iter() + .map(|row| serde_json::Value::Object(row.into_iter().collect())) + .collect() +} + +fn filter_shape(filter: Option<&serde_json::Value>) -> Option { + filter + .map(|value| match value { + serde_json::Value::Object(object) if object.contains_key("where") => "sql_where", + serde_json::Value::Object(_) => "structured_object", + serde_json::Value::Array(_) => "structured_array", + _ => "scalar", + }) + .map(str::to_string) +} - (out, truncated) +fn query_error_problem( + error: temps_query::DataError, + service_id: i32, + entity: &str, + limit: usize, +) -> Problem { + let (status, title, kind, safe_detail) = match &error { + temps_query::DataError::ResultLimitExceeded { .. } => ( + StatusCode::PAYLOAD_TOO_LARGE, + "Query Result Too Large", + "result_limit_exceeded", + Some(error.to_string()), + ), + temps_query::DataError::InvalidQuery(_) => ( + StatusCode::BAD_REQUEST, + "Invalid Query", + "invalid_query", + Some(error.to_string()), + ), + temps_query::DataError::QueryFailed(_) + | temps_query::DataError::BackendQueryFailed { .. } => ( + StatusCode::BAD_REQUEST, + "Query Failed", + "backend_query_failed", + None, + ), + temps_query::DataError::NotFound(_) => { + (StatusCode::NOT_FOUND, "Not Found", "not_found", None) + } + temps_query::DataError::QueryTimeout(_) => ( + StatusCode::GATEWAY_TIMEOUT, + "Query Timed Out", + "query_timeout", + None, + ), + _ => ( + StatusCode::INTERNAL_SERVER_ERROR, + "Query Error", + "internal_query_error", + None, + ), + }; + tracing::warn!( + service_id, + entity, + limit, + error_kind = kind, + "external-service data query failed" + ); + let mut problem = temps_core::problemdetails::new(status).with_title(title); + if let Some(detail) = safe_detail { + problem = problem.with_detail(detail); + } + problem } /// Whether an agent-originated request may read this service's rows. @@ -553,6 +608,7 @@ pub async fn read_entity_rows( .await .ok(); let filters = parse_row_filter(query.filter.as_deref(), filter_schema.as_ref())?; + let audit_filter_shape = filter_shape(filters.as_ref()); let segments: Vec = path_str.split('/').map(String::from).collect(); let path = ContainerPath::new(segments); @@ -567,29 +623,14 @@ pub async fn read_entity_rows( sort_order: query.sort_order, timeout_ms: Some(crate::query_service::effective_timeout_ms(None)), include_nulls: true, + budget: query_budget(is_ai_call), }; let result = app_state .query_service .query_data(service_id, &path, &entity, filters, options) .await - .map_err(|e| { - let (status, title) = match &e { - temps_query::DataError::QueryFailed(_) => (StatusCode::BAD_REQUEST, "Query Failed"), - temps_query::DataError::InvalidQuery(_) => { - (StatusCode::BAD_REQUEST, "Invalid Query") - } - temps_query::DataError::NotFound(_) => (StatusCode::NOT_FOUND, "Not Found"), - temps_query::DataError::QueryTimeout(_) => { - (StatusCode::GATEWAY_TIMEOUT, "Query Timed Out") - } - _ => (StatusCode::INTERNAL_SERVER_ERROR, "Query Error"), - }; - - temps_core::problemdetails::new(status) - .with_title(title) - .with_detail(e.to_string()) - })?; + .map_err(|error| query_error_problem(error, service_id, &entity, limit))?; let total_count = result.stats.total_rows.unwrap_or(result.stats.row_count) as u64; let execution_time_ms = result.stats.execution_ms; @@ -604,12 +645,8 @@ pub async fn read_entity_rows( }) .collect(); - let budget = if is_ai_call { - AI_MAX_RESPONSE_BYTES - } else { - MAX_RESPONSE_BYTES - }; - let (rows, truncated) = take_rows_within_budget(result.rows, budget); + let truncated = result.stats.truncated; + let rows = data_rows_to_json(result.rows); // Record what the model actually read. The `ai_data_access` toggle is // audited, but that only says the door was opened; this says what went @@ -628,7 +665,7 @@ pub async fn read_entity_rows( entity: entity.clone(), returned_rows: rows.len(), truncated, - filter: query.filter.clone(), + filter_shape: audit_filter_shape, }; if let Err(e) = app_state.audit_service.create_audit_log(&audit).await { tracing::error!(service_id, error = %e, "Failed to write AI row-read audit log"); @@ -1364,36 +1401,14 @@ pub async fn query_data( sort_order: request.sort_order, timeout_ms: Some(crate::query_service::effective_timeout_ms(None)), include_nulls: true, + budget: query_budget(false), }; let result = app_state .query_service .query_data(service_id, &path, &entity, request.filters, options) .await - .map_err(|e| { - // Determine if this is a user error (400) or server error (500) - let (status, title) = match &e { - temps_query::DataError::QueryFailed(_) => { - // Query syntax errors are user errors - (StatusCode::BAD_REQUEST, "Query Failed") - } - temps_query::DataError::InvalidQuery(_) => { - (StatusCode::BAD_REQUEST, "Invalid Query") - } - temps_query::DataError::NotFound(_) => (StatusCode::NOT_FOUND, "Not Found"), - temps_query::DataError::QueryTimeout(_) => { - (StatusCode::GATEWAY_TIMEOUT, "Query Timed Out") - } - _ => { - // Other errors are server errors - (StatusCode::INTERNAL_SERVER_ERROR, "Query Error") - } - }; - - temps_core::problemdetails::new(status) - .with_title(title) - .with_detail(e.to_string()) // Use to_string() instead of format! to avoid extra nesting - })?; + .map_err(|error| query_error_problem(error, service_id, &entity, request.limit))?; let total_count = result.stats.total_rows.unwrap_or(result.stats.row_count) as u64; let execution_time_ms = result.stats.execution_ms; @@ -1410,7 +1425,8 @@ pub async fn query_data( // This route is not reachable by the agent (it is absent from the write // allowlist), so the human budget always applies. - let (rows, truncated) = take_rows_within_budget(result.rows, MAX_RESPONSE_BYTES); + let truncated = result.stats.truncated; + let rows = data_rows_to_json(result.rows); let response = QueryDataResponse { fields, @@ -1488,19 +1504,23 @@ pub async fn download_object( // Set response headers let mut headers = axum::http::HeaderMap::new(); - headers.insert( - header::CONTENT_TYPE, - content_type - .unwrap_or_else(|| "application/octet-stream".to_string()) - .parse() - .unwrap(), - ); - headers.insert( - header::CONTENT_DISPOSITION, - format!("attachment; filename=\"{}\"", entity) - .parse() - .unwrap(), - ); + let content_type = content_type.unwrap_or_else(|| "application/octet-stream".to_string()); + let content_type = axum::http::HeaderValue::from_str(&content_type).map_err(|_| { + temps_core::problemdetails::new(StatusCode::INTERNAL_SERVER_ERROR) + .with_title("Invalid Object Content Type") + .with_detail("The object provider returned an invalid content type") + })?; + let content_disposition = axum::http::HeaderValue::from_str(&format!( + "attachment; filename=\"{}\"", + entity.replace(['"', '\\'], "_") + )) + .map_err(|_| { + temps_core::problemdetails::new(StatusCode::BAD_REQUEST) + .with_title("Invalid Object Name") + .with_detail("The object name cannot be represented in a download header") + })?; + headers.insert(header::CONTENT_TYPE, content_type); + headers.insert(header::CONTENT_DISPOSITION, content_disposition); Ok((headers, body)) } @@ -1772,45 +1792,39 @@ mod tests { assert_eq!(effective_row_offset(MAX_OFFSET), MAX_OFFSET); } - fn row_of_size(bytes: usize) -> temps_query::DataRow { - let mut row = temps_query::DataRow::new(); - row.insert("blob".to_string(), serde_json::json!("x".repeat(bytes))); - row - } - #[test] - fn take_rows_within_budget_stops_before_blowing_the_budget() { - // MAX_ROWS assumes rows are small. A bytea/jsonb column holding an - // upload or a session blob is megabytes on its own, so a page of them - // is the same OOM the row clamp exists to prevent, reached along the - // axis the row clamp does not measure. - let rows: Vec<_> = (0..10).map(|_| row_of_size(1000)).collect(); - let (kept, truncated) = take_rows_within_budget(rows, 3_000); - - assert!(truncated, "should have reported truncation"); - assert!(kept.len() < 10, "should have dropped rows"); - assert!(!kept.is_empty(), "should have kept what fit"); + fn agent_budget_is_stricter_than_human_budget() { + let human = query_budget(false); + let agent = query_budget(true); + assert_eq!(human.max_bytes, MAX_RESPONSE_BYTES); + assert_eq!(agent.max_bytes, AI_MAX_RESPONSE_BYTES); + assert!(agent.max_cell_bytes < human.max_cell_bytes); + assert!(agent.max_value_elements_per_row < human.max_value_elements_per_row); } #[test] - fn take_rows_within_budget_reports_complete_pages_as_complete() { - // `truncated` drives whether a caller believes the table ended here, so - // a false positive is as bad as a false negative. - let rows: Vec<_> = (0..3).map(|_| row_of_size(10)).collect(); - let (kept, truncated) = take_rows_within_budget(rows, MAX_RESPONSE_BYTES); - - assert_eq!(kept.len(), 3); - assert!(!truncated); + fn filter_shape_discards_email_and_token_literals() { + let filter = serde_json::json!({ + "where": "email = 'alice@example.com' AND token = 'tok_live_secret'" + }); + let shape = filter_shape(Some(&filter)).expect("filter shape"); + assert_eq!(shape, "sql_where"); + assert!(!shape.contains("alice@example.com")); + assert!(!shape.contains("tok_live_secret")); } #[test] - fn take_rows_within_budget_always_emits_at_least_one_row() { - // A single row bigger than the entire budget must not produce an empty - // page: that reads as "no results", and no amount of paging gets past - // it. Emit the one row and flag truncation instead. - let (kept, _) = take_rows_within_budget(vec![row_of_size(4096)], 16); - - assert_eq!(kept.len(), 1, "must not return an unpageable empty page"); + fn query_error_problem_does_not_echo_backend_diagnostics() { + let secret = "alice@example.com tok_live_secret"; + let problem = query_error_problem( + temps_query::DataError::QueryFailed(format!("syntax near {secret}")), + 7, + "users", + 100, + ); + let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); + assert!(!serialized.contains("alice@example.com")); + assert!(!serialized.contains("tok_live_secret")); } #[test] diff --git a/crates/temps-providers/src/mariadb_query.rs b/crates/temps-providers/src/mariadb_query.rs index 9909bc667..3512a082f 100644 --- a/crates/temps-providers/src/mariadb_query.rs +++ b/crates/temps-providers/src/mariadb_query.rs @@ -1,12 +1,13 @@ use async_trait::async_trait; -use base64::Engine; -use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow}; -use sqlx::{Column, Row, TypeInfo}; +use futures::TryStreamExt; +use sqlx::mysql::{MySqlPool, MySqlPoolOptions}; +use sqlx::Row; use std::collections::HashMap; use temps_query::{ - Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, - DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, - Introspect, QueryOptions, QueryResult, QuerySchemaProvider, QueryStats, Queryable, Result, + BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, + DataError, DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, + FieldType, Introspect, QueryOptions, QueryResult, QuerySchemaProvider, QueryStats, Queryable, + Result, }; use tracing::{debug, error, warn}; @@ -131,83 +132,6 @@ impl MariaDbSource { _ => FieldType::String, } } - - fn row_to_datarow(row: &MySqlRow) -> Result { - let mut data_row = HashMap::new(); - for (idx, column) in row.columns().iter().enumerate() { - let value = Self::extract_value(row, idx)?; - data_row.insert(column.name().to_string(), value); - } - Ok(data_row) - } - - fn extract_value(row: &MySqlRow, idx: usize) -> Result { - let column = &row.columns()[idx]; - let type_name = column.type_info().name().to_ascii_lowercase(); - - let value = match type_name.as_str() { - "bool" | "boolean" => row - .try_get::, _>(idx) - .ok() - .flatten() - .map(serde_json::Value::Bool) - .unwrap_or(serde_json::Value::Null), - "tinyint" | "smallint" | "mediumint" | "int" | "integer" | "year" => row - .try_get::, _>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::Number(v.into())) - .unwrap_or(serde_json::Value::Null), - "bigint" => row - .try_get::, _>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::Number(v.into())) - .or_else(|| { - row.try_get::, _>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::Number(v.into())) - }) - .unwrap_or(serde_json::Value::Null), - "float" => row - .try_get::, _>(idx) - .ok() - .flatten() - .and_then(|v| serde_json::Number::from_f64(v as f64)) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null), - "double" | "real" => row - .try_get::, _>(idx) - .ok() - .flatten() - .and_then(serde_json::Number::from_f64) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null), - "json" => row - .try_get::, _>(idx) - .ok() - .flatten() - .and_then(|v| serde_json::from_str(&v).ok()) - .unwrap_or(serde_json::Value::Null), - "binary" | "varbinary" | "tinyblob" | "blob" | "mediumblob" | "longblob" => row - .try_get::>, _>(idx) - .ok() - .flatten() - .map(|v| { - serde_json::Value::String(base64::engine::general_purpose::STANDARD.encode(v)) - }) - .unwrap_or(serde_json::Value::Null), - _ => row - .try_get::, _>(idx) - .ok() - .flatten() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - }; - - Ok(value) - } } #[async_trait] @@ -572,10 +496,12 @@ impl Queryable for MariaDbSource { ) -> Result { let database_name = database_from_path(container_path, &self.database_name)?; validate_identifier("table", entity_name)?; + let schema = self.get_schema(container_path, entity_name).await?; let start = std::time::Instant::now(); + let projection = mariadb_json_projection(&schema); let mut sql = format!( - "SELECT * FROM {}.{}", + "SELECT {projection} AS __temps_payload FROM {}.{}", quote_identifier(database_name), quote_identifier(entity_name) ); @@ -604,8 +530,12 @@ impl Queryable for MariaDbSource { let limit = options.limit.unwrap_or(100); let offset = options.offset.unwrap_or(0); sql.push_str(" LIMIT ? OFFSET ?"); + let sql = with_wire_row_budget(&sql); - debug!("Executing MariaDB query: {}", sql); + debug!( + entity = entity_name, + limit, offset, "executing MariaDB data query" + ); // SECURITY / AVAILABILITY: bound the query server-side. // @@ -626,20 +556,65 @@ impl Queryable for MariaDbSource { })?; apply_statement_timeout(&mut conn, timeout_ms, database_name).await; - let rows = sqlx::query(&sql) + let mut stream = sqlx::query(&sql) + .bind(options.budget.max_bytes as u64) .bind(limit as i64) .bind(offset as i64) - .fetch_all(&mut *conn) - .await - .map_err(|e| { - error!("MariaDB query failed: {}", e); - DataError::QueryFailed(format!("{}\n\nQuery: {}", e, sql)) + .fetch(&mut *conn); + let mut bounded = BoundedRows::new(options.budget); + while let Some(row) = stream.try_next().await.map_err(|_error| { + error!(entity = entity_name, limit, "MariaDB row stream failed"); + DataError::BackendQueryFailed { + backend: "MariaDB", + entity: entity_name.to_string(), + } + })? { + let observed = row.try_get::("__temps_size").map_err(|_error| { + error!( + entity = entity_name, + limit, "MariaDB bounded row size decode failed" + ); + DataError::BackendQueryFailed { + backend: "MariaDB", + entity: entity_name.to_string(), + } })?; + let payload = row + .try_get::, _>("__temps_row") + .map_err(|_error| { + error!( + entity = entity_name, + limit, "MariaDB bounded row decode failed" + ); + DataError::BackendQueryFailed { + backend: "MariaDB", + entity: entity_name.to_string(), + } + })? + .ok_or_else(|| DataError::ResultLimitExceeded { + entity: entity_name.to_string(), + limit_kind: "wire_row_bytes", + limit: options.budget.max_bytes, + observed: usize::try_from(observed).unwrap_or(usize::MAX), + })?; + let data_row = serde_json::from_str::(&payload).map_err(|_error| { + error!( + entity = entity_name, + limit, "MariaDB bounded JSON row decode failed" + ); + DataError::BackendQueryFailed { + backend: "MariaDB", + entity: entity_name.to_string(), + } + })?; + if !bounded.push(entity_name, data_row)? { + break; + } + } + drop(stream); drop(conn); - let data_rows: Result> = rows.iter().map(Self::row_to_datarow).collect(); - let data_rows = data_rows?; - let schema = self.get_schema(container_path, entity_name).await?; + let (data_rows, truncated) = bounded.into_parts(); let row_count = data_rows.len(); Ok(QueryResult { @@ -651,6 +626,7 @@ impl Queryable for MariaDbSource { execution_ms: start.elapsed().as_millis() as u64, has_more: row_count >= limit, next_cursor: None, + truncated, }, }) } @@ -826,7 +802,43 @@ fn database_from_path<'a>( } fn quote_identifier(value: &str) -> String { - format!("`{}`", value) + format!("`{}`", value.replace('`', "``")) +} + +/// Build one JSON value per source row so MariaDB can measure it before the +/// driver sees any user-controlled cell contents. Binary values retain the +/// previous API representation (base64 text). +fn mariadb_json_projection(schema: &DatasetSchema) -> String { + let entries = schema.fields.iter().flat_map(|field| { + let key_hex = field + .name + .as_bytes() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let key = format!("CONVERT(X'{key_hex}' USING utf8mb4)"); + let column = quote_identifier(&field.name); + let value = if field.field_type == FieldType::Bytes { + format!("TO_BASE64({column})") + } else { + column + }; + [key, value] + }); + + format!("JSON_OBJECT({})", entries.collect::>().join(", ")) +} + +/// Apply the per-row response ceiling in the same statement that reads the +/// row. Oversized JSON is replaced by NULL plus its length, so it never crosses +/// the database wire into the control plane and there is no check/read race. +fn with_wire_row_budget(sql: &str) -> String { + format!( + "SELECT CASE WHEN OCTET_LENGTH(__temps_payload) <= ? \ + THEN __temps_payload ELSE NULL END AS __temps_row, \ + OCTET_LENGTH(__temps_payload) AS __temps_size \ + FROM ({sql}) AS __temps_bounded" + ) } fn validate_identifier(label: &str, value: &str) -> Result<()> { @@ -1191,6 +1203,38 @@ pub(crate) fn is_mariadb_compatible_image(image: &str) -> bool { mod tests { use super::*; + #[test] + fn generated_query_guards_encoded_row_before_wire_transfer() { + let schema = DatasetSchema { + fields: vec![ + FieldDef { + name: "display_name".to_string(), + field_type: FieldType::String, + nullable: false, + description: None, + }, + FieldDef { + name: "avatar".to_string(), + field_type: FieldType::Bytes, + nullable: true, + description: None, + }, + ], + partitions: None, + primary_key: None, + }; + let projection = mariadb_json_projection(&schema); + let sql = with_wire_row_budget(&format!( + "SELECT {projection} AS __temps_payload FROM `app`.`users` LIMIT ? OFFSET ?" + )); + + assert!(sql.contains("OCTET_LENGTH(__temps_payload) <= ?")); + assert!(sql.contains("THEN __temps_payload ELSE NULL")); + assert!(sql.contains("JSON_OBJECT(")); + assert!(sql.contains("TO_BASE64(`avatar`)")); + assert!(sql.contains("FROM `app`.`users` LIMIT ? OFFSET ?")); + } + fn assert_where_rejected(clause: &str) { assert!( validate_where_clause(clause).is_err(), diff --git a/crates/temps-query-mongodb/src/lib.rs b/crates/temps-query-mongodb/src/lib.rs index c2d2463f1..4ce836804 100644 --- a/crates/temps-query-mongodb/src/lib.rs +++ b/crates/temps-query-mongodb/src/lib.rs @@ -34,8 +34,8 @@ use mongodb::{ }; use std::collections::HashMap; use temps_query::{ - Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, - DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, Result, + BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, + DataError, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, Result, }; use tracing::{debug, error}; @@ -278,7 +278,9 @@ impl MongoDBSource { row_count: None, // Would require counting documents size_bytes: None, schema: None, - metadata: Some(serde_json::to_value(metadata_map).unwrap()), + metadata: Some(serde_json::Value::Object( + metadata_map.into_iter().collect(), + )), } }) .collect(); @@ -361,7 +363,9 @@ impl MongoDBSource { row_count: doc_count.map(|c| c as usize), size_bytes, schema, - metadata: Some(serde_json::to_value(metadata_map).unwrap()), + metadata: Some(serde_json::Value::Object( + metadata_map.into_iter().collect(), + )), }) } @@ -630,8 +634,6 @@ impl temps_query::Queryable for MongoDBSource { Document::new() }; - debug!("MongoDB filter: {:?}", filter_doc); - // Apply pagination let limit = options.limit.unwrap_or(100) as i64; let skip = options.offset.unwrap_or(0) as u64; @@ -648,8 +650,8 @@ impl temps_query::Queryable for MongoDBSource { }; debug!( - "MongoDB query: filter={:?}, limit={}, skip={}, sort={:?}", - filter_doc, limit, skip, sort_doc + entity = entity_name, + limit, skip, "executing MongoDB data query" ); let start_time = std::time::Instant::now(); @@ -671,20 +673,34 @@ impl temps_query::Queryable for MongoDBSource { .skip(skip) .max_time(max_time) .await - .map_err(|e| { - error!("MongoDB query failed: {}", e); - DataError::QueryFailed(format!("MongoDB query failed: {}", e)) + .map_err(|_error| { + error!(entity = entity_name, limit, "MongoDB query failed"); + DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + } })?; - // Collect results - let mut rows = Vec::new(); - while cursor.advance().await.map_err(|e| { - error!("Failed to iterate MongoDB cursor: {}", e); - DataError::QueryFailed(format!("Failed to iterate results: {}", e)) + // Decode one cursor document at a time into the shared bounded + // collector. MongoDB itself caps one BSON document at 16 MiB; the + // tighter query budget rejects it before it can accumulate with peers. + let mut bounded = BoundedRows::new(options.budget); + while cursor.advance().await.map_err(|_error| { + error!(entity = entity_name, limit, "MongoDB cursor failed"); + DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + } })? { - let doc = cursor.deserialize_current().map_err(|e| { - error!("Failed to deserialize MongoDB document: {}", e); - DataError::QueryFailed(format!("Failed to deserialize document: {}", e)) + let doc = cursor.deserialize_current().map_err(|_error| { + error!( + entity = entity_name, + limit, "MongoDB document decode failed" + ); + DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + } })?; // Convert Document to HashMap for DataRow @@ -696,8 +712,11 @@ impl temps_query::Queryable for MongoDBSource { row_map.insert(key, json_value); } } - rows.push(row_map); + if !bounded.push(entity_name, row_map)? { + break; + } } + let (rows, truncated) = bounded.into_parts(); // Get total count (expensive, but needed for pagination). // @@ -709,9 +728,12 @@ impl temps_query::Queryable for MongoDBSource { .count_documents(filter_doc) .max_time(max_time) .await - .map_err(|e| { - error!("Failed to count MongoDB documents: {}", e); - DataError::QueryFailed(format!("Failed to count documents: {}", e)) + .map_err(|_error| { + error!(entity = entity_name, limit, "MongoDB count failed"); + DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + } })?; let execution_time = start_time.elapsed(); @@ -769,6 +791,7 @@ impl temps_query::Queryable for MongoDBSource { execution_ms: execution_time.as_millis() as u64, has_more, next_cursor: None, // MongoDB uses offset-based pagination, not cursors + truncated, }, }) } diff --git a/crates/temps-query-postgres/Cargo.toml b/crates/temps-query-postgres/Cargo.toml index 629490ae0..ff507bfab 100644 --- a/crates/temps-query-postgres/Cargo.toml +++ b/crates/temps-query-postgres/Cargo.toml @@ -19,6 +19,7 @@ tracing = "0.1" thiserror = "2.0" chrono = "0.4" uuid = "1.0" +futures-util = "0.3" [dev-dependencies] testcontainers = { workspace = true } diff --git a/crates/temps-query-postgres/src/lib.rs b/crates/temps-query-postgres/src/lib.rs index aa4679ae4..7e30eeef7 100644 --- a/crates/temps-query-postgres/src/lib.rs +++ b/crates/temps-query-postgres/src/lib.rs @@ -3,14 +3,15 @@ //! Implements DataSource, Introspect, and Queryable traits for PostgreSQL. use async_trait::async_trait; +use futures_util::{pin_mut, TryStreamExt}; use std::collections::HashMap; use std::sync::Arc; use temps_query::{ - Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, - DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, + BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, + DataError, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, Introspect, QueryOptions, QueryResult, QueryStats, Queryable, Result, }; -use tokio_postgres::{Client, NoTls, Row}; +use tokio_postgres::{types::ToSql, Client, NoTls}; use tokio_postgres_rustls::MakeRustlsConnect; use tracing::{debug, error, warn}; @@ -32,6 +33,21 @@ fn escape_ident(name: &str) -> String { name.replace('"', "\"\"") } +/// Keep an individual PostgreSQL row bounded before it crosses the wire. +/// `CASE` returns only NULL plus the observed size for oversized rows, so the +/// client driver never allocates their text/blob/json contents. The projection +/// and size decision are part of the same statement, avoiding a TOCTOU race. +fn with_wire_row_budget(sql: &str, max_bytes: usize) -> String { + format!( + "SELECT CASE WHEN OCTET_LENGTH(__temps_payload) <= {max_bytes} \ + THEN __temps_payload::jsonb ELSE NULL END AS __temps_row, \ + OCTET_LENGTH(__temps_payload)::bigint AS __temps_size \ + FROM ({sql}) AS __temps_source \ + CROSS JOIN LATERAL (SELECT TO_JSONB(__temps_source)::text AS __temps_payload) \ + AS __temps_encoded" + ) +} + /// A certificate verifier that accepts all server certificates (including self-signed). /// /// SECURITY: this verifies nothing — with it, TLS gives encryption against a @@ -1212,102 +1228,6 @@ impl PostgresSource { _ => FieldType::String, // Default fallback } } - - /// Convert PostgreSQL row to DataRow - fn row_to_datarow(row: &Row) -> Result { - let mut data_row = HashMap::new(); - - for (idx, column) in row.columns().iter().enumerate() { - let name = column.name().to_string(); - let value = Self::extract_value(row, idx)?; - data_row.insert(name, value); - } - - Ok(data_row) - } - - /// Extract value from PostgreSQL row - fn extract_value(row: &Row, idx: usize) -> Result { - let column = &row.columns()[idx]; - let type_name = column.type_().name(); - - let value = match type_name { - "bool" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(serde_json::Value::Bool) - .unwrap_or(serde_json::Value::Null), - - "int2" | "int4" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::Number(v.into())) - .unwrap_or(serde_json::Value::Null), - - "int8" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::Number(v.into())) - .unwrap_or(serde_json::Value::Null), - - "float4" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .and_then(|v| serde_json::Number::from_f64(v as f64)) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null), - - "float8" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .and_then(serde_json::Number::from_f64) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null), - - "varchar" | "text" | "char" | "bpchar" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null), - - "timestamp" | "timestamptz" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::String(v.to_string())) - .unwrap_or(serde_json::Value::Null), - - "json" | "jsonb" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .unwrap_or(serde_json::Value::Null), - - "uuid" => row - .try_get::<_, Option>(idx) - .ok() - .flatten() - .map(|v| serde_json::Value::String(v.to_string())) - .unwrap_or(serde_json::Value::Null), - - _ => { - // Try to get as string for unknown types - row.try_get::<_, Option>(idx) - .ok() - .flatten() - .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null) - } - }; - - Ok(value) - } } #[async_trait] @@ -1952,8 +1872,12 @@ impl Queryable for PostgresSource { let limit = options.limit.unwrap_or(100); let offset = options.offset.unwrap_or(0); sql.push_str(&format!(" LIMIT {} OFFSET {}", limit, offset)); + let sql = with_wire_row_budget(&sql, options.budget.max_bytes); - debug!("Executing query: {}", sql); + debug!( + entity = entity_name, + limit, offset, "executing PostgreSQL data query" + ); // Safety: SQL injection is prevented by validate_sql() for WHERE clauses // and escape_ident() for identifiers. The database user should be read-only @@ -1985,43 +1909,67 @@ impl Queryable for PostgresSource { warn!("Failed to set statement_timeout for data browser query: {e}"); } - let rows = client.query(&sql, &[]).await.map_err(|e| { - error!("PostgreSQL query failed: {}", e); - error!("Failed SQL: {}", sql); - - // Extract detailed error message from PostgreSQL error - let error_msg = if let Some(db_error) = e.as_db_error() { - // Build detailed error message from PostgreSQL error fields - let mut msg = db_error.message().to_string(); - - if let Some(detail) = db_error.detail() { - msg.push_str(&format!("\nDetail: {}", detail)); - } - - if let Some(hint) = db_error.hint() { - msg.push_str(&format!("\nHint: {}", hint)); + let stream = client + .query_raw(&sql, std::iter::empty::<&dyn ToSql>()) + .await + .map_err(|_error| { + error!(entity = entity_name, limit, "PostgreSQL query failed"); + DataError::BackendQueryFailed { + backend: "PostgreSQL", + entity: entity_name.to_string(), } - - if let Some(position) = db_error.position() { - msg.push_str(&format!("\nPosition: {:?}", position)); + })?; + pin_mut!(stream); + let mut bounded = BoundedRows::new(options.budget); + while let Some(row) = stream.try_next().await.map_err(|_error| { + error!(entity = entity_name, limit, "PostgreSQL row stream failed"); + DataError::BackendQueryFailed { + backend: "PostgreSQL", + entity: entity_name.to_string(), + } + })? { + let observed = row.try_get::<_, i64>("__temps_size").map_err(|_error| { + error!( + entity = entity_name, + limit, "PostgreSQL bounded row size decode failed" + ); + DataError::BackendQueryFailed { + backend: "PostgreSQL", + entity: entity_name.to_string(), } - - if let Some(column) = db_error.column() { - msg.push_str(&format!("\nColumn: {}", column)); + })?; + let payload = row + .try_get::<_, Option>("__temps_row") + .map_err(|_error| { + error!( + entity = entity_name, + limit, "PostgreSQL bounded row decode failed" + ); + DataError::BackendQueryFailed { + backend: "PostgreSQL", + entity: entity_name.to_string(), + } + })? + .ok_or_else(|| DataError::ResultLimitExceeded { + entity: entity_name.to_string(), + limit_kind: "wire_row_bytes", + limit: options.budget.max_bytes, + observed: usize::try_from(observed).unwrap_or(usize::MAX), + })?; + let data_row = match payload { + serde_json::Value::Object(values) => values.into_iter().collect(), + _ => { + return Err(DataError::SerializationError(format!( + "PostgreSQL bounded row for entity '{}' was not an object", + entity_name + ))) } - - msg - } else { - // Non-database error (connection error, etc.) - format!("{}", e) }; - - DataError::QueryFailed(format!("{}\n\nQuery: {}", error_msg, sql)) - })?; - - // Convert rows to DataRow - let data_rows: Result> = rows.iter().map(Self::row_to_datarow).collect(); - let data_rows = data_rows?; + if !bounded.push(entity_name, data_row)? { + break; + } + } + let (data_rows, truncated) = bounded.into_parts(); // Get schema from first row or from table schema let schema = self.get_schema(container_path, entity_name).await?; @@ -2040,6 +1988,7 @@ impl Queryable for PostgresSource { execution_ms, has_more: row_count >= limit, next_cursor: None, + truncated, }, }) } @@ -2130,6 +2079,20 @@ impl Queryable for PostgresSource { } } +#[cfg(test)] +mod wire_budget_tests { + use super::with_wire_row_budget; + + #[test] + fn generated_query_guards_encoded_row_before_wire_transfer() { + let sql = with_wire_row_budget("SELECT * FROM public.users", 262_144); + assert!(sql.contains("OCTET_LENGTH(__temps_payload) <= 262144")); + assert!(sql.contains("THEN __temps_payload::jsonb ELSE NULL")); + assert!(sql.contains("TO_JSONB(__temps_source)::text")); + assert!(sql.contains("SELECT * FROM public.users")); + } +} + impl temps_query::QuerySchemaProvider for PostgresSource { fn get_filter_schema(&self) -> serde_json::Value { serde_json::json!({ diff --git a/crates/temps-query-redis/src/lib.rs b/crates/temps-query-redis/src/lib.rs index 36bb20ec9..cd43122cc 100644 --- a/crates/temps-query-redis/src/lib.rs +++ b/crates/temps-query-redis/src/lib.rs @@ -32,9 +32,9 @@ use redis::aio::ConnectionManager; use redis::{AsyncCommands, RedisError}; use std::collections::HashMap; use temps_query::{ - Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, - DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, - QueryOptions, QueryResult, Queryable, Result, + BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, + DataError, DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, + FieldType, QueryOptions, QueryResult, QueryStats, Queryable, Result, }; use tracing::{debug, error}; @@ -43,6 +43,24 @@ pub struct RedisSource { connection: ConnectionManager, } +fn redis_aggregate_limit(options: &QueryOptions) -> usize { + // A JSON array/object consumes at least two structural elements per Redis + // item (container entry + scalar), and zset rows consume more. Dividing by + // four is deliberately conservative and keeps the shared structural + // limiter from accepting a backend page that was already too large. + let structural_limit = (options.budget.max_value_elements_per_row / 4).max(1); + options.limit.unwrap_or(100).min(structural_limit) +} + +fn trim_extra(values: &mut Vec, limit: usize) -> bool { + if values.len() > limit { + values.truncate(limit); + true + } else { + false + } +} + impl RedisSource { /// Create a new Redis data source /// @@ -205,7 +223,12 @@ impl RedisSource { } /// Get the value of a specific key - async fn get_key_value(&self, db: i32, key: &str) -> Result { + async fn get_key_value( + &self, + db: i32, + key: &str, + options: &QueryOptions, + ) -> Result<(DataRow, bool)> { let mut conn = self.get_db_connection(db).await?; debug!("Getting value for key '{}' in database {}", key, db); @@ -239,37 +262,102 @@ impl RedisSource { DataError::QueryFailed(format!("Failed to get key TTL: {}", e)) })?; + // Reject a huge value before asking Redis to send it. Redis strings and + // collection members may each be hundreds of MiB; paging the collection + // alone does not protect the control plane from one giant member. + let memory_usage: Option = redis::cmd("MEMORY") + .arg("USAGE") + .arg(key) + .query_async(&mut conn) + .await + .map_err(|e: RedisError| { + error!(key_type, error = %e, "failed to inspect Redis value size"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + if let Some(observed) = memory_usage { + if observed > options.budget.max_cell_bytes { + return Err(DataError::ResultLimitExceeded { + entity: key.to_string(), + limit_kind: "redis_value_bytes", + limit: options.budget.max_cell_bytes, + observed, + }); + } + } + + let offset = options.offset.unwrap_or(0); + let aggregate_limit = redis_aggregate_limit(options); + let start = isize::try_from(offset).unwrap_or(isize::MAX); + let end = start.saturating_add(isize::try_from(aggregate_limit).unwrap_or(isize::MAX)); + let mut truncated = false; + // Get value based on type let value = match key_type.as_str() { "string" => { - let v: String = conn.get(key).await.map_err(|e: RedisError| { - error!("Failed to get string value: {}", e); - DataError::QueryFailed(format!("Failed to get string value: {}", e)) - })?; - serde_json::Value::String(v) + let max_string_bytes = options.budget.max_cell_bytes.saturating_sub(2).max(1); + let end = i64::try_from(max_string_bytes.saturating_sub(1)).unwrap_or(i64::MAX); + let v: Vec = redis::cmd("GETRANGE") + .arg(key) + .arg(0) + .arg(end) + .query_async(&mut conn) + .await + .map_err(|e: RedisError| { + error!("Failed to get string value: {}", e); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + serde_json::Value::String(String::from_utf8_lossy(&v).into_owned()) } "list" => { - let v: Vec = conn.lrange(key, 0, -1).await.map_err(|e: RedisError| { - error!("Failed to get list value: {}", e); - DataError::QueryFailed(format!("Failed to get list value: {}", e)) - })?; + let mut v: Vec = + conn.lrange(key, start, end) + .await + .map_err(|e: RedisError| { + error!("Failed to get list value: {}", e); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + truncated = trim_extra(&mut v, aggregate_limit); serde_json::json!(v) } "set" => { - let v: Vec = conn.smembers(key).await.map_err(|e: RedisError| { - error!("Failed to get set value: {}", e); - DataError::QueryFailed(format!("Failed to get set value: {}", e)) - })?; + let (_cursor, mut v): (u64, Vec) = redis::cmd("SSCAN") + .arg(key) + .arg(0) + .arg("COUNT") + .arg(aggregate_limit.saturating_add(1)) + .query_async(&mut conn) + .await + .map_err(|e: RedisError| { + error!(error = %e, "failed to page Redis set"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + truncated = trim_extra(&mut v, aggregate_limit); serde_json::json!(v) } "zset" => { - let v: Vec<(String, f64)> = - conn.zrange_withscores(key, 0, -1) - .await - .map_err(|e: RedisError| { - error!("Failed to get sorted set value: {}", e); - DataError::QueryFailed(format!("Failed to get sorted set value: {}", e)) - })?; + let mut v: Vec<(String, f64)> = conn + .zrange_withscores(key, start, end) + .await + .map_err(|e: RedisError| { + error!("Failed to get sorted set value: {}", e); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + truncated = trim_extra(&mut v, aggregate_limit); serde_json::json!(v .into_iter() .map(|(member, score)| { @@ -278,12 +366,22 @@ impl RedisSource { .collect::>()) } "hash" => { - let v: HashMap = - conn.hgetall(key).await.map_err(|e: RedisError| { - error!("Failed to get hash value: {}", e); - DataError::QueryFailed(format!("Failed to get hash value: {}", e)) + let (_cursor, mut values): (u64, Vec<(String, String)>) = redis::cmd("HSCAN") + .arg(key) + .arg(0) + .arg("COUNT") + .arg(aggregate_limit.saturating_add(1)) + .query_async(&mut conn) + .await + .map_err(|e: RedisError| { + error!(error = %e, "failed to page Redis hash"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } })?; - serde_json::json!(v) + truncated = trim_extra(&mut values, aggregate_limit); + serde_json::json!(values.into_iter().collect::>()) } "stream" => { // For streams, just indicate it's a stream - full stream reading is complex @@ -301,7 +399,7 @@ impl RedisSource { row.insert("ttl".to_string(), serde_json::Value::Number(ttl.into())); row.insert("value".to_string(), value); - Ok(row) + Ok((row, truncated)) } /// Get information about a specific key @@ -383,7 +481,9 @@ impl RedisSource { row_count: Some(1), size_bytes: None, schema: Some(schema), - metadata: Some(serde_json::to_value(metadata_map).unwrap()), + metadata: Some(serde_json::Value::Object( + metadata_map.into_iter().collect(), + )), }) } } @@ -578,7 +678,7 @@ impl Queryable for RedisSource { container_path: &ContainerPath, entity_name: &str, _filters: Option, - _options: QueryOptions, + options: QueryOptions, ) -> Result { let start = std::time::Instant::now(); @@ -604,7 +704,7 @@ impl Queryable for RedisSource { } // Get the key value - let row = self.get_key_value(db_num, entity_name).await?; + let (row, value_truncated) = self.get_key_value(db_num, entity_name, &options).await?; let execution_ms = start.elapsed().as_millis() as u64; // Define schema for the result @@ -639,7 +739,21 @@ impl Queryable for RedisSource { primary_key: Some(vec!["key".to_string()]), }; - Ok(QueryResult::new(schema, vec![row], execution_ms)) + let mut bounded = BoundedRows::new(options.budget); + bounded.push(entity_name, row)?; + let (rows, budget_truncated) = bounded.into_parts(); + Ok(QueryResult { + schema, + stats: QueryStats { + row_count: rows.len(), + total_rows: Some(1), + execution_ms, + has_more: false, + next_cursor: None, + truncated: value_truncated || budget_truncated, + }, + rows, + }) } async fn count( @@ -720,6 +834,7 @@ impl Queryable for RedisSource { #[cfg(test)] mod tests { use super::*; + use temps_query::QueryBudget; #[test] fn test_source_type() { @@ -739,4 +854,21 @@ mod tests { assert!((0..=15).contains(&15)); assert!(!(0..=15).contains(&16)); } + + #[test] + fn redis_aggregate_pages_are_capped_by_structural_budget() { + let options = QueryOptions { + limit: Some(10_000), + budget: QueryBudget { + max_value_elements_per_row: 40, + ..QueryBudget::default() + }, + ..QueryOptions::default() + }; + assert_eq!(redis_aggregate_limit(&options), 10); + + let mut values = vec!["value"; 11]; + assert!(trim_extra(&mut values, redis_aggregate_limit(&options))); + assert_eq!(values.len(), 10); + } } diff --git a/crates/temps-query/src/budget.rs b/crates/temps-query/src/budget.rs new file mode 100644 index 000000000..ad3085e4b --- /dev/null +++ b/crates/temps-query/src/budget.rs @@ -0,0 +1,253 @@ +use crate::{DataError, DataRow, QueryBudget, Result}; + +/// Incremental result collector shared by every query backend. +/// +/// Backends must feed rows into this collector as their driver yields them; +/// collecting driver rows into an intermediate `Vec` defeats the memory bound. +pub struct BoundedRows { + budget: QueryBudget, + rows: Vec, + used_bytes: usize, + used_cells: usize, + truncated: bool, +} + +impl BoundedRows { + pub fn new(budget: QueryBudget) -> Self { + Self { + budget, + rows: Vec::new(), + used_bytes: 0, + used_cells: 0, + truncated: false, + } + } + + /// Add one decoded row. Returns `false` when the page-level byte/cell + /// budget is full and the backend should stop reading its cursor. + /// Oversized individual rows/cells are rejected, including the first row. + pub fn push(&mut self, entity: &str, row: DataRow) -> Result { + if row.len() > self.budget.max_cells_per_row { + return Err(limit_error( + entity, + "cells_per_row", + self.budget.max_cells_per_row, + row.len(), + )); + } + + let mut row_bytes = 2usize; + let mut row_elements = 0usize; + for (name, value) in &row { + let stats = value_stats(value, 0, self.budget.max_nesting_depth).ok_or_else(|| { + limit_error( + entity, + "nesting_depth", + self.budget.max_nesting_depth, + self.budget.max_nesting_depth.saturating_add(1), + ) + })?; + if stats.bytes > self.budget.max_cell_bytes { + return Err(limit_error( + entity, + "cell_bytes", + self.budget.max_cell_bytes, + stats.bytes, + )); + } + row_elements = row_elements.saturating_add(stats.elements); + row_bytes = row_bytes + .saturating_add(json_string_bytes(name)) + .saturating_add(stats.bytes) + .saturating_add(2); + } + if row_elements > self.budget.max_value_elements_per_row { + return Err(limit_error( + entity, + "value_elements_per_row", + self.budget.max_value_elements_per_row, + row_elements, + )); + } + + let next_bytes = self.used_bytes.saturating_add(row_bytes); + let next_cells = self.used_cells.saturating_add(row.len()); + if next_bytes > self.budget.max_bytes || next_cells > self.budget.max_cells { + if self.rows.is_empty() { + let (kind, limit, observed) = if next_bytes > self.budget.max_bytes { + ("response_bytes", self.budget.max_bytes, next_bytes) + } else { + ("response_cells", self.budget.max_cells, next_cells) + }; + return Err(limit_error(entity, kind, limit, observed)); + } + self.truncated = true; + return Ok(false); + } + + self.used_bytes = next_bytes; + self.used_cells = next_cells; + self.rows.push(row); + Ok(true) + } + + pub fn into_parts(self) -> (Vec, bool) { + (self.rows, self.truncated) + } +} + +fn limit_error(entity: &str, limit_kind: &'static str, limit: usize, observed: usize) -> DataError { + DataError::ResultLimitExceeded { + entity: entity.to_string(), + limit_kind, + limit, + observed, + } +} + +#[derive(Clone, Copy)] +struct ValueStats { + bytes: usize, + elements: usize, +} + +fn value_stats(value: &serde_json::Value, depth: usize, max_depth: usize) -> Option { + if depth > max_depth { + return None; + } + let stats = match value { + serde_json::Value::Null => ValueStats { + bytes: 4, + elements: 1, + }, + serde_json::Value::Bool(_) => ValueStats { + bytes: 5, + elements: 1, + }, + serde_json::Value::Number(number) => ValueStats { + bytes: number.to_string().len(), + elements: 1, + }, + serde_json::Value::String(value) => ValueStats { + bytes: json_string_bytes(value), + elements: 1, + }, + serde_json::Value::Array(values) => { + let mut bytes = 2usize; + let mut elements = values.len(); + for value in values { + let child = value_stats(value, depth + 1, max_depth)?; + bytes = bytes.saturating_add(child.bytes).saturating_add(1); + elements = elements.saturating_add(child.elements); + } + ValueStats { bytes, elements } + } + serde_json::Value::Object(values) => { + let mut bytes = 2usize; + let mut elements = values.len(); + for (name, value) in values { + let child = value_stats(value, depth + 1, max_depth)?; + bytes = bytes + .saturating_add(json_string_bytes(name)) + .saturating_add(child.bytes) + .saturating_add(2); + elements = elements.saturating_add(child.elements); + } + ValueStats { bytes, elements } + } + }; + Some(stats) +} + +/// Exact UTF-8 byte length of serde_json's compact string representation, +/// without allocating the encoded string. This accounts for control-character, +/// quote and backslash expansion so the response budget cannot be bypassed by +/// strings whose encoded form is much larger than their source bytes. +fn json_string_bytes(value: &str) -> usize { + value.chars().fold(2usize, |bytes, character| { + let encoded = match character { + '"' | '\\' | '\u{0008}' | '\u{000c}' | '\n' | '\r' | '\t' => 2, + '\u{0000}'..='\u{001f}' => 6, + _ => character.len_utf8(), + }; + bytes.saturating_add(encoded) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn budget(max_bytes: usize) -> QueryBudget { + QueryBudget { + max_bytes, + max_cells: 16, + max_cells_per_row: 8, + max_cell_bytes: max_bytes, + max_value_elements_per_row: 16, + max_nesting_depth: 8, + } + } + + fn row(value: serde_json::Value) -> DataRow { + [("value".to_string(), value)].into_iter().collect() + } + + #[test] + fn oversized_first_row_is_rejected() { + let mut rows = BoundedRows::new(budget(64)); + let error = rows + .push("users", row(serde_json::json!("x".repeat(1_024)))) + .expect_err("oversized first row must never be emitted"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + entity, + limit_kind: "cell_bytes", + .. + } if entity == "users" + )); + } + + #[test] + fn later_row_over_page_budget_truncates_without_adding_it() { + let mut rows = BoundedRows::new(budget(80)); + assert!(rows.push("users", row(serde_json::json!("small"))).unwrap()); + assert!(!rows + .push("users", row(serde_json::json!("y".repeat(70)))) + .unwrap()); + let (rows, truncated) = rows.into_parts(); + assert_eq!(rows.len(), 1); + assert!(truncated); + } + + #[test] + fn aggregate_element_limit_is_enforced() { + let mut rows = BoundedRows::new(budget(4_096)); + let error = rows + .push("cache", row(serde_json::json!(vec![1; 32]))) + .expect_err("aggregate element count must be bounded"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "value_elements_per_row", + .. + } + )); + } + + #[test] + fn escaped_json_bytes_count_toward_the_budget() { + let mut rows = BoundedRows::new(budget(40)); + let error = rows + .push("events", row(serde_json::json!("\u{0000}".repeat(10)))) + .expect_err("encoded control characters must count as six bytes each"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "cell_bytes", + .. + } + )); + } +} diff --git a/crates/temps-query/src/error.rs b/crates/temps-query/src/error.rs index c0661de42..3a244f754 100644 --- a/crates/temps-query/src/error.rs +++ b/crates/temps-query/src/error.rs @@ -43,6 +43,27 @@ pub enum DataError { #[error("Serialization error: {0}")] SerializationError(String), + /// A backend result exceeded a server-defined memory/shape budget. + #[error( + "Query result for entity '{entity}' exceeded {limit_kind} limit {limit} \ + (observed at least {observed})" + )] + ResultLimitExceeded { + entity: String, + limit_kind: &'static str, + limit: usize, + observed: usize, + }, + + /// A backend query failed. The source error is logged separately and is + /// intentionally absent here because database diagnostics may echo filter + /// literals or raw SQL. + #[error("{backend} query failed for entity '{entity}'")] + BackendQueryFailed { + backend: &'static str, + entity: String, + }, + /// Permission denied #[error("Permission denied: {0}")] PermissionDenied(String), diff --git a/crates/temps-query/src/lib.rs b/crates/temps-query/src/lib.rs index 2c6f24def..0780a5268 100644 --- a/crates/temps-query/src/lib.rs +++ b/crates/temps-query/src/lib.rs @@ -64,12 +64,14 @@ //! - `temps-query-redis` - Redis implementation (future) //! - `temps-query-mongodb` - MongoDB implementation (future) +mod budget; pub mod error; pub mod registry; pub mod traits; pub mod types; // Re-export commonly used items +pub use budget::BoundedRows; pub use error::{DataError, Result}; pub use registry::{ConnectionConfig, DataSourceFactory, QueryRegistry}; pub use traits::{ @@ -79,5 +81,5 @@ pub use traits::{ pub use types::{ Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataRow, DatabaseInfo, DatasetSchema, EntityCountHint, EntityInfo, EntityRef, FieldDef, FieldType, - NamespaceInfo, NamespaceRef, QueryOptions, QueryResult, QueryStats, + NamespaceInfo, NamespaceRef, QueryBudget, QueryOptions, QueryResult, QueryStats, }; diff --git a/crates/temps-query/src/types.rs b/crates/temps-query/src/types.rs index 48eaf32a5..50a1d1389 100644 --- a/crates/temps-query/src/types.rs +++ b/crates/temps-query/src/types.rs @@ -167,6 +167,33 @@ pub struct QueryOptions { pub timeout_ms: Option, /// Include null values pub include_nulls: bool, + /// Server-controlled memory and structural budget. Callers must not expose + /// these values as unconstrained request parameters. + #[serde(default)] + pub budget: QueryBudget, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct QueryBudget { + pub max_bytes: usize, + pub max_cells: usize, + pub max_cells_per_row: usize, + pub max_cell_bytes: usize, + pub max_value_elements_per_row: usize, + pub max_nesting_depth: usize, +} + +impl Default for QueryBudget { + fn default() -> Self { + Self { + max_bytes: 8 * 1024 * 1024, + max_cells: 16_384, + max_cells_per_row: 256, + max_cell_bytes: 1024 * 1024, + max_value_elements_per_row: 10_000, + max_nesting_depth: 64, + } + } } impl Default for QueryOptions { @@ -179,6 +206,7 @@ impl Default for QueryOptions { sort_order: Some("asc".to_string()), timeout_ms: Some(30000), include_nulls: true, + budget: QueryBudget::default(), } } } @@ -199,6 +227,8 @@ pub struct QueryStats { pub has_more: bool, /// Next cursor for pagination (if applicable) pub next_cursor: Option, + /// True when a backend stopped reading because the result budget filled. + pub truncated: bool, } /// Result of executing a query @@ -227,6 +257,7 @@ impl QueryResult { execution_ms, has_more, next_cursor: None, + truncated: false, }, } } @@ -251,6 +282,7 @@ impl QueryResult { execution_ms, has_more, next_cursor, + truncated: false, }, } } From 5ca86340cd5a99b71a34cf08822e3efa45f832d9 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 13:51:22 +0200 Subject: [PATCH 04/10] fix(security): reject oversized rows before encoding --- .../src/handlers/query_handlers.rs | 59 +++- crates/temps-providers/src/mariadb_query.rs | 247 ++++++++++++---- crates/temps-query-mongodb/src/lib.rs | 118 +++++++- crates/temps-query-postgres/src/lib.rs | 272 ++++++++++++++++-- 4 files changed, 600 insertions(+), 96 deletions(-) diff --git a/crates/temps-providers/src/handlers/query_handlers.rs b/crates/temps-providers/src/handlers/query_handlers.rs index feaa50993..c179e43e5 100644 --- a/crates/temps-providers/src/handlers/query_handlers.rs +++ b/crates/temps-providers/src/handlers/query_handlers.rs @@ -340,13 +340,13 @@ fn query_error_problem( StatusCode::PAYLOAD_TOO_LARGE, "Query Result Too Large", "result_limit_exceeded", - Some(error.to_string()), + Some("The query result exceeds the configured response limits".to_string()), ), temps_query::DataError::InvalidQuery(_) => ( StatusCode::BAD_REQUEST, "Invalid Query", "invalid_query", - Some(error.to_string()), + Some("The query parameters are invalid for this data source".to_string()), ), temps_query::DataError::QueryFailed(_) | temps_query::DataError::BackendQueryFailed { .. } => ( @@ -520,13 +520,12 @@ fn parse_row_filter( match serde_json::from_str::(raw) { Ok(value) => Ok(Some(value)), - Err(e) => { + Err(_error) => { let mut problem = temps_core::problemdetails::new(StatusCode::BAD_REQUEST) .with_title("Invalid Filter") - .with_detail(format!( - "The `filter` parameter must be JSON matching this service's filter schema, \ - but it failed to parse: {e}. Received: {raw}" - )); + .with_detail( + "The `filter` parameter must be valid JSON matching this service's filter schema", + ); if let Some(schema) = filter_schema { problem = problem.with_value("expected_filter_schema", schema.clone()); } @@ -1827,6 +1826,40 @@ mod tests { assert!(!serialized.contains("tok_live_secret")); } + #[test] + fn invalid_query_problem_does_not_echo_submitted_values() { + let problem = query_error_problem( + temps_query::DataError::InvalidQuery( + "invalid email alice@example.com with tok_live_secret".to_string(), + ), + 7, + "users", + 100, + ); + let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); + assert!(!serialized.contains("alice@example.com")); + assert!(!serialized.contains("tok_live_secret")); + assert!(serialized.contains("query parameters are invalid")); + } + + #[test] + fn result_limit_problem_does_not_echo_sensitive_entity_name() { + let problem = query_error_problem( + temps_query::DataError::ResultLimitExceeded { + entity: "session:tok_live_secret:alice@example.com".to_string(), + limit_kind: "wire_cell_bytes", + limit: 64, + observed: 1_024, + }, + 7, + "session:tok_live_secret:alice@example.com", + 1, + ); + let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); + assert!(!serialized.contains("alice@example.com")); + assert!(!serialized.contains("tok_live_secret")); + } + #[test] fn entity_names_are_user_data_for_key_value_and_object_stores() { // `list_entities` is allowlisted for the agent as schema navigation. @@ -1894,6 +1927,18 @@ mod tests { assert_eq!(err.status_code, StatusCode::BAD_REQUEST); } + #[test] + fn malformed_filter_problem_does_not_echo_email_or_token() { + let err = parse_row_filter( + Some(r#"{"email":"alice@example.com","token":"tok_live_secret""#), + None, + ) + .expect_err("malformed secret-bearing JSON must be rejected"); + let serialized = serde_json::to_string(&err.body).expect("problem body serializes"); + assert!(!serialized.contains("alice@example.com")); + assert!(!serialized.contains("tok_live_secret")); + } + #[test] fn parse_row_filter_echoes_expected_schema_so_callers_can_self_correct() { let schema = serde_json::json!({ diff --git a/crates/temps-providers/src/mariadb_query.rs b/crates/temps-providers/src/mariadb_query.rs index 3512a082f..f4323725f 100644 --- a/crates/temps-providers/src/mariadb_query.rs +++ b/crates/temps-providers/src/mariadb_query.rs @@ -6,8 +6,8 @@ use std::collections::HashMap; use temps_query::{ BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, DataRow, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, - FieldType, Introspect, QueryOptions, QueryResult, QuerySchemaProvider, QueryStats, Queryable, - Result, + FieldType, Introspect, QueryBudget, QueryOptions, QueryResult, QuerySchemaProvider, QueryStats, + Queryable, Result, }; use tracing::{debug, error, warn}; @@ -132,6 +132,45 @@ impl MariaDbSource { _ => FieldType::String, } } + + async fn query_columns( + &self, + database_name: &str, + entity_name: &str, + ) -> Result> { + let rows = sqlx::query( + "SELECT COLUMN_NAME, DATA_TYPE FROM information_schema.COLUMNS \ + WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? ORDER BY ORDINAL_POSITION", + ) + .bind(database_name) + .bind(entity_name) + .fetch_all(&self.pool) + .await + .map_err(|_error| { + DataError::SchemaError(format!( + "Failed to inspect MariaDB query columns for '{}.{}'", + database_name, entity_name + )) + })?; + rows.into_iter() + .map(|row| { + Ok(MariaQueryColumn { + name: row.try_get("COLUMN_NAME").map_err(|_error| { + DataError::SchemaError(format!( + "Failed to decode a column name for '{}.{}'", + database_name, entity_name + )) + })?, + data_type: row.try_get("DATA_TYPE").map_err(|_error| { + DataError::SchemaError(format!( + "Failed to decode a column type for '{}.{}'", + database_name, entity_name + )) + })?, + }) + }) + .collect() + } } #[async_trait] @@ -497,11 +536,11 @@ impl Queryable for MariaDbSource { let database_name = database_from_path(container_path, &self.database_name)?; validate_identifier("table", entity_name)?; let schema = self.get_schema(container_path, entity_name).await?; + let columns = self.query_columns(database_name, entity_name).await?; let start = std::time::Instant::now(); - let projection = mariadb_json_projection(&schema); let mut sql = format!( - "SELECT {projection} AS __temps_payload FROM {}.{}", + "SELECT * FROM {}.{}", quote_identifier(database_name), quote_identifier(entity_name) ); @@ -530,7 +569,7 @@ impl Queryable for MariaDbSource { let limit = options.limit.unwrap_or(100); let offset = options.offset.unwrap_or(0); sql.push_str(" LIMIT ? OFFSET ?"); - let sql = with_wire_row_budget(&sql); + let sql = with_wire_row_budget(&sql, &columns, options.budget)?; debug!( entity = entity_name, @@ -557,7 +596,6 @@ impl Queryable for MariaDbSource { apply_statement_timeout(&mut conn, timeout_ms, database_name).await; let mut stream = sqlx::query(&sql) - .bind(options.budget.max_bytes as u64) .bind(limit as i64) .bind(offset as i64) .fetch(&mut *conn); @@ -579,6 +617,12 @@ impl Queryable for MariaDbSource { entity: entity_name.to_string(), } })?; + let observed_cell = row + .try_get::("__temps_max_cell") + .map_err(|_error| DataError::BackendQueryFailed { + backend: "MariaDB", + entity: entity_name.to_string(), + })?; let payload = row .try_get::, _>("__temps_row") .map_err(|_error| { @@ -591,11 +635,28 @@ impl Queryable for MariaDbSource { entity: entity_name.to_string(), } })? - .ok_or_else(|| DataError::ResultLimitExceeded { - entity: entity_name.to_string(), - limit_kind: "wire_row_bytes", - limit: options.budget.max_bytes, - observed: usize::try_from(observed).unwrap_or(usize::MAX), + .ok_or_else(|| { + let observed_cell = usize::try_from(observed_cell).unwrap_or(usize::MAX); + let (limit_kind, limit, observed) = + if observed_cell > options.budget.max_cell_bytes { + ( + "wire_cell_bytes", + options.budget.max_cell_bytes, + observed_cell, + ) + } else { + ( + "wire_row_bytes", + options.budget.max_bytes, + usize::try_from(observed).unwrap_or(usize::MAX), + ) + }; + DataError::ResultLimitExceeded { + entity: entity_name.to_string(), + limit_kind, + limit, + observed, + } })?; let data_row = serde_json::from_str::(&payload).map_err(|_error| { error!( @@ -805,23 +866,29 @@ fn quote_identifier(value: &str) -> String { format!("`{}`", value.replace('`', "``")) } -/// Build one JSON value per source row so MariaDB can measure it before the -/// driver sees any user-controlled cell contents. Binary values retain the -/// previous API representation (base64 text). -fn mariadb_json_projection(schema: &DatasetSchema) -> String { - let entries = schema.fields.iter().flat_map(|field| { - let key_hex = field +#[derive(Clone, Debug)] +struct MariaQueryColumn { + name: String, + data_type: String, +} + +fn mariadb_json_projection(columns: &[MariaQueryColumn]) -> String { + let entries = columns.iter().flat_map(|column| { + let key_hex = column .name .as_bytes() .iter() .map(|byte| format!("{byte:02x}")) .collect::(); let key = format!("CONVERT(X'{key_hex}' USING utf8mb4)"); - let column = quote_identifier(&field.name); - let value = if field.field_type == FieldType::Bytes { - format!("TO_BASE64({column})") + let column_ref = format!("__temps_source.{}", quote_identifier(&column.name)); + let value = if matches!( + column.data_type.as_str(), + "binary" | "varbinary" | "tinyblob" | "blob" | "mediumblob" | "longblob" + ) { + format!("TO_BASE64({column_ref})") } else { - column + column_ref }; [key, value] }); @@ -829,16 +896,62 @@ fn mariadb_json_projection(schema: &DatasetSchema) -> String { format!("JSON_OBJECT({})", entries.collect::>().join(", ")) } -/// Apply the per-row response ceiling in the same statement that reads the -/// row. Oversized JSON is replaced by NULL plus its length, so it never crosses -/// the database wire into the control plane and there is no check/read race. -fn with_wire_row_budget(sql: &str) -> String { - format!( - "SELECT CASE WHEN OCTET_LENGTH(__temps_payload) <= ? \ - THEN __temps_payload ELSE NULL END AS __temps_row, \ - OCTET_LENGTH(__temps_payload) AS __temps_size \ - FROM ({sql}) AS __temps_bounded" - ) +fn mariadb_column_admission(column: &MariaQueryColumn) -> Result { + let value = format!("__temps_source.{}", quote_identifier(&column.name)); + let estimate = match column.data_type.as_str() { + "binary" | "varbinary" | "tinyblob" | "blob" | "mediumblob" | "longblob" => { + format!("OCTET_LENGTH({value}) * 5 + 8") + } + "char" | "varchar" | "tinytext" | "text" | "mediumtext" | "longtext" | "enum" | "set" => { + format!("OCTET_LENGTH({value}) * 6 + 8") + } + "json" => format!("JSON_STORAGE_SIZE({value}) * 6 + 8"), + "bool" | "boolean" | "tinyint" | "smallint" | "mediumint" | "int" | "integer" + | "bigint" | "float" | "double" | "real" | "decimal" | "numeric" | "date" | "datetime" + | "timestamp" | "time" | "year" => "128".to_string(), + unsupported => { + return Err(DataError::OperationNotSupported(format!( + "MariaDB column '{}' uses unsupported type '{}'", + column.name, unsupported + ))) + } + }; + Ok(format!("COALESCE({estimate}, 4)")) +} + +/// Admission expressions execute before the JSON constructor. Rejected rows +/// return only conservative byte metadata, never the original values. +fn with_wire_row_budget( + sql: &str, + columns: &[MariaQueryColumn], + budget: QueryBudget, +) -> Result { + let estimates = columns + .iter() + .map(mariadb_column_admission) + .collect::>>()?; + let max_cell = if estimates.is_empty() { + "0".to_string() + } else { + format!("GREATEST({})", estimates.join(", ")) + }; + let key_overhead = columns.iter().fold(2usize, |total, column| { + total.saturating_add(column.name.len().saturating_mul(6).saturating_add(4)) + }); + let row_size = if estimates.is_empty() { + key_overhead.to_string() + } else { + format!("{key_overhead} + {}", estimates.join(" + ")) + }; + let projection = mariadb_json_projection(columns); + + Ok(format!( + "SELECT CASE WHEN {max_cell} <= {} AND {row_size} <= {} \ + THEN {projection} ELSE NULL END AS __temps_row, \ + {row_size} AS __temps_size, {max_cell} AS __temps_max_cell \ + FROM ({sql}) AS __temps_source", + budget.max_cell_bytes, budget.max_bytes + )) } fn validate_identifier(label: &str, value: &str) -> Result<()> { @@ -1205,36 +1318,58 @@ mod tests { #[test] fn generated_query_guards_encoded_row_before_wire_transfer() { - let schema = DatasetSchema { - fields: vec![ - FieldDef { - name: "display_name".to_string(), - field_type: FieldType::String, - nullable: false, - description: None, - }, - FieldDef { - name: "avatar".to_string(), - field_type: FieldType::Bytes, - nullable: true, - description: None, - }, - ], - partitions: None, - primary_key: None, + let columns = vec![ + MariaQueryColumn { + name: "display_name".to_string(), + data_type: "longtext".to_string(), + }, + MariaQueryColumn { + name: "avatar".to_string(), + data_type: "longblob".to_string(), + }, + MariaQueryColumn { + name: "settings".to_string(), + data_type: "json".to_string(), + }, + ]; + let budget = QueryBudget { + max_bytes: 262_144, + max_cell_bytes: 65_536, + ..QueryBudget::default() }; - let projection = mariadb_json_projection(&schema); - let sql = with_wire_row_budget(&format!( - "SELECT {projection} AS __temps_payload FROM `app`.`users` LIMIT ? OFFSET ?" - )); - - assert!(sql.contains("OCTET_LENGTH(__temps_payload) <= ?")); - assert!(sql.contains("THEN __temps_payload ELSE NULL")); + let sql = with_wire_row_budget( + "SELECT * FROM `app`.`users` LIMIT ? OFFSET ?", + &columns, + budget, + ) + .expect("supported schema should build a bounded query"); + + assert!(sql.contains("OCTET_LENGTH(__temps_source.`display_name`) * 6")); + assert!(sql.contains("OCTET_LENGTH(__temps_source.`avatar`) * 5")); + assert!(sql.contains("JSON_STORAGE_SIZE(__temps_source.`settings`) * 6")); + assert!(sql.contains("<= 65536")); + assert!(sql.contains("<= 262144")); + assert!(sql.contains("THEN JSON_OBJECT(")); assert!(sql.contains("JSON_OBJECT(")); - assert!(sql.contains("TO_BASE64(`avatar`)")); + assert_eq!(sql.matches("JSON_OBJECT(").count(), 1); + assert!(sql.contains("TO_BASE64(__temps_source.`avatar`)")); assert!(sql.contains("FROM `app`.`users` LIMIT ? OFFSET ?")); } + #[test] + fn unsupported_mariadb_types_are_rejected_before_query_execution() { + let error = with_wire_row_budget( + "SELECT * FROM `app`.`places`", + &[MariaQueryColumn { + name: "shape".to_string(), + data_type: "geometry".to_string(), + }], + QueryBudget::default(), + ) + .expect_err("unknown encodings cannot be safely admitted"); + assert!(matches!(error, DataError::OperationNotSupported(_))); + } + fn assert_where_rejected(clause: &str) { assert!( validate_where_clause(clause).is_err(), diff --git a/crates/temps-query-mongodb/src/lib.rs b/crates/temps-query-mongodb/src/lib.rs index 4ce836804..7827e9e27 100644 --- a/crates/temps-query-mongodb/src/lib.rs +++ b/crates/temps-query-mongodb/src/lib.rs @@ -28,7 +28,7 @@ use async_trait::async_trait; use mongodb::{ - bson::{doc, Document}, + bson::{doc, Bson, Document}, options::ClientOptions, Client, }; @@ -101,6 +101,43 @@ const ALLOWED_FILTER_OPERATORS: [&str; 15] = [ /// a shallow check. Honest filters are one or two levels. const MAX_FILTER_DEPTH: usize = 8; +fn bounded_document_pipeline( + filter: Document, + sort: Document, + skip: u64, + limit: i64, + max_document_bytes: usize, +) -> Vec { + let byte_limit = i64::try_from(max_document_bytes).unwrap_or(i64::MAX); + vec![ + doc! { "$match": filter }, + doc! { "$sort": sort }, + doc! { "$skip": i64::try_from(skip).unwrap_or(i64::MAX) }, + doc! { "$limit": limit }, + doc! { + "$replaceWith": { + "$let": { + "vars": { "temps_size": { "$bsonSize": "$$ROOT" } }, + "in": { + "$cond": [ + { "$lte": ["$$temps_size", byte_limit] }, + { + "__temps_admitted": true, + "__temps_size": "$$temps_size", + "__temps_doc": "$$ROOT" + }, + { + "__temps_admitted": false, + "__temps_size": "$$temps_size" + } + ] + } + } + } + }, + ] +} + /// Server-side ceiling for the standalone `count`, which carries no /// `QueryOptions` and therefore no caller deadline. Matches the SQL backends. const COUNT_TIMEOUT_MS: u64 = 10_000; @@ -644,7 +681,9 @@ impl temps_query::Queryable for MongoDBSource { Some("desc") | Some("DESC") => -1, _ => 1, }; - doc! { sort_by: sort_order } + [(sort_by.clone(), Bson::Int32(sort_order))] + .into_iter() + .collect() } else { doc! { "_id": 1 } // Default sort by _id ascending }; @@ -665,12 +704,21 @@ impl temps_query::Queryable for MongoDBSource { // every skipped document regardless of `limit`. let max_time = std::time::Duration::from_millis(options.timeout_ms.unwrap_or(30_000)); - // Execute query + // `$bsonSize` is evaluated by mongod before the conditional projection. + // Using the cell ceiling as the document ceiling is deliberately + // conservative: no individual nested value can cross the wire above + // the per-cell budget, even though MongoDB documents are dynamic. + let wire_document_limit = options.budget.max_bytes.min(options.budget.max_cell_bytes); + let pipeline = bounded_document_pipeline( + filter_doc.clone(), + sort_doc, + skip, + limit, + wire_document_limit, + ); let mut cursor = collection - .find(filter_doc.clone()) - .sort(sort_doc) - .limit(limit) - .skip(skip) + .aggregate(pipeline) + .batch_size(1) .max_time(max_time) .await .map_err(|_error| { @@ -692,7 +740,7 @@ impl temps_query::Queryable for MongoDBSource { entity: entity_name.to_string(), } })? { - let doc = cursor.deserialize_current().map_err(|_error| { + let mut envelope = cursor.deserialize_current().map_err(|_error| { error!( entity = entity_name, limit, "MongoDB document decode failed" @@ -703,7 +751,37 @@ impl temps_query::Queryable for MongoDBSource { } })?; - // Convert Document to HashMap for DataRow + let observed = match envelope.remove("__temps_size") { + Some(Bson::Int32(value)) => usize::try_from(value).unwrap_or(usize::MAX), + Some(Bson::Int64(value)) => usize::try_from(value).unwrap_or(usize::MAX), + _ => { + return Err(DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + }) + } + }; + let admitted = envelope + .remove("__temps_admitted") + .and_then(|value| value.as_bool()) + .unwrap_or(false); + if !admitted { + return Err(DataError::ResultLimitExceeded { + entity: entity_name.to_string(), + limit_kind: "wire_document_bytes", + limit: wire_document_limit, + observed, + }); + } + let doc = envelope + .remove("__temps_doc") + .and_then(|value| value.as_document().cloned()) + .ok_or_else(|| DataError::BackendQueryFailed { + backend: "MongoDB", + entity: entity_name.to_string(), + })?; + + // Convert the admitted Document to HashMap for DataRow. let mut row_map = std::collections::HashMap::new(); for (key, value) in doc { // Convert BSON to serde_json::Value @@ -876,6 +954,28 @@ impl temps_query::Queryable for MongoDBSource { mod tests { use super::*; + #[test] + fn aggregation_sizes_before_conditionally_projecting_document() { + let pipeline = bounded_document_pipeline( + doc! { "status": "active" }, + doc! { "created_at": -1 }, + 25, + 10, + 65_536, + ); + let encoded = mongodb::bson::serialize_to_bson(&pipeline) + .expect("test pipeline should serialize") + .to_string(); + + assert!(encoded.contains("$bsonSize")); + assert!(encoded.contains("$cond")); + assert!(encoded.contains("__temps_admitted")); + assert!(encoded.contains("__temps_doc")); + assert!(encoded.contains("65536")); + assert_eq!(pipeline[2], doc! { "$skip": 25_i64 }); + assert_eq!(pipeline[3], doc! { "$limit": 10_i64 }); + } + #[test] fn test_source_type() { assert_eq!("mongodb", "mongodb"); diff --git a/crates/temps-query-postgres/src/lib.rs b/crates/temps-query-postgres/src/lib.rs index 7e30eeef7..d301643da 100644 --- a/crates/temps-query-postgres/src/lib.rs +++ b/crates/temps-query-postgres/src/lib.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use temps_query::{ BoundedRows, Capability, ContainerCapabilities, ContainerInfo, ContainerPath, ContainerType, DataError, DataSource, DatasetSchema, EntityCountHint, EntityInfo, FieldDef, FieldType, - Introspect, QueryOptions, QueryResult, QueryStats, Queryable, Result, + Introspect, QueryBudget, QueryOptions, QueryResult, QueryStats, Queryable, Result, }; use tokio_postgres::{types::ToSql, Client, NoTls}; use tokio_postgres_rustls::MakeRustlsConnect; @@ -33,19 +33,83 @@ fn escape_ident(name: &str) -> String { name.replace('"', "\"\"") } -/// Keep an individual PostgreSQL row bounded before it crosses the wire. -/// `CASE` returns only NULL plus the observed size for oversized rows, so the -/// client driver never allocates their text/blob/json contents. The projection -/// and size decision are part of the same statement, avoiding a TOCTOU race. -fn with_wire_row_budget(sql: &str, max_bytes: usize) -> String { - format!( - "SELECT CASE WHEN OCTET_LENGTH(__temps_payload) <= {max_bytes} \ - THEN __temps_payload::jsonb ELSE NULL END AS __temps_row, \ - OCTET_LENGTH(__temps_payload)::bigint AS __temps_size \ +#[derive(Clone, Debug)] +struct PgQueryColumn { + name: String, + data_type: String, +} + +fn pg_column_admission(column: &PgQueryColumn, row_budget: usize) -> Result { + let value = format!("__temps_source.\"{}\"", escape_ident(&column.name)); + let expression = match column.data_type.as_str() { + "character varying" | "character" | "text" => { + format!("COALESCE(OCTET_LENGTH({value})::bigint * 6 + 2, 4)") + } + "bytea" => format!("COALESCE(OCTET_LENGTH({value})::bigint * 2 + 8, 4)"), + "json" => format!("COALESCE(OCTET_LENGTH({value}::text)::bigint * 6 + 8, 4)"), + "jsonb" | "ARRAY" => format!( + "CASE WHEN {value} IS NULL THEN 4 WHEN PG_COLUMN_COMPRESSION({value}) IS NOT NULL \ + THEN {} ELSE PG_COLUMN_SIZE({value})::bigint * 8 + 64 END", + row_budget.saturating_add(1) + ), + "boolean" + | "smallint" + | "integer" + | "bigint" + | "real" + | "double precision" + | "numeric" + | "decimal" + | "date" + | "timestamp without time zone" + | "timestamp with time zone" + | "uuid" => { + format!("COALESCE(PG_COLUMN_SIZE({value})::bigint * 8 + 64, 4)") + } + unsupported => { + return Err(DataError::OperationNotSupported(format!( + "PostgreSQL column '{}' uses unsupported type '{}'", + column.name, unsupported + ))) + } + }; + Ok(expression) +} + +/// Admit rows from non-materializing per-column upper bounds. `TO_JSONB` is +/// located only in the admitted CASE arm, so rejected values are never encoded. +fn with_wire_row_budget( + sql: &str, + columns: &[PgQueryColumn], + budget: QueryBudget, +) -> Result { + let estimates = columns + .iter() + .map(|column| pg_column_admission(column, budget.max_bytes)) + .collect::>>()?; + let max_cell = if estimates.is_empty() { + "0".to_string() + } else { + format!("GREATEST({})", estimates.join(", ")) + }; + let key_overhead = columns.iter().fold(2usize, |total, column| { + total.saturating_add(column.name.len().saturating_mul(6).saturating_add(4)) + }); + let row_size = if estimates.is_empty() { + key_overhead.to_string() + } else { + format!("{key_overhead} + {}", estimates.join(" + ")) + }; + + Ok(format!( + "SELECT CASE WHEN __temps_max_cell <= {} AND __temps_row_size <= {} \ + THEN TO_JSONB(__temps_source) ELSE NULL END AS __temps_row, \ + __temps_row_size AS __temps_size, __temps_max_cell \ FROM ({sql}) AS __temps_source \ - CROSS JOIN LATERAL (SELECT TO_JSONB(__temps_source)::text AS __temps_payload) \ - AS __temps_encoded" - ) + CROSS JOIN LATERAL (SELECT {row_size}::bigint AS __temps_row_size, \ + {max_cell}::bigint AS __temps_max_cell) AS __temps_admission", + budget.max_cell_bytes, budget.max_bytes + )) } /// A certificate verifier that accepts all server certificates (including self-signed). @@ -1228,6 +1292,34 @@ impl PostgresSource { _ => FieldType::String, // Default fallback } } + + async fn query_columns( + &self, + schema_name: &str, + entity_name: &str, + ) -> Result> { + let rows = self + .client + .query( + "SELECT column_name, data_type FROM information_schema.columns \ + WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position", + &[&schema_name, &entity_name], + ) + .await + .map_err(|_error| { + DataError::SchemaError(format!( + "Failed to inspect PostgreSQL query columns for '{}.{}'", + schema_name, entity_name + )) + })?; + Ok(rows + .into_iter() + .map(|row| PgQueryColumn { + name: row.get(0), + data_type: row.get(1), + }) + .collect()) + } } #[async_trait] @@ -1819,6 +1911,7 @@ impl Queryable for PostgresSource { } let schema_name = &container_path.segments[1]; + let columns = self.query_columns(schema_name, entity_name).await?; let start = std::time::Instant::now(); @@ -1872,7 +1965,7 @@ impl Queryable for PostgresSource { let limit = options.limit.unwrap_or(100); let offset = options.offset.unwrap_or(0); sql.push_str(&format!(" LIMIT {} OFFSET {}", limit, offset)); - let sql = with_wire_row_budget(&sql, options.budget.max_bytes); + let sql = with_wire_row_budget(&sql, &columns, options.budget)?; debug!( entity = entity_name, @@ -1938,6 +2031,12 @@ impl Queryable for PostgresSource { entity: entity_name.to_string(), } })?; + let observed_cell = row + .try_get::<_, i64>("__temps_max_cell") + .map_err(|_error| DataError::BackendQueryFailed { + backend: "PostgreSQL", + entity: entity_name.to_string(), + })?; let payload = row .try_get::<_, Option>("__temps_row") .map_err(|_error| { @@ -1950,11 +2049,28 @@ impl Queryable for PostgresSource { entity: entity_name.to_string(), } })? - .ok_or_else(|| DataError::ResultLimitExceeded { - entity: entity_name.to_string(), - limit_kind: "wire_row_bytes", - limit: options.budget.max_bytes, - observed: usize::try_from(observed).unwrap_or(usize::MAX), + .ok_or_else(|| { + let observed_cell = usize::try_from(observed_cell).unwrap_or(usize::MAX); + let (limit_kind, limit, observed) = + if observed_cell > options.budget.max_cell_bytes { + ( + "wire_cell_bytes", + options.budget.max_cell_bytes, + observed_cell, + ) + } else { + ( + "wire_row_bytes", + options.budget.max_bytes, + usize::try_from(observed).unwrap_or(usize::MAX), + ) + }; + DataError::ResultLimitExceeded { + entity: entity_name.to_string(), + limit_kind, + limit, + observed, + } })?; let data_row = match payload { serde_json::Value::Object(values) => values.into_iter().collect(), @@ -2081,16 +2197,54 @@ impl Queryable for PostgresSource { #[cfg(test)] mod wire_budget_tests { - use super::with_wire_row_budget; + use super::{with_wire_row_budget, PgQueryColumn}; + use temps_query::QueryBudget; #[test] fn generated_query_guards_encoded_row_before_wire_transfer() { - let sql = with_wire_row_budget("SELECT * FROM public.users", 262_144); - assert!(sql.contains("OCTET_LENGTH(__temps_payload) <= 262144")); - assert!(sql.contains("THEN __temps_payload::jsonb ELSE NULL")); - assert!(sql.contains("TO_JSONB(__temps_source)::text")); + let columns = vec![ + PgQueryColumn { + name: "bio".to_string(), + data_type: "text".to_string(), + }, + PgQueryColumn { + name: "payload".to_string(), + data_type: "jsonb".to_string(), + }, + ]; + let budget = QueryBudget { + max_bytes: 262_144, + max_cell_bytes: 65_536, + ..QueryBudget::default() + }; + let sql = with_wire_row_budget("SELECT * FROM public.users", &columns, budget) + .expect("supported schema should build a bounded query"); + assert!(sql.contains("OCTET_LENGTH(__temps_source.\"bio\")::bigint * 6")); + assert!(sql.contains("PG_COLUMN_COMPRESSION(__temps_source.\"payload\")")); + assert!(sql.contains("__temps_max_cell <= 65536")); + assert!(sql.contains("__temps_row_size <= 262144")); + assert!(sql.contains("THEN TO_JSONB(__temps_source) ELSE NULL")); + assert!(!sql.contains("TO_JSONB(__temps_source)::text")); + assert_eq!(sql.matches("TO_JSONB(__temps_source)").count(), 1); assert!(sql.contains("SELECT * FROM public.users")); } + + #[test] + fn unknown_types_are_rejected_before_query_execution() { + let error = with_wire_row_budget( + "SELECT * FROM public.widgets", + &[PgQueryColumn { + name: "shape".to_string(), + data_type: "USER-DEFINED".to_string(), + }], + QueryBudget::default(), + ) + .expect_err("unknown output functions have no safe size upper bound"); + assert!(matches!( + error, + temps_query::DataError::OperationNotSupported(_) + )); + } } impl temps_query::QuerySchemaProvider for PostgresSource { @@ -2773,6 +2927,76 @@ mod tests { drop(source); } + #[tokio::test] + async fn oversized_first_row_is_rejected_by_real_postgres_query() { + let container = match GenericImage::new("postgres", "18-alpine") + .with_exposed_port(ContainerPort::Tcp(5432)) + .with_wait_for(WaitFor::message_on_stderr( + "database system is ready to accept connections", + )) + .with_env_var("POSTGRES_DB", "postgres") + .with_env_var("POSTGRES_USER", "postgres") + .with_env_var("POSTGRES_HOST_AUTH_METHOD", "trust") + .start() + .await + { + Ok(container) => container, + Err(error) => { + eprintln!("Docker unavailable; skipping PostgreSQL row-budget test: {error}"); + return; + } + }; + let host = container + .get_host() + .await + .expect("started PostgreSQL container must expose its host") + .to_string(); + let port = container + .get_host_port_ipv4(5432) + .await + .expect("started PostgreSQL container must expose port 5432"); + let source = match PostgresSource::connect(&host, port, "postgres", "", "postgres").await { + Ok(source) => source, + Err(error) => { + eprintln!("PostgreSQL container unavailable after startup; skipping: {error}"); + return; + } + }; + source + .client + .batch_execute( + "CREATE TABLE oversized_rows (id bigint, payload text); \ + INSERT INTO oversized_rows VALUES (1, repeat('x', 2097152));", + ) + .await + .expect("oversized row fixture should be created"); + let budget = QueryBudget { + max_bytes: 128 * 1024, + max_cell_bytes: 64 * 1024, + ..QueryBudget::default() + }; + let error = source + .query( + &ContainerPath::from_slice(&["postgres", "public"]), + "oversized_rows", + None, + QueryOptions { + limit: Some(1), + budget, + ..QueryOptions::default() + }, + ) + .await + .expect_err("the first oversized row must be rejected"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "wire_cell_bytes", + .. + } + )); + } + #[tokio::test] async fn test_quoted_function_identifier_rejected_by_real_query_paths() { let container = match GenericImage::new("postgres", "18-alpine") From 77d20572b0033ff1ccbbc4a6d3ade48e2521ff6e Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 14:11:12 +0200 Subject: [PATCH 05/10] fix(security): redact audited data identifiers --- crates/temps-providers/src/handlers/audit.rs | 65 ++++++++++++--- .../src/handlers/query_handlers.rs | 79 +++++++++++++------ crates/temps-query-postgres/src/lib.rs | 31 +++++++- 3 files changed, 141 insertions(+), 34 deletions(-) diff --git a/crates/temps-providers/src/handlers/audit.rs b/crates/temps-providers/src/handlers/audit.rs index 352c036d0..44ffa51f4 100644 --- a/crates/temps-providers/src/handlers/audit.rs +++ b/crates/temps-providers/src/handlers/audit.rs @@ -99,16 +99,35 @@ impl AuditOperation for AiDataAccessChangedAudit { /// a suspected prompt injection the operator's first question is "what did the /// model see?", and without this record there is nothing to answer it with. /// -/// Deliberately records the location and shape of the read (service, container -/// path, entity, row count) and never the values themselves — an audit log that -/// copied the rows would just be a second place the same secrets live. +/// Deliberately records only stable IDs and fixed categories. Service names, +/// container paths and entity names are omitted because key-value/object +/// backends routinely store emails, tokens and filenames in those identifiers. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AiBackendCategory { + Relational, + Document, + KeyValue, + ObjectStore, + Unknown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AiEntityCategory { + Table, + Collection, + Key, + Object, + Unknown, +} + #[derive(Debug, Clone, Serialize)] pub struct AiRowsReadAudit { pub context: AuditContext, pub service_id: i32, - pub service_name: String, - pub container_path: String, - pub entity: String, + pub backend_category: AiBackendCategory, + pub entity_category: AiEntityCategory, pub returned_rows: usize, pub truncated: bool, /// Bounded structural category only; literal filter values are never @@ -478,9 +497,8 @@ mod tests { user_agent: "test-agent".to_string(), }, service_id: 7, - service_name: "postgres".to_string(), - container_path: "app/public".to_string(), - entity: "users".to_string(), + backend_category: AiBackendCategory::Relational, + entity_category: AiEntityCategory::Table, returned_rows: 1, truncated: false, filter_shape: Some("sql_where".to_string()), @@ -491,4 +509,33 @@ mod tests { assert!(!serialized.contains("alice@example.com")); assert!(!serialized.contains("tok_live_secret")); } + + #[test] + fn ai_row_read_audit_omits_secret_bearing_identifiers() { + // These are the raw values available at the handler boundary. The + // audit type has no fields capable of accepting them. + let service_name = "service-alice@example.com"; + let container_path = "bucket/tok_live_secret"; + let entity = "session:alice@example.com:tok_live_secret"; + let audit = AiRowsReadAudit { + context: AuditContext { + user_id: 42, + ip_address: None, + user_agent: "test-agent".to_string(), + }, + service_id: 7, + backend_category: AiBackendCategory::KeyValue, + entity_category: AiEntityCategory::Key, + returned_rows: 1, + truncated: false, + filter_shape: Some("structured_object".to_string()), + }; + let serialized = AuditOperation::serialize(&audit).expect("audit serializes"); + + for secret_identifier in [service_name, container_path, entity] { + assert!(!serialized.contains(secret_identifier)); + } + assert!(serialized.contains("key_value")); + assert!(serialized.contains("entity_category")); + } } diff --git a/crates/temps-providers/src/handlers/query_handlers.rs b/crates/temps-providers/src/handlers/query_handlers.rs index c179e43e5..3d83f1cd5 100644 --- a/crates/temps-providers/src/handlers/query_handlers.rs +++ b/crates/temps-providers/src/handlers/query_handlers.rs @@ -12,7 +12,9 @@ use temps_core::problemdetails::Problem; use temps_query::{ContainerInfo, ContainerPath, EntityInfo, QueryBudget, QueryOptions}; use utoipa::ToSchema; -use super::audit::{AiDataAccessChangedAudit, AiRowsReadAudit}; +use super::audit::{ + AiBackendCategory, AiDataAccessChangedAudit, AiEntityCategory, AiRowsReadAudit, +}; use super::types::AppState; // ============================================================================ @@ -329,12 +331,7 @@ fn filter_shape(filter: Option<&serde_json::Value>) -> Option { .map(str::to_string) } -fn query_error_problem( - error: temps_query::DataError, - service_id: i32, - entity: &str, - limit: usize, -) -> Problem { +fn query_error_problem(error: temps_query::DataError, service_id: i32, limit: usize) -> Problem { let (status, title, kind, safe_detail) = match &error { temps_query::DataError::ResultLimitExceeded { .. } => ( StatusCode::PAYLOAD_TOO_LARGE, @@ -373,7 +370,6 @@ fn query_error_problem( }; tracing::warn!( service_id, - entity, limit, error_kind = kind, "external-service data query failed" @@ -418,6 +414,20 @@ fn entity_names_are_user_data(service_type: &str) -> bool { ) } +/// Fixed audit categories derived only from an allowlisted engine type. +/// Unknown future engines fail closed to non-identifying categories. +fn ai_read_audit_categories(service_type: &str) -> (AiBackendCategory, AiEntityCategory) { + match service_type.trim().to_ascii_lowercase().as_str() { + "postgres" | "postgresql" | "mysql" | "mariadb" => { + (AiBackendCategory::Relational, AiEntityCategory::Table) + } + "mongodb" => (AiBackendCategory::Document, AiEntityCategory::Collection), + "redis" => (AiBackendCategory::KeyValue, AiEntityCategory::Key), + "s3" | "rustfs" => (AiBackendCategory::ObjectStore, AiEntityCategory::Object), + _ => (AiBackendCategory::Unknown, AiEntityCategory::Unknown), + } +} + /// Enforce the `ai_data_access` opt-in on an endpoint that returns *entity /// names*, for engines where those names are user data. /// @@ -463,14 +473,15 @@ async fn apply_entity_name_gate( /// A no-op for human callers — their authorization is `ExternalServicesRead`, /// checked by the caller before this runs. /// -/// Returns the service name for an agent call that passed the gate, so the -/// caller can name the service in the audit record without a second lookup. +/// Returns fixed, non-identifying backend/entity categories for an agent call +/// that passed the gate. User-controlled service and entity names never leave +/// this function for audit serialization. async fn enforce_ai_data_access( app_state: &AppState, service_id: i32, is_ai_call: bool, what: &str, -) -> Result, Problem> { +) -> Result, Problem> { if !is_ai_call { return Ok(None); } @@ -486,7 +497,7 @@ async fn enforce_ai_data_access( })?; if ai_may_read_rows(is_ai_call, service.ai_data_access) { - return Ok(Some(service.name.clone())); + return Ok(Some(ai_read_audit_categories(&service.service_type))); } Err(temps_core::problemdetails::new(StatusCode::FORBIDDEN) @@ -598,7 +609,7 @@ pub async fn read_entity_rows( // projects, so there is no single owning project to scope to, and the // `ai_data_access` opt-in below is the intended boundary — an operator // enables row access per service, whatever project the chat is about. - let ai_service_name = + let ai_audit_categories = enforce_ai_data_access(&app_state, service_id, is_ai_call, "row data").await?; let filter_schema = app_state @@ -629,7 +640,7 @@ pub async fn read_entity_rows( .query_service .query_data(service_id, &path, &entity, filters, options) .await - .map_err(|error| query_error_problem(error, service_id, &entity, limit))?; + .map_err(|error| query_error_problem(error, service_id, limit))?; let total_count = result.stats.total_rows.unwrap_or(result.stats.row_count) as u64; let execution_time_ms = result.stats.execution_ms; @@ -649,9 +660,9 @@ pub async fn read_entity_rows( // Record what the model actually read. The `ai_data_access` toggle is // audited, but that only says the door was opened; this says what went - // through it. Location and shape only — never the values, or the audit log - // becomes a second copy of the same secrets. - if let Some(service_name) = ai_service_name { + // through it. Stable ID, fixed categories and shape only — names and values + // can both contain secrets, so the audit type cannot represent either. + if let Some((backend_category, entity_category)) = ai_audit_categories { let audit = AiRowsReadAudit { context: temps_core::audit::AuditContext { user_id: auth.user_id(), @@ -659,9 +670,8 @@ pub async fn read_entity_rows( user_agent: metadata.user_agent.clone(), }, service_id, - service_name, - container_path: path_str.clone(), - entity: entity.clone(), + backend_category, + entity_category, returned_rows: rows.len(), truncated, filter_shape: audit_filter_shape, @@ -1407,7 +1417,7 @@ pub async fn query_data( .query_service .query_data(service_id, &path, &entity, request.filters, options) .await - .map_err(|error| query_error_problem(error, service_id, &entity, request.limit))?; + .map_err(|error| query_error_problem(error, service_id, request.limit))?; let total_count = result.stats.total_rows.unwrap_or(result.stats.row_count) as u64; let execution_time_ms = result.stats.execution_ms; @@ -1818,7 +1828,6 @@ mod tests { let problem = query_error_problem( temps_query::DataError::QueryFailed(format!("syntax near {secret}")), 7, - "users", 100, ); let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); @@ -1833,7 +1842,6 @@ mod tests { "invalid email alice@example.com with tok_live_secret".to_string(), ), 7, - "users", 100, ); let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); @@ -1852,7 +1860,6 @@ mod tests { observed: 1_024, }, 7, - "session:tok_live_secret:alice@example.com", 1, ); let serialized = serde_json::to_string(&problem.body).expect("problem body serializes"); @@ -1886,6 +1893,30 @@ mod tests { assert!(!entity_names_are_user_data(" PostgreSQL ")); } + #[test] + fn ai_audit_categories_are_fixed_and_unknown_backends_fail_closed() { + assert_eq!( + ai_read_audit_categories("postgres"), + (AiBackendCategory::Relational, AiEntityCategory::Table) + ); + assert_eq!( + ai_read_audit_categories("mongodb"), + (AiBackendCategory::Document, AiEntityCategory::Collection) + ); + assert_eq!( + ai_read_audit_categories("redis"), + (AiBackendCategory::KeyValue, AiEntityCategory::Key) + ); + assert_eq!( + ai_read_audit_categories("s3"), + (AiBackendCategory::ObjectStore, AiEntityCategory::Object) + ); + assert_eq!( + ai_read_audit_categories("future-secret-bearing-engine"), + (AiBackendCategory::Unknown, AiEntityCategory::Unknown) + ); + } + #[test] fn entity_names_are_user_data_defaults_closed_for_unknown_engines() { // Secure-by-default: a backend added later is gated until someone has diff --git a/crates/temps-query-postgres/src/lib.rs b/crates/temps-query-postgres/src/lib.rs index d301643da..1425633ce 100644 --- a/crates/temps-query-postgres/src/lib.rs +++ b/crates/temps-query-postgres/src/lib.rs @@ -47,11 +47,17 @@ fn pg_column_admission(column: &PgQueryColumn, row_budget: usize) -> Result format!("COALESCE(OCTET_LENGTH({value})::bigint * 2 + 8, 4)"), "json" => format!("COALESCE(OCTET_LENGTH({value}::text)::bigint * 6 + 8, 4)"), - "jsonb" | "ARRAY" => format!( + "jsonb" => format!( "CASE WHEN {value} IS NULL THEN 4 WHEN PG_COLUMN_COMPRESSION({value}) IS NOT NULL \ THEN {} ELSE PG_COLUMN_SIZE({value})::bigint * 8 + 64 END", row_budget.saturating_add(1) ), + "ARRAY" => { + return Err(DataError::OperationNotSupported( + "PostgreSQL array columns are not supported by the bounded data browser" + .to_string(), + )) + } "boolean" | "smallint" | "integer" @@ -2245,6 +2251,29 @@ mod wire_budget_tests { temps_query::DataError::OperationNotSupported(_) )); } + + #[test] + fn every_array_type_is_rejected_before_json_query_generation() { + for name in [ + "token_tok_live_secret", + "large_numeric_array", + "custom_type_array", + ] { + let error = with_wire_row_budget( + "SELECT * FROM public.array_payloads", + &[PgQueryColumn { + name: name.to_string(), + data_type: "ARRAY".to_string(), + }], + QueryBudget::default(), + ) + .expect_err("array output expansion has no generic stored-size bound"); + let diagnostic = error.to_string(); + assert!(diagnostic.contains("array columns are not supported")); + assert!(!diagnostic.contains("tok_live_secret")); + assert!(!diagnostic.contains("TO_JSONB")); + } + } } impl temps_query::QuerySchemaProvider for PostgresSource { From 791c660aeeaef300e1c18a04aa06b6ee81c24da4 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 14:26:04 +0200 Subject: [PATCH 06/10] fix(cli): sanitize untrusted terminal output --- .../temps-cli/src/commands/data/index.test.ts | 10 +++ apps/temps-cli/src/commands/data/index.ts | 68 ++++++++++++------- apps/temps-cli/src/ui/table.test.ts | 17 +++++ apps/temps-cli/src/ui/table.ts | 18 +++-- apps/temps-cli/src/ui/terminal.test.ts | 22 ++++++ apps/temps-cli/src/ui/terminal.ts | 33 +++++++++ 6 files changed, 139 insertions(+), 29 deletions(-) create mode 100644 apps/temps-cli/src/ui/table.test.ts create mode 100644 apps/temps-cli/src/ui/terminal.test.ts create mode 100644 apps/temps-cli/src/ui/terminal.ts diff --git a/apps/temps-cli/src/commands/data/index.test.ts b/apps/temps-cli/src/commands/data/index.test.ts index 9ecdd6f93..91a7fdcda 100644 --- a/apps/temps-cli/src/commands/data/index.test.ts +++ b/apps/temps-cli/src/commands/data/index.test.ts @@ -58,6 +58,16 @@ describe('cell', () => { expect(cell(0)).toBe('0') expect(cell(false)).toBe('false') }) + + test('strips terminal control sequences before truncating', () => { + const malicious = + '\x1b[31mred\x1b[0m\x1b]8;;https://attacker.example\x1b\\link\x1b]8;;\x1b\\' + + '\x1b]52;c;dG9rX2xpdmVfc2VjcmV0\x07\nforged' + const out = cell(malicious, 80) + expect(out).toBe('redlink forged') + expect(out).not.toContain('\x1b') + expect(out).not.toContain('\n') + }) }) describe('formatBytes', () => { diff --git a/apps/temps-cli/src/commands/data/index.ts b/apps/temps-cli/src/commands/data/index.ts index c985f3490..790982db6 100644 --- a/apps/temps-cli/src/commands/data/index.ts +++ b/apps/temps-cli/src/commands/data/index.ts @@ -19,6 +19,7 @@ import type { } from '../../api/types.gen.js' import { withSpinner } from '../../ui/spinner.js' import { printTable, type TableColumn } from '../../ui/table.js' +import { sanitizeTerminalText } from '../../ui/terminal.js' import { newline, header, @@ -132,9 +133,9 @@ async function resolveService( ) if (!match) { - const available = services.map((s) => s.name).join(', ') + const available = services.map((s) => sanitizeTerminalText(s.name)).join(', ') throw new Error( - `Service "${nameOrId}" not found. Available: ${available || '(none)'}`, + `Service "${sanitizeTerminalText(nameOrId)}" not found. Available: ${available || '(none)'}`, ) } return { id: match.id, name: match.name, service_type: match.service_type } @@ -149,7 +150,8 @@ async function resolveService( */ export function cell(value: unknown, maxLength = 40): string { if (value === null || value === undefined) return colors.dim('null') - const text = typeof value === 'string' ? value : JSON.stringify(value) + const raw = typeof value === 'string' ? value : JSON.stringify(value) + const text = sanitizeTerminalText(raw) return text.length > maxLength ? `${text.slice(0, maxLength - 1)}…` : text } @@ -180,8 +182,8 @@ export function validateFilter(raw: string | undefined, service: string): string } catch (e) { throw new Error( `--filter must be JSON, but failed to parse: ${(e as Error).message}\n` + - ` Received: ${raw}\n` + - ` Run "temps data info ${service}" to see this backend's filter schema.`, + ` Received: ${sanitizeTerminalText(raw)}\n` + + ` Run "temps data info ${sanitizeTerminalText(service)}" to see this backend's filter schema.`, ) } } @@ -214,7 +216,9 @@ async function infoCmd( } newline() - header(`${icons.info} ${service.name} (${service.service_type})`) + header( + `${icons.info} ${sanitizeTerminalText(service.name)} (${sanitizeTerminalText(service.service_type)})`, + ) newline() if (!support?.supported) { @@ -235,7 +239,7 @@ async function infoCmd( ? 'holds containers' : 'leaf' console.log( - ` ${colors.muted(`level ${level.level}`)} ${colors.bold(level.container_type)} ${colors.dim(`— ${holds}`)}`, + ` ${colors.muted(`level ${level.level}`)} ${colors.bold(sanitizeTerminalText(level.container_type))} ${colors.dim(`— ${holds}`)}`, ) } newline() @@ -246,7 +250,7 @@ async function infoCmd( newline() } - info(`Next: temps data containers ${service.name}`) + info(`Next: temps data containers ${sanitizeTerminalText(service.name)}`) newline() } @@ -283,9 +287,11 @@ async function containersCmd( return } - const scope = options.path ? ` under ${options.path}` : '' + const scope = options.path ? ` under ${sanitizeTerminalText(options.path)}` : '' newline() - header(`${icons.folder} Containers in ${service.name}${scope} (${rows.length})`) + header( + `${icons.folder} Containers in ${sanitizeTerminalText(service.name)}${scope} (${rows.length})`, + ) if (rows.length === 0) { info('No containers found.') @@ -312,11 +318,14 @@ async function containersCmd( // common thing to get wrong, so show a real one rather than a placeholder. const first = rows[0] if (first) { - const nextPath = options.path ? `${options.path}/${first.name}` : first.name + const nextPath = sanitizeTerminalText( + options.path ? `${options.path}/${first.name}` : first.name, + ) + const safeServiceName = sanitizeTerminalText(service.name) if (first.can_contain_entities) { - info(`Next: temps data tables ${service.name} --path ${nextPath}`) + info(`Next: temps data tables ${safeServiceName} --path ${nextPath}`) } else if (first.can_contain_containers) { - info(`Next: temps data containers ${service.name} --path ${nextPath}`) + info(`Next: temps data containers ${safeServiceName} --path ${nextPath}`) } } newline() @@ -349,10 +358,15 @@ async function tablesCmd( } newline() - header(`${icons.folder} ${options.path} in ${service.name} (${entities.length})`) + header( + `${icons.folder} ${sanitizeTerminalText(options.path)} in ${sanitizeTerminalText(service.name)} (${entities.length})`, + ) if (entities.length === 0) { - info('No entities found. Check the path with: temps data containers ' + service.name) + info( + 'No entities found. Check the path with: temps data containers ' + + sanitizeTerminalText(service.name), + ) newline() return } @@ -373,7 +387,9 @@ async function tablesCmd( } const firstEntity = entities[0] if (firstEntity) { - info(`Next: temps data rows ${service.name} ${firstEntity.name} --path ${options.path}`) + info( + `Next: temps data rows ${sanitizeTerminalText(service.name)} ${sanitizeTerminalText(firstEntity.name)} --path ${sanitizeTerminalText(options.path)}`, + ) } newline() } @@ -403,7 +419,9 @@ async function schemaCmd( } newline() - header(`${icons.info} ${options.path}/${entity} (${info_?.entity_type ?? 'entity'})`) + header( + `${icons.info} ${sanitizeTerminalText(options.path)}/${sanitizeTerminalText(entity)} (${sanitizeTerminalText(info_?.entity_type ?? 'entity')})`, + ) newline() if (info_?.row_count !== null && info_?.row_count !== undefined) { keyValue('Rows', String(info_.row_count)) @@ -474,7 +492,7 @@ async function rowsCmd( newline() header( - `${icons.info} ${options.path}/${entity} — ${result?.returned_count ?? rows.length} of ${result?.total_count ?? '?'} rows`, + `${icons.info} ${sanitizeTerminalText(options.path)}/${sanitizeTerminalText(entity)} — ${result?.returned_count ?? rows.length} of ${result?.total_count ?? '?'} rows`, ) if (rows.length === 0) { @@ -543,7 +561,7 @@ async function aiAccessCmd( } newline() - header(`${icons.info} AI data access — ${service.name}`) + header(`${icons.info} AI data access — ${sanitizeTerminalText(service.name)}`) newline() keyValue( 'Built-in assistant may read rows', @@ -552,8 +570,8 @@ async function aiAccessCmd( newline() info( current?.enabled - ? `Disable with: temps data ai-access ${service.name} --disable` - : `Enable with: temps data ai-access ${service.name} --enable`, + ? `Disable with: temps data ai-access ${sanitizeTerminalText(service.name)} --disable` + : `Enable with: temps data ai-access ${sanitizeTerminalText(service.name)} --enable`, ) info( colors.dim( @@ -587,14 +605,18 @@ async function aiAccessCmd( newline() if (enabled) { - success(`The built-in AI assistant can now read rows from ${service.name}.`) + success( + `The built-in AI assistant can now read rows from ${sanitizeTerminalText(service.name)}.`, + ) newline() warning( 'Rows are sent to your configured AI provider. If this service stores password', ) warning('hashes, tokens or personal data, that data leaves your infrastructure.') } else { - success(`The built-in AI assistant can no longer read rows from ${service.name}.`) + success( + `The built-in AI assistant can no longer read rows from ${sanitizeTerminalText(service.name)}.`, + ) info('Table and column names remain readable.') } newline() diff --git a/apps/temps-cli/src/ui/table.test.ts b/apps/temps-cli/src/ui/table.test.ts new file mode 100644 index 000000000..eb2f424d3 --- /dev/null +++ b/apps/temps-cli/src/ui/table.test.ts @@ -0,0 +1,17 @@ +import { describe, expect, test } from 'bun:test' +import { createTable } from './table.js' + +describe('createTable', () => { + test('sanitizes headers and values before terminal rendering', () => { + const rendered = createTable( + [{ name: '\x1b]52;c;dG9rX2xpdmVfc2VjcmV0\x07safe\nrow' }], + [{ header: '\x1b[31mName\x1b[0m', key: 'name' }], + { style: 'minimal' }, + ) + + expect(rendered).toContain('Name') + expect(rendered).toContain('safe row') + expect(rendered).not.toContain(']52;') + expect(rendered).not.toContain('\x1b[31m') + }) +}) diff --git a/apps/temps-cli/src/ui/table.ts b/apps/temps-cli/src/ui/table.ts index 8b097bafa..7be3a5160 100644 --- a/apps/temps-cli/src/ui/table.ts +++ b/apps/temps-cli/src/ui/table.ts @@ -1,6 +1,7 @@ import Table from 'cli-table3' import chalk from 'chalk' import { colors } from './output.js' +import { sanitizeTerminalText } from './terminal.js' export interface TableColumn { header: string @@ -121,7 +122,7 @@ export function createTable( const preset = stylePresets[options.style ?? 'default'] const table = new Table({ - head: columns.map((col) => colors.bold(col.header)), + head: columns.map((col) => colors.bold(sanitizeTerminalText(col.header))), colAligns: columns.map((col) => col.align ?? 'left'), colWidths: columns.map((col) => col.width ?? null), ...preset, @@ -139,7 +140,8 @@ export function createTable( value = '' } - let strValue = value === null || value === undefined ? '' : String(value) + let strValue = + value === null || value === undefined ? '' : sanitizeTerminalText(value) if (col.color) { strValue = col.color(strValue, item) @@ -177,8 +179,11 @@ export function detailsTable( }) for (const [key, value] of Object.entries(details)) { - const displayValue = value === null || value === undefined ? colors.muted('not set') : String(value) - table.push([colors.muted(key), displayValue]) + const displayValue = + value === null || value === undefined + ? colors.muted('not set') + : sanitizeTerminalText(value) + table.push([colors.muted(sanitizeTerminalText(key)), displayValue]) } console.log(table.toString()) @@ -188,6 +193,7 @@ export function detailsTable( * Status badge formatter */ export function statusBadge(status: string): string { + const safeStatus = sanitizeTerminalText(status) const statusColors: Record string> = { running: chalk.green, active: chalk.green, @@ -210,6 +216,6 @@ export function statusBadge(status: string): string { cancelled: chalk.red, } - const colorFn = statusColors[status.toLowerCase()] ?? chalk.white - return colorFn(`● ${status}`) + const colorFn = statusColors[safeStatus.toLowerCase()] ?? chalk.white + return colorFn(`● ${safeStatus}`) } diff --git a/apps/temps-cli/src/ui/terminal.test.ts b/apps/temps-cli/src/ui/terminal.test.ts new file mode 100644 index 000000000..d3293550a --- /dev/null +++ b/apps/temps-cli/src/ui/terminal.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, test } from 'bun:test' +import { sanitizeTerminalText } from './terminal.js' + +describe('sanitizeTerminalText', () => { + test('strips CSI colour and cursor sequences', () => { + expect(sanitizeTerminalText('\x1b[31mforged\x1b[0m')).toBe('forged') + expect(sanitizeTerminalText('before\x1b[2Jafter')).toBe('beforeafter') + }) + + test('strips OSC 8 links and OSC 52 clipboard writes', () => { + const link = '\x1b]8;;https://attacker.example\x1b\\click\x1b]8;;\x1b\\' + const clipboard = '\x1b]52;c;dG9rX2xpdmVfc2VjcmV0\x07visible' + expect(sanitizeTerminalText(link)).toBe('click') + expect(sanitizeTerminalText(clipboard)).toBe('visible') + }) + + test('collapses newlines and strips C0, C1, and bidi overrides', () => { + expect(sanitizeTerminalText('one\r\ntwo\nthree\x00\x9b31m\u202E')).toBe( + 'one two three', + ) + }) +}) diff --git a/apps/temps-cli/src/ui/terminal.ts b/apps/temps-cli/src/ui/terminal.ts new file mode 100644 index 000000000..c235decb1 --- /dev/null +++ b/apps/temps-cli/src/ui/terminal.ts @@ -0,0 +1,33 @@ +/** + * Make untrusted text safe for human-readable terminal output. + * + * JSON output must not use this helper: JSON.stringify already escapes control + * characters and callers expect the original data. This is for strings that + * will be written directly to a terminal, where ANSI/OSC sequences can alter + * the clipboard, links, title, colours, cursor position, or visible history. + */ +export function sanitizeTerminalText(value: unknown): string { + let text = String(value) + + // OSC sequences end with BEL or ST. Handle both 7-bit (ESC ]) and 8-bit + // (C1 OSC) forms, including an unterminated sequence at end-of-input. + text = text.replace( + /\x1B\](?:[^\x07\x1B]|\x1B(?!\\))*(?:\x07|\x1B\\|$)/g, + '', + ) + text = text.replace(/\x9D[^\x07\x9C]*(?:\x07|\x9C|$)/g, '') + + // CSI covers colour, cursor movement, erasure, and terminal mode changes. + text = text.replace(/(?:\x1B\[|\x9B)[0-?]*[ -/]*[@-~]/g, '') + + // Strip remaining two-byte/intermediate ESC sequences and lone ESC bytes. + text = text.replace(/\x1B[ -/]*[@-~]?/g, '') + + // A value must never create a second terminal line. Preserve word separation + // while removing the remaining C0/C1 controls and DEL. + text = text.replace(/\r\n?|\n/g, ' ') + text = text.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]/g, '') + + // Bidirectional overrides can visually reorder commands and identifiers. + return text.replace(/[\u202A-\u202E\u2066-\u2069]/g, '') +} From 799920f1072d116e89e50cf7ae7c7993e45e9a99 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Thu, 6 Aug 2026 14:35:33 +0200 Subject: [PATCH 07/10] fix(cli): sanitize invalid filter errors --- .../temps-cli/src/commands/data/index.test.ts | 20 ++++++++++++++++++- apps/temps-cli/src/commands/data/index.ts | 7 +++---- apps/temps-cli/src/ui/terminal.test.ts | 6 +++++- apps/temps-cli/src/ui/terminal.ts | 7 ++++--- 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/apps/temps-cli/src/commands/data/index.test.ts b/apps/temps-cli/src/commands/data/index.test.ts index 91a7fdcda..71189aa87 100644 --- a/apps/temps-cli/src/commands/data/index.test.ts +++ b/apps/temps-cli/src/commands/data/index.test.ts @@ -17,7 +17,7 @@ describe('validateFilter', () => { // envelope. Sending it would let the server drop the filter and return a // full unfiltered table, which reads as a correct answer. expect(() => validateFilter("plan = 'pro'", 'my-db')).toThrow( - /--filter must be JSON/, + /--filter must be valid JSON/, ) }) @@ -30,6 +30,24 @@ describe('validateFilter', () => { test('echoes what was received so a typo is visible', () => { expect(() => validateFilter('{oops}', 'my-db')).toThrow(/\{oops\}/) }) + + test('does not expose parser or terminal controls for a malicious filter', () => { + const payload = + '{"where":"\x1b]52;c;dG9rX2xpdmVfc2VjcmV0\x07\x1b[31m\n\u202E"' + let message = '' + try { + validateFilter(payload, 'service\x1b]8;;https://attacker.example\x1b\\') + } catch (error) { + message = (error as Error).message + } + + expect(message).toContain('--filter must be valid JSON') + expect(message).not.toContain('\x1b') + expect(message).not.toContain('\x07') + expect(message).not.toContain('\n') + expect(message).not.toContain('\u202E') + expect(message).not.toContain('Unexpected') + }) }) describe('cell', () => { diff --git a/apps/temps-cli/src/commands/data/index.ts b/apps/temps-cli/src/commands/data/index.ts index 790982db6..b662d31ed 100644 --- a/apps/temps-cli/src/commands/data/index.ts +++ b/apps/temps-cli/src/commands/data/index.ts @@ -179,11 +179,10 @@ export function validateFilter(raw: string | undefined, service: string): string try { JSON.parse(raw) return raw - } catch (e) { + } catch { throw new Error( - `--filter must be JSON, but failed to parse: ${(e as Error).message}\n` + - ` Received: ${sanitizeTerminalText(raw)}\n` + - ` Run "temps data info ${sanitizeTerminalText(service)}" to see this backend's filter schema.`, + `--filter must be valid JSON. Received: ${sanitizeTerminalText(raw)}. ` + + `Run "temps data info ${sanitizeTerminalText(service)}" to see this backend's filter schema.`, ) } } diff --git a/apps/temps-cli/src/ui/terminal.test.ts b/apps/temps-cli/src/ui/terminal.test.ts index d3293550a..4fead1895 100644 --- a/apps/temps-cli/src/ui/terminal.test.ts +++ b/apps/temps-cli/src/ui/terminal.test.ts @@ -15,7 +15,11 @@ describe('sanitizeTerminalText', () => { }) test('collapses newlines and strips C0, C1, and bidi overrides', () => { - expect(sanitizeTerminalText('one\r\ntwo\nthree\x00\x9b31m\u202E')).toBe( + expect( + sanitizeTerminalText( + 'one\r\ntwo\tthree\x00\x9b31m\u061C\u200E\u200F\u202E', + ), + ).toBe( 'one two three', ) }) diff --git a/apps/temps-cli/src/ui/terminal.ts b/apps/temps-cli/src/ui/terminal.ts index c235decb1..780ea886f 100644 --- a/apps/temps-cli/src/ui/terminal.ts +++ b/apps/temps-cli/src/ui/terminal.ts @@ -25,9 +25,10 @@ export function sanitizeTerminalText(value: unknown): string { // A value must never create a second terminal line. Preserve word separation // while removing the remaining C0/C1 controls and DEL. - text = text.replace(/\r\n?|\n/g, ' ') + text = text.replace(/\r\n?|\n|\t/g, ' ') text = text.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]/g, '') - // Bidirectional overrides can visually reorder commands and identifiers. - return text.replace(/[\u202A-\u202E\u2066-\u2069]/g, '') + // Bidirectional marks, isolates, and overrides can visually reorder commands + // and identifiers even though they do not alter the underlying byte order. + return text.replace(/[\u061C\u200E\u200F\u202A-\u202E\u2066-\u2069]/g, '') } From 5258a228dba2d290577f88558a768d61ca9b0121 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Fri, 7 Aug 2026 14:25:11 +0200 Subject: [PATCH 08/10] fix(security): resolve PR review findings --- Cargo.lock | 3 + apps/temps-cli/openapi-ts.config.ts | 6 +- apps/temps-cli/src/commands/context/index.ts | 12 +- apps/temps-cli/src/commands/data/index.ts | 7 +- apps/temps-cli/src/commands/flags/index.ts | 4 +- .../temps-cli/src/commands/instances/index.ts | 6 +- .../src/commands/notifications/index.ts | 4 +- .../temps-cli/src/commands/providers/index.ts | 6 +- .../src/commands/services/restore.ts | 19 +- apps/temps-cli/src/ui/table.test.ts | 11 + crates/temps-auth/src/audit.rs | 5 + crates/temps-auth/src/temps_middleware.rs | 29 +- crates/temps-cli/src/commands/migrate.rs | 4 + crates/temps-cli/src/commands/serve/mod.rs | 8 + crates/temps-database/Cargo.toml | 1 + crates/temps-database/src/connection.rs | 109 +++++++ crates/temps-database/src/lib.rs | 27 +- ...00001_index_permission_denied_retention.rs | 52 ---- crates/temps-migrations/src/migration/mod.rs | 2 - crates/temps-providers/src/mariadb_query.rs | 100 ++++++- crates/temps-query-mongodb/Cargo.toml | 3 + crates/temps-query-mongodb/src/lib.rs | 83 ++++++ crates/temps-query-postgres/src/lib.rs | 53 +++- crates/temps-query-redis/Cargo.toml | 3 + crates/temps-query-redis/src/lib.rs | 267 ++++++++++++++---- web/src/components/audit/AuditLogItem.test.ts | 26 +- web/src/pages/AuditLogs.test.ts | 10 +- web/src/pages/AuditLogs.tsx | 34 +-- 28 files changed, 718 insertions(+), 176 deletions(-) delete mode 100644 crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs diff --git a/Cargo.lock b/Cargo.lock index 416d2af4a..0b47ff883 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10778,6 +10778,7 @@ dependencies = [ "once_cell", "sea-orm", "sea-orm-migration", + "sqlx 0.8.6", "temps-core", "temps-entities", "temps-migrations", @@ -12295,6 +12296,7 @@ dependencies = [ "serde", "serde_json", "temps-query", + "testcontainers", "thiserror 2.0.19", "tokio", "tracing", @@ -12331,6 +12333,7 @@ dependencies = [ "serde", "serde_json", "temps-query", + "testcontainers", "thiserror 2.0.19", "tokio", "tracing", diff --git a/apps/temps-cli/openapi-ts.config.ts b/apps/temps-cli/openapi-ts.config.ts index c17bf158e..af54e3d93 100644 --- a/apps/temps-cli/openapi-ts.config.ts +++ b/apps/temps-cli/openapi-ts.config.ts @@ -3,11 +3,9 @@ import { defineConfig } from '@hey-api/openapi-ts' export default defineConfig({ input: 'openapi.json', // input: 'http://localhost:3000/api-docs/openapi.json', - output: { - path: 'src/api', - }, - client: '@hey-api/client-fetch', + output: 'src/api', plugins: [ + '@hey-api/client-fetch', '@hey-api/sdk', '@hey-api/typescript', ], diff --git a/apps/temps-cli/src/commands/context/index.ts b/apps/temps-cli/src/commands/context/index.ts index 5fcaa1608..236c19423 100644 --- a/apps/temps-cli/src/commands/context/index.ts +++ b/apps/temps-cli/src/commands/context/index.ts @@ -87,24 +87,28 @@ async function listAction(options: { json?: boolean }): Promise { const columns: TableColumn[] = [ { header: '', - accessor: (c) => (isActiveRow(c) ? colors.success('●') : colors.muted('○')), + accessor: (c) => (isActiveRow(c) ? '●' : '○'), + color: (value, c) => isActiveRow(c) ? colors.success(value) : colors.muted(value), }, { header: 'Name', key: 'name', color: (v) => colors.bold(String(v)) }, { header: 'URL', key: 'url', color: (v) => colors.primary(String(v)) }, { header: 'Email', - accessor: (c) => c.email || colors.muted('-'), + accessor: (c) => c.email || '-', + color: (value, c) => c.email ? value : colors.muted(value), }, { header: 'Key', - accessor: (c) => (c.keyPrefix ? colors.muted(`${c.keyPrefix}…`) : colors.muted('-')), + accessor: (c) => (c.keyPrefix ? `${c.keyPrefix}…` : '-'), + color: (value) => colors.muted(value), }, { header: 'Expires', accessor: (c) => c.expiresAt ? new Date(c.expiresAt).toISOString().split('T')[0] ?? '-' - : colors.muted('-'), + : '-', + color: (value, c) => c.expiresAt ? value : colors.muted(value), }, ] diff --git a/apps/temps-cli/src/commands/data/index.ts b/apps/temps-cli/src/commands/data/index.ts index b662d31ed..bf8587a64 100644 --- a/apps/temps-cli/src/commands/data/index.ts +++ b/apps/temps-cli/src/commands/data/index.ts @@ -308,7 +308,9 @@ async function containersCmd( ? (c.entity_type_label ?? 'entities') : c.can_contain_containers ? (c.child_container_type ?? 'containers') - : colors.dim('—'), + : '—', + color: (value, c) => + c.can_contain_entities || c.can_contain_containers ? value : colors.dim(value), }, ] printTable(rows, columns, { style: 'minimal' }) @@ -375,7 +377,8 @@ async function tablesCmd( { header: 'Type', key: 'entity_type' }, { header: 'Rows', - accessor: (e) => (e.row_count === null || e.row_count === undefined ? colors.dim('—') : String(e.row_count)), + accessor: (e) => (e.row_count === null || e.row_count === undefined ? '—' : String(e.row_count)), + color: (value, e) => e.row_count === null || e.row_count === undefined ? colors.dim(value) : value, }, { header: 'Size', accessor: (e) => formatBytes(e.size_bytes) }, ] diff --git a/apps/temps-cli/src/commands/flags/index.ts b/apps/temps-cli/src/commands/flags/index.ts index b4df9c56e..f0886bc1c 100644 --- a/apps/temps-cli/src/commands/flags/index.ts +++ b/apps/temps-cli/src/commands/flags/index.ts @@ -279,7 +279,7 @@ async function listFlagsCmd(options: { ? { header: options.environment!, accessor: (f) => effectiveValue(f, environmentId) } : { header: 'Default', accessor: (f) => formatValue(f.default_value) }, { header: 'Client', accessor: (f) => (f.client_visible ? 'yes' : 'no') }, - { header: 'Status', accessor: (f) => (f.archived_at ? colors.dim('archived') : 'active') }, + { header: 'Status', accessor: (f) => (f.archived_at ? 'archived' : 'active'), color: (value, f) => f.archived_at ? colors.dim(value) : value }, ] printTable(flags, columns, { style: 'minimal' }) @@ -337,7 +337,7 @@ async function getFlagCmd( flag.environments, [ { header: 'Environment ID', accessor: (e) => String(e.environment_id) }, - { header: 'Enabled', accessor: (e) => (e.enabled ? 'yes' : colors.warning('no (kill switch)')) }, + { header: 'Enabled', accessor: (e) => (e.enabled ? 'yes' : 'no (kill switch)'), color: (value, e) => e.enabled ? value : colors.warning(value) }, { header: 'Value', accessor: (e) => formatValue(e.value) }, ], { style: 'minimal' }, diff --git a/apps/temps-cli/src/commands/instances/index.ts b/apps/temps-cli/src/commands/instances/index.ts index 348cf62d8..73ec7cace 100644 --- a/apps/temps-cli/src/commands/instances/index.ts +++ b/apps/temps-cli/src/commands/instances/index.ts @@ -51,13 +51,15 @@ async function listAction(options: { json?: boolean }): Promise { const columns: TableColumn[] = [ { header: '', - accessor: (i) => i.isDefault ? colors.success('●') : colors.muted('○'), + accessor: (i) => i.isDefault ? '●' : '○', + color: (value, i) => i.isDefault ? colors.success(value) : colors.muted(value), }, { header: 'Name', key: 'name', color: (v) => colors.bold(v) }, { header: 'URL', key: 'url', color: (v) => colors.primary(v) }, { header: 'Email', - accessor: (i) => i.email ?? colors.muted('-'), + accessor: (i) => i.email ?? '-', + color: (value, i) => i.email ? value : colors.muted(value), }, ] diff --git a/apps/temps-cli/src/commands/notifications/index.ts b/apps/temps-cli/src/commands/notifications/index.ts index c10feb4e7..65b6208c9 100644 --- a/apps/temps-cli/src/commands/notifications/index.ts +++ b/apps/temps-cli/src/commands/notifications/index.ts @@ -10,7 +10,7 @@ import { testNotificationProvider as testProvider2, updateNotificationProvider as updateProvider2, updateSlackProvider, - updateEmailProvider, + updateNotificationEmailProvider, } from '../../api/sdk.gen.js' import type { NotificationProviderResponse } from '../../api/types.gen.js' import { withSpinner } from '../../ui/spinner.js' @@ -567,7 +567,7 @@ async function updateEmailProviderAction( } | null const updated = await withSpinner('Updating Email provider...', async () => { - const { data, error } = await updateEmailProvider({ + const { data, error } = await updateNotificationEmailProvider({ client, path: { id }, body: { diff --git a/apps/temps-cli/src/commands/providers/index.ts b/apps/temps-cli/src/commands/providers/index.ts index 25812df3a..e51b3d933 100644 --- a/apps/temps-cli/src/commands/providers/index.ts +++ b/apps/temps-cli/src/commands/providers/index.ts @@ -961,9 +961,9 @@ async function syncConnectionAction(options: IdOptions): Promise { return data }) - success(`Synced ${result?.total_count ?? 0} repositories for connection ${id}`) - if (result?.synced_at) { - info(`Synced at: ${result.synced_at}`) + success(`Repository sync started for connection ${result?.connection_id ?? id}`) + if (result?.started_at) { + info(`Started at: ${result.started_at}`) } } diff --git a/apps/temps-cli/src/commands/services/restore.ts b/apps/temps-cli/src/commands/services/restore.ts index 443bc460b..a293c48c8 100644 --- a/apps/temps-cli/src/commands/services/restore.ts +++ b/apps/temps-cli/src/commands/services/restore.ts @@ -150,14 +150,16 @@ async function listBackupsAction(options: ListBackupsOptions): Promise { accessor: (r) => typeof r.size_bytes === 'number' && r.size_bytes != null ? formatBytes(r.size_bytes) - : colors.muted('—'), + : '—', + color: (value, r) => typeof r.size_bytes === 'number' && r.size_bytes != null ? value : colors.muted(value), }, { header: 'Location', accessor: (r) => (r.location ?? '').startsWith('s3://') - ? colors.success('WAL-G') - : colors.muted('legacy'), + ? 'WAL-G' + : 'legacy', + color: (value, r) => (r.location ?? '').startsWith('s3://') ? colors.success(value) : colors.muted(value), }, { header: 'Created', @@ -358,17 +360,18 @@ async function listRunsAction(options: RestoreRunsOptions): Promise { { header: 'Phase', accessor: (r) => r.phase }, { header: 'Status', - accessor: (r) => - statusBadge( - r.status === 'completed' ? 'active' : r.status === 'failed' ? 'inactive' : 'pending', - ), + accessor: (r) => r.status, + color: (_value, r) => statusBadge( + r.status === 'completed' ? 'active' : r.status === 'failed' ? 'inactive' : 'pending', + ), }, { header: 'Target', accessor: (r) => r.target_service_id != null ? `#${r.target_service_id} (${r.target_service_name ?? ''})` - : colors.muted('—'), + : '—', + color: (value, r) => r.target_service_id != null ? value : colors.muted(value), }, { header: 'Started', diff --git a/apps/temps-cli/src/ui/table.test.ts b/apps/temps-cli/src/ui/table.test.ts index eb2f424d3..8778596fe 100644 --- a/apps/temps-cli/src/ui/table.test.ts +++ b/apps/temps-cli/src/ui/table.test.ts @@ -14,4 +14,15 @@ describe('createTable', () => { expect(rendered).not.toContain(']52;') expect(rendered).not.toContain('\x1b[31m') }) + + test('applies trusted column styling after sanitizing accessor output', () => { + const rendered = createTable( + [{ status: '\x1b]52;c;Zm9yZ2Vk\x07inactive' }], + [{ header: 'Status', accessor: (row) => row.status, color: (value) => `\x1b[33m${value}\x1b[0m` }], + { style: 'minimal' }, + ) + + expect(rendered).toContain('\x1b[33minactive\x1b[0m') + expect(rendered).not.toContain(']52;') + }) }) diff --git a/crates/temps-auth/src/audit.rs b/crates/temps-auth/src/audit.rs index d522aa9be..d00a638c1 100644 --- a/crates/temps-auth/src/audit.rs +++ b/crates/temps-auth/src/audit.rs @@ -705,6 +705,10 @@ pub struct PermissionDeniedAudit { /// per-window persistence budgets. pub suppressed_by_budget: bool, pub ip_address: Option, + /// Persisted through [`AuditOperation::user_agent`] in the dedicated audit + /// column. Skipping it here avoids duplicating attacker-controlled origin + /// metadata in the JSON payload. + #[serde(skip_serializing)] pub user_agent: String, } @@ -913,5 +917,6 @@ mod failure_audit_tests { assert!(json.contains("\"attempt_count\":4")); assert!(json.contains("\"credential_id\":17")); assert!(!json.contains("token_name")); + assert!(!json.contains("test-agent")); } } diff --git a/crates/temps-auth/src/temps_middleware.rs b/crates/temps-auth/src/temps_middleware.rs index bd6a4d3cd..b55990a33 100644 --- a/crates/temps-auth/src/temps_middleware.rs +++ b/crates/temps-auth/src/temps_middleware.rs @@ -25,6 +25,11 @@ use crate::{ apikey_service::ApiKeyService, auth_service::AuthService, deployment_token_service::DeploymentTokenValidationService, user_service::UserService, }; + +/// Permission-denial events are attacker-triggerable. Keep the queued and +/// persisted origin metadata small even when the HTTP stack accepts a large +/// request header block. +const MAX_AUDIT_USER_AGENT_BYTES: usize = 512; use temps_core::CookieCrypto; /// Authentication middleware that implements TempsMiddleware @@ -219,7 +224,7 @@ impl AuthMiddleware { .map(|metadata| { ( Some(metadata.ip_address.clone()), - metadata.user_agent.clone(), + bounded_audit_user_agent(&metadata.user_agent), ) }) .unwrap_or_else(|| (None, "unknown".to_string())); @@ -260,6 +265,18 @@ impl AuthMiddleware { } } +fn bounded_audit_user_agent(user_agent: &str) -> String { + if user_agent.len() <= MAX_AUDIT_USER_AGENT_BYTES { + return user_agent.to_string(); + } + + let mut end = MAX_AUDIT_USER_AGENT_BYTES; + while !user_agent.is_char_boundary(end) { + end -= 1; + } + user_agent[..end].to_string() +} + fn safe_principal(auth: &crate::context::AuthContext) -> SafePrincipal { let (source, credential_id) = match &auth.source { crate::context::AuthSource::Session { .. } => (AuthSourceKind::Session, None), @@ -362,6 +379,16 @@ mod tests { ) } + #[test] + fn denial_user_agent_is_utf8_safely_byte_bounded() { + let oversized = format!("{}{}", "a".repeat(MAX_AUDIT_USER_AGENT_BYTES - 1), "éé"); + let bounded = bounded_audit_user_agent(&oversized); + + assert!(bounded.len() <= MAX_AUDIT_USER_AGENT_BYTES); + assert!(bounded.is_char_boundary(bounded.len())); + assert_eq!(bounded, "a".repeat(MAX_AUDIT_USER_AGENT_BYTES - 1)); + } + #[test] fn marked_guard_403_becomes_an_audit_event() { let response = temps_core::error_builder::ErrorBuilder::new(StatusCode::FORBIDDEN) diff --git a/crates/temps-cli/src/commands/migrate.rs b/crates/temps-cli/src/commands/migrate.rs index 335044e06..137a21ed2 100644 --- a/crates/temps-cli/src/commands/migrate.rs +++ b/crates/temps-cli/src/commands/migrate.rs @@ -173,6 +173,10 @@ impl MigrateCommand { ); } + // Non-blocking indexes cannot run in SeaORM's migration transaction. + // The explicit migrate command is the ideal place to wait for them. + temps_database::run_post_migration_indexes(&db).await?; + // Continuous-aggregate backfill is idempotent and safe to run here // (the operator is already waiting on this command). if let Err(e) = temps_database::run_post_migration_backfill(&db).await { diff --git a/crates/temps-cli/src/commands/serve/mod.rs b/crates/temps-cli/src/commands/serve/mod.rs index b8ac83935..8a85d44c5 100644 --- a/crates/temps-cli/src/commands/serve/mod.rs +++ b/crates/temps-cli/src/commands/serve/mod.rs @@ -307,6 +307,14 @@ impl ServeCommand { { let backfill_db = db.clone(); rt.spawn(async move { + if let Err(e) = + temps_database::run_post_migration_indexes(backfill_db.as_ref()).await + { + tracing::warn!( + "Post-migration index build failed (will retry on next startup): {}", + e + ); + } if let Err(e) = temps_database::run_post_migration_backfill(backfill_db.as_ref()).await { diff --git a/crates/temps-database/Cargo.toml b/crates/temps-database/Cargo.toml index e7e884c8f..d069a9f58 100644 --- a/crates/temps-database/Cargo.toml +++ b/crates/temps-database/Cargo.toml @@ -14,6 +14,7 @@ temps-entities = { path = "../temps-entities" } temps-migrations = { path = "../temps-migrations" } sea-orm = { workspace = true, features = ["runtime-tokio-native-tls", "sqlx-postgres"] } sea-orm-migration = { workspace = true } +sqlx = { workspace = true } tokio = { workspace = true } uuid = { workspace = true } testcontainers = { workspace = true } diff --git a/crates/temps-database/src/connection.rs b/crates/temps-database/src/connection.rs index 382f452ea..e67492783 100644 --- a/crates/temps-database/src/connection.rs +++ b/crates/temps-database/src/connection.rs @@ -547,6 +547,115 @@ pub async fn get_pending_migration_names(db: &DbConnection) -> ServiceResult ServiceResult<()> { + let pool = db.get_postgres_connection_pool(); + let mut connection = pool.acquire().await.map_err(|error| { + ServiceError::Database(format!( + "Failed to acquire database connection for post-migration indexes: {error}" + )) + })?; + + sqlx::query("SET lock_timeout = '5s'") + .execute(&mut *connection) + .await + .map_err(|error| { + ServiceError::Database(format!( + "Failed to set lock timeout for permission-denied retention index: {error}" + )) + })?; + if let Err(error) = sqlx::query("SET statement_timeout = '10min'") + .execute(&mut *connection) + .await + { + let _ = sqlx::query("RESET lock_timeout") + .execute(&mut *connection) + .await; + return Err(ServiceError::Database(format!( + "Failed to set statement timeout for permission-denied retention index: {error}" + ))); + } + + let maintenance_result: ServiceResult<()> = async { + let index_is_valid = sqlx::query_scalar::<_, bool>( + "SELECT index.indisvalid \ + FROM pg_index AS index \ + WHERE index.indexrelid = \ + to_regclass('idx_audit_logs_permission_denied_retention')", + ) + .fetch_optional(&mut *connection) + .await + .map_err(|error| { + ServiceError::Database(format!( + "Failed to inspect permission-denied retention index: {error}" + )) + })?; + + // A cancelled/failed concurrent build leaves an invalid index behind. + // `IF NOT EXISTS` would otherwise accept it forever and prevent retries. + if index_is_valid == Some(false) { + sqlx::query( + "DROP INDEX CONCURRENTLY IF EXISTS idx_audit_logs_permission_denied_retention", + ) + .execute(&mut *connection) + .await + .map_err(|error| { + ServiceError::Database(format!( + "Failed to remove invalid permission-denied retention index: {error}" + )) + })?; + } + + sqlx::query( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS \ + idx_audit_logs_permission_denied_retention \ + ON audit_logs (audit_date ASC, id ASC) \ + WHERE operation_type = 'PERMISSION_DENIED'", + ) + .execute(&mut *connection) + .await + .map_err(|error| { + ServiceError::Database(format!( + "Failed to build concurrent permission-denied retention index \ + '{PERMISSION_DENIED_RETENTION_INDEX}': {error}" + )) + })?; + + Ok(()) + } + .await; + + // Do not leak maintenance-specific session settings back into the pool. + let reset_lock_result = sqlx::query("RESET lock_timeout") + .execute(&mut *connection) + .await; + let reset_statement_result = sqlx::query("RESET statement_timeout") + .execute(&mut *connection) + .await; + + maintenance_result?; + reset_lock_result.map_err(|error| { + ServiceError::Database(format!( + "Built permission-denied retention index but failed to reset lock timeout: {error}" + )) + })?; + reset_statement_result.map_err(|error| { + ServiceError::Database(format!( + "Built permission-denied retention index but failed to reset statement timeout: {error}" + )) + })?; + + Ok(()) +} + /// Run post-migration backfill for continuous aggregates. /// /// `CALL refresh_continuous_aggregate()` cannot run inside a transaction block, diff --git a/crates/temps-database/src/lib.rs b/crates/temps-database/src/lib.rs index fd619d271..bc933cf85 100644 --- a/crates/temps-database/src/lib.rs +++ b/crates/temps-database/src/lib.rs @@ -8,8 +8,8 @@ pub use approx_count::{approximate_row_count, count_for_pagination, CountKind}; pub use connection::{ cancel_migration_backend, connect_for_migrate, connect_without_migrations, establish_connection, get_pending_migration_names, run_migrations, run_migrations_reported, - run_migrations_streaming, run_post_migration_backfill, DbConnection, MigrationProgress, - MigrationRunReport, MigrationStepResult, + run_migrations_streaming, run_post_migration_backfill, run_post_migration_indexes, + DbConnection, MigrationProgress, MigrationRunReport, MigrationStepResult, }; // Export test utilities for use by other crates in their tests @@ -107,7 +107,7 @@ mod tests { // Retry connection setup let mut retries = 5; - let _connection = loop { + let connection = loop { match establish_connection(&database_url).await { Ok(conn) => break conn, Err(e) if retries > 0 => { @@ -124,6 +124,27 @@ mod tests { } }; + // Heavy audit index creation is deliberately outside the migration + // transaction. Prove the post-migration path is idempotent and creates + // the expected partial concurrent index against a real PostgreSQL. + run_post_migration_indexes(connection.as_ref()).await?; + run_post_migration_indexes(connection.as_ref()).await?; + let index = connection + .query_one(sea_orm::Statement::from_string( + sea_orm::DatabaseBackend::Postgres, + "SELECT indexdef FROM pg_indexes \ + WHERE indexname = 'idx_audit_logs_permission_denied_retention'" + .to_owned(), + )) + .await? + .ok_or_else(|| anyhow::anyhow!("permission-denied retention index was not created"))?; + let index_definition: String = index.try_get("", "indexdef")?; + assert!(index_definition.contains("audit_date")); + // PostgreSQL may render the predicate with type casts and additional + // parentheses, so assert its semantic pieces instead of one formatting. + assert!(index_definition.contains("operation_type")); + assert!(index_definition.contains("PERMISSION_DENIED")); + // If we get here, migrations ran successfully and connection is established println!("✅ Database connection with migrations established successfully"); diff --git a/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs b/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs deleted file mode 100644 index be081563d..000000000 --- a/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs +++ /dev/null @@ -1,52 +0,0 @@ -//! Adds the partial index used by the bounded permission-denial retention -//! worker. Without it, each small delete batch could scan the entire audit -//! history as the table grows, defeating the bounded-maintenance contract. - -use sea_orm_migration::prelude::*; - -const INDEX_NAME: &str = "idx_audit_logs_permission_denied_retention"; - -#[derive(DeriveMigrationName)] -pub struct Migration; - -#[async_trait::async_trait] -impl MigrationTrait for Migration { - async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { - // Migrations run transactionally, so PostgreSQL does not permit - // CONCURRENTLY. Bound both lock acquisition and build time rather than - // allowing startup to hang indefinitely on a busy audit table. - manager - .get_connection() - .execute_unprepared("SET LOCAL lock_timeout = '5s'") - .await?; - manager - .get_connection() - .execute_unprepared("SET LOCAL statement_timeout = '30s'") - .await?; - manager - .get_connection() - .execute_unprepared(&format!( - "CREATE INDEX IF NOT EXISTS {INDEX_NAME} \ - ON audit_logs (audit_date ASC, id ASC) \ - WHERE operation_type = 'PERMISSION_DENIED'" - )) - .await?; - Ok(()) - } - - async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { - manager - .get_connection() - .execute_unprepared("SET LOCAL lock_timeout = '5s'") - .await?; - manager - .get_connection() - .execute_unprepared("SET LOCAL statement_timeout = '30s'") - .await?; - manager - .get_connection() - .execute_unprepared(&format!("DROP INDEX IF EXISTS {INDEX_NAME}")) - .await?; - Ok(()) - } -} diff --git a/crates/temps-migrations/src/migration/mod.rs b/crates/temps-migrations/src/migration/mod.rs index a139d9a74..9ad65cdd3 100644 --- a/crates/temps-migrations/src/migration/mod.rs +++ b/crates/temps-migrations/src/migration/mod.rs @@ -176,7 +176,6 @@ mod m20260804_000001_add_ai_data_access_to_external_services; mod m20260804_000001_add_must_change_password_to_users; pub mod m20260805_000001_index_normalized_managed_domains; pub mod m20260806_000001_sandbox_workspace_lifecycle; -mod m20260806_000001_index_permission_denied_retention; pub struct Migrator; @@ -368,7 +367,6 @@ impl MigratorTrait for Migrator { ), Box::new(m20260805_000001_index_normalized_managed_domains::Migration), Box::new(m20260806_000001_sandbox_workspace_lifecycle::Migration), - Box::new(m20260806_000001_index_permission_denied_retention::Migration), ] } } diff --git a/crates/temps-providers/src/mariadb_query.rs b/crates/temps-providers/src/mariadb_query.rs index f4323725f..eeefca3c7 100644 --- a/crates/temps-providers/src/mariadb_query.rs +++ b/crates/temps-providers/src/mariadb_query.rs @@ -607,22 +607,30 @@ impl Queryable for MariaDbSource { entity: entity_name.to_string(), } })? { - let observed = row.try_get::("__temps_size").map_err(|_error| { + let observed = row.try_get::("__temps_size").map_err(|error| { error!( entity = entity_name, - limit, "MariaDB bounded row size decode failed" + limit, + error = %error, + "MariaDB bounded row size decode failed" ); DataError::BackendQueryFailed { backend: "MariaDB", entity: entity_name.to_string(), } })?; - let observed_cell = row - .try_get::("__temps_max_cell") - .map_err(|_error| DataError::BackendQueryFailed { + let observed_cell = row.try_get::("__temps_max_cell").map_err(|error| { + error!( + entity = entity_name, + limit, + error = %error, + "MariaDB bounded cell size decode failed" + ); + DataError::BackendQueryFailed { backend: "MariaDB", entity: entity_name.to_string(), - })?; + } + })?; let payload = row .try_get::, _>("__temps_row") .map_err(|_error| { @@ -1315,6 +1323,11 @@ pub(crate) fn is_mariadb_compatible_image(image: &str) -> bool { #[cfg(test)] mod tests { use super::*; + use testcontainers::{ + core::{ContainerPort, WaitFor}, + runners::AsyncRunner, + GenericImage, ImageExt, + }; #[test] fn generated_query_guards_encoded_row_before_wire_transfer() { @@ -1370,6 +1383,81 @@ mod tests { assert!(matches!(error, DataError::OperationNotSupported(_))); } + #[tokio::test] + async fn real_mariadb_enforces_wire_budget_and_preserves_rows() -> anyhow::Result<()> { + let container = match GenericImage::new("mariadb", "11.4") + .with_exposed_port(ContainerPort::Tcp(3306)) + .with_wait_for(WaitFor::message_on_stderr("ready for connections")) + .with_env_var("MARIADB_ROOT_PASSWORD", "test") + .with_env_var("MARIADB_DATABASE", "app") + .start() + .await + { + Ok(container) => container, + Err(error) => { + eprintln!("Skipping Docker-dependent MariaDB budget test: {error}"); + return Ok(()); + } + }; + let host = container.get_host().await?.to_string(); + let port = container.get_host_port_ipv4(3306).await?; + let mut source = None; + for _ in 0..20 { + match MariaDbSource::connect(&host, port, "root", "test", "app").await { + Ok(connected) => { + source = Some(connected); + break; + } + Err(_) => tokio::time::sleep(std::time::Duration::from_millis(250)).await, + } + } + let source = source.ok_or_else(|| anyhow::anyhow!("MariaDB did not become reachable"))?; + sqlx::query("CREATE TABLE rows_budget (id BIGINT PRIMARY KEY, payload LONGTEXT NOT NULL)") + .execute(&source.pool) + .await?; + sqlx::query("INSERT INTO rows_budget VALUES (1, 'safe'), (2, REPEAT('x', 100000))") + .execute(&source.pool) + .await?; + + let path = ContainerPath::from_slice(&["app"]); + let safe = source + .query( + &path, + "rows_budget", + Some(serde_json::json!({"where": "id = 1"})), + QueryOptions::default(), + ) + .await?; + assert_eq!(safe.rows[0]["payload"], serde_json::json!("safe")); + + let error = source + .query( + &path, + "rows_budget", + Some(serde_json::json!({"where": "id = 2"})), + QueryOptions { + limit: Some(1), + budget: QueryBudget { + max_bytes: 16 * 1024, + max_cell_bytes: 8 * 1024, + ..QueryBudget::default() + }, + ..QueryOptions::default() + }, + ) + .await + .expect_err("oversized MariaDB row must be rejected before JSON transfer"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "wire_cell_bytes", + .. + } + )); + + Ok(()) + } + fn assert_where_rejected(clause: &str) { assert!( validate_where_clause(clause).is_err(), diff --git a/crates/temps-query-mongodb/Cargo.toml b/crates/temps-query-mongodb/Cargo.toml index 8e44a3c88..7c7400267 100644 --- a/crates/temps-query-mongodb/Cargo.toml +++ b/crates/temps-query-mongodb/Cargo.toml @@ -31,3 +31,6 @@ async-trait = "0.1" # Date/time chrono = { version = "0.4", features = ["serde"] } + +[dev-dependencies] +testcontainers = { workspace = true } diff --git a/crates/temps-query-mongodb/src/lib.rs b/crates/temps-query-mongodb/src/lib.rs index 7827e9e27..2d5089eb1 100644 --- a/crates/temps-query-mongodb/src/lib.rs +++ b/crates/temps-query-mongodb/src/lib.rs @@ -953,6 +953,12 @@ impl temps_query::Queryable for MongoDBSource { #[cfg(test)] mod tests { use super::*; + use temps_query::{QueryBudget, QueryOptions, Queryable}; + use testcontainers::{ + core::{ContainerPort, WaitFor}, + runners::AsyncRunner, + GenericImage, + }; #[test] fn aggregation_sizes_before_conditionally_projecting_document() { @@ -976,6 +982,83 @@ mod tests { assert_eq!(pipeline[3], doc! { "$limit": 10_i64 }); } + #[tokio::test] + async fn real_mongodb_enforces_wire_budget_and_preserves_documents() -> anyhow::Result<()> { + let container = match GenericImage::new("mongo", "8") + .with_exposed_port(ContainerPort::Tcp(27017)) + .with_wait_for(WaitFor::message_on_stdout("Waiting for connections")) + .start() + .await + { + Ok(container) => container, + Err(error) => { + eprintln!("Skipping Docker-dependent MongoDB budget test: {error}"); + return Ok(()); + } + }; + let host = container.get_host().await?; + let port = container.get_host_port_ipv4(27017).await?; + let url = format!("mongodb://{host}:{port}/app"); + let mut source = None; + for _ in 0..20 { + match MongoDBSource::new_scoped(&url, Some("app")).await { + Ok(connected) => { + source = Some(connected); + break; + } + Err(_) => tokio::time::sleep(std::time::Duration::from_millis(250)).await, + } + } + let source = source.ok_or_else(|| anyhow::anyhow!("MongoDB did not become reachable"))?; + source + .client + .database("app") + .collection::("rows_budget") + .insert_many([ + doc! { "row_id": 1_i32, "payload": "safe" }, + doc! { "row_id": 2_i32, "payload": "x".repeat(100_000) }, + ]) + .await?; + + let path = ContainerPath::from_slice(&["app"]); + let safe = source + .query( + &path, + "rows_budget", + Some(serde_json::json!({"row_id": 1})), + QueryOptions::default(), + ) + .await?; + assert_eq!(safe.rows[0]["payload"], serde_json::json!("safe")); + + let error = source + .query( + &path, + "rows_budget", + Some(serde_json::json!({"row_id": 2})), + QueryOptions { + limit: Some(1), + budget: QueryBudget { + max_bytes: 16 * 1024, + max_cell_bytes: 8 * 1024, + ..QueryBudget::default() + }, + ..QueryOptions::default() + }, + ) + .await + .expect_err("oversized MongoDB document must be rejected before transfer"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "wire_document_bytes", + .. + } + )); + + Ok(()) + } + #[test] fn test_source_type() { assert_eq!("mongodb", "mongodb"); diff --git a/crates/temps-query-postgres/src/lib.rs b/crates/temps-query-postgres/src/lib.rs index 1425633ce..5f9e4bec8 100644 --- a/crates/temps-query-postgres/src/lib.rs +++ b/crates/temps-query-postgres/src/lib.rs @@ -52,12 +52,11 @@ fn pg_column_admission(column: &PgQueryColumn, row_budget: usize) -> Result { - return Err(DataError::OperationNotSupported( - "PostgreSQL array columns are not supported by the bounded data browser" - .to_string(), - )) - } + "ARRAY" => format!( + "CASE WHEN {value} IS NULL THEN 4 WHEN PG_COLUMN_COMPRESSION({value}) IS NOT NULL \ + THEN {} ELSE PG_COLUMN_SIZE({value})::bigint * 8 + 64 END", + row_budget.saturating_add(1) + ), "boolean" | "smallint" | "integer" @@ -175,6 +174,9 @@ impl rustls::client::danger::ServerCertVerifier for AcceptAllVerifier { pub struct PostgresSource { client: Arc, database_name: String, + /// `statement_timeout` is session-scoped. Serialize the SET + query pair so + /// concurrent row/count calls cannot inherit one another's timeout. + query_timeout_lock: Arc>, } /// True when `input` begins with `keyword` on a token boundary. @@ -372,6 +374,7 @@ impl PostgresSource { Ok(Self { client: Arc::new(client), database_name: database.to_string(), + query_timeout_lock: Arc::new(tokio::sync::Mutex::new(())), }) } @@ -1982,6 +1985,7 @@ impl Queryable for PostgresSource { // and escape_ident() for identifiers. The database user should be read-only // as defense-in-depth. let client = &self.client; + let _timeout_guard = self.query_timeout_lock.lock().await; // SECURITY / AVAILABILITY: bound the query server-side before running it. // @@ -2146,6 +2150,7 @@ impl Queryable for PostgresSource { } let client = &self.client; + let _timeout_guard = self.query_timeout_lock.lock().await; // SECURITY: bound this the same way `query` is bounded. // @@ -2253,13 +2258,13 @@ mod wire_budget_tests { } #[test] - fn every_array_type_is_rejected_before_json_query_generation() { + fn arrays_use_compression_aware_admission_before_json_generation() { for name in [ "token_tok_live_secret", "large_numeric_array", "custom_type_array", ] { - let error = with_wire_row_budget( + let sql = with_wire_row_budget( "SELECT * FROM public.array_payloads", &[PgQueryColumn { name: name.to_string(), @@ -2267,11 +2272,10 @@ mod wire_budget_tests { }], QueryBudget::default(), ) - .expect_err("array output expansion has no generic stored-size bound"); - let diagnostic = error.to_string(); - assert!(diagnostic.contains("array columns are not supported")); - assert!(!diagnostic.contains("tok_live_secret")); - assert!(!diagnostic.contains("TO_JSONB")); + .expect("array columns should remain browsable through bounded admission"); + assert!(sql.contains("PG_COLUMN_COMPRESSION")); + assert!(sql.contains("PG_COLUMN_SIZE")); + assert_eq!(sql.matches("TO_JSONB(__temps_source)").count(), 1); } } } @@ -3024,6 +3028,29 @@ mod tests { .. } )); + + source + .client + .batch_execute( + "CREATE TABLE array_rows (id bigint, tags text[]); \ + INSERT INTO array_rows VALUES (1, ARRAY['alpha', 'beta']);", + ) + .await + .expect("array fixture should be created"); + let result = source + .query( + &ContainerPath::from_slice(&["postgres", "public"]), + "array_rows", + None, + QueryOptions { + limit: Some(1), + budget, + ..QueryOptions::default() + }, + ) + .await + .expect("bounded PostgreSQL arrays should remain browsable"); + assert_eq!(result.rows[0]["tags"], serde_json::json!(["alpha", "beta"])); } #[tokio::test] diff --git a/crates/temps-query-redis/Cargo.toml b/crates/temps-query-redis/Cargo.toml index 842dfd1e3..d03172af6 100644 --- a/crates/temps-query-redis/Cargo.toml +++ b/crates/temps-query-redis/Cargo.toml @@ -28,3 +28,6 @@ async-trait = "0.1" # Date/time chrono = { version = "0.4", features = ["serde"] } + +[dev-dependencies] +testcontainers = { workspace = true } diff --git a/crates/temps-query-redis/src/lib.rs b/crates/temps-query-redis/src/lib.rs index cd43122cc..038d23fdc 100644 --- a/crates/temps-query-redis/src/lib.rs +++ b/crates/temps-query-redis/src/lib.rs @@ -62,6 +62,96 @@ fn trim_extra(values: &mut Vec, limit: usize) -> bool { } impl RedisSource { + async fn scan_set_page( + conn: &mut ConnectionManager, + key: &str, + offset: usize, + limit: usize, + ) -> Result<(Vec, bool)> { + let mut cursor = 0_u64; + let mut skipped = 0_usize; + let mut values = Vec::with_capacity(limit.saturating_add(1)); + let count_hint = limit.saturating_add(1).clamp(1, 512); + + loop { + let (next_cursor, batch): (u64, Vec) = redis::cmd("SSCAN") + .arg(key) + .arg(cursor) + .arg("COUNT") + .arg(count_hint) + .query_async(conn) + .await + .map_err(|error: RedisError| { + error!(error = %error, "failed to page Redis set"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + + for value in batch { + if skipped < offset { + skipped += 1; + } else if values.len() <= limit { + values.push(value); + } + } + cursor = next_cursor; + if values.len() > limit || cursor == 0 { + break; + } + } + + let has_more = values.len() > limit || cursor != 0; + values.truncate(limit); + Ok((values, has_more)) + } + + async fn scan_hash_page( + conn: &mut ConnectionManager, + key: &str, + offset: usize, + limit: usize, + ) -> Result<(Vec<(String, String)>, bool)> { + let mut cursor = 0_u64; + let mut skipped = 0_usize; + let mut values = Vec::with_capacity(limit.saturating_add(1)); + let count_hint = limit.saturating_add(1).clamp(1, 512); + + loop { + let (next_cursor, batch): (u64, Vec<(String, String)>) = redis::cmd("HSCAN") + .arg(key) + .arg(cursor) + .arg("COUNT") + .arg(count_hint) + .query_async(conn) + .await + .map_err(|error: RedisError| { + error!(error = %error, "failed to page Redis hash"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + + for value in batch { + if skipped < offset { + skipped += 1; + } else if values.len() <= limit { + values.push(value); + } + } + cursor = next_cursor; + if values.len() > limit || cursor == 0 { + break; + } + } + + let has_more = values.len() > limit || cursor != 0; + values.truncate(limit); + Ok((values, has_more)) + } + /// Create a new Redis data source /// /// # Arguments @@ -228,7 +318,7 @@ impl RedisSource { db: i32, key: &str, options: &QueryOptions, - ) -> Result<(DataRow, bool)> { + ) -> Result<(DataRow, bool, Option)> { let mut conn = self.get_db_connection(db).await?; debug!("Getting value for key '{}' in database {}", key, db); @@ -262,33 +352,42 @@ impl RedisSource { DataError::QueryFailed(format!("Failed to get key TTL: {}", e)) })?; - // Reject a huge value before asking Redis to send it. Redis strings and - // collection members may each be hundreds of MiB; paging the collection - // alone does not protect the control plane from one giant member. - let memory_usage: Option = redis::cmd("MEMORY") - .arg("USAGE") - .arg(key) - .query_async(&mut conn) - .await - .map_err(|e: RedisError| { - error!(key_type, error = %e, "failed to inspect Redis value size"); - DataError::BackendQueryFailed { - backend: "Redis", - entity: key.to_string(), - } - })?; - if let Some(observed) = memory_usage { + // A string is one response cell, so reject it by logical length before + // transfer. Collection allocation is deliberately not compared with a + // per-cell limit: large collections remain safely browsable in pages. + if key_type == "string" { + let observed: usize = redis::cmd("STRLEN") + .arg(key) + .query_async(&mut conn) + .await + .map_err(|e: RedisError| { + error!(key_type, error = %e, "failed to inspect Redis string size"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; if observed > options.budget.max_cell_bytes { return Err(DataError::ResultLimitExceeded { entity: key.to_string(), - limit_kind: "redis_value_bytes", + limit_kind: "cell_bytes", limit: options.budget.max_cell_bytes, observed, }); } } - let offset = options.offset.unwrap_or(0); + // Redis SCAN cursors cannot represent an exact item boundary because + // COUNT is only a hint. Expose a stable logical offset cursor and replay + // SCAN to that boundary, which avoids dropping the remainder of a batch. + let offset = match options.cursor.as_deref() { + Some(cursor) => cursor.parse::().map_err(|_| { + DataError::InvalidQuery(format!( + "Invalid Redis value cursor '{cursor}': expected a non-negative offset" + )) + })?, + None => options.offset.unwrap_or(0), + }; let aggregate_limit = redis_aggregate_limit(options); let start = isize::try_from(offset).unwrap_or(isize::MAX); let end = start.saturating_add(isize::try_from(aggregate_limit).unwrap_or(isize::MAX)); @@ -329,21 +428,9 @@ impl RedisSource { serde_json::json!(v) } "set" => { - let (_cursor, mut v): (u64, Vec) = redis::cmd("SSCAN") - .arg(key) - .arg(0) - .arg("COUNT") - .arg(aggregate_limit.saturating_add(1)) - .query_async(&mut conn) - .await - .map_err(|e: RedisError| { - error!(error = %e, "failed to page Redis set"); - DataError::BackendQueryFailed { - backend: "Redis", - entity: key.to_string(), - } - })?; - truncated = trim_extra(&mut v, aggregate_limit); + let (v, has_more) = + Self::scan_set_page(&mut conn, key, offset, aggregate_limit).await?; + truncated = has_more; serde_json::json!(v) } "zset" => { @@ -366,21 +453,9 @@ impl RedisSource { .collect::>()) } "hash" => { - let (_cursor, mut values): (u64, Vec<(String, String)>) = redis::cmd("HSCAN") - .arg(key) - .arg(0) - .arg("COUNT") - .arg(aggregate_limit.saturating_add(1)) - .query_async(&mut conn) - .await - .map_err(|e: RedisError| { - error!(error = %e, "failed to page Redis hash"); - DataError::BackendQueryFailed { - backend: "Redis", - entity: key.to_string(), - } - })?; - truncated = trim_extra(&mut values, aggregate_limit); + let (values, has_more) = + Self::scan_hash_page(&mut conn, key, offset, aggregate_limit).await?; + truncated = has_more; serde_json::json!(values.into_iter().collect::>()) } "stream" => { @@ -399,7 +474,8 @@ impl RedisSource { row.insert("ttl".to_string(), serde_json::Value::Number(ttl.into())); row.insert("value".to_string(), value); - Ok((row, truncated)) + let next_cursor = truncated.then(|| offset.saturating_add(aggregate_limit).to_string()); + Ok((row, truncated, next_cursor)) } /// Get information about a specific key @@ -704,7 +780,8 @@ impl Queryable for RedisSource { } // Get the key value - let (row, value_truncated) = self.get_key_value(db_num, entity_name, &options).await?; + let (row, value_truncated, next_cursor) = + self.get_key_value(db_num, entity_name, &options).await?; let execution_ms = start.elapsed().as_millis() as u64; // Define schema for the result @@ -748,8 +825,8 @@ impl Queryable for RedisSource { row_count: rows.len(), total_rows: Some(1), execution_ms, - has_more: false, - next_cursor: None, + has_more: value_truncated, + next_cursor, truncated: value_truncated || budget_truncated, }, rows, @@ -833,8 +910,11 @@ impl Queryable for RedisSource { #[cfg(test)] mod tests { + use std::collections::HashSet; + use super::*; use temps_query::QueryBudget; + use testcontainers::{core::WaitFor, runners::AsyncRunner, GenericImage}; #[test] fn test_source_type() { @@ -871,4 +951,85 @@ mod tests { assert!(trim_extra(&mut values, redis_aggregate_limit(&options))); assert_eq!(values.len(), 10); } + + #[tokio::test] + async fn large_set_remains_pageable_by_offset() -> anyhow::Result<()> { + let container = match GenericImage::new("redis", "7-alpine") + .with_exposed_port(6379_u16.into()) + .with_wait_for(WaitFor::message_on_stdout("Ready to accept connections")) + .start() + .await + { + Ok(container) => container, + Err(error) => { + eprintln!("Skipping Docker-dependent Redis paging test: {error}"); + return Ok(()); + } + }; + let host = container.get_host().await?; + let port = container.get_host_port_ipv4(6379).await?; + let source = RedisSource::new(&format!("redis://{host}:{port}")).await?; + let mut connection = source.get_db_connection(0).await?; + + for chunk_start in (0..4_000).step_by(200) { + let mut pipeline = redis::pipe(); + for index in chunk_start..chunk_start + 200 { + pipeline + .cmd("SADD") + .arg("large-set") + .arg(format!("member-{index:04}-{}", "x".repeat(300))) + .ignore(); + } + pipeline.query_async::<()>(&mut connection).await?; + } + let allocated: usize = redis::cmd("MEMORY") + .arg("USAGE") + .arg("large-set") + .query_async(&mut connection) + .await?; + assert!(allocated > QueryBudget::default().max_cell_bytes); + + let path = ContainerPath::from_slice(&["0"]); + let page = |offset| QueryOptions { + limit: Some(2), + offset: Some(offset), + ..QueryOptions::default() + }; + let first = source.query(&path, "large-set", None, page(0)).await?; + let second = source.query(&path, "large-set", None, page(2)).await?; + let cursor_page = source + .query( + &path, + "large-set", + None, + QueryOptions { + limit: Some(2), + cursor: first.stats.next_cursor.clone(), + ..QueryOptions::default() + }, + ) + .await?; + + let members = |result: &QueryResult| -> HashSet { + result.rows[0]["value"] + .as_array() + .expect("set value is an array") + .iter() + .filter_map(|value| value.as_str().map(str::to_string)) + .collect() + }; + let first_members = members(&first); + let second_members = members(&second); + let cursor_members = members(&cursor_page); + assert_eq!(first_members.len(), 2); + assert_eq!(second_members.len(), 2); + assert!(first_members.is_disjoint(&second_members)); + assert_eq!(second_members, cursor_members); + assert!(first.stats.has_more); + assert_eq!(first.stats.next_cursor.as_deref(), Some("2")); + assert!(first.stats.truncated); + assert!(second.stats.truncated); + + Ok(()) + } } diff --git a/web/src/components/audit/AuditLogItem.test.ts b/web/src/components/audit/AuditLogItem.test.ts index d8648d6e1..8bd771904 100644 --- a/web/src/components/audit/AuditLogItem.test.ts +++ b/web/src/components/audit/AuditLogItem.test.ts @@ -1,7 +1,9 @@ import { describe, expect, test } from 'bun:test' +import { createElement } from 'react' +import { renderToStaticMarkup } from 'react-dom/server' import { describePermissionDenial } from '@/lib/permission-denial-display' -import { categorize } from './AuditLogItem' +import { AuditLogItemRow, categorize } from './AuditLogItem' describe('permission-denial audit presentation', () => { test('categorizes permission denials as authentication events', () => { @@ -22,4 +24,26 @@ describe('permission-denial audit presentation', () => { test('handles missing optional denial metadata', () => { expect(describePermissionDenial()).toBe('Denied a request') }) + + test('renders a permission-denial row through the component dispatch', () => { + const rendered = renderToStaticMarkup( + createElement(AuditLogItemRow, { + id: 1, + operation_type: 'PERMISSION_DENIED', + audit_date: 0, + data: { + method: 'DELETE', + route: '/projects/{project_id}', + auth_source: 'api_key', + attempt_count: 3, + }, + }) + ) + + expect(rendered).toContain( + 'Denied DELETE /projects/{project_id} for api key (3 attempts)' + ) + expect(rendered).toContain('PERMISSION_DENIED') + expect(rendered).toContain('Auth') + }) }) diff --git a/web/src/pages/AuditLogs.test.ts b/web/src/pages/AuditLogs.test.ts index ba9c4aaae..ab6d95b9e 100644 --- a/web/src/pages/AuditLogs.test.ts +++ b/web/src/pages/AuditLogs.test.ts @@ -1,12 +1,18 @@ import { describe, expect, test } from 'bun:test' -import { PERMISSION_DENIED_FILTER } from '@/lib/audit-operation-filters' +import { buildOperationOptions } from './AuditLogs' describe('audit operation filters', () => { test('offers permission denials in the authentication group', () => { - expect(PERMISSION_DENIED_FILTER).toEqual({ + expect( + buildOperationOptions().find( + (option) => option.value === 'PERMISSION_DENIED' + ) + ).toEqual({ value: 'PERMISSION_DENIED', label: 'Permission Denied', + group: 'Authentication', + keywords: 'PERMISSION_DENIED', }) }) }) diff --git a/web/src/pages/AuditLogs.tsx b/web/src/pages/AuditLogs.tsx index d93455aef..9511ee413 100644 --- a/web/src/pages/AuditLogs.tsx +++ b/web/src/pages/AuditLogs.tsx @@ -301,6 +301,23 @@ const OPERATION_GROUPS: OperationGroup[] = [ }, ] +export function buildOperationOptions(): SearchableSelectOption[] { + const options: SearchableSelectOption[] = [ + { value: ALL_FILTER, label: 'All types' }, + ] + for (const group of OPERATION_GROUPS) { + for (const operation of group.operations) { + options.push({ + value: operation.value, + label: operation.label, + group: group.label, + keywords: operation.value, + }) + } + } + return options +} + const ALL_FILTER = '__all__' export function AuditLogs() { @@ -347,22 +364,7 @@ export function AuditLogs() { const hasFilters = !!dateRange || operation !== ALL_FILTER || selectedUserId !== ALL_FILTER - const operationOptions = useMemo(() => { - const opts: SearchableSelectOption[] = [ - { value: ALL_FILTER, label: 'All types' }, - ] - for (const group of OPERATION_GROUPS) { - for (const op of group.operations) { - opts.push({ - value: op.value, - label: op.label, - group: group.label, - keywords: op.value, - }) - } - } - return opts - }, []) + const operationOptions = useMemo(buildOperationOptions, []) const userOptions = useMemo(() => { const opts: SearchableSelectOption[] = [ From da3314c796dabcbbd048e8e8f471f4906687709d Mon Sep 17 00:00:00 2001 From: David Viejo Date: Fri, 7 Aug 2026 14:47:53 +0200 Subject: [PATCH 09/10] fix(security): enforce pre-wire Redis budgets --- crates/temps-cli/src/commands/migrate.rs | 11 + crates/temps-cli/src/commands/serve/mod.rs | 37 +- crates/temps-database/src/connection.rs | 4 + ...00001_index_permission_denied_retention.rs | 21 + crates/temps-migrations/src/migration/mod.rs | 2 + crates/temps-query-redis/src/lib.rs | 388 +++++++++++++----- 6 files changed, 356 insertions(+), 107 deletions(-) create mode 100644 crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs diff --git a/crates/temps-cli/src/commands/migrate.rs b/crates/temps-cli/src/commands/migrate.rs index 137a21ed2..ff3e71def 100644 --- a/crates/temps-cli/src/commands/migrate.rs +++ b/crates/temps-cli/src/commands/migrate.rs @@ -90,6 +90,17 @@ impl MigrateCommand { "{}", "✓ Database already up to date — no migrations to apply.".green() ); + if !self.dry_run { + // Non-transactional maintenance must remain retryable even + // after every schema migration is already recorded. + temps_database::run_post_migration_indexes(&db).await?; + if let Err(e) = temps_database::run_post_migration_backfill(&db).await { + info!( + "Post-migration backfill skipped/failed (refresh policy will catch up): {e}" + ); + } + println!("{}", "✓ Post-migration maintenance complete.".green()); + } return Ok(()); } diff --git a/crates/temps-cli/src/commands/serve/mod.rs b/crates/temps-cli/src/commands/serve/mod.rs index 8a85d44c5..e377597f1 100644 --- a/crates/temps-cli/src/commands/serve/mod.rs +++ b/crates/temps-cli/src/commands/serve/mod.rs @@ -300,21 +300,34 @@ impl ServeCommand { let rt = tokio::runtime::Runtime::new()?; - // Backfill TimescaleDB continuous aggregates on this long-lived runtime, - // detached. `establish_connection` no longer runs this (it would block - // startup on a slow `CALL`); it's idempotent and the refresh policy - // catches up regardless, so it must not gate the proxy bind. + // Run non-transactional indexes and the TimescaleDB backfill on this + // long-lived runtime, detached. Index creation retries with capped + // backoff until the retention query has its supporting index; the + // idempotent backfill remains best-effort and never gates proxy bind. { - let backfill_db = db.clone(); + let index_db = db.clone(); rt.spawn(async move { - if let Err(e) = - temps_database::run_post_migration_indexes(backfill_db.as_ref()).await - { - tracing::warn!( - "Post-migration index build failed (will retry on next startup): {}", - e - ); + let mut retry_delay = std::time::Duration::from_secs(5); + loop { + match temps_database::run_post_migration_indexes(index_db.as_ref()).await { + Ok(()) => break, + Err(e) => { + tracing::warn!( + "Post-migration index build failed; retrying in {:?}: {}", + retry_delay, + e + ); + tokio::time::sleep(retry_delay).await; + retry_delay = retry_delay + .saturating_mul(2) + .min(std::time::Duration::from_secs(300)); + } + } } + }); + + let backfill_db = db.clone(); + rt.spawn(async move { if let Err(e) = temps_database::run_post_migration_backfill(backfill_db.as_ref()).await { diff --git a/crates/temps-database/src/connection.rs b/crates/temps-database/src/connection.rs index e67492783..ba1223f88 100644 --- a/crates/temps-database/src/connection.rs +++ b/crates/temps-database/src/connection.rs @@ -563,6 +563,10 @@ pub async fn run_post_migration_indexes(db: &DatabaseConnection) -> ServiceResul "Failed to acquire database connection for post-migration indexes: {error}" )) })?; + // Session timeouts cannot be reset asynchronously from Drop if this future + // is cancelled. Treat this as a disposable maintenance connection so no + // cancellation or RESET failure can leak its settings back into the pool. + connection.close_on_drop(); sqlx::query("SET lock_timeout = '5s'") .execute(&mut *connection) diff --git a/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs b/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs new file mode 100644 index 000000000..de1050b0d --- /dev/null +++ b/crates/temps-migrations/src/migration/m20260806_000001_index_permission_denied_retention.rs @@ -0,0 +1,21 @@ +//! Compatibility identity for early builds of PR #413. +//! +//! The index itself is built outside the migration transaction with +//! `CREATE INDEX CONCURRENTLY`; retaining this no-op migration lets databases +//! that briefly applied the original transactional migration continue to boot. + +use sea_orm_migration::prelude::*; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, _manager: &SchemaManager) -> Result<(), DbErr> { + Ok(()) + } + + async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> { + Ok(()) + } +} diff --git a/crates/temps-migrations/src/migration/mod.rs b/crates/temps-migrations/src/migration/mod.rs index 9ad65cdd3..9cdb99dde 100644 --- a/crates/temps-migrations/src/migration/mod.rs +++ b/crates/temps-migrations/src/migration/mod.rs @@ -175,6 +175,7 @@ mod m20260803_000002_add_step_up_expires_at_to_sessions; mod m20260804_000001_add_ai_data_access_to_external_services; mod m20260804_000001_add_must_change_password_to_users; pub mod m20260805_000001_index_normalized_managed_domains; +mod m20260806_000001_index_permission_denied_retention; pub mod m20260806_000001_sandbox_workspace_lifecycle; pub struct Migrator; @@ -367,6 +368,7 @@ impl MigratorTrait for Migrator { ), Box::new(m20260805_000001_index_normalized_managed_domains::Migration), Box::new(m20260806_000001_sandbox_workspace_lifecycle::Migration), + Box::new(m20260806_000001_index_permission_denied_retention::Migration), ] } } diff --git a/crates/temps-query-redis/src/lib.rs b/crates/temps-query-redis/src/lib.rs index 038d23fdc..56c88b58e 100644 --- a/crates/temps-query-redis/src/lib.rs +++ b/crates/temps-query-redis/src/lib.rs @@ -43,6 +43,10 @@ pub struct RedisSource { connection: ConnectionManager, } +/// Bounds the work performed by the atomic SSCAN/HSCAN admission scripts so a +/// data-browser request cannot monopolize the linked Redis server. +const MAX_REDIS_VALUE_OFFSET: usize = 10_000; + fn redis_aggregate_limit(options: &QueryOptions) -> usize { // A JSON array/object consumes at least two structural elements per Redis // item (container entry + scalar), and zset rows consume more. Dividing by @@ -52,59 +56,70 @@ fn redis_aggregate_limit(options: &QueryOptions) -> usize { options.limit.unwrap_or(100).min(structural_limit) } -fn trim_extra(values: &mut Vec, limit: usize) -> bool { - if values.len() > limit { - values.truncate(limit); - true - } else { - false +impl RedisSource { + fn wire_budget_error(key: &str, limit: usize, observed: usize) -> DataError { + DataError::ResultLimitExceeded { + entity: key.to_string(), + limit_kind: "cell_bytes", + limit, + observed, + } } -} -impl RedisSource { async fn scan_set_page( conn: &mut ConnectionManager, key: &str, offset: usize, limit: usize, + max_cell_bytes: usize, ) -> Result<(Vec, bool)> { - let mut cursor = 0_u64; - let mut skipped = 0_usize; - let mut values = Vec::with_capacity(limit.saturating_add(1)); - let count_hint = limit.saturating_add(1).clamp(1, 512); - - loop { - let (next_cursor, batch): (u64, Vec) = redis::cmd("SSCAN") - .arg(key) - .arg(cursor) - .arg("COUNT") - .arg(count_hint) - .query_async(conn) + const SCRIPT: &str = r#" +local cursor = '0' +local skipped = 0 +local values = {} +local estimated = 2 +local has_more = 0 +repeat + local page = redis.call('SSCAN', KEYS[1], cursor, 'COUNT', math.min(tonumber(ARGV[2]) + 1, 512)) + cursor = page[1] + for _, value in ipairs(page[2]) do + if skipped < tonumber(ARGV[1]) then + skipped = skipped + 1 + elseif #values >= tonumber(ARGV[2]) then + has_more = 1 + break + else + local next_estimate = estimated + string.len(value) * 6 + 8 + if next_estimate > tonumber(ARGV[3]) then + return {0, next_estimate, 0, {}} + end + table.insert(values, value) + estimated = next_estimate + end + end +until has_more == 1 or cursor == '0' +if cursor ~= '0' then has_more = 1 end +return {1, estimated, has_more, values} +"#; + let (admitted, observed, has_more, values): (i64, usize, i64, Vec) = + redis::Script::new(SCRIPT) + .key(key) + .arg(offset) + .arg(limit) + .arg(max_cell_bytes) + .invoke_async(conn) .await .map_err(|error: RedisError| { - error!(error = %error, "failed to page Redis set"); + error!(error = %error, "failed to safely page Redis set"); DataError::BackendQueryFailed { backend: "Redis", entity: key.to_string(), } })?; - - for value in batch { - if skipped < offset { - skipped += 1; - } else if values.len() <= limit { - values.push(value); - } - } - cursor = next_cursor; - if values.len() > limit || cursor == 0 { - break; - } + if admitted == 0 { + return Err(Self::wire_budget_error(key, max_cell_bytes, observed)); } - - let has_more = values.len() > limit || cursor != 0; - values.truncate(limit); - Ok((values, has_more)) + Ok((values, has_more != 0)) } async fn scan_hash_page( @@ -112,44 +127,151 @@ impl RedisSource { key: &str, offset: usize, limit: usize, + max_cell_bytes: usize, ) -> Result<(Vec<(String, String)>, bool)> { - let mut cursor = 0_u64; - let mut skipped = 0_usize; - let mut values = Vec::with_capacity(limit.saturating_add(1)); - let count_hint = limit.saturating_add(1).clamp(1, 512); + const SCRIPT: &str = r#" +local cursor = '0' +local skipped = 0 +local values = {} +local estimated = 2 +local has_more = 0 +repeat + local page = redis.call('HSCAN', KEYS[1], cursor, 'COUNT', math.min(tonumber(ARGV[2]) + 1, 512)) + cursor = page[1] + for index = 1, #page[2], 2 do + if skipped < tonumber(ARGV[1]) then + skipped = skipped + 1 + elseif (#values / 2) >= tonumber(ARGV[2]) then + has_more = 1 + break + else + local field = page[2][index] + local value = page[2][index + 1] + local next_estimate = estimated + (string.len(field) + string.len(value)) * 6 + 16 + if next_estimate > tonumber(ARGV[3]) then + return {0, next_estimate, 0, {}} + end + table.insert(values, field) + table.insert(values, value) + estimated = next_estimate + end + end +until has_more == 1 or cursor == '0' +if cursor ~= '0' then has_more = 1 end +return {1, estimated, has_more, values} +"#; + let (admitted, observed, has_more, flat_values): (i64, usize, i64, Vec) = + redis::Script::new(SCRIPT) + .key(key) + .arg(offset) + .arg(limit) + .arg(max_cell_bytes) + .invoke_async(conn) + .await + .map_err(|error: RedisError| { + error!(error = %error, "failed to safely page Redis hash"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + if admitted == 0 { + return Err(Self::wire_budget_error(key, max_cell_bytes, observed)); + } + let values = flat_values + .chunks_exact(2) + .map(|pair| (pair[0].clone(), pair[1].clone())) + .collect(); + Ok((values, has_more != 0)) + } - loop { - let (next_cursor, batch): (u64, Vec<(String, String)>) = redis::cmd("HSCAN") - .arg(key) - .arg(cursor) - .arg("COUNT") - .arg(count_hint) - .query_async(conn) + async fn list_page( + conn: &mut ConnectionManager, + key: &str, + offset: usize, + limit: usize, + max_cell_bytes: usize, + ) -> Result<(Vec, bool)> { + const SCRIPT: &str = r#" +local values = redis.call('LRANGE', KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[1]) + tonumber(ARGV[2])) +local has_more = (#values > tonumber(ARGV[2])) and 1 or 0 +if has_more == 1 then table.remove(values) end +local estimated = 2 +for _, value in ipairs(values) do + estimated = estimated + string.len(value) * 6 + 8 + if estimated > tonumber(ARGV[3]) then return {0, estimated, 0, {}} end +end +return {1, estimated, has_more, values} +"#; + let (admitted, observed, has_more, values): (i64, usize, i64, Vec) = + redis::Script::new(SCRIPT) + .key(key) + .arg(offset) + .arg(limit) + .arg(max_cell_bytes) + .invoke_async(conn) .await .map_err(|error: RedisError| { - error!(error = %error, "failed to page Redis hash"); + error!(error = %error, "failed to safely page Redis list"); DataError::BackendQueryFailed { backend: "Redis", entity: key.to_string(), } })?; + if admitted == 0 { + return Err(Self::wire_budget_error(key, max_cell_bytes, observed)); + } + Ok((values, has_more != 0)) + } - for value in batch { - if skipped < offset { - skipped += 1; - } else if values.len() <= limit { - values.push(value); + async fn zset_page( + conn: &mut ConnectionManager, + key: &str, + offset: usize, + limit: usize, + max_cell_bytes: usize, + ) -> Result<(Vec<(String, f64)>, bool)> { + const SCRIPT: &str = r#" +local values = redis.call('ZRANGE', KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[1]) + tonumber(ARGV[2]), 'WITHSCORES') +local has_more = ((#values / 2) > tonumber(ARGV[2])) and 1 or 0 +if has_more == 1 then table.remove(values); table.remove(values) end +local estimated = 2 +for index = 1, #values, 2 do + estimated = estimated + (string.len(values[index]) + string.len(values[index + 1])) * 6 + 24 + if estimated > tonumber(ARGV[3]) then return {0, estimated, 0, {}} end +end +return {1, estimated, has_more, values} +"#; + let (admitted, observed, has_more, flat_values): (i64, usize, i64, Vec) = + redis::Script::new(SCRIPT) + .key(key) + .arg(offset) + .arg(limit) + .arg(max_cell_bytes) + .invoke_async(conn) + .await + .map_err(|error: RedisError| { + error!(error = %error, "failed to safely page Redis sorted set"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), + } + })?; + if admitted == 0 { + return Err(Self::wire_budget_error(key, max_cell_bytes, observed)); + } + let mut values = Vec::with_capacity(flat_values.len() / 2); + for pair in flat_values.chunks_exact(2) { + let score = pair[1].parse::().map_err(|error| { + error!(error = %error, "failed to decode Redis sorted-set score"); + DataError::BackendQueryFailed { + backend: "Redis", + entity: key.to_string(), } - } - cursor = next_cursor; - if values.len() > limit || cursor == 0 { - break; - } + })?; + values.push((pair[0].clone(), score)); } - - let has_more = values.len() > limit || cursor != 0; - values.truncate(limit); - Ok((values, has_more)) + Ok((values, has_more != 0)) } /// Create a new Redis data source @@ -388,9 +510,12 @@ impl RedisSource { })?, None => options.offset.unwrap_or(0), }; + if offset > MAX_REDIS_VALUE_OFFSET { + return Err(DataError::InvalidQuery(format!( + "Redis value offset {offset} exceeds maximum {MAX_REDIS_VALUE_OFFSET}" + ))); + } let aggregate_limit = redis_aggregate_limit(options); - let start = isize::try_from(offset).unwrap_or(isize::MAX); - let end = start.saturating_add(isize::try_from(aggregate_limit).unwrap_or(isize::MAX)); let mut truncated = false; // Get value based on type @@ -414,37 +539,39 @@ impl RedisSource { serde_json::Value::String(String::from_utf8_lossy(&v).into_owned()) } "list" => { - let mut v: Vec = - conn.lrange(key, start, end) - .await - .map_err(|e: RedisError| { - error!("Failed to get list value: {}", e); - DataError::BackendQueryFailed { - backend: "Redis", - entity: key.to_string(), - } - })?; - truncated = trim_extra(&mut v, aggregate_limit); + let (v, has_more) = Self::list_page( + &mut conn, + key, + offset, + aggregate_limit, + options.budget.max_cell_bytes, + ) + .await?; + truncated = has_more; serde_json::json!(v) } "set" => { - let (v, has_more) = - Self::scan_set_page(&mut conn, key, offset, aggregate_limit).await?; + let (v, has_more) = Self::scan_set_page( + &mut conn, + key, + offset, + aggregate_limit, + options.budget.max_cell_bytes, + ) + .await?; truncated = has_more; serde_json::json!(v) } "zset" => { - let mut v: Vec<(String, f64)> = conn - .zrange_withscores(key, start, end) - .await - .map_err(|e: RedisError| { - error!("Failed to get sorted set value: {}", e); - DataError::BackendQueryFailed { - backend: "Redis", - entity: key.to_string(), - } - })?; - truncated = trim_extra(&mut v, aggregate_limit); + let (v, has_more) = Self::zset_page( + &mut conn, + key, + offset, + aggregate_limit, + options.budget.max_cell_bytes, + ) + .await?; + truncated = has_more; serde_json::json!(v .into_iter() .map(|(member, score)| { @@ -453,8 +580,14 @@ impl RedisSource { .collect::>()) } "hash" => { - let (values, has_more) = - Self::scan_hash_page(&mut conn, key, offset, aggregate_limit).await?; + let (values, has_more) = Self::scan_hash_page( + &mut conn, + key, + offset, + aggregate_limit, + options.budget.max_cell_bytes, + ) + .await?; truncated = has_more; serde_json::json!(values.into_iter().collect::>()) } @@ -946,10 +1079,6 @@ mod tests { ..QueryOptions::default() }; assert_eq!(redis_aggregate_limit(&options), 10); - - let mut values = vec!["value"; 11]; - assert!(trim_extra(&mut values, redis_aggregate_limit(&options))); - assert_eq!(values.len(), 10); } #[tokio::test] @@ -1030,6 +1159,75 @@ mod tests { assert!(first.stats.truncated); assert!(second.stats.truncated); + let offset_error = source + .query( + &path, + "large-set", + None, + QueryOptions { + cursor: Some((MAX_REDIS_VALUE_OFFSET + 1).to_string()), + ..QueryOptions::default() + }, + ) + .await + .expect_err("Redis aggregate script work must be offset-bounded"); + assert!(matches!(offset_error, DataError::InvalidQuery(_))); + + let oversized = "y".repeat(100_000); + redis::cmd("RPUSH") + .arg("oversized-list") + .arg(&oversized) + .query_async::<()>(&mut connection) + .await?; + redis::cmd("SADD") + .arg("oversized-set") + .arg(&oversized) + .query_async::<()>(&mut connection) + .await?; + redis::cmd("ZADD") + .arg("oversized-zset") + .arg(1) + .arg(&oversized) + .query_async::<()>(&mut connection) + .await?; + redis::cmd("HSET") + .arg("oversized-hash") + .arg("field") + .arg(&oversized) + .query_async::<()>(&mut connection) + .await?; + + for key in [ + "oversized-list", + "oversized-set", + "oversized-zset", + "oversized-hash", + ] { + let error = source + .query( + &path, + key, + None, + QueryOptions { + limit: Some(2), + budget: QueryBudget { + max_cell_bytes: 8 * 1024, + ..QueryBudget::default() + }, + ..QueryOptions::default() + }, + ) + .await + .expect_err("oversized Redis aggregate element must be rejected in Lua"); + assert!(matches!( + error, + DataError::ResultLimitExceeded { + limit_kind: "cell_bytes", + .. + } + )); + } + Ok(()) } } From 9ff7df5d34a520b8c494b85f977c222e619f0137 Mon Sep 17 00:00:00 2001 From: David Viejo Date: Fri, 7 Aug 2026 15:01:58 +0200 Subject: [PATCH 10/10] fix(redis): bound aggregate admission work --- .../src/permission_denial_recorder.rs | 15 +- crates/temps-query-redis/src/lib.rs | 161 ++++++++++++------ 2 files changed, 124 insertions(+), 52 deletions(-) diff --git a/crates/temps-auth/src/permission_denial_recorder.rs b/crates/temps-auth/src/permission_denial_recorder.rs index 086aa080f..4fe44f57a 100644 --- a/crates/temps-auth/src/permission_denial_recorder.rs +++ b/crates/temps-auth/src/permission_denial_recorder.rs @@ -384,6 +384,7 @@ mod tests { #[derive(Default)] struct RecordingLogger { records: Mutex>, + user_agents: Mutex>, } #[async_trait::async_trait] @@ -395,6 +396,10 @@ mod tests { .lock() .map_err(|_| anyhow::anyhow!("test logger lock poisoned"))? .push(value); + self.user_agents + .lock() + .map_err(|_| anyhow::anyhow!("test user-agent lock poisoned"))? + .push(operation.user_agent().to_string()); Ok(()) } } @@ -451,7 +456,10 @@ mod tests { assert_eq!(records[0]["auth_source"], "api_key"); assert_eq!(records[0]["credential_id"], 9); assert_eq!(records[0]["ip_address"], serde_json::Value::Null); - assert_eq!(records[0]["user_agent"], MIXED_VALUE); + assert_eq!( + logger.user_agents.lock().expect("test user-agent lock")[0], + MIXED_VALUE + ); assert_eq!(records[0]["multiple_origins"], true); assert!(records[0].get("key_name").is_none()); assert!(records[0].get("token_name").is_none()); @@ -470,7 +478,10 @@ mod tests { let records = logger.records.lock().expect("test logger lock"); assert_eq!(records.len(), 1); assert_eq!(records[0]["ip_address"], serde_json::Value::Null); - assert_eq!(records[0]["user_agent"], MIXED_VALUE); + assert_eq!( + logger.user_agents.lock().expect("test user-agent lock")[0], + MIXED_VALUE + ); assert_eq!(records[0]["multiple_origins"], true); } diff --git a/crates/temps-query-redis/src/lib.rs b/crates/temps-query-redis/src/lib.rs index 56c88b58e..0235f8974 100644 --- a/crates/temps-query-redis/src/lib.rs +++ b/crates/temps-query-redis/src/lib.rs @@ -45,15 +45,29 @@ pub struct RedisSource { /// Bounds the work performed by the atomic SSCAN/HSCAN admission scripts so a /// data-browser request cannot monopolize the linked Redis server. -const MAX_REDIS_VALUE_OFFSET: usize = 10_000; +const MAX_REDIS_VALUE_OFFSET: usize = 1_000; +const MAX_REDIS_AGGREGATE_PAGE_ITEMS: usize = 100; -fn redis_aggregate_limit(options: &QueryOptions) -> usize { +fn redis_aggregate_limit(options: &QueryOptions) -> Result { // A JSON array/object consumes at least two structural elements per Redis // item (container entry + scalar), and zset rows consume more. Dividing by // four is deliberately conservative and keeps the shared structural // limiter from accepting a backend page that was already too large. let structural_limit = (options.budget.max_value_elements_per_row / 4).max(1); - options.limit.unwrap_or(100).min(structural_limit) + let requested = options.limit.unwrap_or(100); + if requested == 0 { + return Err(DataError::InvalidQuery( + "Redis aggregate page limit must be at least 1".to_string(), + )); + } + Ok(requested + .min(structural_limit) + .min(MAX_REDIS_AGGREGATE_PAGE_ITEMS)) +} + +fn redis_next_value_cursor(offset: usize, limit: usize, truncated: bool) -> Option { + let next = offset.saturating_add(limit); + (truncated && next.saturating_add(limit) <= MAX_REDIS_VALUE_OFFSET).then(|| next.to_string()) } impl RedisSource { @@ -79,8 +93,13 @@ local skipped = 0 local values = {} local estimated = 2 local has_more = 0 +local encoding = redis.call('OBJECT', 'ENCODING', KEYS[1]) +if encoding == 'listpack' or encoding == 'ziplist' then + local allocated = redis.call('MEMORY', 'USAGE', KEYS[1]) or 0 + if allocated > tonumber(ARGV[3]) then return {0, allocated, 0, {}} end +end repeat - local page = redis.call('SSCAN', KEYS[1], cursor, 'COUNT', math.min(tonumber(ARGV[2]) + 1, 512)) + local page = redis.call('SSCAN', KEYS[1], cursor, 'COUNT', 1) cursor = page[1] for _, value in ipairs(page[2]) do if skipped < tonumber(ARGV[1]) then @@ -135,8 +154,13 @@ local skipped = 0 local values = {} local estimated = 2 local has_more = 0 +local encoding = redis.call('OBJECT', 'ENCODING', KEYS[1]) +if encoding == 'listpack' or encoding == 'ziplist' then + local allocated = redis.call('MEMORY', 'USAGE', KEYS[1]) or 0 + if allocated > tonumber(ARGV[3]) then return {0, allocated, 0, {}} end +end repeat - local page = redis.call('HSCAN', KEYS[1], cursor, 'COUNT', math.min(tonumber(ARGV[2]) + 1, 512)) + local page = redis.call('HSCAN', KEYS[1], cursor, 'COUNT', 1) cursor = page[1] for index = 1, #page[2], 2 do if skipped < tonumber(ARGV[1]) then @@ -193,13 +217,16 @@ return {1, estimated, has_more, values} max_cell_bytes: usize, ) -> Result<(Vec, bool)> { const SCRIPT: &str = r#" -local values = redis.call('LRANGE', KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[1]) + tonumber(ARGV[2])) -local has_more = (#values > tonumber(ARGV[2])) and 1 or 0 -if has_more == 1 then table.remove(values) end +local values = {} local estimated = 2 -for _, value in ipairs(values) do +local has_more = 0 +for index = 0, tonumber(ARGV[2]) do + local value = redis.call('LINDEX', KEYS[1], tonumber(ARGV[1]) + index) + if not value then break end + if index >= tonumber(ARGV[2]) then has_more = 1; break end estimated = estimated + string.len(value) * 6 + 8 if estimated > tonumber(ARGV[3]) then return {0, estimated, 0, {}} end + table.insert(values, value) end return {1, estimated, has_more, values} "#; @@ -232,13 +259,17 @@ return {1, estimated, has_more, values} max_cell_bytes: usize, ) -> Result<(Vec<(String, f64)>, bool)> { const SCRIPT: &str = r#" -local values = redis.call('ZRANGE', KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[1]) + tonumber(ARGV[2]), 'WITHSCORES') -local has_more = ((#values / 2) > tonumber(ARGV[2])) and 1 or 0 -if has_more == 1 then table.remove(values); table.remove(values) end +local values = {} local estimated = 2 -for index = 1, #values, 2 do - estimated = estimated + (string.len(values[index]) + string.len(values[index + 1])) * 6 + 24 +local has_more = 0 +for index = 0, tonumber(ARGV[2]) do + local pair = redis.call('ZRANGE', KEYS[1], tonumber(ARGV[1]) + index, tonumber(ARGV[1]) + index, 'WITHSCORES') + if #pair == 0 then break end + if index >= tonumber(ARGV[2]) then has_more = 1; break end + estimated = estimated + (string.len(pair[1]) + string.len(pair[2])) * 6 + 24 if estimated > tonumber(ARGV[3]) then return {0, estimated, 0, {}} end + table.insert(values, pair[1]) + table.insert(values, pair[2]) end return {1, estimated, has_more, values} "#; @@ -515,7 +546,8 @@ return {1, estimated, has_more, values} "Redis value offset {offset} exceeds maximum {MAX_REDIS_VALUE_OFFSET}" ))); } - let aggregate_limit = redis_aggregate_limit(options); + let aggregate_limit = redis_aggregate_limit(options)?; + let wire_cap = options.budget.max_cell_bytes.min(options.budget.max_bytes); let mut truncated = false; // Get value based on type @@ -539,38 +571,20 @@ return {1, estimated, has_more, values} serde_json::Value::String(String::from_utf8_lossy(&v).into_owned()) } "list" => { - let (v, has_more) = Self::list_page( - &mut conn, - key, - offset, - aggregate_limit, - options.budget.max_cell_bytes, - ) - .await?; + let (v, has_more) = + Self::list_page(&mut conn, key, offset, aggregate_limit, wire_cap).await?; truncated = has_more; serde_json::json!(v) } "set" => { - let (v, has_more) = Self::scan_set_page( - &mut conn, - key, - offset, - aggregate_limit, - options.budget.max_cell_bytes, - ) - .await?; + let (v, has_more) = + Self::scan_set_page(&mut conn, key, offset, aggregate_limit, wire_cap).await?; truncated = has_more; serde_json::json!(v) } "zset" => { - let (v, has_more) = Self::zset_page( - &mut conn, - key, - offset, - aggregate_limit, - options.budget.max_cell_bytes, - ) - .await?; + let (v, has_more) = + Self::zset_page(&mut conn, key, offset, aggregate_limit, wire_cap).await?; truncated = has_more; serde_json::json!(v .into_iter() @@ -580,14 +594,8 @@ return {1, estimated, has_more, values} .collect::>()) } "hash" => { - let (values, has_more) = Self::scan_hash_page( - &mut conn, - key, - offset, - aggregate_limit, - options.budget.max_cell_bytes, - ) - .await?; + let (values, has_more) = + Self::scan_hash_page(&mut conn, key, offset, aggregate_limit, wire_cap).await?; truncated = has_more; serde_json::json!(values.into_iter().collect::>()) } @@ -607,7 +615,7 @@ return {1, estimated, has_more, values} row.insert("ttl".to_string(), serde_json::Value::Number(ttl.into())); row.insert("value".to_string(), value); - let next_cursor = truncated.then(|| offset.saturating_add(aggregate_limit).to_string()); + let next_cursor = redis_next_value_cursor(offset, aggregate_limit, truncated); Ok((row, truncated, next_cursor)) } @@ -915,6 +923,7 @@ impl Queryable for RedisSource { // Get the key value let (row, value_truncated, next_cursor) = self.get_key_value(db_num, entity_name, &options).await?; + let has_more = next_cursor.is_some(); let execution_ms = start.elapsed().as_millis() as u64; // Define schema for the result @@ -958,7 +967,7 @@ impl Queryable for RedisSource { row_count: rows.len(), total_rows: Some(1), execution_ms, - has_more: value_truncated, + has_more, next_cursor, truncated: value_truncated || budget_truncated, }, @@ -1078,7 +1087,22 @@ mod tests { }, ..QueryOptions::default() }; - assert_eq!(redis_aggregate_limit(&options), 10); + assert_eq!(redis_aggregate_limit(&options).expect("valid limit"), 10); + + let zero = QueryOptions { + limit: Some(0), + ..QueryOptions::default() + }; + assert!(matches!( + redis_aggregate_limit(&zero), + Err(DataError::InvalidQuery(_)) + )); + + assert_eq!(redis_next_value_cursor(0, 2, true).as_deref(), Some("2")); + assert_eq!( + redis_next_value_cursor(MAX_REDIS_VALUE_OFFSET - 100, 100, true), + None + ); } #[tokio::test] @@ -1173,6 +1197,43 @@ mod tests { .expect_err("Redis aggregate script work must be offset-bounded"); assert!(matches!(offset_error, DataError::InvalidQuery(_))); + let page_budget_error = source + .query( + &path, + "large-set", + None, + QueryOptions { + limit: Some(2), + budget: QueryBudget { + max_bytes: 64, + max_cell_bytes: 8 * 1024, + ..QueryBudget::default() + }, + ..QueryOptions::default() + }, + ) + .await + .expect_err("Redis admission must enforce the smaller page budget pre-wire"); + assert!(matches!( + page_budget_error, + DataError::ResultLimitExceeded { + limit_kind: "cell_bytes", + limit: 64, + .. + } + )); + + redis::cmd("HSET") + .arg("compact-hash") + .arg("field") + .arg("safe") + .query_async::<()>(&mut connection) + .await?; + let compact_hash = source + .query(&path, "compact-hash", None, QueryOptions::default()) + .await?; + assert_eq!(compact_hash.rows[0]["value"]["field"], "safe"); + let oversized = "y".repeat(100_000); redis::cmd("RPUSH") .arg("oversized-list")