From fcc8616492ae063f22bd57be2838ef4ada7e8443 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 12:28:21 +0800 Subject: [PATCH 01/21] =?UTF-8?q?feat(audit):=20phase=201=20=E2=80=94=20pe?= =?UTF-8?q?r-user=20attribution,=20security-event=20stream,=20admin-op=20a?= =?UTF-8?q?udit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add owner to keys and a user_id/request_id/created_at dimension to the ledger (key.owner authoritative, else request metadata: x-gw-user / OpenAI user / Anthropic metadata.user_id). Per-user usage query + /admin/usage/users. A content-safety event stream (no prompt text) with /admin/audit/events, and an admin-operation audit trail with /admin/audit/ops. Keystore list paginated. --- crates/config/src/lib.rs | 5 + crates/dag/src/executor.rs | 1 + crates/dag/src/nodes.rs | 9 + crates/handler/src/lib.rs | 58 ++++ crates/models/src/request.rs | 7 + crates/state/src/keystore.rs | 29 +- crates/state/src/lib.rs | 15 +- crates/state/src/store.rs | 584 ++++++++++++++++++++++++++++++++++- crates/views/src/lib.rs | 333 +++++++++++++++++++- 9 files changed, 997 insertions(+), 44 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 5cf88ab..30034a7 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -64,6 +64,11 @@ pub struct AkConf { /// Tenant this key belongs to; empty = the implicit `default` tenant. #[serde(default)] pub tenant: String, + /// End user this key is issued to, when one key = one user (the enterprise + /// per-employee model); `None` when the key is shared and per-user + /// attribution rides on request metadata instead. + #[serde(default)] + pub owner: Option, pub qps: f64, pub daily_token_quota: i64, /// tokens-per-minute window limit; None = unlimited. diff --git a/crates/dag/src/executor.rs b/crates/dag/src/executor.rs index 7f9b3fc..e6e0412 100644 --- a/crates/dag/src/executor.rs +++ b/crates/dag/src/executor.rs @@ -137,6 +137,7 @@ mod tests { ak: "t".into(), product: "demo".into(), tenant: "default".into(), + owner: None, qps: 10.0, daily_token_quota: 1000, tokens_per_minute: None, diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 996cc46..5ec4d08 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -583,6 +583,13 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> let requested = param .and_then(|p| p.fallback_from.as_deref()) .unwrap_or(served); + // effective user: the key's owner wins (authoritative), else request metadata + let user_id = ctx + .ak + .owner + .as_deref() + .or(ctx.request.user_id.as_deref()) + .unwrap_or_default(); let record = admission::settle_and_bill( ctx.state.governance.as_ref(), ctx.state.store.as_ref(), @@ -592,6 +599,8 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> ak: &ctx.ak.ak, product: &ctx.ak.product, tenant: &ctx.ak.tenant, + user_id, + request_id: &ctx.request.request_id, requested_model: requested, served_model: served, protocol: param.map(|p| p.protocol.as_str()).unwrap_or_default(), diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 7cab625..8f6aff4 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -19,6 +19,57 @@ use gw_state::{AkInfo, GatewayState, SharedConfig}; pub use offline::{BatchItem, OfflineHandler}; +use std::sync::atomic::{AtomicU64, Ordering}; + +static REQ_SEQ: AtomicU64 = AtomicU64::new(0); + +/// Record a content-safety outcome (no prompt text) against this request's +/// key/user/tenant. Best-effort: a store failure is logged, never fails the +/// request. Surface = the request protocol, or "batch" for offline items. +async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i64) { + let user_id = ctx + .ak + .owner + .as_deref() + .or(ctx.request.user_id.as_deref()) + .unwrap_or_default() + .to_owned(); + let surface = if ctx.request.is_online { + ctx.request + .model_param_v2 + .as_ref() + .map(|p| p.protocol.as_str()) + .unwrap_or_default() + .to_owned() + } else { + "batch".to_owned() + }; + let event = gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id, + tenant: ctx.ak.tenant.clone(), + surface, + rule: rule.to_owned(), + action: action.to_owned(), + hits, + }; + if let Err(e) = ctx.state.store.security_event_add(&event).await { + tracing::warn!(error = %e, "security event write failed"); + } +} + +/// A per-request correlation id: `req--`, time-sortable and +/// unique within the process (the seq disambiguates same-millisecond requests). +pub fn new_request_id() -> String { + let ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis()) + .unwrap_or(0); + format!("req-{ms}-{}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)) +} + /// Runs one request through the plugin pre-stage, the DAG, and the plugin post-stage. #[derive(Clone)] pub struct OnlineHandler { @@ -67,6 +118,9 @@ impl OnlineHandler { /// Run one request: plugin pre → DAG (4 layers) → plugin post. /// The returned context carries the outcome, decision log, billing effects. pub async fn run(&self, mut request: GatewayRequest, ak: AkInfo) -> GResult { + if request.request_id.is_empty() { + request.request_id = new_request_id(); + } // one consistent snapshot for the whole request let snap = self.config.load(); // outbound DLP is a response-buffering boundary: a masked span can @@ -101,6 +155,7 @@ impl OnlineHandler { block, ..Default::default() }); + emit_security_event(&ctx, "blocklist", "block", 1).await; return Ok(ctx); } let redacted = plugins::dlp_redact_request(&snap.cfg.security, &mut request); @@ -114,6 +169,7 @@ impl OnlineHandler { ); if redacted > 0 { ctx.decide("dlp", format!("redacted {redacted} span(s) inbound")); + emit_security_event(&ctx, "dlp", "redact", redacted as i64).await; } // a panicking node must refund too, not leak the reserves; unwind-safe — @@ -161,6 +217,7 @@ impl OnlineHandler { }; if redacted_out > 0 { ctx.decide("dlp", format!("redacted {redacted_out} span(s) outbound")); + emit_security_event(&ctx, "dlp", "redact_out", redacted_out as i64).await; } Ok(ctx) } @@ -558,6 +615,7 @@ mod tests { ak: "ak-t1".into(), product: "p".into(), tenant: "t1".into(), + owner: None, qps: 10.0, daily_token_quota: 100_000, tokens_per_minute: None, diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index b9525e8..80f9a61 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -12,6 +12,13 @@ pub struct GatewayRequest { pub model_param_v2: Option, pub ak: String, pub is_online: bool, + /// End-user attribution from request metadata (OpenAI `user` / Anthropic + /// `metadata.user_id` / `x-gw-user`); only trusted when the key is shared — + /// a key's own `owner` overrides it at billing. + pub user_id: Option, + /// Correlation id assigned at ingress; joins access log, ledger, and audit + /// events for one request. Empty until the handler stamps it. + pub request_id: String, /// When set, a streaming-capable engine forwards chunks here as they arrive /// instead of buffering; the bounded channel is the backpressure seam. pub stream_tx: Option>, diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index 477be82..1b418e3 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -31,8 +31,9 @@ pub trait KeyStore: Send + Sync + std::fmt::Debug { async fn patch(&self, ak: &str, patch: &KeyPatch) -> GResult>; /// Remove a key regardless of source; whether it existed. async fn revoke(&self, ak: &str) -> GResult; - /// Every key, sorted by ak for stable listings. - async fn list(&self) -> GResult>; + /// A page of keys, sorted by ak (stable). `offset`/`limit` bound the scan so + /// a fleet key table with millions of rows never loads whole. + async fn list(&self, offset: usize, limit: usize) -> GResult>; /// Re-apply the config file's key set, leaving admin-created keys untouched. async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> GResult<()>; } @@ -67,11 +68,16 @@ impl PostgresKeyStore { expires_at_epoch_secs BIGINT, banned BOOLEAN NOT NULL DEFAULT FALSE, model_quotas TEXT NOT NULL DEFAULT '{}', + owner TEXT, source TEXT NOT NULL DEFAULT 'admin')", ) .execute(&pool) .await .map_err(|e| crate::sqlx_err("create access_keys schema", e))?; + sqlx::query("ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS owner TEXT") + .execute(&pool) + .await + .map_err(|e| crate::sqlx_err("migrate access_keys owner", e))?; Ok(Self { pool, cache: moka::sync::Cache::builder() @@ -91,7 +97,7 @@ impl PostgresKeyStore { async fn fetch(&self, ak: &str) -> Result, sqlx::Error> { let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys WHERE ak = $1", + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys WHERE ak = $1", ) .bind(ak) .fetch_optional(&self.pool) @@ -139,7 +145,7 @@ impl KeyStore for PostgresKeyStore { .map_err(|e| crate::sqlx_err("begin patch", e))?; let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys WHERE ak = $1 FOR UPDATE", ) .bind(ak) @@ -181,11 +187,14 @@ impl KeyStore for PostgresKeyStore { Ok(n > 0) } - async fn list(&self) -> GResult> { + async fn list(&self, offset: usize, limit: usize) -> GResult> { let rows = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys ORDER BY ak", + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys + ORDER BY ak LIMIT $1 OFFSET $2", ) + .bind(limit.min(i64::MAX as usize) as i64) + .bind(offset.min(i64::MAX as usize) as i64) .fetch_all(&self.pool) .await .map_err(|e| crate::sqlx_err("list keys", e))?; @@ -229,14 +238,15 @@ async fn upsert( let quotas = serde_json::to_string(&*info.model_quotas).unwrap_or_else(|_| "{}".into()); sqlx::query( "INSERT INTO access_keys (ak, product, tenant, qps, daily_token_quota, - tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, source) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, owner, source) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) ON CONFLICT (ak) DO UPDATE SET product = EXCLUDED.product, tenant = EXCLUDED.tenant, qps = EXCLUDED.qps, daily_token_quota = EXCLUDED.daily_token_quota, tokens_per_minute = EXCLUDED.tokens_per_minute, expires_at_epoch_secs = EXCLUDED.expires_at_epoch_secs, banned = EXCLUDED.banned, model_quotas = EXCLUDED.model_quotas, + owner = EXCLUDED.owner, source = CASE WHEN access_keys.source = 'config' AND EXCLUDED.source = 'admin' THEN 'config' ELSE EXCLUDED.source END", ) @@ -249,6 +259,7 @@ async fn upsert( .bind(info.expires_at_epoch_secs) .bind(info.banned) .bind("as) + .bind(&info.owner) .bind(source_str(source)) .execute(exec) .await @@ -268,6 +279,7 @@ fn row_to_info(row: &sqlx::postgres::PgRow) -> AkInfo { model_quotas: std::sync::Arc::new( serde_json::from_str(row.get::<&str, _>(8)).unwrap_or_default(), ), + owner: row.get(9), } } @@ -287,6 +299,7 @@ mod tests { ak: ak.into(), product: "p".into(), tenant: "default".into(), + owner: None, qps, daily_token_quota: 10, tokens_per_minute: None, diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 734b4a6..c5caac6 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -35,6 +35,9 @@ pub struct AkInfo { pub product: String, /// Tenant this key belongs to (`gw_config::DEFAULT_TENANT` when undeclared). pub tenant: String, + /// End user this key is issued to; `Some` = authoritative user attribution + /// (one key per user), `None` = shared key, attribute from request metadata. + pub owner: Option, pub qps: f64, pub daily_token_quota: i64, /// tokens-per-minute window cap; None = unlimited. @@ -97,6 +100,7 @@ impl From<&gw_config::AkConf> for AkInfo { ak: k.ak.clone(), product: k.product.clone(), tenant: k.tenant.clone(), + owner: k.owner.clone(), qps: k.qps, daily_token_quota: k.daily_token_quota, tokens_per_minute: k.tokens_per_minute, @@ -165,11 +169,11 @@ impl AkAuth { Some(e.0.clone()) } - /// Every key in the table, sorted by ak for stable listings. - pub fn list(&self) -> Vec { + /// A page of keys, sorted by ak (stable), `offset..offset+limit`. + pub fn list(&self, offset: usize, limit: usize) -> Vec { let mut keys: Vec = self.keys.iter().map(|e| e.value().0.clone()).collect(); keys.sort_by(|a, b| a.ak.cmp(&b.ak)); - keys + keys.into_iter().skip(offset).take(limit).collect() } /// Remove a key regardless of source; returns whether it existed. @@ -205,8 +209,8 @@ impl KeyStore for AkAuth { async fn revoke(&self, ak: &str) -> gw_models::GResult { Ok(AkAuth::revoke(self, ak)) } - async fn list(&self) -> gw_models::GResult> { - Ok(AkAuth::list(self)) + async fn list(&self, offset: usize, limit: usize) -> gw_models::GResult> { + Ok(AkAuth::list(self, offset, limit)) } async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> gw_models::GResult<()> { AkAuth::reload_config_keys(self, keys); @@ -850,6 +854,7 @@ mod tests { ak: ak.into(), product: "p".into(), tenant: "default".into(), + owner: None, qps: 1.0, daily_token_quota: 10, tokens_per_minute: None, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4f83d15..99c4031 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -33,6 +33,16 @@ pub struct BillingRecord { pub ak: String, pub product: String, pub tenant: String, + /// Effective end user: the key's `owner` if set, else request metadata; empty + /// when neither is present. The precise per-user billing dimension. + #[serde(default)] + pub user_id: String, + /// Ingress correlation id, joins this row to the access log and audit events. + #[serde(default)] + pub request_id: String, + /// Unix seconds the call was billed — the billing-period axis. + #[serde(default)] + pub created_at_epoch_secs: i64, /// Public model the caller requested. pub model: String, /// Model that actually served (differs from `model` after a quota fallback). @@ -125,6 +135,10 @@ pub struct BillingInput<'a> { pub ak: &'a str, pub product: &'a str, pub tenant: &'a str, + /// Effective end user (key owner else request metadata); empty when neither. + pub user_id: &'a str, + /// Ingress correlation id for this request. + pub request_id: &'a str, /// Public model the caller requested (accrues the per-(AK, model) counter). pub requested_model: &'a str, /// Model that actually served — charged at its price (may differ on fallback). @@ -167,6 +181,9 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi ak: b.ak.to_owned(), product: b.product.to_owned(), tenant: b.tenant.to_owned(), + user_id: b.user_id.to_owned(), + request_id: b.request_id.to_owned(), + created_at_epoch_secs: crate::epoch_secs(), model: b.requested_model.to_owned(), served_model: b.served_model.to_owned(), protocol: b.protocol.to_owned(), @@ -193,6 +210,54 @@ pub struct UsageRow { pub vendor_cost_micros: i64, } +/// One row of the per-(user, model) usage rollup over a billing period. +#[derive(Debug, Clone, serde::Serialize)] +pub struct UserUsageRow { + pub user_id: String, + pub model: String, + pub requests: i64, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub total_tokens: i64, + pub cost_micros: i64, + pub vendor_cost_micros: i64, +} + +/// A content-safety outcome, recorded WITHOUT the offending text — only which +/// key/user/rule fired and what the gateway did, so hits are queryable per +/// ak/tenant without retaining prompt content. +#[derive(Debug, Clone, serde::Serialize)] +pub struct SecurityEvent { + pub created_at_epoch_secs: i64, + pub request_id: String, + pub ak: String, + pub user_id: String, + pub tenant: String, + /// Which surface: chat/messages/responses/realtime/… + pub surface: String, + /// The rule family that fired: "blocklist" | "dlp" | a recognizer name. + pub rule: String, + /// What the gateway did: "block" | "redact" | "flag". + pub action: String, + pub hits: i64, +} + +/// One admin-plane mutation, recorded with who/what/when for compliance. +#[derive(Debug, Clone, serde::Serialize)] +pub struct AdminAudit { + pub created_at_epoch_secs: i64, + /// The admin identity: "global" or the tenant name whose token was used. + pub actor: String, + /// The presented scope: "global" | "tenant". + pub scope: String, + /// The mutation: "key_create" | "key_patch" | "key_delete" | "config_publish" | "reload". + pub action: String, + /// The object acted on (an ak, a config version, …). + pub target: String, + pub summary: String, + pub source_ip: String, +} + #[async_trait::async_trait] pub trait Store: Send + Sync + std::fmt::Debug { async fn ledger_add(&self, r: &BillingRecord) -> GResult<()>; @@ -201,6 +266,30 @@ pub trait Store: Send + Sync + std::fmt::Debug { async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)>; /// Usage rolled up by (tenant, requested model), sorted. async fn ledger_usage(&self, tenant: Option<&str>) -> GResult>; + /// Precise per-user cost over `[since, until]` (unix secs), grouped by + /// (user, requested model); optional tenant/user filter. The billing-period + /// query behind per-user invoicing. + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult>; + + /// Append a content-safety event (no prompt text retained). + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()>; + /// The most recent `limit` security events, newest first; optional tenant filter. + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult>; + + /// Append an admin-operation audit entry. + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()>; + /// The most recent `limit` admin-audit entries, newest first. + async fn admin_audit_list(&self, limit: usize) -> GResult>; /// Store `content` under a fresh id owned by `tenant`; returns the metadata. async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult; @@ -285,6 +374,8 @@ pub trait Store: Send + Sync + std::fmt::Debug { #[derive(Debug, Default)] pub struct MemoryStore { records: Mutex>, + sec_events: Mutex>, + audit: Mutex>, files: DashMap, jobs: DashMap, seq: AtomicUsize, @@ -362,6 +453,91 @@ impl Store for MemoryStore { Ok(rollup.into_values().collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + let records = self + .records + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut rollup: std::collections::BTreeMap<(String, String), UserUsageRow> = + std::collections::BTreeMap::new(); + for r in records.iter() { + if tenant.is_some_and(|t| t != r.tenant) + || user.is_some_and(|u| u != r.user_id) + || r.created_at_epoch_secs < since + || r.created_at_epoch_secs > until + { + continue; + } + let e = rollup + .entry((r.user_id.clone(), r.model.clone())) + .or_insert_with(|| UserUsageRow { + user_id: r.user_id.clone(), + model: r.model.clone(), + requests: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + cost_micros: 0, + vendor_cost_micros: 0, + }); + e.requests += 1; + e.prompt_tokens = e.prompt_tokens.saturating_add(r.prompt_tokens); + e.completion_tokens = e.completion_tokens.saturating_add(r.completion_tokens); + e.total_tokens = e.total_tokens.saturating_add(r.total_tokens); + e.cost_micros = e.cost_micros.saturating_add(r.cost_micros); + e.vendor_cost_micros = e.vendor_cost_micros.saturating_add(r.vendor_cost_micros); + } + Ok(rollup.into_values().collect()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + self.sec_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(e.clone()); + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let events = self + .sec_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(events + .iter() + .rev() + .filter(|e| tenant.is_none_or(|t| t == e.tenant)) + .take(limit) + .cloned() + .collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + self.audit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(e.clone()); + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let audit = self + .audit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(audit.iter().rev().take(limit).cloned().collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let id = format!( "file-local-{}", @@ -451,6 +627,9 @@ where cost_micros: row.get(10), vendor_cost_micros: row.get(11), ptu_spillover: row.get(12), + user_id: row.get(13), + request_id: row.get(14), + created_at_epoch_secs: row.get(15), } } @@ -489,6 +668,63 @@ where } } +fn user_usage_row<'r, R>(row: &'r R) -> UserUsageRow +where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, +{ + UserUsageRow { + user_id: row.get(0), + model: row.get(1), + requests: row.get(2), + prompt_tokens: row.get(3), + completion_tokens: row.get(4), + total_tokens: row.get(5), + cost_micros: row.get(6), + vendor_cost_micros: row.get(7), + } +} + +fn security_event_row<'r, R>(row: &'r R) -> SecurityEvent +where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, +{ + SecurityEvent { + created_at_epoch_secs: row.get(0), + request_id: row.get(1), + ak: row.get(2), + user_id: row.get(3), + tenant: row.get(4), + surface: row.get(5), + rule: row.get(6), + action: row.get(7), + hits: row.get(8), + } +} + +fn admin_audit_row<'r, R>(row: &'r R) -> AdminAudit +where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, +{ + AdminAudit { + created_at_epoch_secs: row.get(0), + actor: row.get(1), + scope: row.get(2), + action: row.get(3), + target: row.get(4), + summary: row.get(5), + source_ip: row.get(6), + } +} + /// SQLite-backed store (WAL): ledger, files, and batch jobs in one database /// file; ids derive from rowids so they stay unique across restarts. #[derive(Debug)] @@ -524,7 +760,10 @@ impl SqliteStore { prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL, total_tokens INTEGER NOT NULL, cost_micros INTEGER NOT NULL, vendor_cost_micros INTEGER NOT NULL DEFAULT 0, - ptu_spillover INTEGER NOT NULL DEFAULT 0)", + ptu_spillover INTEGER NOT NULL DEFAULT 0, + user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', + created_at_epoch_secs INTEGER NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS files ( n INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -536,6 +775,17 @@ impl SqliteStore { "CREATE TABLE IF NOT EXISTS batch_results ( batch_id TEXT NOT NULL, idx INTEGER NOT NULL, ok INTEGER NOT NULL, message TEXT NOT NULL, total_tokens INTEGER NOT NULL)", + "CREATE TABLE IF NOT EXISTS security_events ( + n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + surface TEXT NOT NULL DEFAULT '', rule TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', hits INTEGER NOT NULL DEFAULT 0)", + "CREATE TABLE IF NOT EXISTS admin_audit ( + n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, + actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', + summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", ] { sqlx::query(ddl) .execute(&pool) @@ -547,6 +797,9 @@ impl SqliteStore { "ALTER TABLE billing ADD COLUMN tenant TEXT NOT NULL DEFAULT 'default'", "ALTER TABLE billing ADD COLUMN served_model TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN vendor_cost_micros INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN user_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN request_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN created_at_epoch_secs INTEGER NOT NULL DEFAULT 0", // back-fill pre-tenant rows to an unmatchable '' tenant (fail closed) "ALTER TABLE files ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", "ALTER TABLE batches ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", @@ -576,8 +829,8 @@ impl Store for SqliteStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&r.ak) .bind(&r.product) @@ -592,6 +845,9 @@ impl Store for SqliteStore { .bind(r.cost_micros) .bind(r.vendor_cost_micros) .bind(r.ptu_spillover) + .bind(&r.user_id) + .bind(&r.request_id) + .bind(r.created_at_epoch_secs) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -618,7 +874,7 @@ impl Store for SqliteStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs FROM billing ORDER BY n DESC LIMIT ?", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -655,6 +911,99 @@ impl Store for SqliteStore { Ok(rows.iter().map(usage_row).collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + let rows = sqlx::query( + "SELECT user_id, model, COUNT(*), SUM(prompt_tokens), SUM(completion_tokens), + SUM(total_tokens), SUM(cost_micros), SUM(vendor_cost_micros) + FROM billing + WHERE (?1 IS NULL OR tenant = ?1) AND (?2 IS NULL OR user_id = ?2) + AND created_at_epoch_secs BETWEEN ?3 AND ?4 + GROUP BY user_id, model ORDER BY user_id, model", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("roll up user usage", e))?; + Ok(rows.iter().map(user_usage_row).collect()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + sqlx::query( + "INSERT INTO security_events (created_at_epoch_secs, request_id, ak, user_id, + tenant, surface, rule, action, hits) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.request_id) + .bind(&e.ak) + .bind(&e.user_id) + .bind(&e.tenant) + .bind(&e.surface) + .bind(&e.rule) + .bind(&e.action) + .bind(e.hits) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert security event", err))?; + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits FROM security_events + WHERE (?1 IS NULL OR tenant = ?1) ORDER BY n DESC LIMIT ?2", + ) + .bind(tenant) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read security events", e))?; + Ok(rows.iter().map(security_event_row).collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + sqlx::query( + "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, + summary, source_ip) VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.actor) + .bind(&e.scope) + .bind(&e.action) + .bind(&e.target) + .bind(&e.summary) + .bind(&e.source_ip) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert admin audit", err))?; + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, actor, scope, action, target, summary, source_ip + FROM admin_audit ORDER BY n DESC LIMIT ?", + ) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read admin audit", e))?; + Ok(rows.iter().map(admin_audit_row).collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); // SQLite serializes writers, so the MAX(n)+1 subselect is atomic with the insert @@ -816,7 +1165,21 @@ impl PostgresStore { prompt_tokens BIGINT NOT NULL, completion_tokens BIGINT NOT NULL, total_tokens BIGINT NOT NULL, cost_micros BIGINT NOT NULL, vendor_cost_micros BIGINT NOT NULL DEFAULT 0, - ptu_spillover BOOLEAN NOT NULL DEFAULT FALSE)", + ptu_spillover BOOLEAN NOT NULL DEFAULT FALSE, + user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', + created_at_epoch_secs BIGINT NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", + "CREATE TABLE IF NOT EXISTS security_events ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + surface TEXT NOT NULL DEFAULT '', rule TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', hits BIGINT NOT NULL DEFAULT 0)", + "CREATE TABLE IF NOT EXISTS admin_audit ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', + summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", "CREATE TABLE IF NOT EXISTS files ( n BIGSERIAL PRIMARY KEY, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -848,13 +1211,17 @@ impl PostgresStore { .await .map_err(|e| crate::sqlx_err("create postgres schema", e))?; } - sqlx::query( - "ALTER TABLE billing ADD COLUMN IF NOT EXISTS - vendor_cost_micros BIGINT NOT NULL DEFAULT 0", - ) - .execute(&pool) - .await - .map_err(|e| crate::sqlx_err("migrate postgres billing schema", e))?; + for ddl in [ + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS vendor_cost_micros BIGINT NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS request_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS created_at_epoch_secs BIGINT NOT NULL DEFAULT 0", + ] { + sqlx::query(ddl) + .execute(&pool) + .await + .map_err(|e| crate::sqlx_err("migrate postgres billing schema", e))?; + } Ok(Self { pool, ledger_max_rows, @@ -869,8 +1236,8 @@ impl Store for PostgresStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)", ) .bind(&r.ak) .bind(&r.product) @@ -885,6 +1252,9 @@ impl Store for PostgresStore { .bind(r.cost_micros) .bind(r.vendor_cost_micros) .bind(r.ptu_spillover) + .bind(&r.user_id) + .bind(&r.request_id) + .bind(r.created_at_epoch_secs) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -911,7 +1281,7 @@ impl Store for PostgresStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs FROM billing ORDER BY n DESC LIMIT $1", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -955,6 +1325,101 @@ impl Store for PostgresStore { Ok(rows.iter().map(usage_row).collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + let rows = sqlx::query( + "SELECT user_id, model, COUNT(*), + SUM(prompt_tokens)::BIGINT, SUM(completion_tokens)::BIGINT, + SUM(total_tokens)::BIGINT, SUM(cost_micros)::BIGINT, SUM(vendor_cost_micros)::BIGINT + FROM billing + WHERE ($1::text IS NULL OR tenant = $1) AND ($2::text IS NULL OR user_id = $2) + AND created_at_epoch_secs BETWEEN $3 AND $4 + GROUP BY user_id, model ORDER BY user_id, model", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("roll up user usage", e))?; + Ok(rows.iter().map(user_usage_row).collect()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + sqlx::query( + "INSERT INTO security_events (created_at_epoch_secs, request_id, ak, user_id, + tenant, surface, rule, action, hits) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.request_id) + .bind(&e.ak) + .bind(&e.user_id) + .bind(&e.tenant) + .bind(&e.surface) + .bind(&e.rule) + .bind(&e.action) + .bind(e.hits) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert security event", err))?; + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits FROM security_events + WHERE ($1::text IS NULL OR tenant = $1) ORDER BY n DESC LIMIT $2", + ) + .bind(tenant) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read security events", e))?; + Ok(rows.iter().map(security_event_row).collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + sqlx::query( + "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, + summary, source_ip) VALUES ($1, $2, $3, $4, $5, $6, $7)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.actor) + .bind(&e.scope) + .bind(&e.action) + .bind(&e.target) + .bind(&e.summary) + .bind(&e.source_ip) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert admin audit", err))?; + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, actor, scope, action, target, summary, source_ip + FROM admin_audit ORDER BY n DESC LIMIT $1", + ) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read admin audit", e))?; + Ok(rows.iter().map(admin_audit_row).collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); // consume the sequence explicitly — concurrent writers race a MAX(n)+1 subselect @@ -1271,6 +1736,8 @@ mod tests { ak: "k", product: "demo", tenant: "default", + user_id: "u1", + request_id: "req-1", requested_model: "gpt-4o", served_model: "gpt-4o", protocol: "openai-chat", @@ -1292,6 +1759,9 @@ mod tests { ak: "ak-t".into(), product: "p".into(), tenant: "default".into(), + user_id: "u1".into(), + request_id: "req-1".into(), + created_at_epoch_secs: 1_000, model: model.into(), served_model: model.into(), protocol: "openai-chat".into(), @@ -1364,6 +1834,90 @@ mod tests { exercise(&MemoryStore::default()).await; } + async fn exercise_audit(store: &dyn Store) { + let mut a = record("m1"); + a.user_id = "alice".into(); + a.created_at_epoch_secs = 500; + let mut b = record("m1"); + b.user_id = "bob".into(); + b.created_at_epoch_secs = 1_500; + store.ledger_add(&a).await.unwrap(); + store.ledger_add(&b).await.unwrap(); + + let all = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!(all.len(), 2, "two users"); + let alice = store + .usage_by_user(Some("default"), Some("alice"), 0, i64::MAX) + .await + .unwrap(); + assert_eq!(alice.len(), 1); + assert_eq!(alice[0].user_id, "alice"); + assert_eq!(alice[0].total_tokens, 8); + let windowed = store + .usage_by_user(None, None, 1_000, i64::MAX) + .await + .unwrap(); + assert_eq!(windowed.len(), 1, "only bob is in the window"); + assert_eq!(windowed[0].user_id, "bob"); + + store + .security_event_add(&SecurityEvent { + created_at_epoch_secs: 10, + request_id: "req-9".into(), + ak: "ak-t".into(), + user_id: "alice".into(), + tenant: "default".into(), + surface: "openai-chat".into(), + rule: "blocklist".into(), + action: "block".into(), + hits: 1, + }) + .await + .unwrap(); + let events = store.security_events(Some("default"), 10).await.unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].rule, "blocklist"); + assert!( + store + .security_events(Some("ghost"), 10) + .await + .unwrap() + .is_empty(), + "tenant filter excludes others" + ); + + store + .admin_audit_add(&AdminAudit { + created_at_epoch_secs: 20, + actor: "global".into(), + scope: "global".into(), + action: "key_create".into(), + target: "ak-new".into(), + summary: "tenant=t1".into(), + source_ip: "10.0.0.1".into(), + }) + .await + .unwrap(); + let audit = store.admin_audit_list(10).await.unwrap(); + assert_eq!(audit.len(), 1); + assert_eq!(audit[0].action, "key_create"); + assert_eq!(audit[0].target, "ak-new"); + } + + #[tokio::test] + async fn memory_audit_roundtrip() { + exercise_audit(&MemoryStore::default()).await; + } + + #[tokio::test] + async fn sqlite_audit_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("audit.db").to_str().unwrap()) + .await + .unwrap(); + exercise_audit(&store).await; + } + #[tokio::test] async fn batch_result_is_first_writer_wins() { let store = MemoryStore::default(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 44fe90f..cee5dbb 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -32,6 +32,7 @@ use gw_state::{AkInfo, GatewayState}; use serde_json::{Value, json}; const LEDGER_PAGE_DEFAULT: usize = 100; +const KEY_PAGE_DEFAULT: usize = 200; const STREAM_CHANNEL_CAP: usize = 64; /// Per-turn token reserve against the AK daily quota; settled to actuals at billing. const REALTIME_TURN_RESERVE: i64 = 1_000; @@ -120,6 +121,9 @@ pub fn app(state: AppState) -> Router { .route("/admin/config", axum::routing::put(admin_config_put)) .route("/admin/keys", post(admin_key_create).get(admin_key_list)) .route("/admin/usage", get(admin_usage)) + .route("/admin/usage/users", get(admin_usage_users)) + .route("/admin/audit/events", get(admin_security_events)) + .route("/admin/audit/ops", get(admin_audit_ops)) .route( "/admin/keys/{ak}", axum::routing::patch(admin_key_patch).delete(admin_key_delete), @@ -248,6 +252,8 @@ struct RealtimeAdmit { /// Tokens reserved in the AK TPM window; `None` when the key has no TPM cap. tpm_reserved: Option, at: i64, + /// Per-turn correlation id for the ledger row. + request_id: String, snap: Arc, } @@ -310,6 +316,7 @@ async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result "prompt").increment(pt.max(0) as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ct.max(0) as u64); tracing::info!( target: "access", surface, + request_id = %ctx.request.request_id, ak = %ctx.ak.ak, product = %ctx.ak.product, + user_id, model = %model, protocol = mt, account, @@ -822,6 +842,18 @@ async fn accounts(State(s): State) -> Json { Json(resp) } +/// End-user attribution hint from `x-gw-user` — the stable user id a trusted +/// backend (the toC case) sets. Surfaces fall back to the body's own field +/// (OpenAI `user` / Anthropic `metadata.user_id`). A key's `owner` overrides +/// both at billing, so this is only trusted for shared keys. +fn user_header(headers: &HeaderMap) -> Option { + headers + .get("x-gw-user") + .and_then(|v| v.to_str().ok()) + .map(str::to_owned) + .filter(|s| !s.is_empty()) +} + /// AK auth: `Authorization: Bearer ` or `x-api-key: `. The error is /// `(status, message)` so each surface can shape it to its own wire dialect. async fn authenticate(s: &AppState, headers: &HeaderMap) -> Result { @@ -880,6 +912,67 @@ impl AdminScope { AdminScope::Tenant(t) => t == tenant, } } + + /// (actor, scope) for the audit trail: who the token spoke for. + fn audit_identity(&self) -> (&str, &'static str) { + match self { + AdminScope::Global => ("global", "global"), + AdminScope::Tenant(t) => (t.as_str(), "tenant"), + } + } +} + +/// Record a realtime blocklist block to the security-event stream (parity with +/// the REST surfaces). Per-frame DLP redactions are not recorded — too hot. +async fn emit_rt_block(s: &AppState, ak: &AkInfo) { + let event = gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: String::new(), + ak: ak.ak.clone(), + user_id: ak.owner.clone().unwrap_or_default(), + tenant: ak.tenant.clone(), + surface: "realtime".to_owned(), + rule: "blocklist".to_owned(), + action: "block".to_owned(), + hits: 1, + }; + let _ = s.handler.state().store.security_event_add(&event).await; +} + +/// The caller IP for the audit trail, from the LB's forwarding headers. +fn source_ip(headers: &HeaderMap) -> String { + headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.split(',').next()) + .or_else(|| headers.get("x-real-ip").and_then(|v| v.to_str().ok())) + .map(|s| s.trim().to_owned()) + .unwrap_or_default() +} + +/// Record one admin-plane mutation to the audit trail (who/what/when/where). +/// Best-effort: a store failure is logged, never fails the operation. +async fn audit_admin( + s: &AppState, + scope: &AdminScope, + headers: &HeaderMap, + action: &str, + target: &str, + summary: String, +) { + let (actor, scope_kind) = scope.audit_identity(); + let entry = gw_state::AdminAudit { + created_at_epoch_secs: gw_state::epoch_secs(), + actor: actor.to_owned(), + scope: scope_kind.to_owned(), + action: action.to_owned(), + target: target.to_owned(), + summary, + source_ip: source_ip(headers), + }; + if let Err(e) = s.handler.state().store.admin_audit_add(&entry).await { + tracing::warn!(error = %e, action, "admin audit write failed"); + } } /// Admin gate: the global token is checked first (a colliding tenant token @@ -940,7 +1033,7 @@ async fn scoped_key( /// The admin surfaces' public view of a key — one shape for PATCH and GET. fn ak_public_json(k: &AkInfo) -> Value { json!({ - "ak": k.ak, "product": k.product, "tenant": k.tenant, + "ak": k.ak, "product": k.product, "tenant": k.tenant, "owner": k.owner, "qps": k.qps, "daily_token_quota": k.daily_token_quota, "tokens_per_minute": k.tokens_per_minute, "expires_at_epoch_secs": k.expires_at_epoch_secs, @@ -989,6 +1082,15 @@ async fn admin_reload(State(s): State, headers: HeaderMap) -> Response accounts = cfg.accounts.len(), "config reloaded" ); + audit_admin( + &s, + &AdminScope::Global, + &headers, + "reload", + "", + String::new(), + ) + .await; ( StatusCode::OK, Json(json!({ @@ -1040,6 +1142,7 @@ async fn admin_key_create( ak: ak.to_owned(), product: product.to_owned(), tenant: tenant.to_owned(), + owner: body["owner"].as_str().map(str::to_owned), qps: body["qps"].as_f64().unwrap_or(0.0), daily_token_quota: body["daily_token_quota"].as_i64().unwrap_or(0), tokens_per_minute: body["tokens_per_minute"].as_i64(), @@ -1065,6 +1168,15 @@ async fn admin_key_create( { return gateway_error(e); } + audit_admin( + &s, + &scope, + &headers, + "key_create", + ak, + format!("tenant={tenant}"), + ) + .await; ( StatusCode::CREATED, Json(json!({ "ak": ak, "status": "created" })), @@ -1102,7 +1214,10 @@ async fn admin_key_patch( let patched = s.handler.state().auth.patch(&ak, &patch).await; match patched { Err(e) => gateway_error(e), - Ok(Some(info)) => (StatusCode::OK, Json(ak_public_json(&info))).into_response(), + Ok(Some(info)) => { + audit_admin(&s, &scope, &headers, "key_patch", &ak, String::new()).await; + (StatusCode::OK, Json(ak_public_json(&info))).into_response() + } Ok(None) => error_response(404, format!("key {ak} not found")), } } @@ -1122,11 +1237,14 @@ async fn admin_key_delete( } match s.handler.state().auth.revoke(&ak).await { Err(e) => gateway_error(e), - Ok(true) => ( - StatusCode::OK, - Json(json!({ "ak": ak, "status": "revoked" })), - ) - .into_response(), + Ok(true) => { + audit_admin(&s, &scope, &headers, "key_delete", &ak, String::new()).await; + ( + StatusCode::OK, + Json(json!({ "ak": ak, "status": "revoked" })), + ) + .into_response() + } Ok(false) => error_response(404, format!("key {ak} not found")), } } @@ -1151,11 +1269,22 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S Err(e) => return gateway_error(e), }; match s.reload().await { - Ok(()) => ( - StatusCode::OK, - Json(json!({ "status": "published", "version": version })), - ) - .into_response(), + Ok(()) => { + audit_admin( + &s, + &AdminScope::Global, + &headers, + "config_publish", + &version.to_string(), + String::new(), + ) + .await; + ( + StatusCode::OK, + Json(json!({ "status": "published", "version": version })), + ) + .into_response() + } Err(e) => error_response( 500, format!("published v{version} but local reload failed: {e}"), @@ -1163,13 +1292,23 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S } } -/// GET /admin/keys — the key table, scoped: a tenant admin sees only its own keys. -async fn admin_key_list(State(s): State, headers: HeaderMap) -> Response { +/// GET /admin/keys?offset=&limit= — a page of the key table, scoped: a tenant +/// admin sees only its own keys. Paginated so a fleet key table never loads whole. +async fn admin_key_list( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { let scope = match admin_auth(&s, &headers) { Ok(scope) => scope, Err(r) => return r, }; - let listed = match s.handler.state().auth.list().await { + let offset = q.get("offset").and_then(|v| v.parse().ok()).unwrap_or(0); + let limit = q + .get("limit") + .and_then(|v| v.parse().ok()) + .unwrap_or(KEY_PAGE_DEFAULT); + let listed = match s.handler.state().auth.list(offset, limit).await { Ok(v) => v, Err(e) => return gateway_error(e), }; @@ -1178,7 +1317,7 @@ async fn admin_key_list(State(s): State, headers: HeaderMap) -> Respon .filter(|k| scope.covers(&k.tenant)) .map(|k| ak_public_json(&k)) .collect(); - let mut resp = json!({ "count": keys.len() }); + let mut resp = json!({ "count": keys.len(), "offset": offset }); resp["keys"] = Value::Array(keys); Json(resp).into_response() } @@ -1211,6 +1350,96 @@ async fn admin_usage( Json(json!({ "usage": usage })).into_response() } +/// GET /admin/usage/users?user=&since=&until= — precise per-user cost over a +/// billing period, grouped by (user, requested model). Tenant-scoped like +/// [`admin_usage`]; `since`/`until` are unix seconds (default: all time). +async fn admin_usage_users( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let tenant = match &scope { + AdminScope::Tenant(t) => Some(t.clone()), + AdminScope::Global => q.get("tenant").cloned(), + }; + let since = q.get("since").and_then(|v| v.parse().ok()).unwrap_or(0); + let until = q + .get("until") + .and_then(|v| v.parse().ok()) + .unwrap_or(i64::MAX); + let usage = match s + .handler + .state() + .store + .usage_by_user( + tenant.as_deref(), + q.get("user").map(String::as_str), + since, + until, + ) + .await + { + Ok(rows) => rows, + Err(e) => return gateway_error(e), + }; + Json(json!({ "usage": usage })).into_response() +} + +/// GET /admin/audit/events?limit= — content-safety hits (no prompt text), newest +/// first. Tenant-scoped: a tenant admin sees only its own tenant's events. +async fn admin_security_events( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let tenant = match &scope { + AdminScope::Tenant(t) => Some(t.clone()), + AdminScope::Global => q.get("tenant").cloned(), + }; + let limit = q + .get("limit") + .and_then(|v| v.parse().ok()) + .unwrap_or(LEDGER_PAGE_DEFAULT); + match s + .handler + .state() + .store + .security_events(tenant.as_deref(), limit) + .await + { + Ok(events) => Json(json!({ "events": events })).into_response(), + Err(e) => gateway_error(e), + } +} + +/// GET /admin/audit/ops?limit= — admin-operation audit trail, newest first. +/// Global admin only (the trail spans all tenants). +async fn admin_audit_ops( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + if let Err(r) = require_global_admin(&s, &headers) { + return r; + } + let limit = q + .get("limit") + .and_then(|v| v.parse().ok()) + .unwrap_or(LEDGER_PAGE_DEFAULT); + match s.handler.state().store.admin_audit_list(limit).await { + Ok(entries) => Json(json!({ "entries": entries })).into_response(), + Err(e) => gateway_error(e), + } +} + fn gateway_error(e: GatewayError) -> Response { let code = StatusCode::from_u16(e.http_status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); // OpenAI's error schema types `code` as string-or-null, never a number. @@ -1345,6 +1574,7 @@ async fn chat_completions( ); param.typed = Some(typed); param.raw = Value::Object(body.extra); + let user_id = user_header(&headers).or_else(|| param.raw["user"].as_str().map(str::to_owned)); let request = GatewayRequest { is_online: true, @@ -1352,6 +1582,7 @@ async fn chat_completions( ak: ak.ak.clone(), message: messages, model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -1675,6 +1906,8 @@ async fn messages( ModelParamV2::with_name(gw_consts::Protocol::AnthropicMessages, body.model.clone()); param.typed = Some(typed); param.raw = Value::Object(body.extra); + let user_id = user_header(&headers) + .or_else(|| param.raw["metadata"]["user_id"].as_str().map(str::to_owned)); let request = GatewayRequest { is_online: true, @@ -1693,6 +1926,7 @@ async fn messages( }) .collect(), model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -1957,6 +2191,7 @@ async fn run_family( mt: gw_consts::Protocol, typed: TypedParams, messages: Vec, + user_id: Option, ) -> Result { let mut param = ModelParamV2::with_name(mt, model); param.typed = Some(typed); @@ -1965,6 +2200,7 @@ async fn run_family( ak: ak.ak.clone(), message: messages, model_param_v2: Some(param), + user_id, ..Default::default() }; match run_pipeline(s, request, ak).await { @@ -2019,6 +2255,7 @@ async fn completions( gw_consts::Protocol::Completions, typed, vec![ChatMsg::text("user", prompt)], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2066,6 +2303,7 @@ async fn responses( return error_response(400, "input is required"); } let stream = body["stream"].as_bool().unwrap_or(false); + let user_id = user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)); let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, model.clone()); param.raw = body; let request = GatewayRequest { @@ -2073,6 +2311,7 @@ async fn responses( stream, ak: ak.ak.clone(), model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -2216,6 +2455,7 @@ async fn embeddings( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2255,6 +2495,7 @@ async fn images_generations( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2297,6 +2538,7 @@ async fn images_edits( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2336,6 +2578,7 @@ async fn audio_speech( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2391,6 +2634,7 @@ async fn audio_transcriptions( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), ) .await { @@ -2621,6 +2865,63 @@ mod tests { assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); } + #[tokio::test] + async fn user_attribution_and_request_id_land_in_the_ledger() { + let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let store = app_state.handler.state().store.clone(); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer ak-demo-123") + .header("x-gw-user", "user-42") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = app(app_state).oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_, records) = store.ledger_snapshot(10).await.unwrap(); + let row = records.last().expect("a ledger row"); + assert_eq!(row.user_id, "user-42", "x-gw-user attributed the cost"); + assert!(!row.request_id.is_empty(), "request_id stamped"); + assert!(row.created_at_epoch_secs > 0, "created_at stamped"); + let by_user = store + .usage_by_user(None, Some("user-42"), 0, i64::MAX) + .await + .unwrap(); + assert_eq!(by_user.len(), 1); + assert!(by_user[0].total_tokens > 0); + } + + #[tokio::test] + async fn dlp_hit_is_recorded_as_a_security_event() { + let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + assert!(cfg.security.dlp_redact, "embedded config has DLP on"); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let store = app_state.handler.state().store.clone(); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer ak-demo-123") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"mail me at a@b.com"}]}"#, + )) + .unwrap(); + let resp = app(app_state).oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let events = store.security_events(None, 10).await.unwrap(); + assert!( + events.iter().any(|e| e.rule == "dlp"), + "an inbound PII redaction was recorded, no prompt text stored" + ); + } + #[tokio::test] async fn chat_non_stream_ok() { let resp = test_app() From 7e503cfd4bc0eee29855c5e360f0ebb6343b2e96 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 12:35:08 +0800 Subject: [PATCH 02/21] review: apply phase-1 simplify findings (dedup user-id resolution, admin query helpers, comment trim) --- Cargo.lock | 2 + Cargo.toml | 4 +- crates/config/Cargo.toml | 2 + crates/config/src/lib.rs | 130 ++++++++++++++++++++++++++++++---- crates/dag/src/context.rs | 10 +++ crates/dag/src/nodes.rs | 22 +++--- crates/handler/src/lib.rs | 8 +-- crates/handler/src/plugins.rs | 2 + crates/state/src/lib.rs | 3 +- crates/state/src/store.rs | 1 - crates/views/src/lib.rs | 74 +++++++++---------- 11 files changed, 183 insertions(+), 75 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 56b3fdf..07bdc40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -749,9 +749,11 @@ name = "gw-config" version = "0.1.0" dependencies = [ "gw-consts", + "regex", "serde", "serde_yaml", "thiserror", + "tracing", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index e295295..368451a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,10 +76,12 @@ gw-task = { path = "crates/task" } gw-views = { path = "crates/views" } base64 = "0.22" tiktoken-rs = "0.12" -# pure-compute crypto (SigV4 signing, cache keys) — no network capability +# pure-compute crypto (SigV4 signing, cache keys, content-at-rest AEAD) — no network capability sha2 = "0.10" hmac = "0.12" hex = "0.4" +chacha20poly1305 = "0.10" +regex = "1" # Faster iterative builds; keep release optimized. [profile.dev] diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index b716bd0..6d11285 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -11,6 +11,8 @@ gw-consts = { workspace = true } serde = { workspace = true } serde_yaml = { workspace = true } thiserror = { workspace = true } +regex = { workspace = true } +tracing = { workspace = true } [lints] workspace = true diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 30034a7..5c735d9 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -64,9 +64,7 @@ pub struct AkConf { /// Tenant this key belongs to; empty = the implicit `default` tenant. #[serde(default)] pub tenant: String, - /// End user this key is issued to, when one key = one user (the enterprise - /// per-employee model); `None` when the key is shared and per-user - /// attribution rides on request metadata instead. + /// End user this key is issued to (one key = one user); `None` = shared key. #[serde(default)] pub owner: Option, pub qps: f64, @@ -154,15 +152,71 @@ fn default_priority() -> i32 { 1 } -/// Local security policy (rule-based; no cloud security service). +/// What a fired content rule does. `block` denies the request; `flag` lets it +/// through but records the hit; `shadow` is `flag` for a rule under evaluation — +/// same recording, and the caller can tell trial rules apart when auditing. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Action { + #[default] + Block, + Flag, + Shadow, +} + +impl Action { + pub fn as_str(self) -> &'static str { + match self { + Action::Block => "block", + Action::Flag => "flag", + Action::Shadow => "shadow", + } + } +} + +/// A named regex recognizer applied to the same inbound text the blocklist +/// scans. `pattern` is compiled once at config load into [`SecurityConf::regexes`]. +#[derive(Debug, Clone, Deserialize)] +pub struct RegexRule { + pub name: String, + pub pattern: String, + #[serde(default)] + pub action: Action, +} + +/// Local security policy (rule-based; no cloud security service). Lives globally +/// (`security:`) and per-tenant ([`TenantConf::security`]); the tenant's wins +/// whole when present. #[derive(Debug, Clone, Default, Deserialize)] pub struct SecurityConf { /// Blocklist terms; normalized to lower-case (empties dropped) at load. #[serde(default)] pub blocklist: Vec, + /// What a blocklist hit does (default: block). + #[serde(default)] + pub blocklist_action: Action, /// Whether to DLP-redact inbound/outbound content (emails/phone numbers). #[serde(default)] pub dlp_redact: bool, + /// Detect API keys / secrets in inbound text and redact them (a top + /// enterprise ask: stop staff pasting credentials to a model). + #[serde(default)] + pub detect_secrets: bool, + /// Named regex recognizers. + #[serde(default)] + pub regex_rules: Vec, + /// Compiled `regex_rules`, built at config load (rules that fail to compile + /// are dropped with a warning). Never deserialized. + #[serde(skip)] + pub regexes: Vec, +} + +/// A compiled [`RegexRule`], ready to match on the hot path. +#[derive(Debug, Clone)] +pub struct CompiledRule { + pub name: String, + pub action: Action, + pub re: regex::Regex, } /// Account stability policy (in-memory). @@ -272,6 +326,12 @@ pub struct TenantConf { /// Per-model charged-price overrides (else the model's list price applies). #[serde(default)] pub model_prices: std::collections::HashMap, + /// Content-safety policy for this tenant; `None` = use the global `security:`. + #[serde(default)] + pub security: Option, + /// Prompt/response retention for this tenant; `None` = retain nothing. + #[serde(default)] + pub retention: Option, } impl TenantConf { @@ -281,6 +341,26 @@ impl TenantConf { } } +/// How much request/response content a tenant retains, and for how long. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ContentLevel { + /// Store nothing (the default posture). + None, + /// Store the post-DLP redacted text (privacy-preserving audit). + Redacted, + /// Store the raw text (needs `GW_CONTENT_KEY` for at-rest encryption). + Full, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +pub struct RetentionConf { + pub content: ContentLevel, + /// Days after which stored content is purged; 0 = keep until manually purged. + #[serde(default)] + pub days: u32, +} + /// First-class provider preset: `kind` fixes the endpoint, auth style, and /// served wire types, so going live is `kind` + `api_key_env`. #[derive(Debug, Clone, Deserialize)] @@ -469,14 +549,13 @@ impl GatewayConfig { protocols: preset.wires.iter().map(|w| (*w).to_owned()).collect(), }); } - // lower-case once here so security_check needn't rebuild the list per request - self.security.blocklist = self - .security - .blocklist - .iter() - .filter(|w| !w.is_empty()) - .map(|w| w.to_lowercase()) - .collect(); + // normalize the global policy and every tenant override once at load + compile_security(&mut self.security); + for t in &mut self.tenants { + if let Some(sec) = t.security.as_mut() { + compile_security(sec); + } + } Ok(()) } @@ -646,6 +725,33 @@ impl GatewayConfig { } } +/// Normalize a security policy at load: lower-case the blocklist (so scans +/// don't rebuild it per request) and compile the regex rules (dropping any that +/// fail to compile, loudly). +fn compile_security(sec: &mut SecurityConf) { + sec.blocklist = sec + .blocklist + .iter() + .filter(|w| !w.is_empty()) + .map(|w| w.to_lowercase()) + .collect(); + sec.regexes = sec + .regex_rules + .iter() + .filter_map(|r| match regex::Regex::new(&r.pattern) { + Ok(re) => Some(CompiledRule { + name: r.name.clone(), + action: r.action, + re, + }), + Err(e) => { + tracing::warn!(rule = %r.name, error = %e, "dropping uncompilable regex rule"); + None + } + }) + .collect(); +} + /// A token read from the named env var at call time; `None` when the name is /// empty or the var is unset/empty. fn token_from_env(var: &str) -> Option { diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 96be7dc..e6fbb70 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -71,6 +71,16 @@ impl DagContext { self.decisions.push((node, what.into())); } + /// The effective end user: the key's `owner` (authoritative) else request + /// metadata; `""` when neither is present. + pub fn effective_user_id(&self) -> &str { + self.ak + .owner + .as_deref() + .or(self.request.user_id.as_deref()) + .unwrap_or_default() + } + /// The decision trail as `"stage: detail"` lines. pub fn decision_lines(&self) -> impl Iterator + '_ { self.decisions.iter().map(|(n, w)| format!("{n}: {w}")) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 5ec4d08..581b68f 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -583,13 +583,13 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> let requested = param .and_then(|p| p.fallback_from.as_deref()) .unwrap_or(served); - // effective user: the key's owner wins (authoritative), else request metadata - let user_id = ctx - .ak - .owner - .as_deref() - .or(ctx.request.user_id.as_deref()) - .unwrap_or_default(); + // take the reserves into locals first so the whole-ctx borrow that + // `effective_user_id` needs can't clash with these mutable field takes + let (reserved, tpm_reserved, model_quota_key) = ( + ctx.quota_reserved.take().unwrap_or(0), + ctx.tpm_reserved.take(), + ctx.model_quota_key.take(), + ); let record = admission::settle_and_bill( ctx.state.governance.as_ref(), ctx.state.store.as_ref(), @@ -599,7 +599,7 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> ak: &ctx.ak.ak, product: &ctx.ak.product, tenant: &ctx.ak.tenant, - user_id, + user_id: ctx.effective_user_id(), request_id: &ctx.request.request_id, requested_model: requested, served_model: served, @@ -610,10 +610,10 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> total, ptu_spillover, }, - reserved: ctx.quota_reserved.take().unwrap_or(0), - tpm_reserved: ctx.tpm_reserved.take(), + reserved, + tpm_reserved, reserved_at: ctx.quota_at, - model_quota_key: ctx.model_quota_key.take(), + model_quota_key, }, ) .await; diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 8f6aff4..bfa3bbe 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -27,13 +27,7 @@ static REQ_SEQ: AtomicU64 = AtomicU64::new(0); /// key/user/tenant. Best-effort: a store failure is logged, never fails the /// request. Surface = the request protocol, or "batch" for offline items. async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i64) { - let user_id = ctx - .ak - .owner - .as_deref() - .or(ctx.request.user_id.as_deref()) - .unwrap_or_default() - .to_owned(); + let user_id = ctx.effective_user_id().to_owned(); let surface = if ctx.request.is_online { ctx.request .model_param_v2 diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 712cf4a..0aed822 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -297,6 +297,7 @@ mod tests { SecurityConf { blocklist: vec!["forbiddenword".into()], dlp_redact: true, + ..Default::default() } } @@ -317,6 +318,7 @@ mod tests { let s = SecurityConf { blocklist: vec!["forbiddenword".into(), "禁词".into()], dlp_redact: false, + ..Default::default() }; let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "前文 FORBIDDENWORD 后文")], diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index c5caac6..d39ec17 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -35,8 +35,7 @@ pub struct AkInfo { pub product: String, /// Tenant this key belongs to (`gw_config::DEFAULT_TENANT` when undeclared). pub tenant: String, - /// End user this key is issued to; `Some` = authoritative user attribution - /// (one key per user), `None` = shared key, attribute from request metadata. + /// End user this key is issued to (one key = one user); `None` = shared key. pub owner: Option, pub qps: f64, pub daily_token_quota: i64, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 99c4031..4f9a1d4 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -286,7 +286,6 @@ pub trait Store: Send + Sync + std::fmt::Debug { limit: usize, ) -> GResult>; - /// Append an admin-operation audit entry. async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()>; /// The most recent `limit` admin-audit entries, newest first. async fn admin_audit_list(&self, limit: usize) -> GResult>; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index cee5dbb..b5bf2d8 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -752,12 +752,7 @@ fn log_access(surface: &str, ctx: &DagContext, started: Instant) { }) .unwrap_or_default(); let latency = started.elapsed(); - let user_id = ctx - .ak - .owner - .as_deref() - .or(ctx.request.user_id.as_deref()) - .unwrap_or(""); + let user_id = ctx.effective_user_id(); metrics::counter!("gateway_tokens_total", "kind" => "prompt").increment(pt.max(0) as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ct.max(0) as u64); tracing::info!( @@ -842,10 +837,8 @@ async fn accounts(State(s): State) -> Json { Json(resp) } -/// End-user attribution hint from `x-gw-user` — the stable user id a trusted -/// backend (the toC case) sets. Surfaces fall back to the body's own field -/// (OpenAI `user` / Anthropic `metadata.user_id`). A key's `owner` overrides -/// both at billing, so this is only trusted for shared keys. +/// The `x-gw-user` attribution hint; surfaces fall back to the body's own user +/// field. See [`gw_models::GatewayRequest::user_id`] for the trust model. fn user_header(headers: &HeaderMap) -> Option { headers .get("x-gw-user") @@ -920,6 +913,24 @@ impl AdminScope { AdminScope::Tenant(t) => (t.as_str(), "tenant"), } } + + /// The tenant a scoped read is confined to: a tenant admin sees only its + /// own; the global admin may narrow with `?tenant=`. + fn tenant_filter(&self, q: &std::collections::HashMap) -> Option { + match self { + AdminScope::Tenant(t) => Some(t.clone()), + AdminScope::Global => q.get("tenant").cloned(), + } + } +} + +/// A numeric query param, or `default` when absent/unparseable. +fn q_num( + q: &std::collections::HashMap, + key: &str, + default: T, +) -> T { + q.get(key).and_then(|v| v.parse().ok()).unwrap_or(default) } /// Record a realtime blocklist block to the security-event stream (parity with @@ -936,7 +947,9 @@ async fn emit_rt_block(s: &AppState, ak: &AkInfo) { action: "block".to_owned(), hits: 1, }; - let _ = s.handler.state().store.security_event_add(&event).await; + if let Err(e) = s.handler.state().store.security_event_add(&event).await { + tracing::warn!(error = %e, "realtime security event write failed"); + } } /// The caller IP for the audit trail, from the LB's forwarding headers. @@ -1303,11 +1316,8 @@ async fn admin_key_list( Ok(scope) => scope, Err(r) => return r, }; - let offset = q.get("offset").and_then(|v| v.parse().ok()).unwrap_or(0); - let limit = q - .get("limit") - .and_then(|v| v.parse().ok()) - .unwrap_or(KEY_PAGE_DEFAULT); + let offset = q_num(&q, "offset", 0); + let limit = q_num(&q, "limit", KEY_PAGE_DEFAULT); let listed = match s.handler.state().auth.list(offset, limit).await { Ok(v) => v, Err(e) => return gateway_error(e), @@ -1333,10 +1343,7 @@ async fn admin_usage( Ok(scope) => scope, Err(r) => return r, }; - let filter = match &scope { - AdminScope::Tenant(t) => Some(t.clone()), - AdminScope::Global => q.get("tenant").cloned(), - }; + let filter = scope.tenant_filter(&q); let usage = match s .handler .state() @@ -1362,15 +1369,9 @@ async fn admin_usage_users( Ok(scope) => scope, Err(r) => return r, }; - let tenant = match &scope { - AdminScope::Tenant(t) => Some(t.clone()), - AdminScope::Global => q.get("tenant").cloned(), - }; - let since = q.get("since").and_then(|v| v.parse().ok()).unwrap_or(0); - let until = q - .get("until") - .and_then(|v| v.parse().ok()) - .unwrap_or(i64::MAX); + let tenant = scope.tenant_filter(&q); + let since = q_num(&q, "since", 0); + let until = q_num(&q, "until", i64::MAX); let usage = match s .handler .state() @@ -1400,14 +1401,8 @@ async fn admin_security_events( Ok(scope) => scope, Err(r) => return r, }; - let tenant = match &scope { - AdminScope::Tenant(t) => Some(t.clone()), - AdminScope::Global => q.get("tenant").cloned(), - }; - let limit = q - .get("limit") - .and_then(|v| v.parse().ok()) - .unwrap_or(LEDGER_PAGE_DEFAULT); + let tenant = scope.tenant_filter(&q); + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); match s .handler .state() @@ -1430,10 +1425,7 @@ async fn admin_audit_ops( if let Err(r) = require_global_admin(&s, &headers) { return r; } - let limit = q - .get("limit") - .and_then(|v| v.parse().ok()) - .unwrap_or(LEDGER_PAGE_DEFAULT); + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); match s.handler.state().store.admin_audit_list(limit).await { Ok(entries) => Json(json!({ "entries": entries })).into_response(), Err(e) => gateway_error(e), From dbfb25c43e46e5432f4933fa7c5f7ff79072bc4e Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 12:43:26 +0800 Subject: [PATCH 03/21] =?UTF-8?q?feat(security):=20phase=202=20(B1-B3)=20?= =?UTF-8?q?=E2=80=94=20per-tenant=20policy,=20action=20tiers,=20regex+secr?= =?UTF-8?q?et=20recognizers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SecurityConf gains block/flag/shadow action tiers, named regex recognizers, and API-key/secret detection; it now lives per-tenant (TenantConf.security wins over global). The pre-stage records every rule that fires to the security-event stream and denies only on a block-action hit. --- Cargo.lock | 1 + crates/config/src/lib.rs | 13 +++ crates/handler/Cargo.toml | 1 + crates/handler/src/lib.rs | 46 ++++---- crates/handler/src/plugins.rs | 207 +++++++++++++++++++++++++++------- crates/views/src/lib.rs | 15 ++- 6 files changed, 212 insertions(+), 71 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 07bdc40..f6d650e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -815,6 +815,7 @@ dependencies = [ "gw-engines", "gw-models", "gw-state", + "regex", "serde_json", "tokio", "tracing", diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 5c735d9..a6fc0df 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -683,6 +683,19 @@ impl GatewayConfig { Ok(()) } + /// The effective content-safety policy for `tenant`: the tenant's own + /// override when present, else the global `security:`. + pub fn security_for(&self, tenant: &str) -> &SecurityConf { + self.find_tenant(tenant) + .and_then(|t| t.security.as_ref()) + .unwrap_or(&self.security) + } + + /// The retention policy for `tenant`, if it configured one. + pub fn retention_for(&self, tenant: &str) -> Option<&RetentionConf> { + self.find_tenant(tenant).and_then(|t| t.retention.as_ref()) + } + /// Whether `tenant` may call `model`: a declared tenant without an /// allowlist (and the implicit default) allows every model. An undeclared /// non-default tenant fails closed — a runtime key outliving the reload diff --git a/crates/handler/Cargo.toml b/crates/handler/Cargo.toml index 7561dca..6317cdb 100644 --- a/crates/handler/Cargo.toml +++ b/crates/handler/Cargo.toml @@ -17,6 +17,7 @@ gw-dag = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } futures = { workspace = true } +regex = { workspace = true } [dev-dependencies] tokio = { workspace = true } diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index bfa3bbe..da3256d 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -115,25 +115,34 @@ impl OnlineHandler { if request.request_id.is_empty() { request.request_id = new_request_id(); } - // one consistent snapshot for the whole request + // one consistent snapshot for the whole request; the tenant's own + // security policy wins over the global one when it sets one let snap = self.config.load(); + let sec = snap.cfg.security_for(&ak.tenant); // outbound DLP is a response-buffering boundary: a masked span can // straddle deltas, so no engine may stream raw ones — enforced here so // no caller can opt out - let dlp = snap.cfg.security.dlp_redact; + let dlp = sec.dlp_redact; if dlp { request.stream_tx = None; } - // blocklist on the ORIGINAL content, before DLP — else a blocklisted term - // inside a redacted span (a domain in an email) is masked out and slips - if let Some(block) = plugins::security_check(&snap.cfg.security, &mut request) { - let mut ctx = DagContext::new( - snap.cfg.clone(), - snap.state.clone(), - self.transport.clone(), - request, - ak, - ); + // scan the ORIGINAL content, before DLP — else a blocklisted term inside + // a redacted span (a domain in an email) is masked out and slips + let scan = plugins::security_check(sec, &mut request); + + let mut ctx = DagContext::new( + snap.cfg.clone(), + snap.state.clone(), + self.transport.clone(), + request, + ak, + ); + // every rule that fired is recorded (block/flag/shadow alike); only a + // block-action hit denies the request + for hit in &scan.hits { + emit_security_event(&ctx, &hit.rule, hit.action.as_str(), hit.count).await; + } + if let Some(block) = scan.block { ctx.decide( "security_check", format!("blocked (code {})", block.err_code), @@ -149,18 +158,9 @@ impl OnlineHandler { block, ..Default::default() }); - emit_security_event(&ctx, "blocklist", "block", 1).await; return Ok(ctx); } - let redacted = plugins::dlp_redact_request(&snap.cfg.security, &mut request); - - let mut ctx = DagContext::new( - snap.cfg.clone(), - snap.state.clone(), - self.transport.clone(), - request, - ak, - ); + let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { ctx.decide("dlp", format!("redacted {redacted} span(s) inbound")); emit_security_event(&ctx, "dlp", "redact", redacted as i64).await; @@ -199,7 +199,7 @@ impl OnlineHandler { } let redacted_out = if let Some(outcome) = ctx.outcome.as_mut() { - let n = plugins::dlp_redact_response(&snap.cfg.security, &mut outcome.response); + let n = plugins::dlp_redact_response(sec, &mut outcome.response); // raw decoded deltas are pre-redaction; drop them so no downstream // reconstruction can replay unmasked text past the boundary if dlp { diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 0aed822..97cf589 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -4,40 +4,78 @@ //! typed params. The post-stage redacts the outbound message; streaming //! surfaces buffer and replay the redacted text when outbound DLP is on. -use gw_config::SecurityConf; +use gw_config::{Action, SecurityConf}; use gw_models::{Block, GatewayRequest, GatewayResponse}; const BLOCKED_MSG: &str = "this content cannot be answered, please try a different request"; -/// Blocklist check. Returns Block on a hit (block=true implies hit=true). -/// Scans the same inbound text DLP covers, so no surface is a bypass; takes -/// `&mut` only to share the one traversal with the redaction pass. -pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> Option { - if sec.blocklist.is_empty() { - return None; +/// One content rule that fired on a request, for the security-event stream. +pub struct RuleHit { + pub rule: String, + pub action: Action, + pub count: i64, +} + +/// The result of scanning one request: whether to deny it, plus every rule that +/// fired (for the audit stream). A `block`-action hit fills `block`; `flag`/ +/// `shadow` hits only populate `hits`. +#[derive(Default)] +pub struct ScanOutcome { + pub block: Option, + pub hits: Vec, +} + +/// Scan inbound text against the blocklist and the regex recognizers. One +/// traversal of the same fields DLP covers, so no surface is a bypass. `&mut` +/// only to share the traversal; nothing is rewritten here. +pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanOutcome { + if sec.blocklist.is_empty() && sec.regexes.is_empty() { + return ScanOutcome::default(); } - // the shared traversal counts hits; a String is never rewritten here - let mut scan = |s: &mut String| usize::from(blocklist_hit(sec, s)); - let mut blocked = 0; + let mut blocklist_count = 0i64; + let mut regex_counts = vec![0i64; sec.regexes.len()]; + let mut visit = |s: &mut String| { + blocklist_count += i64::from(blocklist_hit(sec, s)); + for (i, r) in sec.regexes.iter().enumerate() { + regex_counts[i] += r.re.find_iter(s).count() as i64; + } + 0 + }; for msg in &mut request.message { - blocked += usize::from(blocklist_hit(sec, &msg.content)); + visit(&mut msg.content); if let Some(parts) = &mut msg.parts { - blocked += walk_json_strings(parts, &mut scan); + walk_json_strings(parts, &mut visit); } } if let Some(param) = request.model_param_v2.as_mut() { - blocked += walk_json_strings(&mut param.raw, &mut scan); + walk_json_strings(&mut param.raw, &mut visit); if let Some(typed) = param.typed.as_mut() { - blocked += for_each_typed_text(typed, &mut scan); + for_each_typed_text(typed, &mut visit); } } - if blocked > 0 { - return Some(Block::blocked( - BLOCKED_MSG, - gw_consts::ErrCode::EMPTY_RESP.value() as i32, - )); + + let mut hits = Vec::new(); + if blocklist_count > 0 { + hits.push(RuleHit { + rule: "blocklist".to_owned(), + action: sec.blocklist_action, + count: blocklist_count, + }); + } + for (r, count) in sec.regexes.iter().zip(regex_counts) { + if count > 0 { + hits.push(RuleHit { + rule: r.name.clone(), + action: r.action, + count, + }); + } } - None + let block = hits + .iter() + .any(|h| h.action == Action::Block) + .then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)); + ScanOutcome { block, hits } } /// Case-insensitive blocklist test; terms are pre-lowercased at config load. @@ -122,46 +160,59 @@ pub fn dlp_redact_realtime_frame(sec: &SecurityConf, frame: &mut serde_json::Val gw_engines::realtime::visit_frame_text(frame, &mut redact_str) } -/// DLP inbound redaction: emails and 11-digit phone numbers. +/// DLP inbound redaction: emails, 11-digit phone numbers, and — when +/// `detect_secrets` is on — API keys / credentials. pub fn dlp_redact_request(sec: &SecurityConf, request: &mut GatewayRequest) -> usize { - if !sec.dlp_redact { + if !sec.dlp_redact && !sec.detect_secrets { return 0; } + let secrets = sec.detect_secrets; + let pii = sec.dlp_redact; + let mut redact_field = |s: &mut String| redact_in_place(s, pii, secrets); let mut hits = 0; for msg in &mut request.message { - hits += redact_str(&mut msg.content); + hits += redact_field(&mut msg.content); // engines forward `parts` (not `content`) when present, so PII must be // scrubbed inside the parts' text blocks too if let Some(parts) = &mut msg.parts { - hits += redact_parts_text(parts); + hits += redact_parts_text(parts, pii, secrets); } } // non-chat surfaces carry user text outside `message` (Responses raw body, // family typed params) — scrub those too or they reach the vendor unredacted if let Some(param) = request.model_param_v2.as_mut() { - hits += walk_json_strings(&mut param.raw, &mut redact_str); + hits += walk_json_strings(&mut param.raw, &mut redact_field); if let Some(typed) = param.typed.as_mut() { - hits += for_each_typed_text(typed, &mut redact_str); + hits += for_each_typed_text(typed, &mut redact_field); } } hits } -/// Redact one string in place; the hit count. -fn redact_str(s: &mut String) -> usize { - match redact(s) { - Some((redacted, n)) => { - *s = redacted; - n - } - None => 0, +/// Redact one string in place (email/phone via `pii`, secrets via `secrets`); +/// the hit count. +fn redact_in_place(s: &mut String, pii: bool, secrets: bool) -> usize { + let mut hits = 0; + if pii && let Some((redacted, n)) = redact(s) { + *s = redacted; + hits += n; + } + if secrets && let Some((redacted, n)) = redact_secrets(s) { + *s = redacted; + hits += n; } + hits +} + +/// Redact one string in place; email/phone only (the outbound/realtime path). +fn redact_str(s: &mut String) -> usize { + redact_in_place(s, true, false) } /// Redact PII inside a multimodal `parts` array's text blocks, in place. /// Deliberately narrower than the blocklist scan: non-text parts (image URLs, /// base64 data) are never rewritten, so a redaction can't corrupt them. -fn redact_parts_text(parts: &mut serde_json::Value) -> usize { +fn redact_parts_text(parts: &mut serde_json::Value, pii: bool, secrets: bool) -> usize { let Some(arr) = parts.as_array_mut() else { return 0; }; @@ -169,15 +220,36 @@ fn redact_parts_text(parts: &mut serde_json::Value) -> usize { for p in arr { if p["type"] == "text" && let Some(t) = p["text"].as_str() - && let Some((redacted, n)) = redact(t) { - p["text"] = serde_json::Value::String(redacted); - hits += n; + let mut s = t.to_owned(); + let n = redact_in_place(&mut s, pii, secrets); + if n > 0 { + p["text"] = serde_json::Value::String(s); + hits += n; + } } } hits } +/// Mask credential shapes (API keys, tokens, private-key headers) with +/// `[REDACTED_SECRET]`. High-precision patterns to avoid mauling normal text. +fn redact_secrets(text: &str) -> Option<(String, usize)> { + static SECRETS: std::sync::LazyLock = std::sync::LazyLock::new(|| { + #[allow(clippy::expect_used)] // a compile-time literal, covered by tests + regex::Regex::new( + r"sk-[A-Za-z0-9_-]{20,}|AKIA[0-9A-Z]{16}|gh[pousr]_[A-Za-z0-9]{36,}|xox[baprs]-[A-Za-z0-9-]{10,}|-----BEGIN [A-Z ]*PRIVATE KEY-----", + ) + .expect("static secret patterns compile") + }); + let mut count = 0; + let out = SECRETS.replace_all(text, |_: ®ex::Captures| { + count += 1; + "[REDACTED_SECRET]" + }); + (count > 0).then(|| (out.into_owned(), count)) +} + /// DLP outbound redaction: the flat `message` plus the structured payloads the /// non-chat surfaces return (`response_v2`, `tool_calls`), so vendor-introduced /// PII can't leak through a field the surface serializes verbatim. @@ -307,10 +379,14 @@ mod tests { message: vec![ChatMsg::text("user", "say ForbiddenWord now")], ..Default::default() }; - let block = security_check(&sec(), &mut req).unwrap(); + let block = security_check(&sec(), &mut req).block.unwrap(); assert!(block.block && block.hit); assert_eq!(block.err_code, 4003); - assert!(security_check(&sec(), &mut GatewayRequest::default()).is_none()); + assert!( + security_check(&sec(), &mut GatewayRequest::default()) + .block + .is_none() + ); } #[test] @@ -324,12 +400,59 @@ mod tests { message: vec![ChatMsg::text("user", "前文 FORBIDDENWORD 后文")], ..Default::default() }; - assert!(security_check(&s, &mut req).is_some()); + assert!(security_check(&s, &mut req).block.is_some()); let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "包含 禁词 的内容")], ..Default::default() }; - assert!(security_check(&s, &mut req).is_some()); + assert!(security_check(&s, &mut req).block.is_some()); + } + + #[test] + fn flag_action_records_a_hit_without_blocking() { + let s = SecurityConf { + blocklist: vec!["watchword".into()], + blocklist_action: Action::Flag, + ..Default::default() + }; + let mut req = GatewayRequest { + message: vec![ChatMsg::text("user", "contains watchword here")], + ..Default::default() + }; + let out = security_check(&s, &mut req); + assert!(out.block.is_none(), "flag does not deny"); + assert_eq!(out.hits.len(), 1); + assert_eq!(out.hits[0].action, Action::Flag); + } + + #[test] + fn regex_rule_blocks_and_redact_secrets_masks() { + let mut s = SecurityConf { + regex_rules: vec![gw_config::RegexRule { + name: "ssn".into(), + pattern: r"\d{3}-\d{2}-\d{4}".into(), + action: Action::Block, + }], + ..Default::default() + }; + // compile_security runs at config load; replicate it for the unit test + s.regexes = vec![gw_config::CompiledRule { + name: "ssn".into(), + action: Action::Block, + re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), + }]; + let mut req = GatewayRequest { + message: vec![ChatMsg::text("user", "my ssn is 123-45-6789")], + ..Default::default() + }; + assert!( + security_check(&s, &mut req).block.is_some(), + "regex Block denies" + ); + + let (masked, n) = redact_secrets("key sk-abcdefghijklmnopqrstuvwxyz012345 end").unwrap(); + assert_eq!(n, 1); + assert!(masked.contains("[REDACTED_SECRET]") && !masked.contains("sk-abc")); } #[test] diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index b5bf2d8..cdd3684 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -403,8 +403,9 @@ async fn realtime_session( .await; continue; }; - // same blocklist + inbound DLP every REST surface runs - let sec = &s.handler.cfg().security; + // same blocklist + inbound DLP every REST surface runs, tenant policy first + let cfg = s.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut ev) { emit_rt_block(&s, &ak).await; let _ = socket @@ -565,8 +566,9 @@ async fn realtime_bridge( Some(Ok(_)) => continue, // ping/pong handled by the ws stacks }; if let Some(mut frame) = frame { - // same blocklist + inbound DLP every REST surface runs - let sec = &s.handler.cfg().security; + // same blocklist + inbound DLP every REST surface runs, tenant policy first + let cfg = s.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut frame) { emit_rt_block(&s, &ak).await; if cl_tx @@ -686,9 +688,10 @@ async fn realtime_bridge( } // outbound DLP, per frame (a span straddling deltas is // beyond a relay that cannot buffer) + let cfg = s.handler.cfg(); if relay && gw_handler::plugins::dlp_redact_realtime_frame( - &s.handler.cfg().security, + cfg.security_for(&ak.tenant), &mut v, ) > 0 { @@ -1632,7 +1635,7 @@ fn spawn_stream_pipeline( started: Instant, ) -> tokio::sync::mpsc::Receiver { let (tx, rx) = tokio::sync::mpsc::channel::(STREAM_CHANNEL_CAP); - let dlp = s.handler.cfg().security.dlp_redact; + let dlp = s.handler.cfg().security_for(&ak.tenant).dlp_redact; if !dlp { request.stream_tx = Some(tx.clone()); } From ae72f1a81c4c18c3e89f0a4651669060353229de Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 12:57:17 +0800 Subject: [PATCH 04/21] =?UTF-8?q?feat(retention):=20phase=202=20(C2/C4/D2)?= =?UTF-8?q?=20=E2=80=94=20gated=20content=20retention,=20purge,=20billing?= =?UTF-8?q?=20precision?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-tenant retention (none/redacted/full) stores prompts/responses in a request_content table, sealed at rest with XChaCha20-Poly1305 under GW_CONTENT_KEY (full refuses to store raw without a key). An hourly purge task deletes expired content. Billing rows carry an 'estimated' flag (aborted-stream estimates), and /admin/usage/users gains CSV export. --- Cargo.lock | 96 ++++++++++++++- crates/dag/src/nodes.rs | 16 ++- crates/handler/src/lib.rs | 106 ++++++++++++++++ crates/server/src/main.rs | 2 + crates/state/Cargo.toml | 3 + crates/state/src/content.rs | 100 ++++++++++++++++ crates/state/src/lib.rs | 2 + crates/state/src/store.rs | 232 ++++++++++++++++++++++++++++++++++-- crates/task/src/lib.rs | 32 ++++- crates/views/src/lib.rs | 29 +++++ 10 files changed, 600 insertions(+), 18 deletions(-) create mode 100644 crates/state/src/content.rs diff --git a/Cargo.lock b/Cargo.lock index f6d650e..1677ee9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,16 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common 0.1.7", + "generic-array", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -241,6 +251,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "chacha20" version = "0.10.1" @@ -252,6 +273,30 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20 0.9.1", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", + "zeroize", +] + [[package]] name = "cmov" version = "0.5.4" @@ -354,6 +399,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] @@ -872,12 +918,15 @@ version = "0.1.0" dependencies = [ "arc-swap", "async-trait", + "base64", + "chacha20poly1305", "dashmap", "futures", "governor", "gw-config", "gw-consts", "gw-models", + "hex", "metrics", "moka", "redis", @@ -1226,6 +1275,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -1466,6 +1524,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "parking" version = "2.2.1" @@ -1513,6 +1577,17 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures 0.2.17", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -1654,7 +1729,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ - "chacha20", + "chacha20 0.10.1", "getrandom 0.4.3", "rand_core 0.10.1", ] @@ -1669,6 +1744,15 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + [[package]] name = "rand_core" version = "0.9.5" @@ -2666,6 +2750,16 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common 0.1.7", + "subtle", +] + [[package]] name = "unsafe-libyaml" version = "0.2.11" diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 581b68f..57db750 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -536,7 +536,7 @@ impl DagNode for CostCalc { }; let ct = enc.encode_len(&resp.message) as i64; ctx.decide("cost_calc", format!("aborted stream, billed {pt}+{ct}")); - return bill(ctx, pt, ct, pt.saturating_add(ct)).await; + return bill(ctx, pt, ct, pt.saturating_add(ct), true).await; } // default rate is 1:1; the formula carries future weighted rates. // saturating sums so a malformed usage subtree can't overflow the totals @@ -564,13 +564,20 @@ impl DagNode for CostCalc { resp.prompt_tokens.saturating_add(resp.completion_tokens), ), }; - bill(ctx, prompt, completion, total).await + bill(ctx, prompt, completion, total, false).await } } /// Settle reserves and write the ledger for one served request via the shared -/// [`admission::settle_and_bill`] orchestration. -async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> GResult<()> { +/// [`admission::settle_and_bill`] orchestration. `estimated` marks a bill from +/// an aborted stream's estimated counts rather than a vendor usage payload. +async fn bill( + ctx: &mut DagContext, + prompt: i64, + completion: i64, + total: i64, + estimated: bool, +) -> GResult<()> { let ptu_spillover = ctx .outcome .as_ref() @@ -609,6 +616,7 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> completion, total, ptu_spillover, + estimated, }, reserved, tpm_reserved, diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index da3256d..193cb8e 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -54,6 +54,89 @@ async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i } } +/// Persist this request's prompt and response per the tenant's retention policy. +/// `full` content is sealed at rest (and refuses to store raw without a key, +/// downgrading to the redacted text); `redacted` stores the post-DLP text, +/// sealed too when a key is present. Best-effort — a store failure is logged. +async fn persist_content( + ctx: &DagContext, + retention: gw_config::RetentionConf, + raw_prompt: Option, + raw_response: Option, +) { + use gw_config::ContentLevel; + let want_full = retention.content == ContentLevel::Full; + let sealed_ok = gw_state::sealing_available(); + // full without a key must NOT land raw — fall back to the redacted text + let store_full = want_full && sealed_ok; + if want_full && !sealed_ok { + tracing::warn!("full retention configured but GW_CONTENT_KEY unset; storing redacted text"); + } + + let now = gw_state::epoch_secs(); + let expires = if retention.days > 0 { + now + retention.days as i64 * 86_400 + } else { + 0 + }; + let redacted_prompt = || { + ctx.request + .message + .iter() + .map(|m| m.content.as_str()) + .collect::>() + .join("\n") + }; + let redacted_response = ctx + .outcome + .as_ref() + .map(|o| o.response.message.clone()) + .unwrap_or_default(); + + let items = [ + ( + "prompt", + if store_full { + raw_prompt.unwrap_or_else(redacted_prompt) + } else { + redacted_prompt() + }, + ), + ( + "response", + if store_full { + raw_response.unwrap_or_else(|| redacted_response.clone()) + } else { + redacted_response.clone() + }, + ), + ]; + for (kind, text) in items { + if text.is_empty() { + continue; + } + // seal whenever a key exists (defense in depth even for redacted text) + let (content, sealed) = match gw_state::content::seal(&text) { + Some(ct) => (ct, true), + None => (text, false), + }; + let record = gw_state::ContentRecord { + created_at_epoch_secs: now, + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id: ctx.effective_user_id().to_owned(), + tenant: ctx.ak.tenant.clone(), + kind: kind.to_owned(), + content, + sealed, + expires_at_epoch_secs: expires, + }; + if let Err(e) = ctx.state.store.content_add(&record).await { + tracing::warn!(error = %e, kind, "content retention write failed"); + } + } +} + /// A per-request correlation id: `req--`, time-sortable and /// unique within the process (the seq disambiguates same-millisecond requests). pub fn new_request_id() -> String { @@ -160,6 +243,21 @@ impl OnlineHandler { }); return Ok(ctx); } + // capture raw prompt before DLP when the tenant retains full content + let retention = snap.cfg.retention_for(&ctx.ak.tenant).copied(); + let full = matches!( + retention, + Some(r) if r.content == gw_config::ContentLevel::Full + ); + let raw_prompt = full.then(|| { + ctx.request + .message + .iter() + .map(|m| m.content.as_str()) + .collect::>() + .join("\n") + }); + let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { ctx.decide("dlp", format!("redacted {redacted} span(s) inbound")); @@ -198,6 +296,11 @@ impl OnlineHandler { outcome.response.model = requested; } + // capture raw response before outbound DLP for full retention + let raw_response = full + .then(|| ctx.outcome.as_ref().map(|o| o.response.message.clone())) + .flatten(); + let redacted_out = if let Some(outcome) = ctx.outcome.as_mut() { let n = plugins::dlp_redact_response(sec, &mut outcome.response); // raw decoded deltas are pre-redaction; drop them so no downstream @@ -209,6 +312,9 @@ impl OnlineHandler { } else { 0 }; + if let Some(r) = retention.filter(|r| r.content != gw_config::ContentLevel::None) { + persist_content(&ctx, r, raw_prompt, raw_response).await; + } if redacted_out > 0 { ctx.decide("dlp", format!("redacted {redacted_out} span(s) outbound")); emit_security_event(&ctx, "dlp", "redact_out", redacted_out as i64).await; diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 5a883db..b8e5c71 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -83,6 +83,7 @@ async fn main() -> anyhow::Result<()> { // daily quota reset; governance is a preserved seam, so this survives reloads let quota_task = gw_task::spawn_quota_reset(state.clone(), gw_task::DAILY); + let purge_task = gw_task::spawn_content_purge(state.clone(), gw_task::PURGE_PERIOD); let distributed_batches = state.store.distributed_batches(); let transport = select_transport()?; @@ -190,6 +191,7 @@ async fn main() -> anyhow::Result<()> { .await?; quota_task.abort(); + purge_task.abort(); tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/state/Cargo.toml b/crates/state/Cargo.toml index 3ad744c..0fe9db1 100644 --- a/crates/state/Cargo.toml +++ b/crates/state/Cargo.toml @@ -23,6 +23,9 @@ serde = { workspace = true } tokio = { workspace = true } futures = { workspace = true } metrics = { workspace = true } +chacha20poly1305 = { workspace = true } +hex = { workspace = true } +base64 = { workspace = true } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/state/src/content.rs b/crates/state/src/content.rs new file mode 100644 index 0000000..3a0e7ca --- /dev/null +++ b/crates/state/src/content.rs @@ -0,0 +1,100 @@ +//! Content-at-rest sealing for the retention store: authenticated encryption +//! (XChaCha20-Poly1305) under a process-wide deployment key from +//! `GW_CONTENT_KEY` (64 hex chars = 32 bytes). No key configured → sealing is +//! unavailable, and `full` retention refuses to store raw content. + +use std::sync::LazyLock; + +use base64::Engine as _; +use chacha20poly1305::aead::{Aead, KeyInit, OsRng}; +use chacha20poly1305::{AeadCore, XChaCha20Poly1305}; + +/// One stored prompt or response, priced out of the ledger's blast radius: +/// content may be ciphertext (sealed) or plaintext (redacted level). +#[derive(Debug, Clone, serde::Serialize)] +pub struct ContentRecord { + pub created_at_epoch_secs: i64, + pub request_id: String, + pub ak: String, + pub user_id: String, + pub tenant: String, + /// "prompt" | "response". + pub kind: String, + /// Sealed (base64 nonce||ciphertext) or plaintext, per [`ContentRecord::sealed`]. + pub content: String, + /// Whether `content` is sealed ciphertext. + pub sealed: bool, + /// Unix seconds after which the purge job deletes this row; 0 = keep. + pub expires_at_epoch_secs: i64, +} + +/// The process-wide content key, loaded once from `GW_CONTENT_KEY`. +static CIPHER: LazyLock> = LazyLock::new(|| { + let raw = std::env::var("GW_CONTENT_KEY").ok()?; + match hex::decode(raw.trim()) { + Ok(bytes) if bytes.len() == 32 => Some(XChaCha20Poly1305::new(bytes.as_slice().into())), + _ => { + tracing::error!("GW_CONTENT_KEY is not 64 hex chars (32 bytes); content sealing off"); + None + } + } +}); + +/// Whether a deployment key is configured (so `full` retention may store raw). +pub fn sealing_available() -> bool { + CIPHER.is_some() +} + +/// Seal `plaintext` into base64(nonce ‖ ciphertext); `None` when no key is set. +pub fn seal(plaintext: &str) -> Option { + seal_with(CIPHER.as_ref()?, plaintext) +} + +/// Reverse [`seal`]; `None` when no key is set or the input is malformed/forged. +pub fn open(sealed: &str) -> Option { + open_with(CIPHER.as_ref()?, sealed) +} + +fn seal_with(cipher: &XChaCha20Poly1305, plaintext: &str) -> Option { + let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng); + let ct = cipher.encrypt(&nonce, plaintext.as_bytes()).ok()?; + let mut out = nonce.to_vec(); + out.extend_from_slice(&ct); + Some(base64::engine::general_purpose::STANDARD.encode(out)) +} + +fn open_with(cipher: &XChaCha20Poly1305, sealed: &str) -> Option { + let raw = base64::engine::general_purpose::STANDARD + .decode(sealed) + .ok()?; + if raw.len() < 24 { + return None; + } + let (nonce, ct) = raw.split_at(24); + let pt = cipher.decrypt(nonce.into(), ct).ok()?; + String::from_utf8(pt).ok() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn seal_open_roundtrips_and_rejects_tamper() { + let cipher = XChaCha20Poly1305::new((&[7u8; 32]).into()); + let sealed = seal_with(&cipher, "credit card 4111 1111 1111 1111").unwrap(); + assert!(!sealed.contains("4111"), "ciphertext hides the plaintext"); + assert_eq!( + open_with(&cipher, &sealed).unwrap(), + "credit card 4111 1111 1111 1111" + ); + let other = XChaCha20Poly1305::new((&[8u8; 32]).into()); + assert!(open_with(&other, &sealed).is_none(), "wrong key can't open"); + let mut tampered = sealed.clone(); + tampered.push('x'); + assert!( + open_with(&cipher, &tampered).is_none(), + "tamper is rejected" + ); + } +} diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index d39ec17..350cfe6 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -15,12 +15,14 @@ use gw_models::Account; pub mod admission; pub mod configstore; +pub mod content; pub mod governance; pub mod health; pub mod keystore; pub mod store; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; +pub use content::{ContentRecord, sealing_available}; pub use governance::{Governance, MemoryGovernance, RedisGovernance}; pub use health::{HealthStore, RedisHealth}; pub use keystore::{KeyStore, PostgresKeyStore}; diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4f9a1d4..b92999c 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -59,6 +59,10 @@ pub struct BillingRecord { /// PTU spilled over to a paygo account (a failover occurred). #[serde(default)] pub ptu_spillover: bool, + /// Token counts were estimated (an aborted stream billed from delivered + /// text), not read from a vendor usage payload. + #[serde(default)] + pub estimated: bool, } /// Offline batch job status. @@ -149,6 +153,8 @@ pub struct BillingInput<'a> { pub completion: i64, pub total: i64, pub ptu_spillover: bool, + /// Counts are estimated (aborted stream), not vendor-reported. + pub estimated: bool, } /// Clamp a metered token count into `[0, MAX_METERED_TOKENS]`. @@ -194,6 +200,7 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi cost_micros: gw_models::cost_micros(prompt, completion, charged), vendor_cost_micros: gw_models::cost_micros(prompt, completion, vendor), ptu_spillover: b.ptu_spillover, + estimated: b.estimated, } } @@ -290,6 +297,14 @@ pub trait Store: Send + Sync + std::fmt::Debug { /// The most recent `limit` admin-audit entries, newest first. async fn admin_audit_list(&self, limit: usize) -> GResult>; + /// Store one retained prompt/response record (per-tenant retention policy). + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()>; + /// Delete content whose `expires_at_epoch_secs` is in `(0, now]`; returns the + /// number deleted. Rows with `expires_at = 0` are kept until manual purge. + async fn content_purge(&self, now_epoch_secs: i64) -> GResult; + /// The retained content for one request (both prompt and response rows). + async fn content_for(&self, request_id: &str) -> GResult>; + /// Store `content` under a fresh id owned by `tenant`; returns the metadata. async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult; async fn file_get(&self, id: &str) -> GResult>; @@ -375,6 +390,7 @@ pub struct MemoryStore { records: Mutex>, sec_events: Mutex>, audit: Mutex>, + content: Mutex>, files: DashMap, jobs: DashMap, seq: AtomicUsize, @@ -537,6 +553,36 @@ impl Store for MemoryStore { Ok(audit.iter().rev().take(limit).cloned().collect()) } + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + self.content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(r.clone()); + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let mut content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let before = content.len(); + content.retain(|r| r.expires_at_epoch_secs == 0 || r.expires_at_epoch_secs > now); + Ok((before - content.len()) as u64) + } + + async fn content_for(&self, request_id: &str) -> GResult> { + let content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(content + .iter() + .filter(|r| r.request_id == request_id) + .cloned() + .collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let id = format!( "file-local-{}", @@ -629,6 +675,7 @@ where user_id: row.get(13), request_id: row.get(14), created_at_epoch_secs: row.get(15), + estimated: row.get(16), } } @@ -724,6 +771,27 @@ where } } +fn content_row<'r, R>(row: &'r R) -> crate::ContentRecord +where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, + bool: sqlx::Decode<'r, R::Database> + sqlx::Type, +{ + crate::ContentRecord { + created_at_epoch_secs: row.get(0), + request_id: row.get(1), + ak: row.get(2), + user_id: row.get(3), + tenant: row.get(4), + kind: row.get(5), + content: row.get(6), + sealed: row.get(7), + expires_at_epoch_secs: row.get(8), + } +} + /// SQLite-backed store (WAL): ledger, files, and batch jobs in one database /// file; ids derive from rowids so they stay unique across restarts. #[derive(Debug)] @@ -761,7 +829,8 @@ impl SqliteStore { vendor_cost_micros INTEGER NOT NULL DEFAULT 0, ptu_spillover INTEGER NOT NULL DEFAULT 0, user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', - created_at_epoch_secs INTEGER NOT NULL DEFAULT 0)", + created_at_epoch_secs INTEGER NOT NULL DEFAULT 0, + estimated INTEGER NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS files ( n INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL, @@ -785,6 +854,14 @@ impl SqliteStore { actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", + "CREATE TABLE IF NOT EXISTS request_content ( + n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + kind TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', + sealed INTEGER NOT NULL DEFAULT 0, expires_at_epoch_secs INTEGER NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS content_expiry_idx ON request_content (expires_at_epoch_secs)", + "CREATE INDEX IF NOT EXISTS content_request_idx ON request_content (request_id)", ] { sqlx::query(ddl) .execute(&pool) @@ -799,6 +876,7 @@ impl SqliteStore { "ALTER TABLE billing ADD COLUMN user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN request_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN created_at_epoch_secs INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN estimated INTEGER NOT NULL DEFAULT 0", // back-fill pre-tenant rows to an unmatchable '' tenant (fail closed) "ALTER TABLE files ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", "ALTER TABLE batches ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", @@ -828,8 +906,9 @@ impl Store for SqliteStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&r.ak) .bind(&r.product) @@ -847,6 +926,7 @@ impl Store for SqliteStore { .bind(&r.user_id) .bind(&r.request_id) .bind(r.created_at_epoch_secs) + .bind(r.estimated) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -873,7 +953,8 @@ impl Store for SqliteStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated FROM billing ORDER BY n DESC LIMIT ?", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -1003,6 +1084,51 @@ impl Store for SqliteStore { Ok(rows.iter().map(admin_audit_row).collect()) } + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + sqlx::query( + "INSERT INTO request_content (created_at_epoch_secs, request_id, ak, user_id, + tenant, kind, content, sealed, expires_at_epoch_secs) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(r.created_at_epoch_secs) + .bind(&r.request_id) + .bind(&r.ak) + .bind(&r.user_id) + .bind(&r.tenant) + .bind(&r.kind) + .bind(&r.content) + .bind(r.sealed) + .bind(r.expires_at_epoch_secs) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("insert content", e))?; + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content + WHERE expires_at_epoch_secs > 0 AND expires_at_epoch_secs <= ?", + ) + .bind(now) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("purge content", e))?; + Ok(r.rows_affected()) + } + + async fn content_for(&self, request_id: &str) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, + sealed, expires_at_epoch_secs FROM request_content WHERE request_id = ? ORDER BY n", + ) + .bind(request_id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read content", e))?; + Ok(rows.iter().map(content_row).collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); // SQLite serializes writers, so the MAX(n)+1 subselect is atomic with the insert @@ -1166,7 +1292,8 @@ impl PostgresStore { vendor_cost_micros BIGINT NOT NULL DEFAULT 0, ptu_spillover BOOLEAN NOT NULL DEFAULT FALSE, user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', - created_at_epoch_secs BIGINT NOT NULL DEFAULT 0)", + created_at_epoch_secs BIGINT NOT NULL DEFAULT 0, + estimated BOOLEAN NOT NULL DEFAULT FALSE)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS security_events ( n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, @@ -1179,6 +1306,14 @@ impl PostgresStore { actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", + "CREATE TABLE IF NOT EXISTS request_content ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + kind TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', + sealed BOOLEAN NOT NULL DEFAULT FALSE, expires_at_epoch_secs BIGINT NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS content_expiry_idx ON request_content (expires_at_epoch_secs)", + "CREATE INDEX IF NOT EXISTS content_request_idx ON request_content (request_id)", "CREATE TABLE IF NOT EXISTS files ( n BIGSERIAL PRIMARY KEY, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -1215,6 +1350,7 @@ impl PostgresStore { "ALTER TABLE billing ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN IF NOT EXISTS request_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN IF NOT EXISTS created_at_epoch_secs BIGINT NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS estimated BOOLEAN NOT NULL DEFAULT FALSE", ] { sqlx::query(ddl) .execute(&pool) @@ -1235,8 +1371,9 @@ impl Store for PostgresStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)", ) .bind(&r.ak) .bind(&r.product) @@ -1254,6 +1391,7 @@ impl Store for PostgresStore { .bind(&r.user_id) .bind(&r.request_id) .bind(r.created_at_epoch_secs) + .bind(r.estimated) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -1280,7 +1418,8 @@ impl Store for PostgresStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated FROM billing ORDER BY n DESC LIMIT $1", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -1419,6 +1558,51 @@ impl Store for PostgresStore { Ok(rows.iter().map(admin_audit_row).collect()) } + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + sqlx::query( + "INSERT INTO request_content (created_at_epoch_secs, request_id, ak, user_id, + tenant, kind, content, sealed, expires_at_epoch_secs) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(r.created_at_epoch_secs) + .bind(&r.request_id) + .bind(&r.ak) + .bind(&r.user_id) + .bind(&r.tenant) + .bind(&r.kind) + .bind(&r.content) + .bind(r.sealed) + .bind(r.expires_at_epoch_secs) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("insert content", e))?; + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content + WHERE expires_at_epoch_secs > 0 AND expires_at_epoch_secs <= $1", + ) + .bind(now) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("purge content", e))?; + Ok(r.rows_affected()) + } + + async fn content_for(&self, request_id: &str) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, + sealed, expires_at_epoch_secs FROM request_content WHERE request_id = $1 ORDER BY n", + ) + .bind(request_id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read content", e))?; + Ok(rows.iter().map(content_row).collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); // consume the sequence explicitly — concurrent writers race a MAX(n)+1 subselect @@ -1745,6 +1929,7 @@ mod tests { completion: i64::MAX, total: i64::MAX, ptu_spillover: false, + estimated: false, }, ); assert_eq!(rec.prompt_tokens, MAX_METERED_TOKENS); @@ -1771,6 +1956,7 @@ mod tests { cost_micros: 42, vendor_cost_micros: 7, ptu_spillover: false, + estimated: false, } } @@ -1908,6 +2094,36 @@ mod tests { exercise_audit(&MemoryStore::default()).await; } + #[tokio::test] + async fn content_retention_stores_and_purges() { + let store = MemoryStore::default(); + let rec = |kind: &str, expires: i64| crate::ContentRecord { + created_at_epoch_secs: 100, + request_id: "req-1".into(), + ak: "ak".into(), + user_id: "u".into(), + tenant: "default".into(), + kind: kind.into(), + content: "hello".into(), + sealed: false, + expires_at_epoch_secs: expires, + }; + store.content_add(&rec("prompt", 200)).await.unwrap(); + store.content_add(&rec("response", 0)).await.unwrap(); + let got = store.content_for("req-1").await.unwrap(); + assert_eq!(got.len(), 2); + + assert_eq!( + store.content_purge(150).await.unwrap(), + 0, + "not yet expired" + ); + assert_eq!(store.content_purge(250).await.unwrap(), 1, "prompt expired"); + let kept = store.content_for("req-1").await.unwrap(); + assert_eq!(kept.len(), 1, "the keep-forever response survives"); + assert_eq!(kept[0].kind, "response"); + } + #[tokio::test] async fn sqlite_audit_roundtrip() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index 42df45b..4ae0d1a 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -1,8 +1,6 @@ -//! Local background tasks. -//! -//! Currently just one pure in-process task: periodic AK daily quota reset. -//! Batch job execution lives in gw-handler::offline (spawned on submit) and -//! needs no separate poller. +//! Local background tasks: periodic AK daily quota reset and retained-content +//! purge. Batch job execution lives in gw-handler::offline (spawned on submit) +//! and needs no separate poller. use std::sync::Arc; use std::time::Duration; @@ -11,6 +9,8 @@ use gw_state::GatewayState; /// The production period: once a day. pub const DAILY: Duration = Duration::from_secs(24 * 60 * 60); +/// Retained content is swept for expiry this often. +pub const PURGE_PERIOD: Duration = Duration::from_secs(60 * 60); /// Spawn the daily quota reset loop. Returns the join handle (abort to stop). /// `period` is configurable so tests don't wait 24h. @@ -29,6 +29,28 @@ pub fn spawn_quota_reset( }) } +/// Spawn the retained-content purge loop: deletes content rows whose retention +/// window has elapsed. Returns the join handle (abort to stop). +pub fn spawn_content_purge( + state: Arc, + period: Duration, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let mut tick = tokio::time::interval(period); + tick.tick().await; + loop { + tick.tick().await; + match state.store.content_purge(gw_state::epoch_secs()).await { + Ok(n) if n > 0 => { + tracing::info!(target: "task", purged = n, "content purge") + } + Ok(_) => {} + Err(e) => tracing::warn!(error = %e, "content purge failed"), + } + } + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index cdd3684..c168fe3 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -362,6 +362,7 @@ async fn bill_realtime_turn( completion: ot, total, ptu_spillover: false, + estimated: false, }, reserved: admit.reserved, tpm_reserved: admit.tpm_reserved, @@ -1390,9 +1391,37 @@ async fn admin_usage_users( Ok(rows) => rows, Err(e) => return gateway_error(e), }; + if q.get("format").map(String::as_str) == Some("csv") { + let mut csv = String::from( + "user_id,model,requests,prompt_tokens,completion_tokens,total_tokens,cost_micros,vendor_cost_micros\n", + ); + for u in &usage { + csv.push_str(&format!( + "{},{},{},{},{},{},{},{}\n", + csv_field(&u.user_id), + csv_field(&u.model), + u.requests, + u.prompt_tokens, + u.completion_tokens, + u.total_tokens, + u.cost_micros, + u.vendor_cost_micros, + )); + } + return ([("content-type", "text/csv")], csv).into_response(); + } Json(json!({ "usage": usage })).into_response() } +/// Quote a CSV field that may contain a comma, quote, or newline (RFC 4180). +fn csv_field(s: &str) -> String { + if s.contains([',', '"', '\n']) { + format!("\"{}\"", s.replace('"', "\"\"")) + } else { + s.to_owned() + } +} + /// GET /admin/audit/events?limit= — content-safety hits (no prompt text), newest /// first. Tenant-scoped: a tenant admin sees only its own tenant's events. async fn admin_security_events( From a096fc11147ecf17809866a260ac627e32305037 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:06:17 +0800 Subject: [PATCH 05/21] feat(phase3): moderation seam + per-user budgets Add a pluggable Moderator trait (AllowModerator default, wired via OnlineHandler::with_moderator, enabled per-tenant by security.moderate with a fail-open/closed posture) and a per-user daily token budget (soft cap keyed by end user, TenantConf.user_daily_token_quota), enforced on REST and realtime. SIEM/OTLP export is the existing access-log + security-event + admin-audit streams. --- crates/config/src/lib.rs | 11 +++ crates/dag/src/nodes.rs | 36 ++++++++- crates/handler/Cargo.toml | 2 +- crates/handler/src/lib.rs | 132 ++++++++++++++++++++++++++++++- crates/handler/src/moderation.rs | 40 ++++++++++ crates/state/src/admission.rs | 53 +++++++++++++ crates/views/src/lib.rs | 15 ++++ 7 files changed, 286 insertions(+), 3 deletions(-) create mode 100644 crates/handler/src/moderation.rs diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index a6fc0df..2bfbd06 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -202,6 +202,13 @@ pub struct SecurityConf { /// enterprise ask: stop staff pasting credentials to a model). #[serde(default)] pub detect_secrets: bool, + /// Route inbound text through the wired external moderator; needs a + /// moderator plugged into the handler (the default one allows everything). + #[serde(default)] + pub moderate: bool, + /// On a moderator error, admit the request (`true`) or deny it (`false`). + #[serde(default)] + pub moderation_fail_open: bool, /// Named regex recognizers. #[serde(default)] pub regex_rules: Vec, @@ -326,6 +333,10 @@ pub struct TenantConf { /// Per-model charged-price overrides (else the model's list price applies). #[serde(default)] pub model_prices: std::collections::HashMap, + /// Per-user daily token budget (a soft cap keyed by end user); `None` = + /// unlimited. Enforced only when the request carries a user attribution. + #[serde(default)] + pub user_daily_token_quota: Option, /// Content-safety policy for this tenant; `None` = use the global `security:`. #[serde(default)] pub security: Option, diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 57db750..018d1a4 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -384,6 +384,31 @@ impl DagNode for AkTpmLimit { } } +/// model_access/user_budget: per-user daily token budget (soft cap), keyed by +/// the effective end user. No-op without a tenant limit or a user attribution. +pub struct UserBudgetGate; + +#[async_trait::async_trait] +impl DagNode for UserBudgetGate { + fn name(&self) -> &'static str { + "user_budget" + } + fn deps(&self) -> &'static [&'static str] { + &["ak_tpm"] + } + async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { + let user = ctx.effective_user_id().to_owned(); + admission::check_user_budget( + ctx.state.governance.as_ref(), + &ctx.cfg, + &ctx.ak.tenant, + &user, + ) + .await + .map_err(limit_denied) + } +} + /// model_access/call_engine: factory dispatch + engine execution + failover. /// On an upstream 5xx the failed account is excluded and reselected once (a /// PTU → paygo spill sets `ptu_spillover`); a second failure propagates as-is. @@ -395,7 +420,7 @@ impl DagNode for CallEngine { "call_engine" } fn deps(&self) -> &'static [&'static str] { - &["ak_tpm"] + &["user_budget"] } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let threshold = ctx.cfg.stability.failure_threshold; @@ -625,6 +650,14 @@ async fn bill( }, ) .await; + admission::consume_user_budget( + ctx.state.governance.as_ref(), + &ctx.cfg, + &ctx.ak.tenant, + ctx.effective_user_id(), + record.total_tokens, + ) + .await; ctx.decide( "cost_calc", format!( @@ -707,6 +740,7 @@ pub fn default_layers() -> Vec { Box::new(ProductQpmLimit), Box::new(ModelQpmLimit), Box::new(AkTpmLimit), + Box::new(UserBudgetGate), Box::new(CallEngine), ], }, diff --git a/crates/handler/Cargo.toml b/crates/handler/Cargo.toml index 6317cdb..9778db7 100644 --- a/crates/handler/Cargo.toml +++ b/crates/handler/Cargo.toml @@ -18,10 +18,10 @@ serde_json = { workspace = true } tokio = { workspace = true } futures = { workspace = true } regex = { workspace = true } +async-trait = { workspace = true } [dev-dependencies] tokio = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } [lints] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 193cb8e..7301bd4 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -3,6 +3,7 @@ //! DAG layers, then the plugin post-stage; `OfflineHandler` reuses the same //! chain for batches. +pub mod moderation; pub mod offline; pub mod plugins; @@ -14,7 +15,7 @@ use gw_config::GatewayConfig; use gw_dag::DagContext; use gw_engines::http_transport::UpstreamPolicy; use gw_engines::{EngineOutcome, SharedTransport}; -use gw_models::{GResult, GatewayError, GatewayRequest, GatewayResponse}; +use gw_models::{Block, GResult, GatewayError, GatewayRequest, GatewayResponse}; use gw_state::{AkInfo, GatewayState, SharedConfig}; pub use offline::{BatchItem, OfflineHandler}; @@ -153,6 +154,7 @@ pub struct OnlineHandler { pub config: SharedConfig, pub transport: SharedTransport, plan: Arc, + moderator: Arc, } impl OnlineHandler { @@ -166,11 +168,19 @@ impl OnlineHandler { config, transport, plan, + moderator: moderation::default_moderator(), }; handler.push_policies(&handler.cfg()); handler } + /// Plug an external content moderator into the pre-stage (enable per tenant + /// via `security.moderate`). + pub fn with_moderator(mut self, moderator: Arc) -> Self { + self.moderator = moderator; + self + } + /// The live config snapshot (cheap atomic load). Introspection surfaces read /// through this so a runtime reload takes effect immediately. pub fn cfg(&self) -> Arc { @@ -243,6 +253,26 @@ impl OnlineHandler { }); return Ok(ctx); } + // external moderation, on the ORIGINAL text (before DLP), when the tenant + // opted in; a real moderator is wired via with_moderator + if sec.moderate + && let Some(block) = self.moderate(&ctx, sec).await + { + ctx.decide("moderation", "denied"); + let response = GatewayResponse { + message: block.message.clone(), + finish_reason: "content_filter".to_owned(), + ..Default::default() + }; + ctx.outcome = Some(EngineOutcome { + response, + http_code: 200, + block, + ..Default::default() + }); + return Ok(ctx); + } + // capture raw prompt before DLP when the tenant retains full content let retention = snap.cfg.retention_for(&ctx.ak.tenant).copied(); let full = matches!( @@ -322,6 +352,38 @@ impl OnlineHandler { Ok(ctx) } + /// Run the wired moderator over the request's inbound text; `Some(Block)` to + /// deny. A moderator error resolves per `moderation_fail_open`. Records a + /// security event on a deny. + async fn moderate(&self, ctx: &DagContext, sec: &gw_config::SecurityConf) -> Option { + let text = ctx + .request + .message + .iter() + .map(|m| m.content.as_str()) + .collect::>() + .join("\n"); + match self.moderator.review(&text).await { + Ok(moderation::Verdict::Allow) => None, + Ok(moderation::Verdict::Deny(reason)) => { + emit_security_event(ctx, "moderation", "block", 1).await; + Some(Block::blocked( + reason, + gw_consts::ErrCode::EMPTY_RESP.value() as i32, + )) + } + Err(e) => { + tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); + (!sec.moderation_fail_open).then(|| { + Block::blocked( + "content moderation is unavailable", + gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, + ) + }) + } + } + } + /// Derive the upstream policies (timeouts/connect-retries) from `cfg` and /// apply them to the transport live. fn push_policies(&self, cfg: &GatewayConfig) { @@ -389,6 +451,74 @@ mod tests { } } + #[derive(Debug)] + struct DenyModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for DenyModerator { + async fn review(&self, _text: &str) -> Result { + Ok(moderation::Verdict::Deny("nope".into())) + } + } + + #[tokio::test] + async fn moderator_denies_when_tenant_enables_it() { + let mut cfg = GatewayConfig::embedded_default().unwrap(); + cfg.security.moderate = true; + let cfg = Arc::new(cfg); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DenyModerator)); + let ctx = h + .run(chat_req("gpt-4o", "hello"), ak(&h).await) + .await + .unwrap(); + let out = ctx.outcome.expect("outcome"); + assert!(out.block.block); + assert_eq!(out.response.finish_reason, "content_filter"); + assert!( + h.state() + .store + .ledger_snapshot(1) + .await + .unwrap() + .1 + .is_empty(), + "a moderated deny bills nothing" + ); + } + + #[tokio::test] + async fn per_user_budget_denies_over_the_cap() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, user_daily_token_quota: 5}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let with_user = |content: &str| GatewayRequest { + is_online: true, + message: vec![ChatMsg::text("user", content)], + model_param_v2: Some(ModelParamV2::with_name(Protocol::OpenaiChat, "gpt-4o")), + user_id: Some("u1".into()), + ..Default::default() + }; + h.run(with_user("first burns the budget"), key.clone()) + .await + .unwrap(); + let err = h + .run(with_user("second is over"), key) + .await + .err() + .expect("second denied by the per-user budget"); + assert_eq!(err.http_status, 429); + } + #[tokio::test] async fn full_pipeline_openai() { let h = handler(); diff --git a/crates/handler/src/moderation.rs b/crates/handler/src/moderation.rs new file mode 100644 index 0000000..d3edd3e --- /dev/null +++ b/crates/handler/src/moderation.rs @@ -0,0 +1,40 @@ +//! The moderation seam: an optional external content-review pass, plugged into +//! the pre-stage. The default [`AllowModerator`] is a deterministic no-op, so +//! the hot path pays nothing until a real moderator is wired in and a tenant +//! turns `moderate` on. External review is latency-bearing and pull-driven, so +//! this ships as a trait, not a built-in integration. + +use std::sync::Arc; + +/// A moderator's decision on one request's text. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Verdict { + Allow, + /// Deny with a user-facing reason. + Deny(String), +} + +/// A pluggable content moderator. `review` sees the request's concatenated +/// inbound text; a real impl calls an external service. An `Err` is a moderator +/// failure — the caller resolves it per its fail-open/closed posture. +#[async_trait::async_trait] +pub trait Moderator: Send + Sync + std::fmt::Debug { + async fn review(&self, text: &str) -> Result; +} + +/// The default: allow everything (a real moderator replaces it via +/// [`crate::OnlineHandler::with_moderator`]). +#[derive(Debug, Default)] +pub struct AllowModerator; + +#[async_trait::async_trait] +impl Moderator for AllowModerator { + async fn review(&self, _text: &str) -> Result { + Ok(Verdict::Allow) + } +} + +/// The default moderator handle. +pub fn default_moderator() -> Arc { + Arc::new(AllowModerator) +} diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index bcc6962..fc97c46 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -25,6 +25,59 @@ pub fn model_quota_key(ak: &str, model: &str) -> String { format!("{ak}|{model}") } +/// The per-user daily-budget counter key, namespaced by tenant so the same +/// user id under two tenants meters separately. +pub fn user_budget_key(tenant: &str, user: &str) -> String { + format!("ub:{tenant}:{user}") +} + +/// Per-user daily token budget (a soft cap): admit while the user is under the +/// tenant's limit. Skipped when the tenant sets no limit or the request carries +/// no user attribution. Consumed at billing via [`consume_user_budget`]. +pub async fn check_user_budget( + gov: &dyn Governance, + cfg: &GatewayConfig, + tenant: &str, + user: &str, +) -> Result<(), String> { + if user.is_empty() { + return Ok(()); + } + let Some(limit) = cfg + .find_tenant(tenant) + .and_then(|t| t.user_daily_token_quota) + else { + return Ok(()); + }; + if gov.quota_check(&user_budget_key(tenant, user), limit).await { + Ok(()) + } else { + Err(format!("daily token budget exhausted for user `{user}`")) + } +} + +/// Accrue actual usage to the per-user daily budget (no-op without a limit or +/// user). Soft cap: check-then-consume, so a burst can overshoot by one turn. +pub async fn consume_user_budget( + gov: &dyn Governance, + cfg: &GatewayConfig, + tenant: &str, + user: &str, + total: i64, +) { + if user.is_empty() || total <= 0 { + return; + } + if cfg + .find_tenant(tenant) + .and_then(|t| t.user_daily_token_quota) + .is_some() + { + gov.quota_consume(&user_budget_key(tenant, user), total) + .await; + } +} + /// The per-(AK, model) daily cap: AK override, else tenant default, else none. pub fn model_quota_limit(cfg: &GatewayConfig, ak: &AkInfo, model: &str) -> Option { ak.model_quotas.get(model).copied().or_else(|| { diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index c168fe3..b75daca 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -295,6 +295,13 @@ async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result "prompt").increment(it as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ot as u64); } From e80b6532fc9b7ce91a7a508de9725734ea2a99a2 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:09:25 +0800 Subject: [PATCH 06/21] review: apply phase-2 simplify findings (zero-alloc part redaction, lazy content capture, shared join) --- crates/handler/src/lib.rs | 61 +++++++++++++++-------------------- crates/handler/src/plugins.rs | 9 ++---- 2 files changed, 28 insertions(+), 42 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 7301bd4..0e4b482 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -80,19 +80,13 @@ async fn persist_content( } else { 0 }; - let redacted_prompt = || { - ctx.request - .message - .iter() - .map(|m| m.content.as_str()) - .collect::>() - .join("\n") + let redacted_prompt = || joined_message_text(&ctx.request); + let redacted_response = || { + ctx.outcome + .as_ref() + .map(|o| o.response.message.clone()) + .unwrap_or_default() }; - let redacted_response = ctx - .outcome - .as_ref() - .map(|o| o.response.message.clone()) - .unwrap_or_default(); let items = [ ( @@ -106,9 +100,9 @@ async fn persist_content( ( "response", if store_full { - raw_response.unwrap_or_else(|| redacted_response.clone()) + raw_response.unwrap_or_else(redacted_response) } else { - redacted_response.clone() + redacted_response() }, ), ]; @@ -138,6 +132,17 @@ async fn persist_content( } } +/// The request's message contents joined with newlines — the text moderation +/// and content retention operate on. +fn joined_message_text(request: &GatewayRequest) -> String { + request + .message + .iter() + .map(|m| m.content.as_str()) + .collect::>() + .join("\n") +} + /// A per-request correlation id: `req--`, time-sortable and /// unique within the process (the seq disambiguates same-millisecond requests). pub fn new_request_id() -> String { @@ -273,20 +278,12 @@ impl OnlineHandler { return Ok(ctx); } - // capture raw prompt before DLP when the tenant retains full content + // capture raw prompt before DLP when the tenant retains full content — + // only worth it if a key exists (full without one downgrades to redacted) let retention = snap.cfg.retention_for(&ctx.ak.tenant).copied(); - let full = matches!( - retention, - Some(r) if r.content == gw_config::ContentLevel::Full - ); - let raw_prompt = full.then(|| { - ctx.request - .message - .iter() - .map(|m| m.content.as_str()) - .collect::>() - .join("\n") - }); + let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) + && gw_state::sealing_available(); + let raw_prompt = capture_raw.then(|| joined_message_text(&ctx.request)); let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { @@ -327,7 +324,7 @@ impl OnlineHandler { } // capture raw response before outbound DLP for full retention - let raw_response = full + let raw_response = capture_raw .then(|| ctx.outcome.as_ref().map(|o| o.response.message.clone())) .flatten(); @@ -356,13 +353,7 @@ impl OnlineHandler { /// deny. A moderator error resolves per `moderation_fail_open`. Records a /// security event on a deny. async fn moderate(&self, ctx: &DagContext, sec: &gw_config::SecurityConf) -> Option { - let text = ctx - .request - .message - .iter() - .map(|m| m.content.as_str()) - .collect::>() - .join("\n"); + let text = joined_message_text(&ctx.request); match self.moderator.review(&text).await { Ok(moderation::Verdict::Allow) => None, Ok(moderation::Verdict::Deny(reason)) => { diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 97cf589..1695eeb 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -219,14 +219,9 @@ fn redact_parts_text(parts: &mut serde_json::Value, pii: bool, secrets: bool) -> let mut hits = 0; for p in arr { if p["type"] == "text" - && let Some(t) = p["text"].as_str() + && let Some(serde_json::Value::String(s)) = p.get_mut("text") { - let mut s = t.to_owned(); - let n = redact_in_place(&mut s, pii, secrets); - if n > 0 { - p["text"] = serde_json::Value::String(s); - hits += n; - } + hits += redact_in_place(s, pii, secrets); } } hits From 92b758f9c5987c3feefba3a003b43c70a900ee9a Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:12:40 +0800 Subject: [PATCH 07/21] docs: document per-tenant audit/security/retention knobs in the embedded config --- conf/gateway.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/conf/gateway.yaml b/conf/gateway.yaml index acdfd49..ede32b6 100644 --- a/conf/gateway.yaml +++ b/conf/gateway.yaml @@ -86,6 +86,17 @@ tenants: fallback_model: gpt-4o-mini model_prices: # e2e: tenant pays a premium over the list price gpt-4o: {input_price_per_1k_micros: 5000, output_price_per_1k_micros: 20000} + # Optional per-tenant knobs (enterprise audit/compliance): + # user_daily_token_quota: 100000 # soft per-end-user daily cap + # retention: {content: redacted, days: 30} # none | redacted | full (full needs GW_CONTENT_KEY) + # security: # overrides the global `security:` for this tenant + # blocklist: ["forbidden"] + # blocklist_action: flag # block | flag | shadow + # dlp_redact: true + # detect_secrets: true # mask API keys / credentials + # moderate: false # route to the wired external moderator + # regex_rules: + # - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} # Public model name → wire protocol, optional provider binding (only that # provider's accounts serve the model), plus demo prices (micros per 1k tokens). From 9a895f557613884b3ed477ce990d42b21a59cab6 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:13:14 +0800 Subject: [PATCH 08/21] docs: per-user billing, enterprise content policy, and audit trails in governance.md --- docs/governance.md | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/docs/governance.md b/docs/governance.md index f9975a1..687fd06 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -96,3 +96,37 @@ redacted text replayed** rather than forwarded token-by-token — DLP trades incremental delivery for a guarantee that no unmasked text reaches the client. Turn `dlp_redact` off to keep incremental delivery; note the embedded demo config ships with it on. + +## Per-user attribution and billing + +Every ledger row carries a `user_id`, `request_id`, and `created_at_epoch_secs`. +The effective user is the key's `owner` (one key = one user, the enterprise +model) if set, else request metadata — the `x-gw-user` header, OpenAI's `user` +field, or Anthropic's `metadata.user_id`. `owner` is authoritative; the +metadata hint is only trusted for shared keys. `GET /admin/usage/users?user=&since=&until=` +returns per-(user, model) cost over a billing period (add `format=csv` for +export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap. + +## Enterprise content policy + +`security:` is global by default; a tenant may override it whole with +`tenants[].security`. Beyond the blocklist and DLP, a policy can carry +`blocklist_action` (`block` denies, `flag` records, `shadow` trials a rule), +named `regex_rules` (each with its own action), and `detect_secrets` (masks API +keys / credentials). Every rule that fires is recorded — without the offending +text — to the security-event stream, queryable at `GET /admin/audit/events`. +Set `moderate: true` to route inbound text through an external moderator wired +into the handler (`moderation_fail_open` picks the posture on a moderator +error). + +## Audit trails + +- **Content-safety events** (`/admin/audit/events`): who/which-rule/what-action, + no prompt text; tenant-scoped. +- **Admin operations** (`/admin/audit/ops`, global admin only): every key CRUD, + config publish, and reload with actor, action, target, and source IP. +- **Content retention** (`tenants[].retention`): `none` (default), `redacted` + (post-DLP text), or `full` (raw). Stored in `request_content`, sealed at rest + with XChaCha20-Poly1305 under `GW_CONTENT_KEY` (64 hex chars); `full` refuses + to store raw without a key and falls back to redacted. `days` sets expiry; an + hourly purge deletes elapsed content. From 045d1496e11ff4c241366a4bf035def544b1b0c7 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:24:18 +0800 Subject: [PATCH 09/21] fix(review): rev-safety + rev-isolation findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - moderation & content retention now cover non-chat surfaces (plugins::inbound_text walks message + parts + raw + typed), closing a Responses/embeddings moderation fail-open. - realtime frames run the full policy (regex + blocklist_action) via realtime_frame_scan, not just the blocklist — WebSocket is not a bypass; per-rule hits recorded. - audit source_ip prefers x-real-ip / rightmost XFF (leftmost is client-spoofable). - CSV export neutralizes spreadsheet formula injection. - /admin/keys pagination filters tenant scope in the store, before paging. - reject ':' in tenant names (per-user budget key aliasing). --- crates/config/src/lib.rs | 16 ++++ crates/handler/src/lib.rs | 17 +--- crates/handler/src/plugins.rs | 157 ++++++++++++++++++++++++++++++++-- crates/state/src/keystore.rs | 19 ++-- crates/state/src/lib.rs | 21 +++-- crates/views/src/lib.rs | 115 ++++++++++++++++++------- 6 files changed, 284 insertions(+), 61 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 2bfbd06..3427340 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -611,6 +611,16 @@ impl GatewayConfig { check_unique("product", self.products.iter().map(|p| p.name.as_str()))?; check_unique("provider", self.providers.iter().map(|p| p.name.as_str()))?; check_unique("tenant", self.tenants.iter().map(|t| t.name.as_str()))?; + // ':' separates the per-user budget counter key `ub:{tenant}:{user}`, so a + // colon in a tenant name could alias another tenant's budget counter + for t in &self.tenants { + if t.name.contains(':') { + return Err(ConfigError::DuplicateName { + kind: "tenant (':' not allowed in name)", + name: t.name.clone(), + }); + } + } // a typo'd tenant would silently fall back to the unrestricted default — reject at load for k in &self.access_keys { if !self.is_known_tenant(&k.tenant) { @@ -1067,6 +1077,12 @@ tenants: [{name: t1}, {name: t1}] Err(ConfigError::DuplicateName { kind: "tenant", .. }) )); + let colon = "listen: {host: h, port: 1}\ntenants: [{name: 'a:b'}]"; + assert!( + GatewayConfig::from_yaml(colon).is_err(), + "a colon in a tenant name is rejected (budget-key aliasing)" + ); + let bad_quota = r#" listen: {host: h, port: 1} models: [{name: m1, protocol: openai-chat}] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 0e4b482..da565d9 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -80,7 +80,7 @@ async fn persist_content( } else { 0 }; - let redacted_prompt = || joined_message_text(&ctx.request); + let redacted_prompt = || plugins::inbound_text(&ctx.request); let redacted_response = || { ctx.outcome .as_ref() @@ -132,17 +132,6 @@ async fn persist_content( } } -/// The request's message contents joined with newlines — the text moderation -/// and content retention operate on. -fn joined_message_text(request: &GatewayRequest) -> String { - request - .message - .iter() - .map(|m| m.content.as_str()) - .collect::>() - .join("\n") -} - /// A per-request correlation id: `req--`, time-sortable and /// unique within the process (the seq disambiguates same-millisecond requests). pub fn new_request_id() -> String { @@ -283,7 +272,7 @@ impl OnlineHandler { let retention = snap.cfg.retention_for(&ctx.ak.tenant).copied(); let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) && gw_state::sealing_available(); - let raw_prompt = capture_raw.then(|| joined_message_text(&ctx.request)); + let raw_prompt = capture_raw.then(|| plugins::inbound_text(&ctx.request)); let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { @@ -353,7 +342,7 @@ impl OnlineHandler { /// deny. A moderator error resolves per `moderation_fail_open`. Records a /// security event on a deny. async fn moderate(&self, ctx: &DagContext, sec: &gw_config::SecurityConf) -> Option { - let text = joined_message_text(&ctx.request); + let text = plugins::inbound_text(&ctx.request); match self.moderator.review(&text).await { Ok(moderation::Verdict::Allow) => None, Ok(moderation::Verdict::Deny(reason)) => { diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 1695eeb..bf3cc2e 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -54,6 +54,13 @@ pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanO } } + tally_scan(sec, blocklist_count, regex_counts) +} + +/// Build a [`ScanOutcome`] from per-rule hit counts: a rule with count > 0 is a +/// hit at its action; any block-action hit denies. Shared by the REST scan and +/// the realtime frame scan so the two apply identical action semantics. +fn tally_scan(sec: &SecurityConf, blocklist_count: i64, regex_counts: Vec) -> ScanOutcome { let mut hits = Vec::new(); if blocklist_count > 0 { hits.push(RuleHit { @@ -78,6 +85,75 @@ pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanO ScanOutcome { block, hits } } +/// All inbound text a request carries — message content, multimodal text +/// parts, the Responses raw body, and the family typed params. The one text +/// view moderation and content retention operate on, so a non-chat surface +/// (input in `raw`/typed with an empty `message`) is not a bypass. The +/// typed-field list mirrors [`for_each_typed_text`]. +pub fn inbound_text(request: &GatewayRequest) -> String { + let mut out = String::new(); + for m in &request.message { + push_text(&mut out, &m.content); + if let Some(serde_json::Value::Array(parts)) = &m.parts { + for p in parts { + if p["type"] == "text" + && let Some(t) = p["text"].as_str() + { + push_text(&mut out, t); + } + } + } + } + if let Some(param) = &request.model_param_v2 { + collect_json_text(¶m.raw, &mut out); + if let Some(typed) = ¶m.typed { + collect_typed_text(typed, &mut out); + } + } + out +} + +fn push_text(out: &mut String, s: &str) { + if !s.is_empty() { + if !out.is_empty() { + out.push('\n'); + } + out.push_str(s); + } +} + +fn collect_json_text(v: &serde_json::Value, out: &mut String) { + match v { + serde_json::Value::String(s) => push_text(out, s), + serde_json::Value::Array(a) => a.iter().for_each(|x| collect_json_text(x, out)), + serde_json::Value::Object(o) => o.values().for_each(|x| collect_json_text(x, out)), + _ => {} + } +} + +fn collect_typed_text(typed: &gw_models::TypedParams, out: &mut String) { + use gw_models::TypedParams as T; + match typed { + T::Chat(p) => { + if let Some(s) = &p.system { + push_text(out, s); + } + if let Some(t) = &p.tools { + collect_json_text(t, out); + } + if let Some(t) = &p.tool_choice { + collect_json_text(t, out); + } + } + T::Embeddings(p) => p.input.iter().for_each(|s| push_text(out, s)), + T::AudioTts(p) => push_text(out, &p.input), + T::Image(p) => push_text(out, &p.prompt), + T::Video(p) => push_text(out, &p.prompt), + T::Search(p) => push_text(out, &p.query), + T::AudioStt(_) => {} + } +} + /// Case-insensitive blocklist test; terms are pre-lowercased at config load. /// ASCII text matches without allocating; non-ASCII falls back to a lowercase copy. fn blocklist_hit(sec: &SecurityConf, text: &str) -> bool { @@ -139,15 +215,24 @@ fn for_each_typed_text( } } -/// Blocklist scan over a realtime frame's text-bearing fields — the same terms -/// every REST surface blocks, so the WebSocket surface is not a bypass. -pub fn realtime_frame_blocked(sec: &SecurityConf, frame: &mut serde_json::Value) -> Option { - if sec.blocklist.is_empty() { - return None; +/// Scan a realtime frame's text-bearing fields against the blocklist AND the +/// regex recognizers, honoring their actions — the same policy every REST +/// surface runs, so the WebSocket surface is not a bypass. Returns the same +/// [`ScanOutcome`] as [`security_check`]: a block-action hit denies the frame. +pub fn realtime_frame_scan(sec: &SecurityConf, frame: &mut serde_json::Value) -> ScanOutcome { + if sec.blocklist.is_empty() && sec.regexes.is_empty() { + return ScanOutcome::default(); } - let hits = - gw_engines::realtime::visit_frame_text(frame, &mut |s| usize::from(blocklist_hit(sec, s))); - (hits > 0).then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)) + let mut blocklist_count = 0i64; + let mut regex_counts = vec![0i64; sec.regexes.len()]; + gw_engines::realtime::visit_frame_text(frame, &mut |s| { + blocklist_count += i64::from(blocklist_hit(sec, s)); + for (i, r) in sec.regexes.iter().enumerate() { + regex_counts[i] += r.re.find_iter(s).count() as i64; + } + 0 + }); + tally_scan(sec, blocklist_count, regex_counts) } /// DLP-redact a realtime frame's text-bearing fields in place; the hit count. @@ -403,6 +488,62 @@ mod tests { assert!(security_check(&s, &mut req).block.is_some()); } + #[test] + fn inbound_text_covers_non_chat_raw_and_typed() { + use gw_models::{ModelParamV2, TypedParams}; + // Responses-style: input rides in raw, message is empty + let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, "m"); + param.raw = serde_json::json!({"input": "secret responses text", "model": "m"}); + let req = GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }; + assert!(inbound_text(&req).contains("secret responses text")); + + // embeddings-style: input rides in typed, message is empty + let mut param = ModelParamV2::with_name(gw_consts::Protocol::Embeddings, "m"); + param.typed = Some(TypedParams::Embeddings(gw_models::EmbeddingParams { + input: vec!["typed embed text".into()], + dimensions: None, + })); + let req = GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }; + assert!(inbound_text(&req).contains("typed embed text")); + } + + #[test] + fn realtime_frame_scan_honors_regex_and_actions() { + let mut s = SecurityConf { + regex_rules: vec![gw_config::RegexRule { + name: "ssn".into(), + pattern: r"\d{3}-\d{2}-\d{4}".into(), + action: Action::Block, + }], + ..Default::default() + }; + s.regexes = vec![gw_config::CompiledRule { + name: "ssn".into(), + action: Action::Block, + re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), + }]; + let mut frame = serde_json::json!({"type":"input_text","text":"my ssn is 123-45-6789"}); + let out = realtime_frame_scan(&s, &mut frame); + assert!(out.block.is_some(), "regex Block denies on realtime too"); + + // flag blocklist: hit recorded, not blocked + let s2 = SecurityConf { + blocklist: vec!["watch".into()], + blocklist_action: Action::Flag, + ..Default::default() + }; + let mut frame = serde_json::json!({"type":"input_text","text":"please watch this"}); + let out = realtime_frame_scan(&s2, &mut frame); + assert!(out.block.is_none(), "flag does not block realtime"); + assert_eq!(out.hits.len(), 1); + } + #[test] fn flag_action_records_a_hit_without_blocking() { let s = SecurityConf { diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index 1b418e3..390df12 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -31,9 +31,12 @@ pub trait KeyStore: Send + Sync + std::fmt::Debug { async fn patch(&self, ak: &str, patch: &KeyPatch) -> GResult>; /// Remove a key regardless of source; whether it existed. async fn revoke(&self, ak: &str) -> GResult; - /// A page of keys, sorted by ak (stable). `offset`/`limit` bound the scan so - /// a fleet key table with millions of rows never loads whole. - async fn list(&self, offset: usize, limit: usize) -> GResult>; + /// A page of keys, sorted by ak (stable), optionally confined to one tenant + /// (a tenant admin's scope). `offset`/`limit` bound the scan so a fleet key + /// table with millions of rows never loads whole — and the tenant filter + /// applies BEFORE paging, so a scoped page isn't emptied by a later filter. + async fn list(&self, tenant: Option<&str>, offset: usize, limit: usize) + -> GResult>; /// Re-apply the config file's key set, leaving admin-created keys untouched. async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> GResult<()>; } @@ -187,12 +190,18 @@ impl KeyStore for PostgresKeyStore { Ok(n > 0) } - async fn list(&self, offset: usize, limit: usize) -> GResult> { + async fn list( + &self, + tenant: Option<&str>, + offset: usize, + limit: usize, + ) -> GResult> { let rows = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys - ORDER BY ak LIMIT $1 OFFSET $2", + WHERE ($1::text IS NULL OR tenant = $1) ORDER BY ak LIMIT $2 OFFSET $3", ) + .bind(tenant) .bind(limit.min(i64::MAX as usize) as i64) .bind(offset.min(i64::MAX as usize) as i64) .fetch_all(&self.pool) diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 350cfe6..1bdc8d0 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -170,9 +170,15 @@ impl AkAuth { Some(e.0.clone()) } - /// A page of keys, sorted by ak (stable), `offset..offset+limit`. - pub fn list(&self, offset: usize, limit: usize) -> Vec { - let mut keys: Vec = self.keys.iter().map(|e| e.value().0.clone()).collect(); + /// A page of keys, sorted by ak (stable), optionally confined to `tenant`, + /// `offset..offset+limit` — the filter applies before paging. + pub fn list(&self, tenant: Option<&str>, offset: usize, limit: usize) -> Vec { + let mut keys: Vec = self + .keys + .iter() + .map(|e| e.value().0.clone()) + .filter(|k| tenant.is_none_or(|t| t == k.tenant)) + .collect(); keys.sort_by(|a, b| a.ak.cmp(&b.ak)); keys.into_iter().skip(offset).take(limit).collect() } @@ -210,8 +216,13 @@ impl KeyStore for AkAuth { async fn revoke(&self, ak: &str) -> gw_models::GResult { Ok(AkAuth::revoke(self, ak)) } - async fn list(&self, offset: usize, limit: usize) -> gw_models::GResult> { - Ok(AkAuth::list(self, offset, limit)) + async fn list( + &self, + tenant: Option<&str>, + offset: usize, + limit: usize, + ) -> gw_models::GResult> { + Ok(AkAuth::list(self, tenant, offset, limit)) } async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> gw_models::GResult<()> { AkAuth::reload_config_keys(self, keys); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index b75daca..53d0427 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -419,11 +419,13 @@ async fn realtime_session( .await; continue; }; - // same blocklist + inbound DLP every REST surface runs, tenant policy first + // same policy (blocklist + regex + actions) + inbound DLP every REST + // surface runs, tenant policy first let cfg = s.handler.cfg(); let sec = cfg.security_for(&ak.tenant); - if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut ev) { - emit_rt_block(&s, &ak).await; + let scan = gw_handler::plugins::realtime_frame_scan(sec, &mut ev); + emit_rt_hits(&s, &ak, &scan.hits).await; + if let Some(block) = scan.block { let _ = socket .send(send(json!({"type":"error","message": block.message}))) .await; @@ -582,11 +584,13 @@ async fn realtime_bridge( Some(Ok(_)) => continue, // ping/pong handled by the ws stacks }; if let Some(mut frame) = frame { - // same blocklist + inbound DLP every REST surface runs, tenant policy first + // same policy (blocklist + regex + actions) + inbound DLP every + // REST surface runs, tenant policy first let cfg = s.handler.cfg(); let sec = cfg.security_for(&ak.tenant); - if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut frame) { - emit_rt_block(&s, &ak).await; + let scan = gw_handler::plugins::realtime_frame_scan(sec, &mut frame); + emit_rt_hits(&s, &ak, &scan.hits).await; + if let Some(block) = scan.block { if cl_tx .send(CMsg::Text(send_err(block.message).to_string().into())) .await @@ -952,32 +956,39 @@ fn q_num( q.get(key).and_then(|v| v.parse().ok()).unwrap_or(default) } -/// Record a realtime blocklist block to the security-event stream (parity with -/// the REST surfaces). Per-frame DLP redactions are not recorded — too hot. -async fn emit_rt_block(s: &AppState, ak: &AkInfo) { - let event = gw_state::SecurityEvent { - created_at_epoch_secs: gw_state::epoch_secs(), - request_id: String::new(), - ak: ak.ak.clone(), - user_id: ak.owner.clone().unwrap_or_default(), - tenant: ak.tenant.clone(), - surface: "realtime".to_owned(), - rule: "blocklist".to_owned(), - action: "block".to_owned(), - hits: 1, - }; - if let Err(e) = s.handler.state().store.security_event_add(&event).await { - tracing::warn!(error = %e, "realtime security event write failed"); +/// Record a realtime frame's content-safety hits to the security-event stream +/// (parity with the REST surfaces). Per-frame DLP redactions are not recorded — +/// too hot. +async fn emit_rt_hits(s: &AppState, ak: &AkInfo, hits: &[gw_handler::plugins::RuleHit]) { + for hit in hits { + let event = gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: String::new(), + ak: ak.ak.clone(), + user_id: ak.owner.clone().unwrap_or_default(), + tenant: ak.tenant.clone(), + surface: "realtime".to_owned(), + rule: hit.rule.clone(), + action: hit.action.as_str().to_owned(), + hits: hit.count, + }; + if let Err(e) = s.handler.state().store.security_event_add(&event).await { + tracing::warn!(error = %e, "realtime security event write failed"); + } } } -/// The caller IP for the audit trail, from the LB's forwarding headers. +/// The caller IP for the audit trail. Prefers `x-real-ip` (the immediate peer +/// the trusted proxy sets) and, failing that, the RIGHTMOST `x-forwarded-for` +/// hop (appended by that proxy) — never the leftmost, which a client can forge. fn source_ip(headers: &HeaderMap) -> String { + if let Some(ip) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) { + return ip.trim().to_owned(); + } headers .get("x-forwarded-for") .and_then(|v| v.to_str().ok()) - .and_then(|v| v.split(',').next()) - .or_else(|| headers.get("x-real-ip").and_then(|v| v.to_str().ok())) + .and_then(|v| v.rsplit(',').next()) .map(|s| s.trim().to_owned()) .unwrap_or_default() } @@ -1337,7 +1348,16 @@ async fn admin_key_list( }; let offset = q_num(&q, "offset", 0); let limit = q_num(&q, "limit", KEY_PAGE_DEFAULT); - let listed = match s.handler.state().auth.list(offset, limit).await { + // filter by the caller's scope IN the store, before paging, so a tenant + // admin's page isn't emptied by a post-hoc filter over the global table + let tenant = scope.tenant_filter(&q); + let listed = match s + .handler + .state() + .auth + .list(tenant.as_deref(), offset, limit) + .await + { Ok(v) => v, Err(e) => return gateway_error(e), }; @@ -1428,12 +1448,24 @@ async fn admin_usage_users( Json(json!({ "usage": usage })).into_response() } -/// Quote a CSV field that may contain a comma, quote, or newline (RFC 4180). +/// A CSV field, RFC-4180 quoted AND neutralized against spreadsheet formula +/// injection: a field opening with a formula trigger (`= + - @` / tab / CR) is +/// prefixed with `'` so Excel/Sheets treat it as text (the value is +/// attacker-controlled — it can carry a user id). fn csv_field(s: &str) -> String { - if s.contains([',', '"', '\n']) { - format!("\"{}\"", s.replace('"', "\"\"")) + let needs_prefix = s + .chars() + .next() + .is_some_and(|c| matches!(c, '=' | '+' | '-' | '@' | '\t' | '\r')); + let body = if needs_prefix { + format!("'{s}") } else { s.to_owned() + }; + if body.contains([',', '"', '\n', '\r']) { + format!("\"{}\"", body.replace('"', "\"\"")) + } else { + body } } @@ -2936,6 +2968,31 @@ mod tests { assert!(by_user[0].total_tokens > 0); } + #[test] + fn source_ip_prefers_trusted_hop_over_spoofable_xff() { + let mut h = HeaderMap::new(); + h.insert("x-forwarded-for", "1.2.3.4, 10.0.0.9".parse().unwrap()); + assert_eq!(source_ip(&h), "10.0.0.9", "rightmost (proxy-appended) hop"); + h.insert("x-real-ip", "10.0.0.5".parse().unwrap()); + assert_eq!(source_ip(&h), "10.0.0.5", "x-real-ip wins over XFF"); + } + + #[test] + fn csv_field_neutralizes_formula_injection() { + assert_eq!(csv_field("alice"), "alice"); + assert_eq!( + csv_field("+cmd"), + "'+cmd", + "formula trigger prefixed with '" + ); + assert_eq!( + csv_field("=SUM(A1,A2)"), + "\"'=SUM(A1,A2)\"", + "prefixed AND quoted (has a comma)" + ); + assert_eq!(csv_field("a,b"), "\"a,b\""); + } + #[tokio::test] async fn dlp_hit_is_recorded_as_a_security_event() { let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); From 5e524131df3aa7c4a31c008a24b0426800c85d2b Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 13:34:53 +0800 Subject: [PATCH 10/21] fix(realtime): resolve x-gw-user for billing, budget, and security events The realtime/WebSocket surface read ak.owner directly, so a shared (ownerless) key attributed every turn to an empty user: per-user cost tracking lost all realtime usage, the per-user daily token budget never fired (a caller could evade the cap by using /v1/realtime), and security events recorded an empty user. Thread the connect-time x-gw-user hint through the realtime path with owner-first resolution, matching the REST effective_user_id. --- crates/views/src/lib.rs | 126 ++++++++++++++++++++++++++++++++-------- 1 file changed, 101 insertions(+), 25 deletions(-) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 53d0427..15a2ade 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -226,8 +226,10 @@ async fn realtime_ws( }; // select "realtime" so subprotocol-offering clients get a valid handshake let ws = ws.protocols(["realtime"]); + // client attribution hint captured at connect (no per-turn body user field) + let hint = user_header(&headers).unwrap_or_default(); if account.endpoint.is_empty() { - ws.on_upgrade(move |socket| realtime_session(socket, s, ak, model, mt, account.name)) + ws.on_upgrade(move |socket| realtime_session(socket, s, ak, model, mt, account.name, hint)) } else if gw_engines::realtime::is_gemini_realtime(&account.provider) { // no pre-generation gate signal in this dialect — refuse rather than bill after the fact error_response( @@ -238,7 +240,7 @@ async fn realtime_ws( ), ) } else { - ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, model, mt, account)) + ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, model, mt, account, hint)) } } @@ -248,6 +250,10 @@ async fn realtime_ws( /// from the admission config when a reload lands mid-turn). struct RealtimeAdmit { ak: AkInfo, + /// Effective attribution user for this turn: the key's owner if set, else + /// the client's connect-time `x-gw-user` hint; empty for an ownerless key + /// with no hint. Captured at admission so billing and budget agree. + user: String, reserved: i64, /// Tokens reserved in the AK TPM window; `None` when the key has no TPM cap. tpm_reserved: Option, @@ -268,6 +274,16 @@ impl RealtimeAdmit { } } +/// The attribution user for a realtime turn: the key's owner is authoritative, +/// falling back to the client's connect-time `x-gw-user` hint (realtime has no +/// per-turn body `user` field). The REST counterpart is `effective_user_id`. +fn realtime_user<'a>(ak: &'a AkInfo, hint: &'a str) -> &'a str { + ak.owner + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or(hint) +} + /// The REST admission chain applied per realtime generation via the shared /// [`admission`] checks, with the key re-fetched each turn so mid-session /// bans/de-entitlements take effect. Two deliberate divergences from the DAG: @@ -275,7 +291,12 @@ impl RealtimeAdmit { /// mid-stream), and the reserve is a fixed turn estimate. Reserves are taken /// last so a denial never leaves one behind; a failed TPM reserve rolls back /// the daily reserve just taken. -async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result { +async fn realtime_gate( + s: &AppState, + ak: &AkInfo, + model: &str, + hint: &str, +) -> Result { let snap = s.handler.config.load(); let (cfg, state) = (&snap.cfg, &snap.state); let ak = match state.auth.authenticate(&ak.ak).await { @@ -295,13 +316,7 @@ async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result Result { - let admit = match realtime_gate(&s, &ak, &model).await { + let admit = match realtime_gate(&s, &ak, &model, &hint).await { Ok(a) => a, Err(denied) => { let _ = socket @@ -518,6 +536,7 @@ async fn realtime_bridge( model: String, mt: gw_consts::Protocol, account: gw_models::Account, + hint: String, ) { use axum::extract::ws::Message as CMsg; use futures::{SinkExt, StreamExt}; @@ -589,7 +608,7 @@ async fn realtime_bridge( let cfg = s.handler.cfg(); let sec = cfg.security_for(&ak.tenant); let scan = gw_handler::plugins::realtime_frame_scan(sec, &mut frame); - emit_rt_hits(&s, &ak, &scan.hits).await; + emit_rt_hits(&s, &ak, &scan.hits, &hint).await; if let Some(block) = scan.block { if cl_tx .send(CMsg::Text(send_err(block.message).to_string().into())) @@ -605,7 +624,7 @@ async fn realtime_bridge( } // gate each generation trigger, not every control frame if is_response_create(&frame) { - match realtime_gate(&s, &ak, &model).await { + match realtime_gate(&s, &ak, &model, &hint).await { Ok(admit) => { pending.push_back(admit); generations += 1; @@ -654,7 +673,7 @@ async fn realtime_bridge( // server-VAD: OpenAI auto-starts a turn with no client // response.create — gate it here like a manual one else if realtime_turn_started(&account.provider, &v) && pending.is_empty() { - match realtime_gate(&s, &ak, &model).await { + match realtime_gate(&s, &ak, &model, &hint).await { Ok(admit) => pending.push_back(admit), Err(denied) => { let _ = up_tx @@ -684,8 +703,10 @@ async fn realtime_bridge( .authenticate(&ak.ak) .await .unwrap_or_else(|| ak.clone()); + let user = realtime_user(&billed, &hint).to_owned(); let unreserved = RealtimeAdmit { ak: billed, + user, reserved: 0, tpm_reserved: None, at: gw_state::epoch_secs(), @@ -959,13 +980,18 @@ fn q_num( /// Record a realtime frame's content-safety hits to the security-event stream /// (parity with the REST surfaces). Per-frame DLP redactions are not recorded — /// too hot. -async fn emit_rt_hits(s: &AppState, ak: &AkInfo, hits: &[gw_handler::plugins::RuleHit]) { +async fn emit_rt_hits( + s: &AppState, + ak: &AkInfo, + hits: &[gw_handler::plugins::RuleHit], + hint: &str, +) { for hit in hits { let event = gw_state::SecurityEvent { created_at_epoch_secs: gw_state::epoch_secs(), request_id: String::new(), ak: ak.ak.clone(), - user_id: ak.owner.clone().unwrap_or_default(), + user_id: realtime_user(ak, hint).to_owned(), tenant: ak.tenant.clone(), surface: "realtime".to_owned(), rule: hit.rule.clone(), @@ -3066,18 +3092,18 @@ mod tests { let gov = || s.handler.state().governance.clone(); let used = || async { gov().quota_used(&ak.ak).await }; - let a1 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a1 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, REALTIME_TURN_RESERVE, "reserved up front"); bill_realtime_turn(&a1, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 30, 70).await; assert_eq!(used().await, 100, "settled to actual (30 + 70)"); - let a2 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a2 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); gov().quota_settle(&a2.ak.ak, -a2.reserved, a2.at).await; assert_eq!(used().await, 100, "dropped turn refunded whole"); - let a3 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a3 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); let ledger_before = s.handler.state().store.ledger_snapshot(1).await.unwrap().0; bill_realtime_turn(&a3, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 0, 0).await; @@ -3103,14 +3129,14 @@ mod tests { .unwrap(); let gov = s.handler.state().governance.clone(); - let a1 = realtime_gate(&s, &ak, "gpt-4o") + let a1 = realtime_gate(&s, &ak, "gpt-4o", "") .await .expect("first admits"); assert_eq!(a1.tpm_reserved, Some(REALTIME_TURN_RESERVE)); let daily_before = gov.quota_used(&ak.ak).await; assert!( - realtime_gate(&s, &ak, "gpt-4o").await.is_err(), + realtime_gate(&s, &ak, "gpt-4o", "").await.is_err(), "second turn denied by the TPM reserve" ); assert_eq!( @@ -3132,7 +3158,7 @@ mod tests { let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); let ak = s.handler.state().auth.authenticate("k-rt").await.unwrap(); - let admit = realtime_gate(&s, &ak, "rt").await.expect("admit"); + let admit = realtime_gate(&s, &ak, "rt", "").await.expect("admit"); s.handler .reload(GatewayConfig::from_yaml(&price(2_000_000)).unwrap()) .await @@ -3146,6 +3172,56 @@ mod tests { ); } + #[tokio::test] + async fn realtime_attributes_user_from_owner_then_header_hint() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt, protocol: realtime, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 1000}]\naccess_keys: [{ak: k-shared, product: p, qps: 10, daily_token_quota: 100000}, {ak: k-owned, product: p, qps: 10, daily_token_quota: 100000, owner: bob}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + + let shared = s + .handler + .state() + .auth + .authenticate("k-shared") + .await + .unwrap(); + let admit = realtime_gate(&s, &shared, "rt", "alice") + .await + .expect("admit"); + assert_eq!(admit.user, "alice", "ownerless key attributes to the hint"); + bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 40, 60).await; + + let owned = s + .handler + .state() + .auth + .authenticate("k-owned") + .await + .unwrap(); + let admit = realtime_gate(&s, &owned, "rt", "mallory") + .await + .expect("admit"); + assert_eq!( + admit.user, "bob", + "owner is authoritative over a spoofed hint" + ); + bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 10, 20).await; + + let (_, records) = s.handler.state().store.ledger_snapshot(2).await.unwrap(); + let users: std::collections::HashSet<&str> = + records.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice"), + "shared-key turn billed to header hint" + ); + assert!(users.contains("bob"), "owned-key turn billed to owner"); + assert!( + !users.contains("mallory"), + "spoofed hint never overrides owner" + ); + } + #[test] fn finish_reason_mapping_both_directions() { assert_eq!(finish_openai("end_turn"), "stop"); From 2b8fee35b6a266ca8469f37f294819b32e93942b Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 15:21:26 +0800 Subject: [PATCH 11/21] refactor(attribution): single AkInfo::attributed_user across REST and realtime effective_user_id (REST) treated Some("") as a set owner and returned empty, while realtime filtered an empty owner and fell back to the hint. An empty string is never an identity; collapse both onto one AkInfo method that treats empty owner as absent, so the two surfaces cannot diverge. --- crates/dag/src/context.rs | 8 +++----- crates/state/src/lib.rs | 31 +++++++++++++++++++++++++++++++ crates/views/src/lib.rs | 18 ++++-------------- 3 files changed, 38 insertions(+), 19 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index e6fbb70..25dcf76 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -72,13 +72,11 @@ impl DagContext { } /// The effective end user: the key's `owner` (authoritative) else request - /// metadata; `""` when neither is present. + /// metadata; `""` when neither is present. Resolution lives on [`AkInfo`] so + /// REST and realtime can't diverge on an empty owner. pub fn effective_user_id(&self) -> &str { self.ak - .owner - .as_deref() - .or(self.request.user_id.as_deref()) - .unwrap_or_default() + .attributed_user(self.request.user_id.as_deref().unwrap_or_default()) } /// The decision trail as `"stage: detail"` lines. diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 1bdc8d0..0c9bd89 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -64,6 +64,18 @@ impl AkInfo { } } + /// The attributed end user: this key's `owner` when set to a non-empty value + /// (authoritative), else the caller-supplied `fallback` — request metadata + /// on REST surfaces, the `x-gw-user` connect hint on realtime. Empty when + /// neither is present. The one resolution every surface shares so an empty + /// owner can't fall back on one surface and swallow attribution on another. + pub fn attributed_user<'a>(&'a self, fallback: &'a str) -> &'a str { + self.owner + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or(fallback) + } + /// Apply a partial quota/lifecycle patch. pub fn apply_patch(&mut self, patch: &KeyPatch) { if let Some(v) = patch.qps { @@ -876,6 +888,25 @@ mod tests { } } + #[test] + fn attributed_user_prefers_nonempty_owner_else_fallback() { + let mut ak = ak_info("k"); + assert_eq!(ak.attributed_user("hint"), "hint", "no owner → fallback"); + ak.owner = Some(String::new()); + assert_eq!( + ak.attributed_user("hint"), + "hint", + "empty owner is not an identity → fallback" + ); + ak.owner = Some("alice".into()); + assert_eq!( + ak.attributed_user("hint"), + "alice", + "owner wins over fallback" + ); + assert_eq!(ak.attributed_user(""), "alice"); + } + #[test] fn rate_limit_qps_change_takes_effect_immediately() { let rl = RateLimiter::default(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 15a2ade..3d4b43e 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -274,16 +274,6 @@ impl RealtimeAdmit { } } -/// The attribution user for a realtime turn: the key's owner is authoritative, -/// falling back to the client's connect-time `x-gw-user` hint (realtime has no -/// per-turn body `user` field). The REST counterpart is `effective_user_id`. -fn realtime_user<'a>(ak: &'a AkInfo, hint: &'a str) -> &'a str { - ak.owner - .as_deref() - .filter(|s| !s.is_empty()) - .unwrap_or(hint) -} - /// The REST admission chain applied per realtime generation via the shared /// [`admission`] checks, with the key re-fetched each turn so mid-session /// bans/de-entitlements take effect. Two deliberate divergences from the DAG: @@ -316,7 +306,7 @@ async fn realtime_gate( admission::check_ak_rate(gov, &ak).await?; admission::check_product_qpm(gov, cfg, &ak.product).await?; admission::check_model_qpm(gov, cfg, model).await?; - admission::check_user_budget(gov, cfg, &ak.tenant, realtime_user(&ak, hint)).await?; + admission::check_user_budget(gov, cfg, &ak.tenant, ak.attributed_user(hint)).await?; if let Some(limit) = admission::model_quota_limit(cfg, &ak, model) && !gov .quota_check(&admission::model_quota_key(&ak.ak, model), limit) @@ -333,7 +323,7 @@ async fn realtime_gate( return Err(denied); } }; - let user = realtime_user(&ak, hint).to_owned(); + let user = ak.attributed_user(hint).to_owned(); Ok(RealtimeAdmit { ak, user, @@ -703,7 +693,7 @@ async fn realtime_bridge( .authenticate(&ak.ak) .await .unwrap_or_else(|| ak.clone()); - let user = realtime_user(&billed, &hint).to_owned(); + let user = billed.attributed_user(&hint).to_owned(); let unreserved = RealtimeAdmit { ak: billed, user, @@ -991,7 +981,7 @@ async fn emit_rt_hits( created_at_epoch_secs: gw_state::epoch_secs(), request_id: String::new(), ak: ak.ak.clone(), - user_id: realtime_user(ak, hint).to_owned(), + user_id: ak.attributed_user(hint).to_owned(), tenant: ak.tenant.clone(), surface: "realtime".to_owned(), rule: hit.rule.clone(), From 94199a9e58f02bfe267d94d668d08c0dd7fbc663 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 15:41:30 +0800 Subject: [PATCH 12/21] fix(security): close realtime secret-redaction, keyless-full retention, and config-publish audit gaps - Realtime DLP now honors detect_secrets: dlp_redact_realtime_frame masked PII only, so a secret in a WS frame reached the model even with detect_secrets on. It now redacts both PII and secrets like REST. - Retention owns its redaction: a keyless 'full' downgrade (and the 'redacted' level) stored plaintext when the tenant forwarded traffic with DLP off. redact_retained always strips PII+secrets for stored text. - Config-publish audit: a published version is the fleet source of truth even if the local reload fails; audit it after publish with the reload outcome instead of only on the success branch. --- crates/handler/src/lib.rs | 15 ++++-- crates/handler/src/plugins.rs | 59 ++++++++++++++++++++++-- crates/views/src/lib.rs | 86 ++++++++++++++++++++++++++++------- 3 files changed, 135 insertions(+), 25 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index da565d9..2788a24 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -80,12 +80,17 @@ async fn persist_content( } else { 0 }; - let redacted_prompt = || plugins::inbound_text(&ctx.request); + // retention owns its redaction: the stored non-full text is always stripped + // of PII/secrets even if the tenant forwards traffic with DLP off, so a + // keyless `full` downgrade or a `redacted` row can't hold raw content + let redacted_prompt = || plugins::redact_retained(&plugins::inbound_text(&ctx.request)); let redacted_response = || { - ctx.outcome - .as_ref() - .map(|o| o.response.message.clone()) - .unwrap_or_default() + plugins::redact_retained( + ctx.outcome + .as_ref() + .map(|o| o.response.message.as_str()) + .unwrap_or_default(), + ) }; let items = [ diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index bf3cc2e..389e232 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -236,13 +236,16 @@ pub fn realtime_frame_scan(sec: &SecurityConf, frame: &mut serde_json::Value) -> } /// DLP-redact a realtime frame's text-bearing fields in place; the hit count. +/// Honors both `dlp_redact` (PII) and `detect_secrets` (credentials), the same +/// as the REST path — a realtime frame must not be a secret-redaction bypass. /// Per-frame best effort: a PII span straddling two deltas is not caught — a /// realtime relay cannot buffer the way the REST stream surfaces do. pub fn dlp_redact_realtime_frame(sec: &SecurityConf, frame: &mut serde_json::Value) -> usize { - if !sec.dlp_redact { + if !sec.dlp_redact && !sec.detect_secrets { return 0; } - gw_engines::realtime::visit_frame_text(frame, &mut redact_str) + let (pii, secrets) = (sec.dlp_redact, sec.detect_secrets); + gw_engines::realtime::visit_frame_text(frame, &mut |s| redact_in_place(s, pii, secrets)) } /// DLP inbound redaction: emails, 11-digit phone numbers, and — when @@ -274,6 +277,16 @@ pub fn dlp_redact_request(sec: &SecurityConf, request: &mut GatewayRequest) -> u hits } +/// A privacy-safe copy of `text` for content retention: PII and secrets are +/// ALWAYS stripped, independent of the tenant's forwarding DLP flags. Retention +/// owns its own redaction so a `redacted` row — or a keyless `full` downgrade — +/// can never persist raw secrets/PII even when inline DLP is disabled. +pub fn redact_retained(text: &str) -> String { + let mut s = text.to_owned(); + redact_in_place(&mut s, true, true); + s +} + /// Redact one string in place (email/phone via `pii`, secrets via `secrets`); /// the hit count. fn redact_in_place(s: &mut String, pii: bool, secrets: bool) -> usize { @@ -289,7 +302,7 @@ fn redact_in_place(s: &mut String, pii: bool, secrets: bool) -> usize { hits } -/// Redact one string in place; email/phone only (the outbound/realtime path). +/// Redact one string in place; email/phone only (the outbound response path). fn redact_str(s: &mut String) -> usize { redact_in_place(s, true, false) } @@ -544,6 +557,46 @@ mod tests { assert_eq!(out.hits.len(), 1); } + #[test] + fn realtime_frame_redacts_secrets_when_detect_secrets_on() { + // detect_secrets on, dlp_redact OFF: a credential in a realtime frame is + // still masked before it reaches the model (the frame is not a bypass) + let s = SecurityConf { + detect_secrets: true, + dlp_redact: false, + ..Default::default() + }; + let mut frame = serde_json::json!({ + "type":"input_text","text":"here is sk-abcdefghijklmnopqrstuvwxyz012345" + }); + let n = dlp_redact_realtime_frame(&s, &mut frame); + assert_eq!(n, 1, "secret masked even with dlp_redact off"); + let text = frame["text"].as_str().unwrap(); + assert!( + text.contains("[REDACTED_SECRET]") && !text.contains("sk-abc"), + "{text}" + ); + + // both flags off: realtime frame untouched + let none = SecurityConf::default(); + let mut frame = + serde_json::json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); + assert_eq!(dlp_redact_realtime_frame(&none, &mut frame), 0); + } + + #[test] + fn redact_retained_strips_pii_and_secrets_unconditionally() { + let out = + redact_retained("mail john.doe@example.com key sk-abcdefghijklmnopqrstuvwxyz012345"); + assert!(out.contains("[REDACTED_EMAIL]"), "{out}"); + assert!(out.contains("[REDACTED_SECRET]"), "{out}"); + assert!( + !out.contains("example.com") && !out.contains("sk-abc"), + "{out}" + ); + assert_eq!(redact_retained("clean text"), "clean text"); + } + #[test] fn flag_action_records_a_hit_without_blocking() { let s = SecurityConf { diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 3d4b43e..238feb5 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1327,23 +1327,29 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S Ok(v) => v, Err(e) => return gateway_error(e), }; - match s.reload().await { - Ok(()) => { - audit_admin( - &s, - &AdminScope::Global, - &headers, - "config_publish", - &version.to_string(), - String::new(), - ) - .await; - ( - StatusCode::OK, - Json(json!({ "status": "published", "version": version })), - ) - .into_response() - } + // publish already made this version the fleet source of truth and notified + // peers; audit it before the local reload can fail, or a failed apply would + // leave the fleet on a config with no publish record + let reload = s.reload().await; + let detail = match &reload { + Ok(()) => String::new(), + Err(e) => format!("local reload failed: {e}"), + }; + audit_admin( + &s, + &AdminScope::Global, + &headers, + "config_publish", + &version.to_string(), + detail, + ) + .await; + match reload { + Ok(()) => ( + StatusCode::OK, + Json(json!({ "status": "published", "version": version })), + ) + .into_response(), Err(e) => error_response( 500, format!("published v{version} but local reload failed: {e}"), @@ -3034,6 +3040,52 @@ mod tests { ); } + #[tokio::test] + async fn full_retention_without_key_never_stores_raw_even_with_dlp_off() { + // full retention + no GW_CONTENT_KEY + tenant DLP fully off: the keyless + // downgrade must still strip secrets, not persist the raw prompt + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + assert!( + !gw_state::sealing_available(), + "test env has no content key" + ); + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let store = app_state.handler.state().store.clone(); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer k1") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"here is sk-abcdefghijklmnopqrstuvwxyz012345"}]}"#, + )) + .unwrap(); + let resp = app(app_state).oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_, rows) = store.ledger_snapshot(1).await.unwrap(); + let stored = store.content_for(&rows[0].request_id).await.unwrap(); + let prompt = stored + .iter() + .find(|c| c.kind == "prompt") + .expect("prompt stored"); + assert!(!prompt.sealed, "no key → unsealed"); + assert!( + prompt.content.contains("[REDACTED_SECRET]"), + "secret masked: {}", + prompt.content + ); + for c in &stored { + assert!( + !c.content.contains("sk-abc"), + "raw secret never persisted: {}", + c.content + ); + } + } + #[tokio::test] async fn chat_non_stream_ok() { let resp = test_app() From 64e7388aab9031686047a6833424cb9d00bd2797 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 15:50:03 +0800 Subject: [PATCH 13/21] fix(batch): persist per-item user attribution so shared-key batches bill and budget correctly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A batch dropped the submitter's user: BatchItem carried only messages and the reconstructed GatewayRequest left user_id empty, so a shared-key batch attributed every item's ledger row to "" and the empty user made the per-user budget a no-op — a per-user cap was evadable via /v1/batches. BatchItem now carries the resolved user (body `user` else `x-gw-user`), persisted in batch_items.user_id so a fleet drainer on another instance still attributes and budgets each item. Owner still overrides at billing. --- crates/handler/src/lib.rs | 55 +++++++++++++++++++++++++++++- crates/handler/src/offline.rs | 33 +++++++----------- crates/models/src/lib.rs | 2 +- crates/models/src/request.rs | 10 ++++++ crates/state/src/store.rs | 64 +++++++++++++++++++++++------------ crates/views/src/lib.rs | 22 ++++++++++-- 6 files changed, 139 insertions(+), 47 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 2788a24..fa310ed 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -18,7 +18,8 @@ use gw_engines::{EngineOutcome, SharedTransport}; use gw_models::{Block, GResult, GatewayError, GatewayRequest, GatewayResponse}; use gw_state::{AkInfo, GatewayState, SharedConfig}; -pub use offline::{BatchItem, OfflineHandler}; +pub use gw_models::BatchItem; +pub use offline::OfflineHandler; use std::sync::atomic::{AtomicU64, Ordering}; @@ -871,9 +872,11 @@ mod tests { let items = vec![ BatchItem { messages: vec![ChatMsg::text("user", "same prompt")], + user: String::new(), }, BatchItem { messages: vec![ChatMsg::text("user", "same prompt")], + user: String::new(), }, ]; let job = off @@ -896,9 +899,11 @@ mod tests { vec![ BatchItem { messages: vec![ChatMsg::text("user", "one")], + user: String::new(), }, BatchItem { messages: vec![ChatMsg::text("user", "two")], + user: String::new(), }, ], ) @@ -915,6 +920,44 @@ mod tests { ); } + #[tokio::test] + async fn batch_attributes_each_item_to_its_user() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let off = OfflineHandler::new(h.clone()); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let job = off + .submit( + key, + "gpt-4o".into(), + vec![ + BatchItem { + messages: vec![ChatMsg::text("user", "for alice")], + user: "alice".into(), + }, + BatchItem { + messages: vec![ChatMsg::text("user", "for bob")], + user: "bob".into(), + }, + ], + ) + .await + .unwrap(); + wait_terminal(&h, &job.id).await; + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + let users: std::collections::HashSet<&str> = + ledger.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice") && users.contains("bob"), + "each shared-key batch item bills to its own user: {users:?}" + ); + } + #[tokio::test] async fn distributed_batch_drained_by_a_separate_handler() { let Ok(url) = std::env::var("GW_TEST_PG_URL") else { @@ -943,9 +986,11 @@ mod tests { vec![ BatchItem { messages: vec![ChatMsg::text("user", "alpha")], + user: "alice".into(), }, BatchItem { messages: vec![ChatMsg::text("user", "beta")], + user: "bob".into(), }, ], ) @@ -974,5 +1019,13 @@ mod tests { let j = completed.expect("drain completed the batch"); assert_eq!(j.results.len(), 2, "both items executed exactly once"); assert!(j.results.iter().all(|r| r.ok && r.total_tokens > 0)); + // per-item attribution survived the enqueue → drain → reconstruct round-trip + let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); + let users: std::collections::HashSet<&str> = + ledger.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice") && users.contains("bob"), + "distributed batch preserved per-item user attribution: {users:?}" + ); } } diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 5636468..a4377a3 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -2,17 +2,11 @@ //! instance's background task; a distributed store (Postgres) only persists the //! items and a fleet drain loop on any instance claims and runs them. -use gw_models::{ChatMsg, GatewayRequest, ModelParamV2}; +use gw_models::{BatchItem, GatewayRequest, ModelParamV2}; use gw_state::{AkInfo, BatchItemResult, BatchJob, BatchStatus}; use crate::OnlineHandler; -/// One batch item: a self-contained message list. -#[derive(Debug, Clone)] -pub struct BatchItem { - pub messages: Vec, -} - /// Batch orchestration built on top of the online handler. #[derive(Clone)] pub struct OfflineHandler { @@ -34,18 +28,19 @@ impl OfflineHandler { items: Vec, ) -> gw_models::GResult { let store = self.online.state().store.clone(); - let msgs: Vec> = items.into_iter().map(|i| i.messages).collect(); if store.distributed_batches() { // atomic: the job becomes claimable only once all items are saved - store.batch_enqueue(&ak.ak, &ak.tenant, &model, &msgs).await + store + .batch_enqueue(&ak.ak, &ak.tenant, &model, &items) + .await } else { let job = store - .batch_create(&ak.ak, &ak.tenant, &model, msgs.len()) + .batch_create(&ak.ak, &ak.tenant, &model, items.len()) .await?; let this = self.clone(); let (id, model) = (job.id.clone(), model.clone()); // claim 0: non-distributed store — no fence, the heartbeat is a no-op - tokio::spawn(async move { this.execute(&id, &ak, &model, msgs, 0).await }); + tokio::spawn(async move { this.execute(&id, &ak, &model, items, 0).await }); Ok(job) } } @@ -54,14 +49,7 @@ impl OfflineHandler { /// the terminal status, heartbeating between items. `claim` is the fence /// token (0 for the in-process path); once a heartbeat reports the batch /// reclaimed, this executor stops rather than double-running items. - async fn execute( - &self, - id: &str, - ak: &AkInfo, - model: &str, - items: Vec>, - claim: i64, - ) { + async fn execute(&self, id: &str, ak: &AkInfo, model: &str, items: Vec, claim: i64) { let store = self.online.state().store.clone(); // the distributed claim already set status=running with the fence bump; // only the in-process path needs this write — unfenced on the distributed @@ -108,7 +96,7 @@ impl OfflineHandler { } }) }; - for (index, messages) in items.into_iter().enumerate() { + for (index, item) in items.into_iter().enumerate() { if lost.load(Relaxed) { break; // reclaimed by another instance; stop running new items } @@ -126,7 +114,10 @@ impl OfflineHandler { let request = GatewayRequest { is_online: false, ak: ak.ak.clone(), - message: messages, + message: item.messages, + // per-item attribution so effective_user_id (owner else this) + // bills and budgets each item to its end user + user_id: (!item.user.is_empty()).then_some(item.user), model_param_v2: Some(ModelParamV2::with_name( gw_consts::Protocol::OpenaiChat, model.to_owned(), diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index f58975b..208e61f 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -20,7 +20,7 @@ pub use params::{ VideoParams, }; pub use request::domain::{Account, ChatMsg}; -pub use request::{GatewayRequest, ModelParamV2}; +pub use request::{BatchItem, GatewayRequest, ModelParamV2}; pub use response::GatewayResponse; pub use response::StreamChunk; pub use token_estimate::{HeuristicEncoder, TokenEncoder, estimate_prompt_tokens}; diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index 80f9a61..e260ffa 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -36,6 +36,16 @@ impl GatewayRequest { } } +/// One queued batch item: a message list plus the client-supplied end-user +/// attribution. `user` is persisted with the item so a distributed drainer on +/// another instance still attributes and budgets it (owner still overrides at +/// billing time); empty when the submitter gave no `user`/`x-gw-user`. +#[derive(Debug, Clone)] +pub struct BatchItem { + pub messages: Vec, + pub user: String, +} + /// Domain types referenced by `GatewayRequest`. Per-vendor long-tail fields /// ride in `raw`/`extra` passthroughs rather than being individually modeled. pub mod domain { diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index b92999c..234d7a2 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -362,12 +362,12 @@ pub trait Store: Send + Sync + std::fmt::Debug { ak: &str, tenant: &str, model: &str, - items: &[Vec], + items: &[gw_models::BatchItem], ) -> GResult { self.batch_create(ak, tenant, model, items.len()).await } /// Load a batch's input items for execution. - async fn batch_load_items(&self, _id: &str) -> GResult>> { + async fn batch_load_items(&self, _id: &str) -> GResult> { Ok(Vec::new()) } /// Claim one pending batch (requeuing stale running ones first); `None` = @@ -1328,6 +1328,8 @@ impl PostgresStore { "CREATE TABLE IF NOT EXISTS batch_items ( batch_id TEXT NOT NULL, idx BIGINT NOT NULL, messages TEXT NOT NULL, PRIMARY KEY (batch_id, idx))", + // per-item end-user attribution so a fleet drainer still bills/budgets it + "ALTER TABLE batch_items ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claimed_at TIMESTAMPTZ", // fence token: bumped on every claim so a reclaimed executor's fenced writes no-op "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claim_seq BIGINT NOT NULL DEFAULT 0", @@ -1803,7 +1805,7 @@ impl Store for PostgresStore { ak: &str, tenant: &str, model: &str, - items: &[Vec], + items: &[gw_models::BatchItem], ) -> GResult { // the batch becomes claimable only once all its items are committed let mut tx = self @@ -1820,15 +1822,18 @@ impl Store for PostgresStore { .fetch_one(&mut *tx) .await .map_err(|e| crate::sqlx_err("insert batch", e))?; - for (idx, msgs) in items.iter().enumerate() { - let json = serde_json::to_string(msgs).unwrap_or_else(|_| "[]".into()); - sqlx::query("INSERT INTO batch_items (batch_id, idx, messages) VALUES ($1, $2, $3)") - .bind(&id) - .bind(idx as i64) - .bind(json) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("save batch item", e))?; + for (idx, item) in items.iter().enumerate() { + let json = serde_json::to_string(&item.messages).unwrap_or_else(|_| "[]".into()); + sqlx::query( + "INSERT INTO batch_items (batch_id, idx, messages, user_id) VALUES ($1, $2, $3, $4)", + ) + .bind(&id) + .bind(idx as i64) + .bind(json) + .bind(&item.user) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("save batch item", e))?; } tx.commit() .await @@ -1844,15 +1849,20 @@ impl Store for PostgresStore { }) } - async fn batch_load_items(&self, id: &str) -> GResult>> { - let rows = sqlx::query("SELECT messages FROM batch_items WHERE batch_id = $1 ORDER BY idx") - .bind(id) - .fetch_all(&self.pool) - .await - .map_err(|e| crate::sqlx_err("load batch items", e))?; + async fn batch_load_items(&self, id: &str) -> GResult> { + let rows = sqlx::query( + "SELECT messages, user_id FROM batch_items WHERE batch_id = $1 ORDER BY idx", + ) + .bind(id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("load batch items", e))?; Ok(rows .iter() - .map(|r| serde_json::from_str(r.get::<&str, _>(0)).unwrap_or_default()) + .map(|r| gw_models::BatchItem { + messages: serde_json::from_str(r.get::<&str, _>(0)).unwrap_or_default(), + user: r.get::(1), + }) .collect()) } @@ -2372,8 +2382,14 @@ mod tests { assert!(store.distributed_batches()); let qmsgs = vec![ - vec![gw_models::ChatMsg::text("user", "one")], - vec![gw_models::ChatMsg::text("user", "two")], + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "one")], + user: "u-one".into(), + }, + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "two")], + user: "u-two".into(), + }, ]; let qjob = store .batch_enqueue("ak-b", "default", "gpt-4o", &qmsgs) @@ -2400,7 +2416,11 @@ mod tests { } let loaded = store.batch_load_items(&qjob.id).await.unwrap(); assert_eq!(loaded.len(), 2); - assert_eq!(loaded[1][0].content, "two"); + assert_eq!(loaded[1].messages[0].content, "two"); + assert_eq!( + loaded[1].user, "u-two", + "per-item user round-trips through pg" + ); let fjob = store .batch_enqueue("ak-f", "default", "gpt-4o", &qmsgs) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 238feb5..b430ca5 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -2769,6 +2769,8 @@ async fn batches_submit( }; let mut model = body["model"].as_str().unwrap_or_default().to_owned(); let mut batch_items = Vec::new(); + // batch-level attribution hint; a per-item body `user` overrides it + let hint = user_header(&headers); if let Some(file_id) = body["input_file_id"].as_str() { let found = s.handler.state().store.file_get(file_id).await; @@ -2790,7 +2792,15 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "input file line missing a messages array"); } - batch_items.push(BatchItem { messages: msgs }); + let user = reqbody["user"] + .as_str() + .or(hint.as_deref()) + .unwrap_or_default() + .to_owned(); + batch_items.push(BatchItem { + messages: msgs, + user, + }); } } else if let Some(items) = body["items"].as_array() { for it in items { @@ -2798,7 +2808,15 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "each item needs a non-empty messages array"); } - batch_items.push(BatchItem { messages: msgs }); + let user = it["user"] + .as_str() + .or(hint.as_deref()) + .unwrap_or_default() + .to_owned(); + batch_items.push(BatchItem { + messages: msgs, + user, + }); } } else { return error_response(400, "either items or input_file_id is required"); From 06667a99a626a72f68a7317f35b5419d3de052dc Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 15:57:27 +0800 Subject: [PATCH 14/21] fix(audit): root admin source IP at the TCP peer, trust proxy headers only when configured MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit source_ip trusted x-real-ip/x-forwarded-for unconditionally, so a client reaching /admin/* directly could forge the audit source IP. It now roots at the real socket peer (served with connect-info) and reads forwarded headers only when the new trust_proxy_headers flag is set — i.e. a trusted proxy fronts the gateway. Default off: the audit records the unspoofable peer. --- crates/config/src/lib.rs | 5 +++ crates/server/src/main.rs | 10 +++-- crates/views/src/lib.rs | 94 ++++++++++++++++++++++++++++++--------- 3 files changed, 85 insertions(+), 24 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 3427340..e439cd4 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -466,6 +466,11 @@ pub struct GatewayConfig { /// Admin surface gate (dynamic config reload / key management). #[serde(default)] pub admin: AdminConf, + /// Trust `x-real-ip` / `x-forwarded-for` for the audit source IP. Off by + /// default: the audit records the real TCP peer, which a client can't forge. + /// Enable only when a trusted proxy fronts the gateway and sets those headers. + #[serde(default)] + pub trust_proxy_headers: bool, /// Stable hash of the source document; see [`Self::generation`]. #[serde(skip)] generation: u64, diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index b8e5c71..77f5106 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -186,9 +186,13 @@ async fn main() -> anyhow::Result<()> { let listener = tokio::net::TcpListener::bind(&addr).await?; tracing::info!("gw listening on http://{addr}"); - axum::serve(listener, router) - .with_graceful_shutdown(shutdown_signal()) - .await?; + // connect-info so the audit trail can root the source IP at the TCP peer + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown_signal()) + .await?; quota_task.abort(); purge_task.abort(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index b430ca5..845a908 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -994,19 +994,45 @@ async fn emit_rt_hits( } } -/// The caller IP for the audit trail. Prefers `x-real-ip` (the immediate peer -/// the trusted proxy sets) and, failing that, the RIGHTMOST `x-forwarded-for` -/// hop (appended by that proxy) — never the leftmost, which a client can forge. -fn source_ip(headers: &HeaderMap) -> String { - if let Some(ip) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) { - return ip.trim().to_owned(); +/// The connecting peer's socket address, read from the connect-info extension. +/// An infallible extractor: `None` when the router is driven without connect +/// info (the test harness), `Some` when served over a real listener. +struct ClientPeer(Option); + +impl axum::extract::FromRequestParts for ClientPeer { + type Rejection = std::convert::Infallible; + + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + _state: &S, + ) -> Result { + Ok(ClientPeer( + parts + .extensions + .get::>() + .map(|ci| ci.0), + )) } - headers - .get("x-forwarded-for") - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.rsplit(',').next()) - .map(|s| s.trim().to_owned()) - .unwrap_or_default() +} + +/// The caller IP for the audit trail. Roots at the real TCP `peer`, which a +/// client cannot forge. Only when `trust_proxy` is set (a trusted proxy fronts +/// the gateway) does it read `x-real-ip`, then the RIGHTMOST `x-forwarded-for` +/// hop (the one that proxy appended) — never the leftmost, which a client forges. +fn source_ip(peer: Option, headers: &HeaderMap, trust_proxy: bool) -> String { + if trust_proxy { + if let Some(ip) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) { + return ip.trim().to_owned(); + } + if let Some(ip) = headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.rsplit(',').next()) + { + return ip.trim().to_owned(); + } + } + peer.map(|p| p.ip().to_string()).unwrap_or_default() } /// Record one admin-plane mutation to the audit trail (who/what/when/where). @@ -1015,6 +1041,7 @@ async fn audit_admin( s: &AppState, scope: &AdminScope, headers: &HeaderMap, + peer: Option, action: &str, target: &str, summary: String, @@ -1027,7 +1054,7 @@ async fn audit_admin( action: action.to_owned(), target: target.to_owned(), summary, - source_ip: source_ip(headers), + source_ip: source_ip(peer, headers, s.handler.cfg().trust_proxy_headers), }; if let Err(e) = s.handler.state().store.admin_audit_add(&entry).await { tracing::warn!(error = %e, action, "admin audit write failed"); @@ -1128,7 +1155,11 @@ fn ct_eq(a: &str, b: &str) -> bool { /// POST /admin/reload — re-read config from source and swap it in atomically; /// governance, store, health, and cache are preserved. -async fn admin_reload(State(s): State, headers: HeaderMap) -> Response { +async fn admin_reload( + State(s): State, + headers: HeaderMap, + ClientPeer(peer): ClientPeer, +) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } @@ -1145,6 +1176,7 @@ async fn admin_reload(State(s): State, headers: HeaderMap) -> Response &s, &AdminScope::Global, &headers, + peer, "reload", "", String::new(), @@ -1170,6 +1202,7 @@ async fn admin_reload(State(s): State, headers: HeaderMap) -> Response async fn admin_key_create( State(s): State, headers: HeaderMap, + ClientPeer(peer): ClientPeer, Json(body): Json, ) -> Response { let scope = match admin_auth(&s, &headers) { @@ -1231,6 +1264,7 @@ async fn admin_key_create( &s, &scope, &headers, + peer, "key_create", ak, format!("tenant={tenant}"), @@ -1247,6 +1281,7 @@ async fn admin_key_create( async fn admin_key_patch( State(s): State, headers: HeaderMap, + ClientPeer(peer): ClientPeer, Path(ak): Path, Json(body): Json, ) -> Response { @@ -1274,7 +1309,7 @@ async fn admin_key_patch( match patched { Err(e) => gateway_error(e), Ok(Some(info)) => { - audit_admin(&s, &scope, &headers, "key_patch", &ak, String::new()).await; + audit_admin(&s, &scope, &headers, peer, "key_patch", &ak, String::new()).await; (StatusCode::OK, Json(ak_public_json(&info))).into_response() } Ok(None) => error_response(404, format!("key {ak} not found")), @@ -1285,6 +1320,7 @@ async fn admin_key_patch( async fn admin_key_delete( State(s): State, headers: HeaderMap, + ClientPeer(peer): ClientPeer, Path(ak): Path, ) -> Response { let scope = match admin_auth(&s, &headers) { @@ -1297,7 +1333,7 @@ async fn admin_key_delete( match s.handler.state().auth.revoke(&ak).await { Err(e) => gateway_error(e), Ok(true) => { - audit_admin(&s, &scope, &headers, "key_delete", &ak, String::new()).await; + audit_admin(&s, &scope, &headers, peer, "key_delete", &ak, String::new()).await; ( StatusCode::OK, Json(json!({ "ak": ak, "status": "revoked" })), @@ -1310,7 +1346,12 @@ async fn admin_key_delete( /// PUT /admin/config — validate, publish to the fleet config store, and reload /// this instance; peers converge via the store's change feed. Global admin only. -async fn admin_config_put(State(s): State, headers: HeaderMap, body: String) -> Response { +async fn admin_config_put( + State(s): State, + headers: HeaderMap, + ClientPeer(peer): ClientPeer, + body: String, +) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } @@ -1339,6 +1380,7 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S &s, &AdminScope::Global, &headers, + peer, "config_publish", &version.to_string(), detail, @@ -3009,12 +3051,22 @@ mod tests { } #[test] - fn source_ip_prefers_trusted_hop_over_spoofable_xff() { + fn source_ip_roots_at_peer_and_ignores_forgeable_headers_untrusted() { + let peer: std::net::SocketAddr = "203.0.113.7:5000".parse().unwrap(); let mut h = HeaderMap::new(); - h.insert("x-forwarded-for", "1.2.3.4, 10.0.0.9".parse().unwrap()); - assert_eq!(source_ip(&h), "10.0.0.9", "rightmost (proxy-appended) hop"); h.insert("x-real-ip", "10.0.0.5".parse().unwrap()); - assert_eq!(source_ip(&h), "10.0.0.5", "x-real-ip wins over XFF"); + h.insert("x-forwarded-for", "1.2.3.4, 10.0.0.9".parse().unwrap()); + // default (untrusted): the client's headers are ignored, the TCP peer wins + assert_eq!(source_ip(Some(peer), &h, false), "203.0.113.7"); + assert_eq!( + source_ip(None, &h, false), + "", + "no peer, no forgeable header" + ); + // trusted proxy: x-real-ip, then the rightmost (proxy-appended) XFF hop + assert_eq!(source_ip(Some(peer), &h, true), "10.0.0.5"); + h.remove("x-real-ip"); + assert_eq!(source_ip(Some(peer), &h, true), "10.0.0.9", "rightmost hop"); } #[test] From 29eac92ffbb6c88d23f5409e4e7431804584c67e Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 15:58:51 +0800 Subject: [PATCH 15/21] docs: per-surface user attribution, retention redaction guarantee, audit source-IP trust --- docs/governance.md | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/docs/governance.md b/docs/governance.md index 687fd06..8e884dd 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -103,7 +103,10 @@ Every ledger row carries a `user_id`, `request_id`, and `created_at_epoch_secs`. The effective user is the key's `owner` (one key = one user, the enterprise model) if set, else request metadata — the `x-gw-user` header, OpenAI's `user` field, or Anthropic's `metadata.user_id`. `owner` is authoritative; the -metadata hint is only trusted for shared keys. `GET /admin/usage/users?user=&since=&until=` +metadata hint is only trusted for shared keys. This holds on every surface — +REST, realtime (the `x-gw-user` hint is captured at WS connect), and batch +(each item's `user`, or the submitter's `x-gw-user`, is persisted with the item +so a fleet drainer still attributes and budgets it). `GET /admin/usage/users?user=&since=&until=` returns per-(user, model) cost over a billing period (add `format=csv` for export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap. @@ -124,9 +127,16 @@ error). - **Content-safety events** (`/admin/audit/events`): who/which-rule/what-action, no prompt text; tenant-scoped. - **Admin operations** (`/admin/audit/ops`, global admin only): every key CRUD, - config publish, and reload with actor, action, target, and source IP. + config publish, and reload with actor, action, target, and source IP. A + config publish is recorded even when the local reload fails, since the version + is already the fleet's source of truth. The source IP is the real TCP peer; + set `trust_proxy_headers: true` (top level) to instead trust `x-real-ip` / + `x-forwarded-for` — do that ONLY behind a proxy that sets them, or a direct + client could forge the recorded IP. - **Content retention** (`tenants[].retention`): `none` (default), `redacted` - (post-DLP text), or `full` (raw). Stored in `request_content`, sealed at rest - with XChaCha20-Poly1305 under `GW_CONTENT_KEY` (64 hex chars); `full` refuses - to store raw without a key and falls back to redacted. `days` sets expiry; an - hourly purge deletes elapsed content. + (PII/secrets stripped), or `full` (raw). Stored in `request_content`, sealed + at rest with XChaCha20-Poly1305 under `GW_CONTENT_KEY` (64 hex chars). + Retention owns its redaction, so `redacted` — and a keyless `full` that falls + back to it — never persists raw secrets/PII even if the tenant forwards + traffic with DLP off. `full` refuses to store raw without a key. `days` sets + expiry; an hourly purge deletes elapsed content. From db7e8196a52e33cfe419c6698eebc7292f7477f0 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 16:44:24 +0800 Subject: [PATCH 16/21] fix(realtime): apply external moderation; freeze audit source IP pre-reload; record inbound realtime DLP events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Realtime bypassed the external moderator: REST runs it in the pre-stage but /v1/realtime only ran blocklist/regex/DLP. A moderated tenant (wired via with_moderator) could be bypassed over the WebSocket. Both realtime paths now call the moderator on inbound frames (new moderate_text seam). - Admin audit source IP is now frozen at request entry, before any config reload, so the op that flips trust_proxy_headers is audited under the policy in effect when it arrived — it can't trust its own forged header. - Inbound realtime DLP redactions are recorded as security events (REST parity); per-delta outbound redactions stay unrecorded (too hot) and are now documented as the one exception. --- Cargo.lock | 1 + crates/handler/src/lib.rs | 51 +++++++++ crates/server/tests/e2e.rs | 68 ++++++++++++ crates/views/Cargo.toml | 1 + crates/views/src/lib.rs | 215 ++++++++++++++++++++++++++++++------- docs/governance.md | 7 ++ 6 files changed, 307 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1677ee9..d2b92d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -951,6 +951,7 @@ dependencies = [ name = "gw-views" version = "0.1.0" dependencies = [ + "async-trait", "axum", "base64", "bytes", diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index fa310ed..315547b 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -344,6 +344,21 @@ impl OnlineHandler { Ok(ctx) } + /// Run the wired moderator over raw text; `Some(reason)` to deny, `None` to + /// allow. A moderator error resolves per `moderation_fail_open`. The seam + /// the realtime surface uses (it has no `DagContext`); the caller records + /// the security event on its own surface. + pub async fn moderate_text(&self, sec: &gw_config::SecurityConf, text: &str) -> Option { + match self.moderator.review(text).await { + Ok(moderation::Verdict::Allow) => None, + Ok(moderation::Verdict::Deny(reason)) => Some(reason), + Err(e) => { + tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); + (!sec.moderation_fail_open).then(|| "content moderation is unavailable".to_owned()) + } + } + } + /// Run the wired moderator over the request's inbound text; `Some(Block)` to /// deny. A moderator error resolves per `moderation_fail_open`. Records a /// security event on a deny. @@ -477,6 +492,42 @@ mod tests { ); } + #[derive(Debug)] + struct ErrModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for ErrModerator { + async fn review(&self, _text: &str) -> Result { + Err("moderator upstream down".into()) + } + } + + #[tokio::test] + async fn moderate_text_allow_deny_and_failure_posture() { + let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let base = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let mut sec = gw_config::SecurityConf::default(); + assert_eq!(base.moderate_text(&sec, "x").await, None, "default allows"); + let deny = base.clone().with_moderator(Arc::new(DenyModerator)); + assert_eq!(deny.moderate_text(&sec, "x").await.as_deref(), Some("nope")); + let err = base.with_moderator(Arc::new(ErrModerator)); + sec.moderation_fail_open = true; + assert_eq!( + err.moderate_text(&sec, "x").await, + None, + "error + fail-open allows" + ); + sec.moderation_fail_open = false; + assert!( + err.moderate_text(&sec, "x").await.is_some(), + "error + fail-closed denies" + ); + } + #[tokio::test] async fn per_user_budget_denies_over_the_cap() { let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, user_daily_token_quota: 5}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index f51e815..31a86eb 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -27,6 +27,74 @@ fn app() -> Router { )) } +#[tokio::test] +async fn admin_audit_freezes_source_ip_before_a_trust_flip() { + const V1: &str = r#" +listen: {host: 127.0.0.1, port: 0} +admin: {token_env: GW_TEST_ADMIN_TOKEN_TRUST} +trust_proxy_headers: false +access_keys: [{ak: ak-x, product: demo, qps: 100, daily_token_quota: 1000000}] +models: [{name: gpt-4o, protocol: openai-chat}] +accounts: [{name: mock-openai-1, provider: openai, protocols: ["openai-chat"]}] +"#; + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { std::env::set_var("GW_TEST_ADMIN_TOKEN_TRUST", "s3cret") }; + let v1 = GatewayConfig::from_yaml(V1).unwrap(); + let v2_yaml = V1.replace("trust_proxy_headers: false", "trust_proxy_headers: true"); + let loader: gw_views::ConfigLoader = Arc::new(move || { + let yaml = v2_yaml.clone(); + Box::pin(async move { GatewayConfig::from_yaml(&yaml).map_err(|e| e.to_string()) }) + as gw_views::ConfigFuture + }); + let state = Arc::new(GatewayState::from_config(&v1)); + let shared = gw_state::SharedConfig::new(Arc::new(v1), state); + let app = gw_views::app(gw_views::AppState::with_config( + shared, + Arc::new(gw_engines::MockTransport), + Some(loader), + )); + + // this reload FLIPS trust_proxy_headers on, while forging x-real-ip on the + // very same request — the op must be audited under the pre-reload policy + let r = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/reload") + .header("authorization", "Bearer s3cret") + .header("x-real-ip", "9.9.9.9") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(r.status(), StatusCode::OK); + + let ops = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/audit/ops") + .header("authorization", "Bearer s3cret") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let j = body_json(ops).await; + let reload = j["entries"] + .as_array() + .unwrap() + .iter() + .find(|e| e["action"] == "reload") + .expect("reload audited"); + assert_ne!( + reload["source_ip"], "9.9.9.9", + "the trust-enabling op must not trust its own forged header" + ); +} + #[tokio::test] async fn admin_reload_is_gated_and_swaps_keys_live() { let r = app() diff --git a/crates/views/Cargo.toml b/crates/views/Cargo.toml index 3dac2cc..b72f0c2 100644 --- a/crates/views/Cargo.toml +++ b/crates/views/Cargo.toml @@ -27,6 +27,7 @@ base64 = { workspace = true } [dev-dependencies] tower = { workspace = true } +async-trait = { workspace = true } [lints] workspace = true diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 845a908..de7a2b0 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -439,7 +439,24 @@ async fn realtime_session( .await; continue; } - gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut ev); + if let Some(reason) = realtime_moderate(&s, sec, &ak, &hint, &mut ev).await { + let _ = socket + .send(send(json!({"type":"error","message": reason}))) + .await; + continue; + } + let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut ev); + if redacted > 0 { + write_rt_event( + &s, + &ak, + ak.attributed_user(&hint), + "dlp", + "redact", + redacted as i64, + ) + .await; + } match ev["type"].as_str().unwrap_or_default() { "input_text" => { let admit = match realtime_gate(&s, &ak, &model, &hint).await { @@ -609,7 +626,20 @@ async fn realtime_bridge( } continue; } - if gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut frame) > 0 { + if let Some(reason) = realtime_moderate(&s, sec, &ak, &hint, &mut frame).await { + if cl_tx + .send(CMsg::Text(send_err(reason).to_string().into())) + .await + .is_err() + { + break; + } + continue; + } + let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut frame); + if redacted > 0 { + write_rt_event(&s, &ak, ak.attributed_user(&hint), "dlp", "redact", redacted as i64) + .await; forward = UMsg::text(frame.to_string()); } // gate each generation trigger, not every control frame @@ -967,31 +997,81 @@ fn q_num( q.get(key).and_then(|v| v.parse().ok()).unwrap_or(default) } +/// Write one realtime security event (`user` already resolved). The shared sink +/// for realtime blocklist/regex hits, moderation denials, and inbound DLP hits. +async fn write_rt_event( + s: &AppState, + ak: &AkInfo, + user: &str, + rule: &str, + action: &str, + hits: i64, +) { + let event = gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: String::new(), + ak: ak.ak.clone(), + user_id: user.to_owned(), + tenant: ak.tenant.clone(), + surface: "realtime".to_owned(), + rule: rule.to_owned(), + action: action.to_owned(), + hits, + }; + if let Err(e) = s.handler.state().store.security_event_add(&event).await { + tracing::warn!(error = %e, "realtime security event write failed"); + } +} + /// Record a realtime frame's content-safety hits to the security-event stream -/// (parity with the REST surfaces). Per-frame DLP redactions are not recorded — -/// too hot. +/// (parity with the REST surfaces). async fn emit_rt_hits( s: &AppState, ak: &AkInfo, hits: &[gw_handler::plugins::RuleHit], hint: &str, ) { + let user = ak.attributed_user(hint).to_owned(); for hit in hits { - let event = gw_state::SecurityEvent { - created_at_epoch_secs: gw_state::epoch_secs(), - request_id: String::new(), - ak: ak.ak.clone(), - user_id: ak.attributed_user(hint).to_owned(), - tenant: ak.tenant.clone(), - surface: "realtime".to_owned(), - rule: hit.rule.clone(), - action: hit.action.as_str().to_owned(), - hits: hit.count, - }; - if let Err(e) = s.handler.state().store.security_event_add(&event).await { - tracing::warn!(error = %e, "realtime security event write failed"); + write_rt_event(s, ak, &user, &hit.rule, hit.action.as_str(), hit.count).await; + } +} + +/// All inbound text a realtime frame carries, for moderation. +fn realtime_frame_text(frame: &mut Value) -> String { + let mut text = String::new(); + gw_engines::realtime::visit_frame_text(frame, &mut |s| { + if !s.is_empty() { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(s); } + 0 + }); + text +} + +/// Moderate a realtime frame's inbound text via the wired moderator — parity +/// with the REST surface, so a moderated tenant can't be bypassed over the +/// WebSocket. `Some(reason)` denies the frame; records a moderation event. +async fn realtime_moderate( + s: &AppState, + sec: &gw_config::SecurityConf, + ak: &AkInfo, + hint: &str, + frame: &mut Value, +) -> Option { + if !sec.moderate { + return None; + } + let text = realtime_frame_text(frame); + if text.is_empty() { + return None; } + let reason = s.handler.moderate_text(sec, &text).await?; + write_rt_event(s, ak, ak.attributed_user(hint), "moderation", "block", 1).await; + Some(reason) } /// The connecting peer's socket address, read from the connect-info extension. @@ -1036,12 +1116,14 @@ fn source_ip(peer: Option, headers: &HeaderMap, trust_prox } /// Record one admin-plane mutation to the audit trail (who/what/when/where). +/// `source` is resolved by the caller at request entry — before any config +/// mutation — so the op that flips `trust_proxy_headers` is audited under the +/// policy in effect when it arrived, not the one it just installed. /// Best-effort: a store failure is logged, never fails the operation. async fn audit_admin( s: &AppState, scope: &AdminScope, - headers: &HeaderMap, - peer: Option, + source: String, action: &str, target: &str, summary: String, @@ -1054,7 +1136,7 @@ async fn audit_admin( action: action.to_owned(), target: target.to_owned(), summary, - source_ip: source_ip(peer, headers, s.handler.cfg().trust_proxy_headers), + source_ip: source, }; if let Err(e) = s.handler.state().store.admin_audit_add(&entry).await { tracing::warn!(error = %e, action, "admin audit write failed"); @@ -1163,6 +1245,8 @@ async fn admin_reload( if let Err(r) = require_global_admin(&s, &headers) { return r; } + // freeze the source IP under the pre-reload policy + let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); match s.reload().await { Ok(()) => { let cfg = s.handler.cfg(); @@ -1172,16 +1256,7 @@ async fn admin_reload( accounts = cfg.accounts.len(), "config reloaded" ); - audit_admin( - &s, - &AdminScope::Global, - &headers, - peer, - "reload", - "", - String::new(), - ) - .await; + audit_admin(&s, &AdminScope::Global, source, "reload", "", String::new()).await; ( StatusCode::OK, Json(json!({ @@ -1209,6 +1284,7 @@ async fn admin_key_create( Ok(scope) => scope, Err(r) => return r, }; + let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); let (Some(ak), Some(product)) = (body["ak"].as_str(), body["product"].as_str()) else { return error_response(400, "ak and product are required"); }; @@ -1263,8 +1339,7 @@ async fn admin_key_create( audit_admin( &s, &scope, - &headers, - peer, + source, "key_create", ak, format!("tenant={tenant}"), @@ -1289,6 +1364,7 @@ async fn admin_key_patch( Ok(scope) => scope, Err(r) => return r, }; + let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); if let Err(r) = scoped_key(&s, &scope, &ak).await { return r; } @@ -1309,7 +1385,7 @@ async fn admin_key_patch( match patched { Err(e) => gateway_error(e), Ok(Some(info)) => { - audit_admin(&s, &scope, &headers, peer, "key_patch", &ak, String::new()).await; + audit_admin(&s, &scope, source, "key_patch", &ak, String::new()).await; (StatusCode::OK, Json(ak_public_json(&info))).into_response() } Ok(None) => error_response(404, format!("key {ak} not found")), @@ -1327,13 +1403,14 @@ async fn admin_key_delete( Ok(scope) => scope, Err(r) => return r, }; + let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); if let Err(r) = scoped_key(&s, &scope, &ak).await { return r; } match s.handler.state().auth.revoke(&ak).await { Err(e) => gateway_error(e), Ok(true) => { - audit_admin(&s, &scope, &headers, peer, "key_delete", &ak, String::new()).await; + audit_admin(&s, &scope, source, "key_delete", &ak, String::new()).await; ( StatusCode::OK, Json(json!({ "ak": ak, "status": "revoked" })), @@ -1355,6 +1432,9 @@ async fn admin_config_put( if let Err(r) = require_global_admin(&s, &headers) { return r; } + // freeze the source IP under the pre-publish policy: this very op may be the + // one enabling trust_proxy_headers, and must not be audited under it + let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); let Some(store) = &s.config_store else { return error_response( 400, @@ -1379,8 +1459,7 @@ async fn admin_config_put( audit_admin( &s, &AdminScope::Global, - &headers, - peer, + source, "config_publish", &version.to_string(), detail, @@ -3156,6 +3235,70 @@ mod tests { } } + #[derive(Debug)] + struct DenyModerator; + + #[async_trait::async_trait] + impl gw_handler::moderation::Moderator for DenyModerator { + async fn review(&self, _text: &str) -> Result { + Ok(gw_handler::moderation::Verdict::Deny( + "blocked by moderator".into(), + )) + } + } + + #[tokio::test] + async fn realtime_moderates_and_records_inbound_dlp() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true, detect_secrets: true}\nmodels: [{name: rt, protocol: realtime}]\naccess_keys: [{ak: k1, product: p, qps: 10, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let handler = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DenyModerator)); + let offline = OfflineHandler::new(handler.clone()); + let app = AppState { + handler, + offline, + loader: None, + config_store: None, + }; + let ak = app.handler.state().auth.authenticate("k1").await.unwrap(); + let cfg = app.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); + + // moderation denies over realtime, not just REST + let mut frame = json!({"type":"input_text","text":"hello there"}); + assert_eq!( + realtime_moderate(&app, sec, &ak, "", &mut frame) + .await + .as_deref(), + Some("blocked by moderator") + ); + // inbound realtime DLP redaction is recorded (the sink both WS paths use) + let mut secret = json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); + let n = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut secret); + assert!(n > 0); + write_rt_event(&app, &ak, ak.attributed_user(""), "dlp", "redact", n as i64).await; + + let events = app + .handler + .state() + .store + .security_events(None, 10) + .await + .unwrap(); + assert!( + events.iter().any(|e| e.rule == "moderation"), + "moderation event" + ); + assert!( + events.iter().any(|e| e.rule == "dlp"), + "inbound realtime DLP event" + ); + } + #[tokio::test] async fn chat_non_stream_ok() { let resp = test_app() diff --git a/docs/governance.md b/docs/governance.md index 8e884dd..598df53 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -122,6 +122,13 @@ Set `moderate: true` to route inbound text through an external moderator wired into the handler (`moderation_fail_open` picks the posture on a moderator error). +The same policy applies on every surface, realtime included: a `/v1/realtime` +WebSocket runs the blocklist, regex rules, DLP redaction, and the external +moderator on inbound frames, so it is not a bypass. One deliberate exception: +per-*delta* outbound DLP redactions on the streaming relay are not written as +individual security events (a store write per token is too hot) — inbound +redactions and every block/flag/moderation hit still are. + ## Audit trails - **Content-safety events** (`/admin/audit/events`): who/which-rule/what-action, From b712d5c8e9275c372eb23594a4ef61c3c2da23be Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 17:14:04 +0800 Subject: [PATCH 17/21] fix(realtime): audit outbound DLP per turn; reconcile the doc exemption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The outbound realtime DLP exemption was broader than documented: every upstream frame was redacted without any security event, while the doc claimed only per-token deltas were exempt. Sum the turn's outbound redactions and record one 'dlp' event at the turn boundary (response.done) — the redaction still applies to every frame, only its audit is aggregated so a store write per streamed token stays off the hot path. Doc reconciled. --- crates/views/src/lib.rs | 25 +++++++++++++++++++++---- docs/governance.md | 9 +++++---- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index de7a2b0..b9dd411 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -597,6 +597,8 @@ async fn realtime_bridge( let mut pending: std::collections::VecDeque = std::collections::VecDeque::new(); // denied server-VAD turn: swallow its upstream frames until its terminal frame let mut suppress = false; + // outbound DLP redactions summed within a turn, recorded once at its boundary + let mut out_redacted = 0i64; loop { tokio::select! { m = cl_rx.next() => { @@ -682,6 +684,7 @@ async fn realtime_bridge( }; let mut relay = true; let mut redacted: Option = None; + let mut turn_ended = false; match frame { Some(mut v) => { if suppress { @@ -746,18 +749,32 @@ async fn realtime_bridge( None => {} } recognized += 1; + turn_ended = true; } // outbound DLP, per frame (a span straddling deltas is // beyond a relay that cannot buffer) let cfg = s.handler.cfg(); - if relay - && gw_handler::plugins::dlp_redact_realtime_frame( + let n = if relay { + gw_handler::plugins::dlp_redact_realtime_frame( cfg.security_for(&ak.tenant), &mut v, - ) > 0 - { + ) + } else { + 0 + }; + if n > 0 { redacted = Some(v.to_string()); } + // a per-token event stream would be too hot: sum the turn's + // outbound redactions and record one event at its boundary + out_redacted += n as i64; + if turn_ended { + if out_redacted > 0 { + write_rt_event(&s, &ak, ak.attributed_user(&hint), "dlp", "redact", out_redacted) + .await; + } + out_redacted = 0; + } } // a denied turn's non-JSON output (e.g. audio deltas) is dropped too None => relay = !suppress, diff --git a/docs/governance.md b/docs/governance.md index 598df53..e643889 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -124,10 +124,11 @@ error). The same policy applies on every surface, realtime included: a `/v1/realtime` WebSocket runs the blocklist, regex rules, DLP redaction, and the external -moderator on inbound frames, so it is not a bypass. One deliberate exception: -per-*delta* outbound DLP redactions on the streaming relay are not written as -individual security events (a store write per token is too hot) — inbound -redactions and every block/flag/moderation hit still are. +moderator on inbound frames, so it is not a bypass. Every hit is audited: inbound +block/flag/moderation and DLP redactions are recorded per frame; outbound DLP +redactions (which stream token-by-token) are summed across the turn and recorded +as one event at the turn boundary — the redaction still applies to every frame, +only its audit is aggregated (a store write per token would be too hot). ## Audit trails From 55e041ec8a3a18e983a692165d433956bef6aab1 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 17:29:24 +0800 Subject: [PATCH 18/21] fix(realtime): flush aborted-turn outbound DLP audit on exit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A turn dropped by an upstream disconnect never hits its response.done boundary, so its accumulated outbound-redaction count was never audited. Flush it on loop exit, alongside the reserve refund — the redaction already applied per frame; this only completes the best-effort audit. --- crates/views/src/lib.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index b9dd411..5da75c3 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -797,6 +797,19 @@ async fn realtime_bridge( for a in pending { a.refund().await; } + // a turn aborted before its boundary (upstream drop) still applied its + // redactions per frame — flush the pending count so the audit isn't lost + if out_redacted > 0 { + write_rt_event( + &s, + &ak, + ak.attributed_user(&hint), + "dlp", + "redact", + out_redacted, + ) + .await; + } if generations > 0 && recognized == 0 { tracing::warn!( account = %account.name, From 3416667c7c07df3008cc0a80d15b9084d074e17a Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 18:35:29 +0800 Subject: [PATCH 19/21] review: apply simplify+code-rs convergence round (policy seams, hot-path allocs, comment trim) --- crates/config/src/lib.rs | 6 +- crates/dag/src/nodes.rs | 6 +- crates/handler/src/lib.rs | 388 +++++++++++++++++----------------- crates/handler/src/offline.rs | 2 - crates/handler/src/plugins.rs | 297 +++++++++++++------------- crates/server/tests/e2e.rs | 2 - crates/state/src/admission.rs | 17 +- crates/state/src/content.rs | 27 ++- crates/state/src/keystore.rs | 7 +- crates/state/src/lib.rs | 8 +- crates/state/src/store.rs | 186 ++++++---------- crates/views/src/lib.rs | 313 ++++++++++++--------------- 12 files changed, 574 insertions(+), 685 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index e439cd4..d6a7b5e 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -198,8 +198,7 @@ pub struct SecurityConf { /// Whether to DLP-redact inbound/outbound content (emails/phone numbers). #[serde(default)] pub dlp_redact: bool, - /// Detect API keys / secrets in inbound text and redact them (a top - /// enterprise ask: stop staff pasting credentials to a model). + /// Detect API keys / credentials in inbound text and redact them. #[serde(default)] pub detect_secrets: bool, /// Route inbound text through the wired external moderator; needs a @@ -616,8 +615,7 @@ impl GatewayConfig { check_unique("product", self.products.iter().map(|p| p.name.as_str()))?; check_unique("provider", self.providers.iter().map(|p| p.name.as_str()))?; check_unique("tenant", self.tenants.iter().map(|t| t.name.as_str()))?; - // ':' separates the per-user budget counter key `ub:{tenant}:{user}`, so a - // colon in a tenant name could alias another tenant's budget counter + // a colon in a tenant name would alias another tenant's `ub:{tenant}:{user}` budget key for t in &self.tenants { if t.name.contains(':') { return Err(ConfigError::DuplicateName { diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 018d1a4..bf0341a 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -397,12 +397,11 @@ impl DagNode for UserBudgetGate { &["ak_tpm"] } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let user = ctx.effective_user_id().to_owned(); admission::check_user_budget( ctx.state.governance.as_ref(), &ctx.cfg, &ctx.ak.tenant, - &user, + ctx.effective_user_id(), ) .await .map_err(limit_denied) @@ -615,8 +614,7 @@ async fn bill( let requested = param .and_then(|p| p.fallback_from.as_deref()) .unwrap_or(served); - // take the reserves into locals first so the whole-ctx borrow that - // `effective_user_id` needs can't clash with these mutable field takes + // locals first: the whole-ctx borrow `effective_user_id` needs can't overlap these takes let (reserved, tpm_reserved, model_quota_key) = ( ctx.quota_reserved.take().unwrap_or(0), ctx.tpm_reserved.take(), diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 315547b..f53595e 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -9,6 +9,7 @@ pub mod plugins; use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use futures::FutureExt; use gw_config::GatewayConfig; @@ -21,133 +22,10 @@ use gw_state::{AkInfo, GatewayState, SharedConfig}; pub use gw_models::BatchItem; pub use offline::OfflineHandler; -use std::sync::atomic::{AtomicU64, Ordering}; +const MODERATION_UNAVAILABLE: &str = "content moderation is unavailable"; static REQ_SEQ: AtomicU64 = AtomicU64::new(0); -/// Record a content-safety outcome (no prompt text) against this request's -/// key/user/tenant. Best-effort: a store failure is logged, never fails the -/// request. Surface = the request protocol, or "batch" for offline items. -async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i64) { - let user_id = ctx.effective_user_id().to_owned(); - let surface = if ctx.request.is_online { - ctx.request - .model_param_v2 - .as_ref() - .map(|p| p.protocol.as_str()) - .unwrap_or_default() - .to_owned() - } else { - "batch".to_owned() - }; - let event = gw_state::SecurityEvent { - created_at_epoch_secs: gw_state::epoch_secs(), - request_id: ctx.request.request_id.clone(), - ak: ctx.ak.ak.clone(), - user_id, - tenant: ctx.ak.tenant.clone(), - surface, - rule: rule.to_owned(), - action: action.to_owned(), - hits, - }; - if let Err(e) = ctx.state.store.security_event_add(&event).await { - tracing::warn!(error = %e, "security event write failed"); - } -} - -/// Persist this request's prompt and response per the tenant's retention policy. -/// `full` content is sealed at rest (and refuses to store raw without a key, -/// downgrading to the redacted text); `redacted` stores the post-DLP text, -/// sealed too when a key is present. Best-effort — a store failure is logged. -async fn persist_content( - ctx: &DagContext, - retention: gw_config::RetentionConf, - raw_prompt: Option, - raw_response: Option, -) { - use gw_config::ContentLevel; - let want_full = retention.content == ContentLevel::Full; - let sealed_ok = gw_state::sealing_available(); - // full without a key must NOT land raw — fall back to the redacted text - let store_full = want_full && sealed_ok; - if want_full && !sealed_ok { - tracing::warn!("full retention configured but GW_CONTENT_KEY unset; storing redacted text"); - } - - let now = gw_state::epoch_secs(); - let expires = if retention.days > 0 { - now + retention.days as i64 * 86_400 - } else { - 0 - }; - // retention owns its redaction: the stored non-full text is always stripped - // of PII/secrets even if the tenant forwards traffic with DLP off, so a - // keyless `full` downgrade or a `redacted` row can't hold raw content - let redacted_prompt = || plugins::redact_retained(&plugins::inbound_text(&ctx.request)); - let redacted_response = || { - plugins::redact_retained( - ctx.outcome - .as_ref() - .map(|o| o.response.message.as_str()) - .unwrap_or_default(), - ) - }; - - let items = [ - ( - "prompt", - if store_full { - raw_prompt.unwrap_or_else(redacted_prompt) - } else { - redacted_prompt() - }, - ), - ( - "response", - if store_full { - raw_response.unwrap_or_else(redacted_response) - } else { - redacted_response() - }, - ), - ]; - for (kind, text) in items { - if text.is_empty() { - continue; - } - // seal whenever a key exists (defense in depth even for redacted text) - let (content, sealed) = match gw_state::content::seal(&text) { - Some(ct) => (ct, true), - None => (text, false), - }; - let record = gw_state::ContentRecord { - created_at_epoch_secs: now, - request_id: ctx.request.request_id.clone(), - ak: ctx.ak.ak.clone(), - user_id: ctx.effective_user_id().to_owned(), - tenant: ctx.ak.tenant.clone(), - kind: kind.to_owned(), - content, - sealed, - expires_at_epoch_secs: expires, - }; - if let Err(e) = ctx.state.store.content_add(&record).await { - tracing::warn!(error = %e, kind, "content retention write failed"); - } - } -} - -/// A per-request correlation id: `req--`, time-sortable and -/// unique within the process (the seq disambiguates same-millisecond requests). -pub fn new_request_id() -> String { - let ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis()) - .unwrap_or(0); - format!("req-{ms}-{}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)) -} - /// Runs one request through the plugin pre-stage, the DAG, and the plugin post-stage. #[derive(Clone)] pub struct OnlineHandler { @@ -208,8 +86,7 @@ impl OnlineHandler { if request.request_id.is_empty() { request.request_id = new_request_id(); } - // one consistent snapshot for the whole request; the tenant's own - // security policy wins over the global one when it sets one + // one consistent snapshot for the whole request let snap = self.config.load(); let sec = snap.cfg.security_for(&ak.tenant); // outbound DLP is a response-buffering boundary: a masked span can @@ -219,8 +96,7 @@ impl OnlineHandler { if dlp { request.stream_tx = None; } - // scan the ORIGINAL content, before DLP — else a blocklisted term inside - // a redacted span (a domain in an email) is masked out and slips + // scan the ORIGINAL content pre-DLP: a blocklisted term inside a redacted span would slip let scan = plugins::security_check(sec, &mut request); let mut ctx = DagContext::new( @@ -230,8 +106,7 @@ impl OnlineHandler { request, ak, ); - // every rule that fired is recorded (block/flag/shadow alike); only a - // block-action hit denies the request + // every fired rule is recorded (block/flag/shadow alike); only a block-action hit denies for hit in &scan.hits { emit_security_event(&ctx, &hit.rule, hit.action.as_str(), hit.count).await; } @@ -240,46 +115,29 @@ impl OnlineHandler { "security_check", format!("blocked (code {})", block.err_code), ); - let response = GatewayResponse { - message: block.message.clone(), - finish_reason: "content_filter".to_owned(), - ..Default::default() - }; - ctx.outcome = Some(EngineOutcome { - response, - http_code: 200, - block, - ..Default::default() - }); + ctx.outcome = Some(content_filter_outcome(block)); return Ok(ctx); } - // external moderation, on the ORIGINAL text (before DLP), when the tenant - // opted in; a real moderator is wired via with_moderator + + // pre-DLP text, computed once for moderation and the retained prompt + let retention = snap + .cfg + .retention_for(&ctx.ak.tenant) + .copied() + .filter(|r| r.content != gw_config::ContentLevel::None); + let inbound = + (sec.moderate || retention.is_some()).then(|| plugins::inbound_text(&mut ctx.request)); + if sec.moderate - && let Some(block) = self.moderate(&ctx, sec).await + && let Some(block) = self + .moderate(&ctx, sec, inbound.as_deref().unwrap_or_default()) + .await { ctx.decide("moderation", "denied"); - let response = GatewayResponse { - message: block.message.clone(), - finish_reason: "content_filter".to_owned(), - ..Default::default() - }; - ctx.outcome = Some(EngineOutcome { - response, - http_code: 200, - block, - ..Default::default() - }); + ctx.outcome = Some(content_filter_outcome(block)); return Ok(ctx); } - // capture raw prompt before DLP when the tenant retains full content — - // only worth it if a key exists (full without one downgrades to redacted) - let retention = snap.cfg.retention_for(&ctx.ak.tenant).copied(); - let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) - && gw_state::sealing_available(); - let raw_prompt = capture_raw.then(|| plugins::inbound_text(&ctx.request)); - let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { ctx.decide("dlp", format!("redacted {redacted} span(s) inbound")); @@ -318,7 +176,9 @@ impl OnlineHandler { outcome.response.model = requested; } - // capture raw response before outbound DLP for full retention + // raw response pre-outbound-DLP, only when full retention can store it (key present) + let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) + && gw_state::sealing_available(); let raw_response = capture_raw .then(|| ctx.outcome.as_ref().map(|o| o.response.message.clone())) .flatten(); @@ -334,8 +194,15 @@ impl OnlineHandler { } else { 0 }; - if let Some(r) = retention.filter(|r| r.content != gw_config::ContentLevel::None) { - persist_content(&ctx, r, raw_prompt, raw_response).await; + if let Some(r) = retention { + persist_content( + &ctx, + r, + capture_raw, + inbound.unwrap_or_default(), + raw_response, + ) + .await; } if redacted_out > 0 { ctx.decide("dlp", format!("redacted {redacted_out} span(s) outbound")); @@ -345,42 +212,53 @@ impl OnlineHandler { } /// Run the wired moderator over raw text; `Some(reason)` to deny, `None` to - /// allow. A moderator error resolves per `moderation_fail_open`. The seam - /// the realtime surface uses (it has no `DagContext`); the caller records - /// the security event on its own surface. + /// allow. The seam the realtime surface uses (it has no `DagContext`); the + /// caller records the security event on its own surface. pub async fn moderate_text(&self, sec: &gw_config::SecurityConf, text: &str) -> Option { - match self.moderator.review(text).await { - Ok(moderation::Verdict::Allow) => None, - Ok(moderation::Verdict::Deny(reason)) => Some(reason), - Err(e) => { - tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); - (!sec.moderation_fail_open).then(|| "content moderation is unavailable".to_owned()) - } + match self.moderation(sec, text).await { + Moderation::Allow => None, + Moderation::Deny(reason) => Some(reason), + Moderation::Unavailable => Some(MODERATION_UNAVAILABLE.to_owned()), } } - /// Run the wired moderator over the request's inbound text; `Some(Block)` to - /// deny. A moderator error resolves per `moderation_fail_open`. Records a - /// security event on a deny. - async fn moderate(&self, ctx: &DagContext, sec: &gw_config::SecurityConf) -> Option { - let text = plugins::inbound_text(&ctx.request); - match self.moderator.review(&text).await { - Ok(moderation::Verdict::Allow) => None, - Ok(moderation::Verdict::Deny(reason)) => { + /// Run the wired moderator over the request's pre-DLP inbound `text`; + /// `Some(Block)` to deny. Records a security event on a moderator deny. + async fn moderate( + &self, + ctx: &DagContext, + sec: &gw_config::SecurityConf, + text: &str, + ) -> Option { + match self.moderation(sec, text).await { + Moderation::Allow => None, + Moderation::Deny(reason) => { emit_security_event(ctx, "moderation", "block", 1).await; Some(Block::blocked( reason, gw_consts::ErrCode::EMPTY_RESP.value() as i32, )) } + Moderation::Unavailable => Some(Block::blocked( + MODERATION_UNAVAILABLE, + gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, + )), + } + } + + /// The one moderator-verdict resolution every surface shares, so the + /// fail-open posture can't drift between REST and realtime. + async fn moderation(&self, sec: &gw_config::SecurityConf, text: &str) -> Moderation { + match self.moderator.review(text).await { + Ok(moderation::Verdict::Allow) => Moderation::Allow, + Ok(moderation::Verdict::Deny(reason)) => Moderation::Deny(reason), Err(e) => { tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); - (!sec.moderation_fail_open).then(|| { - Block::blocked( - "content moderation is unavailable", - gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, - ) - }) + if sec.moderation_fail_open { + Moderation::Allow + } else { + Moderation::Unavailable + } } } } @@ -410,6 +288,137 @@ impl OnlineHandler { } } +/// One resolved moderator verdict: `Unavailable` is a moderator error under a +/// fail-closed posture (fail-open resolves to `Allow`). +enum Moderation { + Allow, + Deny(String), + Unavailable, +} + +/// A per-request correlation id: `req--`, time-sortable and +/// unique within the process (the seq disambiguates same-millisecond requests). +pub fn new_request_id() -> String { + let ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis()) + .unwrap_or(0); + format!("req-{ms}-{}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)) +} + +/// The 200-with-`content_filter` outcome every pre-stage denial returns. +fn content_filter_outcome(block: Block) -> EngineOutcome { + EngineOutcome { + response: GatewayResponse { + message: block.message.clone(), + finish_reason: "content_filter".to_owned(), + ..Default::default() + }, + http_code: 200, + block, + ..Default::default() + } +} + +/// Record a content-safety outcome (no prompt text) against this request's +/// key/user/tenant. Best-effort. Surface = the request protocol, or "batch" +/// for offline items. +async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i64) { + let surface = if ctx.request.is_online { + ctx.request + .model_param_v2 + .as_ref() + .map(|p| p.protocol.as_str()) + .unwrap_or_default() + .to_owned() + } else { + "batch".to_owned() + }; + gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id: ctx.effective_user_id().to_owned(), + tenant: ctx.ak.tenant.clone(), + surface, + rule: rule.to_owned(), + action: action.to_owned(), + hits, + } + .record(ctx.state.store.as_ref()) + .await; +} + +/// Persist this request's prompt and response per the tenant's retention policy. +/// `inbound` is the pre-DLP prompt text; `store_full` (resolved once by the +/// caller: full level AND a content key) stores it sealed, else it is stored +/// PII/secret-stripped — retention owns that redaction, so a row can't hold +/// raw content even with DLP off. Best-effort — a store failure is logged. +async fn persist_content( + ctx: &DagContext, + retention: gw_config::RetentionConf, + store_full: bool, + inbound: String, + raw_response: Option, +) { + if retention.content == gw_config::ContentLevel::Full && !store_full { + tracing::warn!("full retention configured but GW_CONTENT_KEY unset; storing redacted text"); + } + + let now = gw_state::epoch_secs(); + let expires = if retention.days > 0 { + now + retention.days as i64 * 86_400 + } else { + 0 + }; + let redacted_response = || { + plugins::redact_retained( + ctx.outcome + .as_ref() + .map(|o| o.response.message.as_str()) + .unwrap_or_default(), + ) + }; + let prompt = if store_full { + inbound + } else { + plugins::redact_retained(&inbound) + }; + let response = if store_full { + raw_response.unwrap_or_else(redacted_response) + } else { + redacted_response() + }; + + let writes = [("prompt", prompt), ("response", response)] + .into_iter() + .filter(|(_, text)| !text.is_empty()) + .map(|(kind, text)| { + // seal whenever a key exists (defense in depth even for redacted text) + let (content, sealed) = match gw_state::content::seal(&text) { + Some(ct) => (ct, true), + None => (text, false), + }; + let record = gw_state::ContentRecord { + created_at_epoch_secs: now, + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id: ctx.effective_user_id().to_owned(), + tenant: ctx.ak.tenant.clone(), + kind: kind.to_owned(), + content, + sealed, + expires_at_epoch_secs: expires, + }; + async move { + if let Err(e) = ctx.state.store.content_add(&record).await { + tracing::warn!(error = %e, kind, "content retention write failed"); + } + } + }); + futures::future::join_all(writes).await; +} + #[cfg(test)] mod tests { use super::*; @@ -1070,7 +1079,6 @@ mod tests { let j = completed.expect("drain completed the batch"); assert_eq!(j.results.len(), 2, "both items executed exactly once"); assert!(j.results.iter().all(|r| r.ok && r.total_tokens > 0)); - // per-item attribution survived the enqueue → drain → reconstruct round-trip let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); let users: std::collections::HashSet<&str> = ledger.iter().map(|r| r.user_id.as_str()).collect(); diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index a4377a3..81e204b 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -115,8 +115,6 @@ impl OfflineHandler { is_online: false, ak: ak.ak.clone(), message: item.messages, - // per-item attribution so effective_user_id (owner else this) - // bills and budgets each item to its end user user_id: (!item.user.is_empty()).then_some(item.user), model_param_v2: Some(ModelParamV2::with_name( gw_consts::Protocol::OpenaiChat, diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 389e232..795dfd5 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -32,15 +32,8 @@ pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanO if sec.blocklist.is_empty() && sec.regexes.is_empty() { return ScanOutcome::default(); } - let mut blocklist_count = 0i64; - let mut regex_counts = vec![0i64; sec.regexes.len()]; - let mut visit = |s: &mut String| { - blocklist_count += i64::from(blocklist_hit(sec, s)); - for (i, r) in sec.regexes.iter().enumerate() { - regex_counts[i] += r.re.find_iter(s).count() as i64; - } - 0 - }; + let mut counts = ScanCounts::new(sec); + let mut visit = |s: &mut String| counts.visit(s); for msg in &mut request.message { visit(&mut msg.content); if let Some(parts) = &mut msg.parts { @@ -53,61 +46,83 @@ pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanO for_each_typed_text(typed, &mut visit); } } + counts.outcome() +} - tally_scan(sec, blocklist_count, regex_counts) +/// Per-rule hit counters for one scan, shared by the REST request scan and the +/// realtime frame scan so the two apply identical action semantics. +struct ScanCounts<'a> { + sec: &'a SecurityConf, + blocklist: i64, + regex: Vec, } -/// Build a [`ScanOutcome`] from per-rule hit counts: a rule with count > 0 is a -/// hit at its action; any block-action hit denies. Shared by the REST scan and -/// the realtime frame scan so the two apply identical action semantics. -fn tally_scan(sec: &SecurityConf, blocklist_count: i64, regex_counts: Vec) -> ScanOutcome { - let mut hits = Vec::new(); - if blocklist_count > 0 { - hits.push(RuleHit { - rule: "blocklist".to_owned(), - action: sec.blocklist_action, - count: blocklist_count, - }); +impl<'a> ScanCounts<'a> { + fn new(sec: &'a SecurityConf) -> Self { + Self { + sec, + blocklist: 0, + regex: vec![0; sec.regexes.len()], + } } - for (r, count) in sec.regexes.iter().zip(regex_counts) { - if count > 0 { + + fn visit(&mut self, s: &str) -> usize { + self.blocklist += i64::from(blocklist_hit(self.sec, s)); + for (i, r) in self.sec.regexes.iter().enumerate() { + self.regex[i] += r.re.find_iter(s).count() as i64; + } + 0 + } + + /// Fold the counts into a [`ScanOutcome`]: a rule with count > 0 is a hit + /// at its action; any block-action hit denies. + fn outcome(self) -> ScanOutcome { + let mut hits = Vec::new(); + if self.blocklist > 0 { hits.push(RuleHit { - rule: r.name.clone(), - action: r.action, - count, + rule: "blocklist".to_owned(), + action: self.sec.blocklist_action, + count: self.blocklist, }); } + for (r, count) in self.sec.regexes.iter().zip(self.regex) { + if count > 0 { + hits.push(RuleHit { + rule: r.name.clone(), + action: r.action, + count, + }); + } + } + let block = hits + .iter() + .any(|h| h.action == Action::Block) + .then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)); + ScanOutcome { block, hits } } - let block = hits - .iter() - .any(|h| h.action == Action::Block) - .then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)); - ScanOutcome { block, hits } } /// All inbound text a request carries — message content, multimodal text -/// parts, the Responses raw body, and the family typed params. The one text -/// view moderation and content retention operate on, so a non-chat surface -/// (input in `raw`/typed with an empty `message`) is not a bypass. The -/// typed-field list mirrors [`for_each_typed_text`]. -pub fn inbound_text(request: &GatewayRequest) -> String { +/// parts, the Responses raw body, and the family typed params — collected via +/// the same traversals the blocklist scan and DLP run, so the field lists +/// cannot drift apart. The one text view moderation and content retention +/// operate on. `&mut` only to share those traversals; nothing is rewritten. +pub fn inbound_text(request: &mut GatewayRequest) -> String { let mut out = String::new(); - for m in &request.message { - push_text(&mut out, &m.content); - if let Some(serde_json::Value::Array(parts)) = &m.parts { - for p in parts { - if p["type"] == "text" - && let Some(t) = p["text"].as_str() - { - push_text(&mut out, t); - } - } + let mut collect = |s: &mut String| { + push_text(&mut out, s); + 0 + }; + for m in &mut request.message { + collect(&mut m.content); + if let Some(parts) = &mut m.parts { + for_each_part_text(parts, &mut collect); } } - if let Some(param) = &request.model_param_v2 { - collect_json_text(¶m.raw, &mut out); - if let Some(typed) = ¶m.typed { - collect_typed_text(typed, &mut out); + if let Some(param) = request.model_param_v2.as_mut() { + walk_json_strings(&mut param.raw, &mut collect); + if let Some(typed) = param.typed.as_mut() { + for_each_typed_text(typed, &mut collect); } } out @@ -122,38 +137,6 @@ fn push_text(out: &mut String, s: &str) { } } -fn collect_json_text(v: &serde_json::Value, out: &mut String) { - match v { - serde_json::Value::String(s) => push_text(out, s), - serde_json::Value::Array(a) => a.iter().for_each(|x| collect_json_text(x, out)), - serde_json::Value::Object(o) => o.values().for_each(|x| collect_json_text(x, out)), - _ => {} - } -} - -fn collect_typed_text(typed: &gw_models::TypedParams, out: &mut String) { - use gw_models::TypedParams as T; - match typed { - T::Chat(p) => { - if let Some(s) = &p.system { - push_text(out, s); - } - if let Some(t) = &p.tools { - collect_json_text(t, out); - } - if let Some(t) = &p.tool_choice { - collect_json_text(t, out); - } - } - T::Embeddings(p) => p.input.iter().for_each(|s| push_text(out, s)), - T::AudioTts(p) => push_text(out, &p.input), - T::Image(p) => push_text(out, &p.prompt), - T::Video(p) => push_text(out, &p.prompt), - T::Search(p) => push_text(out, &p.query), - T::AudioStt(_) => {} - } -} - /// Case-insensitive blocklist test; terms are pre-lowercased at config load. /// ASCII text matches without allocating; non-ASCII falls back to a lowercase copy. fn blocklist_hit(sec: &SecurityConf, text: &str) -> bool { @@ -215,24 +198,53 @@ fn for_each_typed_text( } } +/// Visit a multimodal `parts` array's text blocks; returns summed hits. +/// Deliberately narrower than the blocklist scan: non-text parts (image URLs, +/// base64 data) are never visited, so a rewrite can't corrupt them. +fn for_each_part_text( + parts: &mut serde_json::Value, + f: &mut impl FnMut(&mut String) -> usize, +) -> usize { + let Some(arr) = parts.as_array_mut() else { + return 0; + }; + let mut hits = 0; + for p in arr { + if p["type"] == "text" + && let Some(serde_json::Value::String(s)) = p.get_mut("text") + { + hits += f(s); + } + } + hits +} + /// Scan a realtime frame's text-bearing fields against the blocklist AND the /// regex recognizers, honoring their actions — the same policy every REST -/// surface runs, so the WebSocket surface is not a bypass. Returns the same -/// [`ScanOutcome`] as [`security_check`]: a block-action hit denies the frame. -pub fn realtime_frame_scan(sec: &SecurityConf, frame: &mut serde_json::Value) -> ScanOutcome { - if sec.blocklist.is_empty() && sec.regexes.is_empty() { - return ScanOutcome::default(); +/// surface runs, so the WebSocket surface is not a bypass. A block-action hit +/// denies the frame. `collect_text` also gathers the frame's text (for +/// moderation) in the same traversal, so a moderated tenant pays one walk per +/// frame, not two. +pub fn realtime_frame_scan( + sec: &SecurityConf, + frame: &mut serde_json::Value, + collect_text: bool, +) -> (ScanOutcome, String) { + let scan_rules = !sec.blocklist.is_empty() || !sec.regexes.is_empty(); + let mut counts = ScanCounts::new(sec); + let mut text = String::new(); + if scan_rules || collect_text { + gw_engines::realtime::visit_frame_text(frame, &mut |s| { + if scan_rules { + counts.visit(s); + } + if collect_text { + push_text(&mut text, s); + } + 0 + }); } - let mut blocklist_count = 0i64; - let mut regex_counts = vec![0i64; sec.regexes.len()]; - gw_engines::realtime::visit_frame_text(frame, &mut |s| { - blocklist_count += i64::from(blocklist_hit(sec, s)); - for (i, r) in sec.regexes.iter().enumerate() { - regex_counts[i] += r.re.find_iter(s).count() as i64; - } - 0 - }); - tally_scan(sec, blocklist_count, regex_counts) + (counts.outcome(), text) } /// DLP-redact a realtime frame's text-bearing fields in place; the hit count. @@ -263,7 +275,7 @@ pub fn dlp_redact_request(sec: &SecurityConf, request: &mut GatewayRequest) -> u // engines forward `parts` (not `content`) when present, so PII must be // scrubbed inside the parts' text blocks too if let Some(parts) = &mut msg.parts { - hits += redact_parts_text(parts, pii, secrets); + hits += for_each_part_text(parts, &mut redact_field); } } // non-chat surfaces carry user text outside `message` (Responses raw body, @@ -307,24 +319,6 @@ fn redact_str(s: &mut String) -> usize { redact_in_place(s, true, false) } -/// Redact PII inside a multimodal `parts` array's text blocks, in place. -/// Deliberately narrower than the blocklist scan: non-text parts (image URLs, -/// base64 data) are never rewritten, so a redaction can't corrupt them. -fn redact_parts_text(parts: &mut serde_json::Value, pii: bool, secrets: bool) -> usize { - let Some(arr) = parts.as_array_mut() else { - return 0; - }; - let mut hits = 0; - for p in arr { - if p["type"] == "text" - && let Some(serde_json::Value::String(s)) = p.get_mut("text") - { - hits += redact_in_place(s, pii, secrets); - } - } - hits -} - /// Mask credential shapes (API keys, tokens, private-key headers) with /// `[REDACTED_SECRET]`. High-precision patterns to avoid mauling normal text. fn redact_secrets(text: &str) -> Option<(String, usize)> { @@ -504,63 +498,68 @@ mod tests { #[test] fn inbound_text_covers_non_chat_raw_and_typed() { use gw_models::{ModelParamV2, TypedParams}; - // Responses-style: input rides in raw, message is empty let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, "m"); param.raw = serde_json::json!({"input": "secret responses text", "model": "m"}); - let req = GatewayRequest { + let mut req = GatewayRequest { model_param_v2: Some(param), ..Default::default() }; - assert!(inbound_text(&req).contains("secret responses text")); + assert!( + inbound_text(&mut req).contains("secret responses text"), + "Responses input rides in raw with an empty message" + ); - // embeddings-style: input rides in typed, message is empty let mut param = ModelParamV2::with_name(gw_consts::Protocol::Embeddings, "m"); param.typed = Some(TypedParams::Embeddings(gw_models::EmbeddingParams { input: vec!["typed embed text".into()], dimensions: None, })); - let req = GatewayRequest { + let mut req = GatewayRequest { model_param_v2: Some(param), ..Default::default() }; - assert!(inbound_text(&req).contains("typed embed text")); + assert!( + inbound_text(&mut req).contains("typed embed text"), + "embeddings input rides in typed with an empty message" + ); } - #[test] - fn realtime_frame_scan_honors_regex_and_actions() { - let mut s = SecurityConf { - regex_rules: vec![gw_config::RegexRule { + fn ssn_block() -> SecurityConf { + SecurityConf { + regexes: vec![gw_config::CompiledRule { name: "ssn".into(), - pattern: r"\d{3}-\d{2}-\d{4}".into(), action: Action::Block, + re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), }], ..Default::default() - }; - s.regexes = vec![gw_config::CompiledRule { - name: "ssn".into(), - action: Action::Block, - re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), - }]; + } + } + + #[test] + fn realtime_frame_scan_honors_regex_and_actions() { + let s = ssn_block(); let mut frame = serde_json::json!({"type":"input_text","text":"my ssn is 123-45-6789"}); - let out = realtime_frame_scan(&s, &mut frame); + let (out, text) = realtime_frame_scan(&s, &mut frame, true); assert!(out.block.is_some(), "regex Block denies on realtime too"); + assert!( + text.contains("123-45-6789"), + "collect_text gathers the frame text in the same walk" + ); - // flag blocklist: hit recorded, not blocked let s2 = SecurityConf { blocklist: vec!["watch".into()], blocklist_action: Action::Flag, ..Default::default() }; let mut frame = serde_json::json!({"type":"input_text","text":"please watch this"}); - let out = realtime_frame_scan(&s2, &mut frame); + let (out, text) = realtime_frame_scan(&s2, &mut frame, false); assert!(out.block.is_none(), "flag does not block realtime"); assert_eq!(out.hits.len(), 1); + assert!(text.is_empty(), "no text collected unless asked"); } #[test] fn realtime_frame_redacts_secrets_when_detect_secrets_on() { - // detect_secrets on, dlp_redact OFF: a credential in a realtime frame is - // still masked before it reaches the model (the frame is not a bypass) let s = SecurityConf { detect_secrets: true, dlp_redact: false, @@ -577,11 +576,14 @@ mod tests { "{text}" ); - // both flags off: realtime frame untouched let none = SecurityConf::default(); let mut frame = serde_json::json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); - assert_eq!(dlp_redact_realtime_frame(&none, &mut frame), 0); + assert_eq!( + dlp_redact_realtime_frame(&none, &mut frame), + 0, + "both flags off leaves the frame untouched" + ); } #[test] @@ -616,20 +618,7 @@ mod tests { #[test] fn regex_rule_blocks_and_redact_secrets_masks() { - let mut s = SecurityConf { - regex_rules: vec![gw_config::RegexRule { - name: "ssn".into(), - pattern: r"\d{3}-\d{2}-\d{4}".into(), - action: Action::Block, - }], - ..Default::default() - }; - // compile_security runs at config load; replicate it for the unit test - s.regexes = vec![gw_config::CompiledRule { - name: "ssn".into(), - action: Action::Block, - re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), - }]; + let s = ssn_block(); let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "my ssn is 123-45-6789")], ..Default::default() diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index 31a86eb..3ce7511 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -54,8 +54,6 @@ accounts: [{name: mock-openai-1, provider: openai, protocols: ["openai-chat"]}] Some(loader), )); - // this reload FLIPS trust_proxy_headers on, while forging x-real-ip on the - // very same request — the op must be audited under the pre-reload policy let r = app .clone() .oneshot( diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index fc97c46..1dcbf05 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -31,6 +31,12 @@ pub fn user_budget_key(tenant: &str, user: &str) -> String { format!("ub:{tenant}:{user}") } +/// The tenant's per-user daily token budget, if it configured one. +pub fn user_budget_limit(cfg: &GatewayConfig, tenant: &str) -> Option { + cfg.find_tenant(tenant) + .and_then(|t| t.user_daily_token_quota) +} + /// Per-user daily token budget (a soft cap): admit while the user is under the /// tenant's limit. Skipped when the tenant sets no limit or the request carries /// no user attribution. Consumed at billing via [`consume_user_budget`]. @@ -43,10 +49,7 @@ pub async fn check_user_budget( if user.is_empty() { return Ok(()); } - let Some(limit) = cfg - .find_tenant(tenant) - .and_then(|t| t.user_daily_token_quota) - else { + let Some(limit) = user_budget_limit(cfg, tenant) else { return Ok(()); }; if gov.quota_check(&user_budget_key(tenant, user), limit).await { @@ -68,11 +71,7 @@ pub async fn consume_user_budget( if user.is_empty() || total <= 0 { return; } - if cfg - .find_tenant(tenant) - .and_then(|t| t.user_daily_token_quota) - .is_some() - { + if user_budget_limit(cfg, tenant).is_some() { gov.quota_consume(&user_budget_key(tenant, user), total) .await; } diff --git a/crates/state/src/content.rs b/crates/state/src/content.rs index 3a0e7ca..721193f 100644 --- a/crates/state/src/content.rs +++ b/crates/state/src/content.rs @@ -9,8 +9,19 @@ use base64::Engine as _; use chacha20poly1305::aead::{Aead, KeyInit, OsRng}; use chacha20poly1305::{AeadCore, XChaCha20Poly1305}; -/// One stored prompt or response, priced out of the ledger's blast radius: -/// content may be ciphertext (sealed) or plaintext (redacted level). +/// The process-wide content key, loaded once from `GW_CONTENT_KEY`. +static CIPHER: LazyLock> = LazyLock::new(|| { + let raw = std::env::var("GW_CONTENT_KEY").ok()?; + match hex::decode(raw.trim()) { + Ok(bytes) if bytes.len() == 32 => Some(XChaCha20Poly1305::new(bytes.as_slice().into())), + _ => { + tracing::error!("GW_CONTENT_KEY is not 64 hex chars (32 bytes); content sealing off"); + None + } + } +}); + +/// One stored prompt or response. #[derive(Debug, Clone, serde::Serialize)] pub struct ContentRecord { pub created_at_epoch_secs: i64, @@ -28,18 +39,6 @@ pub struct ContentRecord { pub expires_at_epoch_secs: i64, } -/// The process-wide content key, loaded once from `GW_CONTENT_KEY`. -static CIPHER: LazyLock> = LazyLock::new(|| { - let raw = std::env::var("GW_CONTENT_KEY").ok()?; - match hex::decode(raw.trim()) { - Ok(bytes) if bytes.len() == 32 => Some(XChaCha20Poly1305::new(bytes.as_slice().into())), - _ => { - tracing::error!("GW_CONTENT_KEY is not 64 hex chars (32 bytes); content sealing off"); - None - } - } -}); - /// Whether a deployment key is configured (so `full` retention may store raw). pub fn sealing_available() -> bool { CIPHER.is_some() diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index 390df12..6277d93 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -31,10 +31,9 @@ pub trait KeyStore: Send + Sync + std::fmt::Debug { async fn patch(&self, ak: &str, patch: &KeyPatch) -> GResult>; /// Remove a key regardless of source; whether it existed. async fn revoke(&self, ak: &str) -> GResult; - /// A page of keys, sorted by ak (stable), optionally confined to one tenant - /// (a tenant admin's scope). `offset`/`limit` bound the scan so a fleet key - /// table with millions of rows never loads whole — and the tenant filter - /// applies BEFORE paging, so a scoped page isn't emptied by a later filter. + /// A page of keys, sorted by ak, optionally confined to one tenant. The + /// tenant filter applies before paging, so a scoped page is never emptied + /// by a later filter; `offset`/`limit` bound the scan. async fn list(&self, tenant: Option<&str>, offset: usize, limit: usize) -> GResult>; /// Re-apply the config file's key set, leaving admin-created keys untouched. diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 0c9bd89..b736a35 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -64,11 +64,9 @@ impl AkInfo { } } - /// The attributed end user: this key's `owner` when set to a non-empty value - /// (authoritative), else the caller-supplied `fallback` — request metadata - /// on REST surfaces, the `x-gw-user` connect hint on realtime. Empty when - /// neither is present. The one resolution every surface shares so an empty - /// owner can't fall back on one surface and swallow attribution on another. + /// The attributed end user: this key's non-empty `owner` (authoritative), + /// else the caller-supplied `fallback` (request metadata / the realtime + /// `x-gw-user` hint). The one resolution every surface shares. pub fn attributed_user<'a>(&'a self, fallback: &'a str) -> &'a str { self.owner .as_deref() diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 234d7a2..08abe30 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -249,6 +249,16 @@ pub struct SecurityEvent { pub hits: i64, } +impl SecurityEvent { + /// Append to `store`. Best-effort: a write failure is logged, never fails + /// the request being audited. + pub async fn record(&self, store: &dyn Store) { + if let Err(e) = store.security_event_add(self).await { + tracing::warn!(error = %e, rule = %self.rule, "security event write failed"); + } + } +} + /// One admin-plane mutation, recorded with who/what/when for compliance. #[derive(Debug, Clone, serde::Serialize)] pub struct AdminAudit { @@ -649,55 +659,63 @@ impl Store for MemoryStore { } } -/// Positional row → record shared by the SQL backends (identical SELECT order). -fn row_to_billing<'r, R>(row: &'r R) -> BillingRecord -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, - bool: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - BillingRecord { - ak: row.get(0), - product: row.get(1), - tenant: row.get(2), - model: row.get(3), - served_model: row.get(4), - protocol: row.get(5), - account: row.get(6), - prompt_tokens: row.get(7), - completion_tokens: row.get(8), - total_tokens: row.get(9), - cost_micros: row.get(10), - vendor_cost_micros: row.get(11), - ptu_spillover: row.get(12), - user_id: row.get(13), - request_id: row.get(14), - created_at_epoch_secs: row.get(15), - estimated: row.get(16), - } +/// Positional row → record mappers shared by the SQL backends: fields decode +/// in the SELECT's column order. One macro so the sqlx trait-bound boilerplate +/// lives once. +macro_rules! row_mapper { + ($name:ident -> $ty:path { $($field:ident),+ $(,)? }) => { + fn $name<'r, R>(row: &'r R) -> $ty + where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, + bool: sqlx::Decode<'r, R::Database> + sqlx::Type, + { + let mut col = 0usize; + $(let $field = row.get(next_col(&mut col));)+ + $ty { $($field),+ } + } + }; } -fn usage_row<'r, R>(row: &'r R) -> UsageRow -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - UsageRow { - tenant: row.get(0), - model: row.get(1), - requests: row.get(2), - prompt_tokens: row.get(3), - completion_tokens: row.get(4), - total_tokens: row.get(5), - cost_micros: row.get(6), - vendor_cost_micros: row.get(7), - } +fn next_col(col: &mut usize) -> usize { + let i = *col; + *col += 1; + i } +row_mapper!(row_to_billing -> BillingRecord { + ak, product, tenant, model, served_model, protocol, account, + prompt_tokens, completion_tokens, total_tokens, cost_micros, + vendor_cost_micros, ptu_spillover, user_id, request_id, + created_at_epoch_secs, estimated, +}); + +row_mapper!(usage_row -> UsageRow { + tenant, model, requests, prompt_tokens, completion_tokens, total_tokens, + cost_micros, vendor_cost_micros, +}); + +row_mapper!(user_usage_row -> UserUsageRow { + user_id, model, requests, prompt_tokens, completion_tokens, total_tokens, + cost_micros, vendor_cost_micros, +}); + +row_mapper!(security_event_row -> SecurityEvent { + created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits, +}); + +row_mapper!(admin_audit_row -> AdminAudit { + created_at_epoch_secs, actor, scope, action, target, summary, source_ip, +}); + +row_mapper!(content_row -> crate::ContentRecord { + created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, + sealed, expires_at_epoch_secs, +}); + fn batch_item_row<'r, R>(row: &'r R) -> BatchItemResult where R: sqlx::Row, @@ -714,84 +732,6 @@ where } } -fn user_usage_row<'r, R>(row: &'r R) -> UserUsageRow -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - UserUsageRow { - user_id: row.get(0), - model: row.get(1), - requests: row.get(2), - prompt_tokens: row.get(3), - completion_tokens: row.get(4), - total_tokens: row.get(5), - cost_micros: row.get(6), - vendor_cost_micros: row.get(7), - } -} - -fn security_event_row<'r, R>(row: &'r R) -> SecurityEvent -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - SecurityEvent { - created_at_epoch_secs: row.get(0), - request_id: row.get(1), - ak: row.get(2), - user_id: row.get(3), - tenant: row.get(4), - surface: row.get(5), - rule: row.get(6), - action: row.get(7), - hits: row.get(8), - } -} - -fn admin_audit_row<'r, R>(row: &'r R) -> AdminAudit -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - AdminAudit { - created_at_epoch_secs: row.get(0), - actor: row.get(1), - scope: row.get(2), - action: row.get(3), - target: row.get(4), - summary: row.get(5), - source_ip: row.get(6), - } -} - -fn content_row<'r, R>(row: &'r R) -> crate::ContentRecord -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, - bool: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - crate::ContentRecord { - created_at_epoch_secs: row.get(0), - request_id: row.get(1), - ak: row.get(2), - user_id: row.get(3), - tenant: row.get(4), - kind: row.get(5), - content: row.get(6), - sealed: row.get(7), - expires_at_epoch_secs: row.get(8), - } -} - /// SQLite-backed store (WAL): ledger, files, and batch jobs in one database /// file; ids derive from rowids so they stay unique across restarts. #[derive(Debug)] diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 5da75c3..c13d273 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -3,6 +3,7 @@ //! structured access-log line per request. use std::convert::Infallible; +use std::fmt::Write as _; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; @@ -427,36 +428,12 @@ async fn realtime_session( .await; continue; }; - // same policy (blocklist + regex + actions) + inbound DLP every REST - // surface runs, tenant policy first - let cfg = s.handler.cfg(); - let sec = cfg.security_for(&ak.tenant); - let scan = gw_handler::plugins::realtime_frame_scan(sec, &mut ev); - emit_rt_hits(&s, &ak, &scan.hits, &hint).await; - if let Some(block) = scan.block { - let _ = socket - .send(send(json!({"type":"error","message": block.message}))) - .await; - continue; - } - if let Some(reason) = realtime_moderate(&s, sec, &ak, &hint, &mut ev).await { + if let Err(reason) = rt_inbound_policy(&s, &ak, &hint, &mut ev).await { let _ = socket .send(send(json!({"type":"error","message": reason}))) .await; continue; } - let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut ev); - if redacted > 0 { - write_rt_event( - &s, - &ak, - ak.attributed_user(&hint), - "dlp", - "redact", - redacted as i64, - ) - .await; - } match ev["type"].as_str().unwrap_or_default() { "input_text" => { let admit = match realtime_gate(&s, &ak, &model, &hint).await { @@ -612,37 +589,22 @@ async fn realtime_bridge( Some(Ok(_)) => continue, // ping/pong handled by the ws stacks }; if let Some(mut frame) = frame { - // same policy (blocklist + regex + actions) + inbound DLP every - // REST surface runs, tenant policy first - let cfg = s.handler.cfg(); - let sec = cfg.security_for(&ak.tenant); - let scan = gw_handler::plugins::realtime_frame_scan(sec, &mut frame); - emit_rt_hits(&s, &ak, &scan.hits, &hint).await; - if let Some(block) = scan.block { - if cl_tx - .send(CMsg::Text(send_err(block.message).to_string().into())) - .await - .is_err() - { - break; + match rt_inbound_policy(&s, &ak, &hint, &mut frame).await { + Err(reason) => { + if cl_tx + .send(CMsg::Text(send_err(reason).to_string().into())) + .await + .is_err() + { + break; + } + continue; } - continue; - } - if let Some(reason) = realtime_moderate(&s, sec, &ak, &hint, &mut frame).await { - if cl_tx - .send(CMsg::Text(send_err(reason).to_string().into())) - .await - .is_err() - { - break; + Ok(redacted) => { + if redacted > 0 { + forward = UMsg::text(frame.to_string()); + } } - continue; - } - let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut frame); - if redacted > 0 { - write_rt_event(&s, &ak, ak.attributed_user(&hint), "dlp", "redact", redacted as i64) - .await; - forward = UMsg::text(frame.to_string()); } // gate each generation trigger, not every control frame if is_response_create(&frame) { @@ -765,14 +727,10 @@ async fn realtime_bridge( if n > 0 { redacted = Some(v.to_string()); } - // a per-token event stream would be too hot: sum the turn's - // outbound redactions and record one event at its boundary + // per-token events would be too hot: sum the turn, record once at its boundary out_redacted += n as i64; if turn_ended { - if out_redacted > 0 { - write_rt_event(&s, &ak, ak.attributed_user(&hint), "dlp", "redact", out_redacted) - .await; - } + flush_rt_out_dlp(&s, &ak, &hint, out_redacted).await; out_redacted = 0; } } @@ -799,17 +757,7 @@ async fn realtime_bridge( } // a turn aborted before its boundary (upstream drop) still applied its // redactions per frame — flush the pending count so the audit isn't lost - if out_redacted > 0 { - write_rt_event( - &s, - &ak, - ak.attributed_user(&hint), - "dlp", - "redact", - out_redacted, - ) - .await; - } + flush_rt_out_dlp(&s, &ak, &hint, out_redacted).await; if generations > 0 && recognized == 0 { tracing::warn!( account = %account.name, @@ -901,10 +849,7 @@ async fn ledger( State(s): State, axum::extract::Query(q): axum::extract::Query>, ) -> Response { - let limit = q - .get("limit") - .and_then(|v| v.parse::().ok()) - .unwrap_or(LEDGER_PAGE_DEFAULT); + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); match s.handler.state().store.ledger_snapshot(limit).await { Ok((count, records)) => Json(json!({ "count": count, "records": records })).into_response(), Err(e) => gateway_error(e), @@ -941,6 +886,14 @@ fn user_header(headers: &HeaderMap) -> Option { .filter(|s| !s.is_empty()) } +/// The one request-metadata attribution precedence the REST surfaces apply: +/// `x-gw-user` header, else the dialect's own user field (OpenAI `user`, +/// Anthropic `metadata.user_id`). Batch items invert it — per-item `user` +/// first — so shared-key batches keep per-item attribution. +fn user_hint(headers: &HeaderMap, field: &Value) -> Option { + user_header(headers).or_else(|| field.as_str().map(str::to_owned)) +} + /// AK auth: `Authorization: Bearer ` or `x-api-key: `. The error is /// `(status, message)` so each surface can shape it to its own wire dialect. async fn authenticate(s: &AppState, headers: &HeaderMap) -> Result { @@ -1037,7 +990,7 @@ async fn write_rt_event( action: &str, hits: i64, ) { - let event = gw_state::SecurityEvent { + gw_state::SecurityEvent { created_at_epoch_secs: gw_state::epoch_secs(), request_id: String::new(), ak: ak.ak.clone(), @@ -1047,10 +1000,51 @@ async fn write_rt_event( rule: rule.to_owned(), action: action.to_owned(), hits, - }; - if let Err(e) = s.handler.state().store.security_event_add(&event).await { - tracing::warn!(error = %e, "realtime security event write failed"); } + .record(s.handler.state().store.as_ref()) + .await; +} + +/// Record a turn's summed outbound DLP redactions as one event; no-op at zero. +async fn flush_rt_out_dlp(s: &AppState, ak: &AkInfo, hint: &str, count: i64) { + if count > 0 { + write_rt_event(s, ak, ak.attributed_user(hint), "dlp", "redact", count).await; + } +} + +/// The full inbound content policy for one realtime frame — the same chain +/// every REST surface runs (scan + hit events, moderation, DLP + event), +/// shared by both WebSocket paths. `Err(reason)` denies the frame; `Ok(n)` is +/// the DLP redaction count (n > 0 means the frame was rewritten). +async fn rt_inbound_policy( + s: &AppState, + ak: &AkInfo, + hint: &str, + frame: &mut Value, +) -> Result { + let cfg = s.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); + let (scan, text) = gw_handler::plugins::realtime_frame_scan(sec, frame, sec.moderate); + emit_rt_hits(s, ak, &scan.hits, hint).await; + if let Some(block) = scan.block { + return Err(block.message); + } + if let Some(reason) = realtime_moderate(s, sec, ak, hint, &text).await { + return Err(reason); + } + let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, frame); + if redacted > 0 { + write_rt_event( + s, + ak, + ak.attributed_user(hint), + "dlp", + "redact", + redacted as i64, + ) + .await; + } + Ok(redacted) } /// Record a realtime frame's content-safety hits to the security-event stream @@ -1061,67 +1055,54 @@ async fn emit_rt_hits( hits: &[gw_handler::plugins::RuleHit], hint: &str, ) { - let user = ak.attributed_user(hint).to_owned(); + let user = ak.attributed_user(hint); for hit in hits { - write_rt_event(s, ak, &user, &hit.rule, hit.action.as_str(), hit.count).await; + write_rt_event(s, ak, user, &hit.rule, hit.action.as_str(), hit.count).await; } } -/// All inbound text a realtime frame carries, for moderation. -fn realtime_frame_text(frame: &mut Value) -> String { - let mut text = String::new(); - gw_engines::realtime::visit_frame_text(frame, &mut |s| { - if !s.is_empty() { - if !text.is_empty() { - text.push('\n'); - } - text.push_str(s); - } - 0 - }); - text -} - -/// Moderate a realtime frame's inbound text via the wired moderator — parity -/// with the REST surface, so a moderated tenant can't be bypassed over the -/// WebSocket. `Some(reason)` denies the frame; records a moderation event. +/// Moderate a realtime frame's inbound `text` (collected by the frame scan) +/// via the wired moderator — parity with the REST surface, so a moderated +/// tenant can't be bypassed over the WebSocket. `Some(reason)` denies the +/// frame; records a moderation event. async fn realtime_moderate( s: &AppState, sec: &gw_config::SecurityConf, ak: &AkInfo, hint: &str, - frame: &mut Value, + text: &str, ) -> Option { - if !sec.moderate { + if !sec.moderate || text.is_empty() { return None; } - let text = realtime_frame_text(frame); - if text.is_empty() { - return None; - } - let reason = s.handler.moderate_text(sec, &text).await?; + let reason = s.handler.moderate_text(sec, text).await?; write_rt_event(s, ak, ak.attributed_user(hint), "moderation", "block", 1).await; Some(reason) } -/// The connecting peer's socket address, read from the connect-info extension. -/// An infallible extractor: `None` when the router is driven without connect -/// info (the test harness), `Some` when served over a real listener. -struct ClientPeer(Option); +/// The caller IP for the admin audit trail, resolved at request entry — before +/// any config mutation the handler performs — so the op that flips +/// `trust_proxy_headers` is audited under the policy in effect when it +/// arrived, not the one it just installed. Empty when the router is driven +/// without connect info (the test harness). +struct AuditSourceIp(String); -impl axum::extract::FromRequestParts for ClientPeer { +impl axum::extract::FromRequestParts for AuditSourceIp { type Rejection = std::convert::Infallible; async fn from_request_parts( parts: &mut axum::http::request::Parts, - _state: &S, + s: &AppState, ) -> Result { - Ok(ClientPeer( - parts - .extensions - .get::>() - .map(|ci| ci.0), - )) + let peer = parts + .extensions + .get::>() + .map(|ci| ci.0); + Ok(AuditSourceIp(source_ip( + peer, + &parts.headers, + s.handler.cfg().trust_proxy_headers, + ))) } } @@ -1146,9 +1127,6 @@ fn source_ip(peer: Option, headers: &HeaderMap, trust_prox } /// Record one admin-plane mutation to the audit trail (who/what/when/where). -/// `source` is resolved by the caller at request entry — before any config -/// mutation — so the op that flips `trust_proxy_headers` is audited under the -/// policy in effect when it arrived, not the one it just installed. /// Best-effort: a store failure is logged, never fails the operation. async fn audit_admin( s: &AppState, @@ -1270,13 +1248,11 @@ fn ct_eq(a: &str, b: &str) -> bool { async fn admin_reload( State(s): State, headers: HeaderMap, - ClientPeer(peer): ClientPeer, + AuditSourceIp(source): AuditSourceIp, ) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } - // freeze the source IP under the pre-reload policy - let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); match s.reload().await { Ok(()) => { let cfg = s.handler.cfg(); @@ -1307,14 +1283,13 @@ async fn admin_reload( async fn admin_key_create( State(s): State, headers: HeaderMap, - ClientPeer(peer): ClientPeer, + AuditSourceIp(source): AuditSourceIp, Json(body): Json, ) -> Response { let scope = match admin_auth(&s, &headers) { Ok(scope) => scope, Err(r) => return r, }; - let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); let (Some(ak), Some(product)) = (body["ak"].as_str(), body["product"].as_str()) else { return error_response(400, "ak and product are required"); }; @@ -1386,7 +1361,7 @@ async fn admin_key_create( async fn admin_key_patch( State(s): State, headers: HeaderMap, - ClientPeer(peer): ClientPeer, + AuditSourceIp(source): AuditSourceIp, Path(ak): Path, Json(body): Json, ) -> Response { @@ -1394,7 +1369,6 @@ async fn admin_key_patch( Ok(scope) => scope, Err(r) => return r, }; - let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); if let Err(r) = scoped_key(&s, &scope, &ak).await { return r; } @@ -1426,14 +1400,13 @@ async fn admin_key_patch( async fn admin_key_delete( State(s): State, headers: HeaderMap, - ClientPeer(peer): ClientPeer, + AuditSourceIp(source): AuditSourceIp, Path(ak): Path, ) -> Response { let scope = match admin_auth(&s, &headers) { Ok(scope) => scope, Err(r) => return r, }; - let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); if let Err(r) = scoped_key(&s, &scope, &ak).await { return r; } @@ -1456,15 +1429,12 @@ async fn admin_key_delete( async fn admin_config_put( State(s): State, headers: HeaderMap, - ClientPeer(peer): ClientPeer, + AuditSourceIp(source): AuditSourceIp, body: String, ) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } - // freeze the source IP under the pre-publish policy: this very op may be the - // one enabling trust_proxy_headers, and must not be audited under it - let source = source_ip(peer, &headers, s.handler.cfg().trust_proxy_headers); let Some(store) = &s.config_store else { return error_response( 400, @@ -1478,9 +1448,7 @@ async fn admin_config_put( Ok(v) => v, Err(e) => return gateway_error(e), }; - // publish already made this version the fleet source of truth and notified - // peers; audit it before the local reload can fail, or a failed apply would - // leave the fleet on a config with no publish record + // audit before the local reload can fail — the published version already leads the fleet let reload = s.reload().await; let detail = match &reload { Ok(()) => String::new(), @@ -1521,8 +1489,7 @@ async fn admin_key_list( }; let offset = q_num(&q, "offset", 0); let limit = q_num(&q, "limit", KEY_PAGE_DEFAULT); - // filter by the caller's scope IN the store, before paging, so a tenant - // admin's page isn't emptied by a post-hoc filter over the global table + // the scope filters in the store before paging, or a tenant admin's page could come back empty let tenant = scope.tenant_filter(&q); let listed = match s .handler @@ -1604,8 +1571,9 @@ async fn admin_usage_users( "user_id,model,requests,prompt_tokens,completion_tokens,total_tokens,cost_micros,vendor_cost_micros\n", ); for u in &usage { - csv.push_str(&format!( - "{},{},{},{},{},{},{},{}\n", + let _ = writeln!( + csv, + "{},{},{},{},{},{},{},{}", csv_field(&u.user_id), csv_field(&u.model), u.requests, @@ -1614,7 +1582,7 @@ async fn admin_usage_users( u.total_tokens, u.cost_micros, u.vendor_cost_micros, - )); + ); } return ([("content-type", "text/csv")], csv).into_response(); } @@ -1818,7 +1786,7 @@ async fn chat_completions( ); param.typed = Some(typed); param.raw = Value::Object(body.extra); - let user_id = user_header(&headers).or_else(|| param.raw["user"].as_str().map(str::to_owned)); + let user_id = user_hint(&headers, ¶m.raw["user"]); let request = GatewayRequest { is_online: true, @@ -2150,8 +2118,7 @@ async fn messages( ModelParamV2::with_name(gw_consts::Protocol::AnthropicMessages, body.model.clone()); param.typed = Some(typed); param.raw = Value::Object(body.extra); - let user_id = user_header(&headers) - .or_else(|| param.raw["metadata"]["user_id"].as_str().map(str::to_owned)); + let user_id = user_hint(&headers, ¶m.raw["metadata"]["user_id"]); let request = GatewayRequest { is_online: true, @@ -2499,7 +2466,7 @@ async fn completions( gw_consts::Protocol::Completions, typed, vec![ChatMsg::text("user", prompt)], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2547,7 +2514,7 @@ async fn responses( return error_response(400, "input is required"); } let stream = body["stream"].as_bool().unwrap_or(false); - let user_id = user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)); + let user_id = user_hint(&headers, &body["user"]); let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, model.clone()); param.raw = body; let request = GatewayRequest { @@ -2699,7 +2666,7 @@ async fn embeddings( gw_consts::Protocol::OpenaiChat, typed, vec![], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2739,7 +2706,7 @@ async fn images_generations( gw_consts::Protocol::OpenaiChat, typed, vec![], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2782,7 +2749,7 @@ async fn images_edits( gw_consts::Protocol::OpenaiChat, typed, vec![], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2822,7 +2789,7 @@ async fn audio_speech( gw_consts::Protocol::OpenaiChat, typed, vec![], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2878,7 +2845,7 @@ async fn audio_transcriptions( gw_consts::Protocol::OpenaiChat, typed, vec![], - user_header(&headers).or_else(|| body["user"].as_str().map(str::to_owned)), + user_hint(&headers, &body["user"]), ) .await { @@ -2922,6 +2889,13 @@ async fn batches_submit( let mut batch_items = Vec::new(); // batch-level attribution hint; a per-item body `user` overrides it let hint = user_header(&headers); + let item_user = |v: &Value| { + v["user"] + .as_str() + .or(hint.as_deref()) + .unwrap_or_default() + .to_owned() + }; if let Some(file_id) = body["input_file_id"].as_str() { let found = s.handler.state().store.file_get(file_id).await; @@ -2943,14 +2917,9 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "input file line missing a messages array"); } - let user = reqbody["user"] - .as_str() - .or(hint.as_deref()) - .unwrap_or_default() - .to_owned(); batch_items.push(BatchItem { messages: msgs, - user, + user: item_user(reqbody), }); } } else if let Some(items) = body["items"].as_array() { @@ -2959,14 +2928,9 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "each item needs a non-empty messages array"); } - let user = it["user"] - .as_str() - .or(hint.as_deref()) - .unwrap_or_default() - .to_owned(); batch_items.push(BatchItem { messages: msgs, - user, + user: item_user(it), }); } } else { @@ -3165,15 +3129,21 @@ mod tests { let mut h = HeaderMap::new(); h.insert("x-real-ip", "10.0.0.5".parse().unwrap()); h.insert("x-forwarded-for", "1.2.3.4, 10.0.0.9".parse().unwrap()); - // default (untrusted): the client's headers are ignored, the TCP peer wins - assert_eq!(source_ip(Some(peer), &h, false), "203.0.113.7"); + assert_eq!( + source_ip(Some(peer), &h, false), + "203.0.113.7", + "untrusted: forgeable headers ignored, the TCP peer wins" + ); assert_eq!( source_ip(None, &h, false), "", "no peer, no forgeable header" ); - // trusted proxy: x-real-ip, then the rightmost (proxy-appended) XFF hop - assert_eq!(source_ip(Some(peer), &h, true), "10.0.0.5"); + assert_eq!( + source_ip(Some(peer), &h, true), + "10.0.0.5", + "trusted proxy: x-real-ip wins" + ); h.remove("x-real-ip"); assert_eq!(source_ip(Some(peer), &h, true), "10.0.0.9", "rightmost hop"); } @@ -3221,8 +3191,6 @@ mod tests { #[tokio::test] async fn full_retention_without_key_never_stores_raw_even_with_dlp_off() { - // full retention + no GW_CONTENT_KEY + tenant DLP fully off: the keyless - // downgrade must still strip secrets, not persist the raw prompt let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; assert!( !gw_state::sealing_available(), @@ -3298,15 +3266,12 @@ mod tests { let cfg = app.handler.cfg(); let sec = cfg.security_for(&ak.tenant); - // moderation denies over realtime, not just REST - let mut frame = json!({"type":"input_text","text":"hello there"}); assert_eq!( - realtime_moderate(&app, sec, &ak, "", &mut frame) + realtime_moderate(&app, sec, &ak, "", "hello there") .await .as_deref(), Some("blocked by moderator") ); - // inbound realtime DLP redaction is recorded (the sink both WS paths use) let mut secret = json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); let n = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut secret); assert!(n > 0); From 0c9ceaee0f3863a6cc9ed5ff16ac414fad4c0fde Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 18:36:20 +0800 Subject: [PATCH 20/21] feat(retention): read retained content back via /admin/audit/content/{request_id} --- crates/views/src/lib.rs | 77 +++++++++++++++++++++++++++++++++++++++-- docs/governance.md | 4 ++- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index c13d273..a93cac4 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -125,6 +125,7 @@ pub fn app(state: AppState) -> Router { .route("/admin/usage/users", get(admin_usage_users)) .route("/admin/audit/events", get(admin_security_events)) .route("/admin/audit/ops", get(admin_audit_ops)) + .route("/admin/audit/content/{request_id}", get(admin_content_get)) .route( "/admin/keys/{ak}", axum::routing::patch(admin_key_patch).delete(admin_key_delete), @@ -1652,6 +1653,48 @@ async fn admin_audit_ops( } } +/// GET /admin/audit/content/{request_id} — the retained prompt/response rows +/// for one request, unsealed when the content key is present (a sealed row +/// without it returns `content: null`). Tenant-scoped like the other reads. +async fn admin_content_get( + State(s): State, + headers: HeaderMap, + Path(request_id): Path, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let rows = match s.handler.state().store.content_for(&request_id).await { + Ok(rows) => rows, + Err(e) => return gateway_error(e), + }; + let entries: Vec = rows + .into_iter() + .filter(|r| scope.covers(&r.tenant)) + .map(|r| { + let content = if r.sealed { + gw_state::content::open(&r.content) + .map(Value::String) + .unwrap_or(Value::Null) + } else { + Value::String(r.content) + }; + json!({ + "created_at_epoch_secs": r.created_at_epoch_secs, + "kind": r.kind, + "ak": r.ak, + "user_id": r.user_id, + "tenant": r.tenant, + "sealed": r.sealed, + "expires_at_epoch_secs": r.expires_at_epoch_secs, + "content": content, + }) + }) + .collect(); + Json(json!({ "request_id": request_id, "entries": entries })).into_response() +} + fn gateway_error(e: GatewayError) -> Response { let code = StatusCode::from_u16(e.http_status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); // OpenAI's error schema types `code` as string-or-null, never a number. @@ -3191,7 +3234,9 @@ mod tests { #[tokio::test] async fn full_retention_without_key_never_stores_raw_even_with_dlp_off() { - let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_CONTENT_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { std::env::set_var("GW_TEST_CONTENT_ADMIN", "s3cret") }; assert!( !gw_state::sealing_available(), "test env has no content key" @@ -3200,6 +3245,7 @@ mod tests { let state = Arc::new(GatewayState::from_config(&cfg)); let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); let store = app_state.handler.state().store.clone(); + let router = app(app_state); let req = Request::builder() .method("POST") .uri("/v1/chat/completions") @@ -3209,7 +3255,7 @@ mod tests { r#"{"model":"gpt-4o","messages":[{"role":"user","content":"here is sk-abcdefghijklmnopqrstuvwxyz012345"}]}"#, )) .unwrap(); - let resp = app(app_state).oneshot(req).await.unwrap(); + let resp = router.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); let (_, rows) = store.ledger_snapshot(1).await.unwrap(); @@ -3231,6 +3277,33 @@ mod tests { c.content ); } + + let read = router + .oneshot( + Request::builder() + .method("GET") + .uri(format!("/admin/audit/content/{}", rows[0].request_id)) + .header("authorization", "Bearer s3cret") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(read.status(), StatusCode::OK); + let j = body_json(read).await; + let entries = j["entries"].as_array().unwrap(); + assert_eq!(entries.len(), 2, "prompt and response rows read back"); + let prompt_entry = entries + .iter() + .find(|e| e["kind"] == "prompt") + .expect("prompt entry"); + assert!( + prompt_entry["content"] + .as_str() + .unwrap() + .contains("[REDACTED_SECRET]"), + "read-back returns the redacted text" + ); } #[derive(Debug)] diff --git a/docs/governance.md b/docs/governance.md index e643889..3fc1d79 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -147,4 +147,6 @@ only its audit is aggregated (a store write per token would be too hot). Retention owns its redaction, so `redacted` — and a keyless `full` that falls back to it — never persists raw secrets/PII even if the tenant forwards traffic with DLP off. `full` refuses to store raw without a key. `days` sets - expiry; an hourly purge deletes elapsed content. + expiry; an hourly purge deletes elapsed content. Read back with + `GET /admin/audit/content/{request_id}` (tenant-scoped; sealed rows are + unsealed when the key is present, else returned as `content: null`). From c810598f903e3c76a2bf6fe4e1bfea48dce12420 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 15 Jul 2026 22:59:16 +0800 Subject: [PATCH 21/21] docs: document per-user attribution, audit trails, content policy, and retention --- README.md | 1 + docs/api.md | 28 +++++++++++++++++++------- docs/configuration.md | 47 +++++++++++++++++++++++++++++++++++++++---- docs/deployment.md | 1 + docs/observability.md | 22 +++++++++++++++----- 5 files changed, 83 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index c80c512..5a64c12 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ key-based auth, quotas, rate limits, failover, and a billing ledger. - **Staged request pipeline** — a 4-layer DAG per request: model resolve / quota / cache lookup → account selection (priority, PTU-first, failover) → rate limits + engine call (retry on upstream 5xx) → usage extraction, billing, cache store - **Governance built in** — access-key auth, daily token quotas, QPS / QPM / TPM limits at key, product, and model level, request-level TTL cache, account cooldown and recovery, DLP redaction and blocklist plugins. Admission reserves then settles, so concurrent requests can't overshoot a quota - **Multi-tenant** — keys carry a tenant; tenants get a pooled QPS bucket, a model entitlement allowlist, per-(key, model) quota defaults with an optional fallback-model degrade, key lifecycle (expiry/ban), and tenant-scoped admin tokens. Billing records charged cost and (optionally) vendor cost per row, so margin is queryable per tenant × model +- **Per-user billing & enterprise audit** — every ledger row attributes to an effective end user (the key's `owner`, else the request's `x-gw-user` / `user` hint) with a `request_id`, so cost rolls up per user (`/admin/usage/users`) and a soft per-user daily budget applies on every surface. Per-tenant content policy adds blocklist action tiers (block / flag / shadow), regex recognizers, secret masking, and an external-moderation seam; every hit is recorded without prompt text. An admin-operation trail (key CRUD / config / reload, with source IP) and optional at-rest content retention complete the audit surfaces (`/admin/audit/*`) - **Fleet-ready** — run N instances behind a load balancer: Postgres shares config (versioned + a change feed), the access-key table, the ledger/files/batches store, and a distributed batch queue any instance drains; Redis shares rate/quota/TPM counters, account health, and optionally the response cache. Single-node stays zero-dependency - **Providers behind traits** — engines talk to upstreams through a `Transport` seam; accounts with a real endpoint go over HTTP (reqwest + rustls), accounts without one are served by a deterministic in-process mock; AWS SigV4 signing included - **Observability built in** — Prometheus `/metrics` (per-route request/status counters, per-pipeline-stage latency, token counters), structured access logs diff --git a/docs/api.md b/docs/api.md index e9fe2a7..2bce4a7 100644 --- a/docs/api.md +++ b/docs/api.md @@ -20,6 +20,11 @@ error shape, so its SDKs can dispatch on it: {"type": "error", "error": {"type": "invalid_request_error", "message": "..."}} ``` +For per-user attribution on a shared key, send `x-gw-user: ` (it also reads +OpenAI's body `user` field and Anthropic's `metadata.user_id`). A key's own +`owner` overrides the hint, so a key issued to one user always bills to that +user. See [Governance](governance.md#per-user-attribution-and-billing). + ## OpenAI-compatible | Method | Path | Notes | @@ -89,10 +94,13 @@ A realtime model bound to an account with a real `endpoint` bridges the session to that vendor's realtime WebSocket: a transparent relay, with the gateway enforcing the same governance chain as the REST path per generation — tenant and AK QPS, product/model QPM, per-(key, model) and daily-token quota, TPM — plus -billing (shared pricing) from the vendor's usage. Content security also applies: -the blocklist gates inbound frames and DLP redacts text fields in both -directions (per frame — a PII span straddling two deltas is beyond a relay -that cannot buffer). Each generation re-checks the +billing (shared pricing) from the vendor's usage. The full content policy also +applies, so the WebSocket is not a bypass: the blocklist, regex recognizers, and +(when enabled) the external moderator gate inbound frames, and DLP — emails, +phone numbers, and credential masking — redacts text fields in both directions +(per frame — a PII span straddling two deltas is beyond a relay that cannot +buffer). Every hit is audited without prompt text; per-user attribution comes +from the `x-gw-user` hint captured at connect. Each generation re-checks the key, so a key banned, expired, or revoked (or a model de-entitled) mid-session stops generating. An endpoint-less account serves a local mock session (OpenAI Realtime event shape) for offline development. @@ -122,14 +130,20 @@ surface on a private network regardless. | POST | `/admin/reload` | re-read config from source and swap it in atomically (global token only) | | PUT | `/admin/config` | validate + publish a new config document to the fleet config store; every instance reloads via the change feed (global token; needs `storage.postgres_url`) | | GET | `/admin/keys` | list keys (a tenant token sees only its own tenant's) | -| POST | `/admin/keys` | create/replace a key: `{ak, product, tenant?, qps, daily_token_quota, tokens_per_minute?, expires_at_epoch_secs?, banned?, model_quotas?}` | +| POST | `/admin/keys` | create/replace a key: `{ak, product, tenant?, owner?, qps, daily_token_quota, tokens_per_minute?, expires_at_epoch_secs?, banned?, model_quotas?}` (`owner` binds the key to one end user — authoritative for attribution) | | PATCH | `/admin/keys/{ak}` | update any of `qps` / `daily_token_quota` / `tokens_per_minute` / `expires_at_epoch_secs` (null clears) / `banned` | | DELETE | `/admin/keys/{ak}` | revoke a key | | GET | `/admin/usage` | ledger rollup by tenant × model (requests, tokens, charged `cost_micros`, `vendor_cost_micros` for margin); `?tenant=` filter for the global token | +| GET | `/admin/usage/users` | per-user cost rollup (user × model) over a billing period: `?since=&until=` (unix secs), `?user=` filter, `?format=csv` export; tenant-scoped | +| GET | `/admin/audit/events` | content-safety hits (blocklist / regex / DLP / moderation) recorded without prompt text; `?limit=`; tenant-scoped | +| GET | `/admin/audit/ops` | admin-operation trail (key CRUD, config publish, reload) with actor, target, and source IP; `?limit=`; global token only | +| GET | `/admin/audit/content/{request_id}` | retained prompt/response for one request, unsealed when `GW_CONTENT_KEY` is set (sealed rows without it return `content: null`); tenant-scoped | Two token tiers: the global token (`admin.token_env`) manages everything; a -tenant's `admin_token_env` token manages only that tenant's keys and usage -(cross-tenant keys answer 404, reload/config-publish answer 403). +tenant's `admin_token_env` token manages only that tenant's keys, usage, and +content-safety events, scoped to its own tenant (cross-tenant keys answer 404; +reload, config-publish, and the cross-tenant `/admin/audit/ops` trail answer +403). A reload rebuilds the AK table (config keys), models, providers, tenants, and accounts while preserving the runtime seams — governance counters, the durable diff --git a/docs/configuration.md b/docs/configuration.md index 8ec8d43..4878b38 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -6,7 +6,9 @@ One YAML file configures the gateway. Resolution order: 2. otherwise the embedded default (the repo's `conf/gateway.yaml`) `GW_HOST` / `GW_PORT` override `listen.host` / `listen.port` at runtime -(the container image sets `GW_HOST=0.0.0.0`). +(the container image sets `GW_HOST=0.0.0.0`). `GW_CONTENT_KEY` (64 hex chars = +32 bytes) is the deployment key that seals retained content at rest; without it, +`full` retention refuses to store raw text and falls back to redacted. ## Sections @@ -47,6 +49,9 @@ access_keys: - ak: ak-demo-123 # bearer / x-api-key value clients send product: demo # product group (for product-level QPM) tenant: acme # optional; absent = the unrestricted `default` tenant + owner: alice # optional; binds the key to one end user (authoritative + # for per-user attribution; a shared key omits it and + # falls back to the request's `x-gw-user` / `user`) qps: 100 # per-key request rate daily_token_quota: 1000000 tokens_per_minute: 600 # optional TPM window limit @@ -69,12 +74,26 @@ tenants: admin_token_env: ACME_ADMIN_TOKEN # optional tenant-scoped /admin token model_prices: # optional per-model charged-price override for this tenant gpt-4o: {input_price_per_1k_micros: 5000, output_price_per_1k_micros: 20000} + user_daily_token_quota: 100000 # optional soft per-end-user daily cap + security: # optional; overrides the global `security:` WHOLE for this tenant + blocklist: ["forbidden"] + blocklist_action: flag # block | flag | shadow + detect_secrets: true + regex_rules: + - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} + retention: # optional prompt/response retention; absent = retain nothing + content: redacted # none | redacted | full (full needs GW_CONTENT_KEY) + days: 30 # purge after N days; 0 = keep until manually purged ``` Keys without a `tenant` join the implicit `default` tenant (no pooled limits, entitled to every model), so a flat config keeps working unchanged. The model catalog (`GET /v1/models`) filters to the caller's entitlement. +`user_daily_token_quota`, `security`, and `retention` are enterprise controls +detailed in [Governance](governance.md); `security` replaces the global policy +outright when present (it is not merged field-by-field). + ### `models` — public model names and dispatch ```yaml @@ -134,9 +153,15 @@ the charged `cost_micros` and margin is queryable per tenant/model via ### `security`, `stability`, `products` ```yaml -security: - dlp_redact: true # redact emails/phone numbers before egress - blocklist: ["badword"] # reject requests containing listed terms +security: # global default; a tenant may override it whole + dlp_redact: true # redact emails/phone numbers, both directions + detect_secrets: true # also mask API keys / credentials in inbound text + blocklist: ["badword"] # reject/flag requests containing listed terms + blocklist_action: block # block (deny) | flag (record) | shadow (trial a rule) + regex_rules: # named recognizers, each with its own action + - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} + moderate: false # route inbound text through the wired external moderator + moderation_fail_open: false # on a moderator error: admit (true) or deny (false) stability: failure_threshold: 3 # consecutive failures before an account cools down @@ -147,6 +172,20 @@ products: qpm: 120 # product-level request rate ``` +Every rule that fires (block / flag / DLP / moderation) is recorded without the +prompt text to the security-event stream (`GET /admin/audit/events`). The same +policy runs on the realtime WebSocket, so it is not a bypass. `moderate` needs a +moderator wired into the handler — the default one allows everything. See +[Governance](governance.md#enterprise-content-policy). + +### Top-level flags + +```yaml +trust_proxy_headers: false # audit source IP: false = the real TCP peer (unforgeable); + # true = trust x-real-ip / rightmost x-forwarded-for hop + # (only behind a proxy that sets them) +``` + ## Observability `GET /metrics` serves the Prometheus registry: `gateway_requests_total` diff --git a/docs/deployment.md b/docs/deployment.md index e95e32c..2269ae6 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -24,6 +24,7 @@ GW_CONFIG=/etc/gateway.yaml ./target/release/gw | `GW_HOST` | override `listen.host` (containers set `0.0.0.0`) | | `GW_PORT` | override `listen.port` | | `GW_TRANSPORT` | `mock` (zero egress) / `http` (no mock) / unset (auto-route) | +| `GW_CONTENT_KEY` | 64 hex chars (32 bytes); seals retained content at rest. Without it, `full` retention stores redacted text instead of raw | | `RUST_LOG` | log level, e.g. `info`, `gw_views=debug` | | provider key vars | named by each account's `api_key_env` | diff --git a/docs/observability.md b/docs/observability.md index 824c81f..c11ce30 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -22,14 +22,26 @@ status codes, protocol/stage names) — no per-key or per-model cardinality. ## Access log One structured line per request goes to stdout (via `tracing`; control level -with `RUST_LOG`), carrying `ak`, `product`, `model`, `protocol`, `account`, -`status`, `prompt_tokens`, `completion_tokens`, `total_tokens`, and -`latency_ms`. +with `RUST_LOG`), carrying `request_id`, `ak`, `product`, `user_id`, `model`, +`protocol`, `account`, `status`, `prompt_tokens`, `completion_tokens`, +`total_tokens`, and `latency_ms`. `request_id` joins the access log to the +ledger row and the audit events for the same request. ## Billing ledger `GET /internal/ledger?limit=N` returns the most recent `N` billing records (newest first); `count` is always the true total, independent of the page size. Records persist when a SQLite store is configured and can be capped with -`storage.ledger_max_rows`. Each record has the access key, product, model, -protocol, account, token counts, cost, and the PTU-spillover flag. +`storage.ledger_max_rows`. Each record carries `request_id`, the access key, +product, `user_id` (effective end user), model, protocol, account, token counts, +cost, `created_at_epoch_secs`, the PTU-spillover flag, and an `estimated` flag +(set when counts came from an aborted stream rather than a vendor usage payload). + +## Audit trails + +Three operator-facing audit surfaces, all under the gated `/admin` prefix and +covered in [Governance](governance.md#audit-trails): `GET /admin/audit/events` +(content-safety hits, no prompt text), `GET /admin/audit/ops` (admin-plane +mutations with source IP), and `GET /admin/usage/users` (per-user cost). Content +retention, when a tenant enables it, is read back via +`GET /admin/audit/content/{request_id}`.