diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index 0b3148b79..8413fd62e 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -3679,6 +3679,7 @@ fn add_variant_titles(doc: &mut Value) { "Azure AI Content Safety Prompt Shield", "Azure AI Content Safety Text Moderation", "Aliyun Text Moderation", + "PII Detection & Redaction", ], ), ( @@ -3803,6 +3804,10 @@ fn add_missing_property_descriptions(doc: &mut Value) { "/components/schemas/Guardrail/oneOf/4/properties/kind", "Guardrail provider type. Use `aliyun_text_moderation` for Aliyun text moderation.", ), + ( + "/components/schemas/Guardrail/oneOf/5/properties/kind", + "Guardrail provider type. Use `pii` for in-process sensitive-data detection and redaction.", + ), ( "/components/schemas/KeywordPattern/oneOf/0/properties/kind", "Pattern type. Use `literal` to match the value as plain text.", diff --git a/crates/aisix-core/src/models/guardrail.rs b/crates/aisix-core/src/models/guardrail.rs index 3bd8aeb59..583d561d4 100644 --- a/crates/aisix-core/src/models/guardrail.rs +++ b/crates/aisix-core/src/models/guardrail.rs @@ -325,6 +325,87 @@ fn default_aliyun_risk_level_threshold() -> String { "high".to_owned() } +/// One built-in detector selection for `kind: "pii"`. The `type` names a +/// detector compiled into the DP (`aisix-guardrails::pii::BUILTIN_DETECTORS`); +/// `action` optionally overrides the guardrail-level `default_action` for +/// this detector only. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct PiiDetectorConfig { + /// Built-in detector id: `email`, `china_mobile`, `china_id_card`, + /// `bank_card`, `us_ssn`, `ip_address`, `api_key`, `jwt`, `private_key`. + #[serde(rename = "type")] + #[schemars(length(min = 1))] + pub detector_type: String, + /// Per-detector action override: `mask` or `block`. Falls back to the + /// guardrail's `default_action` when omitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub action: Option, +} + +/// One operator-supplied regex detector for `kind: "pii"`. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct PiiCustomPattern { + /// Detector name surfaced in the mask token (`[_REDACTED]`), + /// telemetry counts, and block reasons. Never the matched value. + #[schemars(length(min = 1, max = 64))] + pub name: String, + /// Regular expression the DP compiles at chain build. Invalid patterns + /// are logged and the row is skipped (same contract as keyword rules). + #[schemars(length(min = 1))] + pub regex: String, + /// Per-pattern action override: `mask` or `block`. Falls back to the + /// guardrail's `default_action` when omitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub action: Option, +} + +/// Config block for `kind: "pii"`. In-process sensitive-data detection and +/// redaction (#932): built-in detectors + custom regex patterns, each with a +/// `mask` (redact-and-continue) or `block` action, applied to request and/or +/// response text per the row's `hook_point`. +/// +/// `mask` rewrites each matched span to `[_REDACTED]` and lets the +/// request/response continue; `block` rejects with the standard 422 +/// content-filter envelope. Matched values never appear in logs, telemetry, +/// or error envelopes — only detector names and match counts do. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct PiiConfig { + /// Built-in detectors to enable. Unknown detector ids are rejected at + /// build time (row skipped + warn), so a typo can't silently disarm + /// the policy. + #[serde(default)] + pub detectors: Vec, + /// Operator-supplied regex detectors, evaluated after the built-ins. + #[serde(default)] + pub custom_patterns: Vec, + /// Action for detectors that don't set their own: `mask` (default) or + /// `block`. + #[serde(default = "default_pii_action")] + pub default_action: String, + + // --- streaming-output controls (consumed by aisix-proxy) --- + // Masking a streamed response requires the whole response to be held + // back (the mask spans chunk boundaries), so kind=pii always uses the + // buffer_full policy on the output hook. These knobs mirror the ACS/ + // Aliyun buffer_full parameters. + /// Max bytes buffered for a streamed response before `on_buffer_exceeded` applies. + #[serde(default = "default_acs_max_buffer_bytes")] + #[schemars(range(min = 1))] + pub max_buffer_bytes: u64, + /// Buffer-overflow policy. Use `fail_open` to release output unscanned + /// (and unmasked) when the buffer cap is hit; the default `fail_closed` + /// blocks the response instead. + #[serde(default = "default_acs_on_buffer_exceeded")] + pub on_buffer_exceeded: String, +} + +fn default_pii_action() -> String { + "mask".to_owned() +} + fn default_aliyun_window_size() -> u32 { 2_000 } @@ -380,6 +461,10 @@ pub enum GuardrailKind { /// `TextModerationPlus` action on `green-cip..aliyuncs.com`, on /// input and/or output, including streaming output. AliyunTextModeration(AliyunTextModerationConfig), + /// In-process sensitive-data detection + redaction (#932): built-in + /// detectors + custom regex, per-detector `mask`/`block` actions, on + /// input and/or output, including streaming output. Always available. + Pii(PiiConfig), } impl GuardrailKind { @@ -394,6 +479,7 @@ impl GuardrailKind { "azure_content_safety_text_moderation" } GuardrailKind::AliyunTextModeration(_) => "aliyun_text_moderation", + GuardrailKind::Pii(_) => "pii", } } } @@ -903,6 +989,61 @@ mod tests { } } + #[test] + fn pii_kind_parses_with_defaults() { + // cp-api omits unset fields (omitempty); the DP must apply the + // documented defaults so a minimal row still works. + let v = json!({ + "name": "mask-pii", + "kind": "pii", + "detectors": [ + { "type": "email" }, + { "type": "china_id_card", "action": "block" } + ] + }); + let g: Guardrail = serde_json::from_value(v).unwrap(); + match g.config { + GuardrailKind::Pii(ref c) => { + assert_eq!(c.detectors.len(), 2); + assert_eq!(c.detectors[0].detector_type, "email"); + assert_eq!(c.detectors[0].action, None); + assert_eq!(c.detectors[1].action.as_deref(), Some("block")); + assert!(c.custom_patterns.is_empty()); + assert_eq!(c.default_action, "mask"); + assert_eq!(c.max_buffer_bytes, 262_144); + assert_eq!(c.on_buffer_exceeded, "fail_closed"); + } + _ => panic!("expected Pii variant"), + } + } + + #[test] + fn pii_kind_round_trips_custom_patterns_and_overrides() { + let v = json!({ + "name": "mask-pii", + "kind": "pii", + "default_action": "block", + "custom_patterns": [ + { "name": "employee_id", "regex": "\\bEMP-\\d{6}\\b", "action": "mask" } + ], + "max_buffer_bytes": 1024, + "on_buffer_exceeded": "fail_open" + }); + let g: Guardrail = serde_json::from_value(v).unwrap(); + match g.config { + GuardrailKind::Pii(ref c) => { + assert!(c.detectors.is_empty()); + assert_eq!(c.custom_patterns.len(), 1); + assert_eq!(c.custom_patterns[0].name, "employee_id"); + assert_eq!(c.custom_patterns[0].action.as_deref(), Some("mask")); + assert_eq!(c.default_action, "block"); + assert_eq!(c.max_buffer_bytes, 1024); + assert_eq!(c.on_buffer_exceeded, "fail_open"); + } + _ => panic!("expected Pii variant"), + } + } + #[test] fn resource_trait_uses_name_and_guardrails_kind() { let mut g: Guardrail = serde_json::from_value(json!({ diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index acc7a47f5..00f8e82a7 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -39,7 +39,8 @@ pub use guardrail::{ AliyunTextModerationConfig, AppliedGuardrail, AzureContentSafetyConfig, AzureContentSafetyTextModerationConfig, BedrockAWSCredentials, BedrockConfig, BedrockLatencyMode, Guardrail, GuardrailAttachment, GuardrailHookPoint, GuardrailKind, - GuardrailScopeType, KeywordConfig, KeywordPattern, + GuardrailScopeType, KeywordConfig, KeywordPattern, PiiConfig, PiiCustomPattern, + PiiDetectorConfig, }; pub use mcp_server::{McpAuthType, McpServer, McpTransport}; pub use model::{ diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index 61148cb7c..cb5943d89 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -24,7 +24,8 @@ use async_trait::async_trait; use crate::index::{GuardrailIndex, RequestContext, ScopeKind}; use crate::keyword::{KeywordBlocklist, KeywordRule}; -use crate::{Guardrail, GuardrailChain, GuardrailVerdict, StreamOutputPolicy}; +use crate::pii::{builtin_rule, PiiAction, PiiGuardrail, PiiRule}; +use crate::{Guardrail, GuardrailChain, GuardrailVerdict, Redaction, StreamOutputPolicy}; /// A snapshot table's guardrail entries in deterministic chain order: /// `created_at` ascending (RFC3339 strings in a fixed offset compare @@ -211,6 +212,58 @@ fn build_one_inner( }; Ok(Some(Arc::new(blocklist))) } + GuardrailKind::Pii(cfg) => { + if cfg.detectors.is_empty() && cfg.custom_patterns.is_empty() { + return Ok(None); + } + let default_action = + PiiAction::parse(&cfg.default_action).ok_or_else(|| BuildError::InvalidValue { + field: "default_action", + value: cfg.default_action.clone(), + })?; + let mut rules: Vec = + Vec::with_capacity(cfg.detectors.len() + cfg.custom_patterns.len()); + for d in &cfg.detectors { + let action = match d.action.as_deref() { + None => default_action, + Some(s) => PiiAction::parse(s).ok_or_else(|| BuildError::InvalidValue { + field: "detectors[].action", + value: s.to_owned(), + })?, + }; + let rule = builtin_rule(&d.detector_type, action).ok_or_else(|| { + BuildError::InvalidValue { + field: "detectors[].type", + value: d.detector_type.clone(), + } + })?; + rules.push(rule); + } + for p in &cfg.custom_patterns { + let action = match p.action.as_deref() { + None => default_action, + Some(s) => PiiAction::parse(s).ok_or_else(|| BuildError::InvalidValue { + field: "custom_patterns[].action", + value: s.to_owned(), + })?, + }; + let rule = PiiRule::new(p.name.clone(), &p.regex, action, None).map_err(|e| { + BuildError::InvalidRegex { + pattern: p.regex.clone(), + source: e, + } + })?; + rules.push(rule); + } + let on_exceeded_fail_open = cfg.on_buffer_exceeded == "fail_open"; + let g = PiiGuardrail::new( + rules, + row.hook_point, + usize::try_from(cfg.max_buffer_bytes).unwrap_or(usize::MAX), + on_exceeded_fail_open, + ); + Ok(Some(Arc::new(g))) + } #[cfg(feature = "bedrock")] GuardrailKind::Bedrock(cfg) => { // Phase 2: build the AWS-SDK-backed dispatcher. cp-api @@ -299,6 +352,11 @@ enum BuildError { pattern: String, source: regex::Error, }, + /// An enum-ish config field carrying a value the DP doesn't know + /// (unknown built-in detector id, unrecognised action). The row is + /// skipped + warned rather than silently running a weaker policy. + #[error("invalid {field} value {value:?}")] + InvalidValue { field: &'static str, value: String }, /// A guardrail kind whose runtime dispatch was compiled out via /// feature flags (e.g. a pruned build that excluded `--features bedrock` /// or `--features azure-content-safety`). The chain treats the row as @@ -331,6 +389,19 @@ struct MonitorGuardrail { } impl MonitorGuardrail { + /// Log the mask counts a monitor-mode redacting guardrail WOULD have + /// applied. Counts carry detector names only, never matched values. + fn observe_redaction(&self, hook: &'static str, r: Option) { + if let Some(r) = r { + tracing::info!( + guardrail_name = %self.row_name, + hook, + counts = ?r.counts, + "guardrail in monitor mode observed maskable spans; not redacting (enforcement_mode=monitor)", + ); + } + } + fn observe(&self, hook: &'static str, verdict: GuardrailVerdict) -> GuardrailVerdict { match verdict { GuardrailVerdict::Block { reason, .. } => { @@ -372,6 +443,28 @@ impl Guardrail for MonitorGuardrail { fn runs_on_output(&self) -> bool { self.inner.runs_on_output() } + + // Monitor mode observes without modifying: redaction is suppressed the + // same way Block verdicts are downgraded. The would-be mask counts are + // logged so operators can stage a redaction rule and audit its impact + // before enforcing it. + fn redacts_input(&self) -> bool { + false + } + + fn redacts_output(&self) -> bool { + false + } + + fn redact_input_text(&self, text: &str) -> Option { + self.observe_redaction("input", self.inner.redact_input_text(text)); + None + } + + fn redact_output_text(&self, text: &str) -> Option { + self.observe_redaction("output", self.inner.redact_output_text(text)); + None + } } /// `mandatory: true` decorator. A remote guardrail that can't reach its @@ -440,6 +533,24 @@ impl Guardrail for MandatoryGuardrail { fn runs_on_output(&self) -> bool { self.inner.runs_on_output() } + + // Mandatory only rewrites the failure (Bypass) verdict; redaction + // passes straight through to the inner guardrail. + fn redacts_input(&self) -> bool { + self.inner.redacts_input() + } + + fn redacts_output(&self) -> bool { + self.inner.redacts_output() + } + + fn redact_input_text(&self, text: &str) -> Option { + self.inner.redact_input_text(text) + } + + fn redact_output_text(&self, text: &str) -> Option { + self.inner.redact_output_text(text) + } } /// Adapter that wraps a snapshot handle and rebuilds the runtime @@ -538,6 +649,22 @@ impl Guardrail for LiveGuardrailChain { fn runs_on_output(&self) -> bool { self.current().runs_on_output() } + + fn redacts_input(&self) -> bool { + self.current().redacts_input() + } + + fn redacts_output(&self) -> bool { + self.current().redacts_output() + } + + fn redact_input_text(&self, text: &str) -> Option { + self.current().redact_input_text(text) + } + + fn redact_output_text(&self, text: &str) -> Option { + self.current().redact_output_text(text) + } } // --------------------------------------------------------------------------- diff --git a/crates/aisix-guardrails/src/chain.rs b/crates/aisix-guardrails/src/chain.rs index 875b61d7f..f7414171e 100644 --- a/crates/aisix-guardrails/src/chain.rs +++ b/crates/aisix-guardrails/src/chain.rs @@ -13,7 +13,7 @@ use aisix_core::AppliedGuardrail; use aisix_gateway::{ChatFormat, ChatResponse}; use async_trait::async_trait; -use crate::{Guardrail, GuardrailVerdict, StreamOutputPolicy}; +use crate::{Guardrail, GuardrailVerdict, Redaction, StreamOutputPolicy}; /// One chain member: the runtime guardrail plus the operator-facing name /// of the row it was built from. The name is what `Block` verdicts are @@ -204,6 +204,66 @@ impl Guardrail for GuardrailChain { None => GuardrailVerdict::Allow, } } + + fn redacts_input(&self) -> bool { + self.members.iter().any(|m| m.guardrail.redacts_input()) + } + + fn redacts_output(&self) -> bool { + self.members.iter().any(|m| m.guardrail.redacts_output()) + } + + /// Members apply in chain order, each rewriting the previous member's + /// output, so stacked redacting guardrails compose. Counts merge across + /// members. + fn redact_input_text(&self, text: &str) -> Option { + fold_redactions( + text, + self.members + .iter() + .filter_map(|m| m.guardrail.redacts_input().then_some(&m.guardrail)), + true, + ) + } + + fn redact_output_text(&self, text: &str) -> Option { + fold_redactions( + text, + self.members + .iter() + .filter_map(|m| m.guardrail.redacts_output().then_some(&m.guardrail)), + false, + ) + } +} + +/// Fold `text` through each member's redactor, merging counts. `None` +/// when no member changed anything. +fn fold_redactions<'a>( + text: &str, + members: impl Iterator>, + input: bool, +) -> Option { + let mut current: Option = None; + for g in members { + let src = current.as_ref().map_or(text, |r| r.text.as_str()); + let next = if input { + g.redact_input_text(src) + } else { + g.redact_output_text(src) + }; + if let Some(r) = next { + current = Some(match current.take() { + None => r, + Some(mut acc) => { + acc.text = r.text; + Redaction::merge_counts(&mut acc.counts, &r.counts); + acc + } + }); + } + } + current } #[cfg(test)] diff --git a/crates/aisix-guardrails/src/lib.rs b/crates/aisix-guardrails/src/lib.rs index 7a127a4f4..44ed363c6 100644 --- a/crates/aisix-guardrails/src/lib.rs +++ b/crates/aisix-guardrails/src/lib.rs @@ -26,6 +26,7 @@ mod build; mod chain; mod index; mod keyword; +mod pii; #[cfg(feature = "azure-content-safety")] mod prompt_shield; #[cfg(feature = "azure-content-safety")] @@ -72,6 +73,7 @@ pub(crate) fn message_scan_text(m: &ChatMessage) -> String { pub fn supported_kinds() -> &'static [&'static str] { &[ "keyword", + "pii", #[cfg(feature = "azure-content-safety")] "azure_content_safety", #[cfg(feature = "azure-content-safety")] @@ -93,6 +95,7 @@ pub use build::{ pub use chain::GuardrailChain; pub use index::{GuardrailIndex, RequestContext}; pub use keyword::{KeywordBlocklist, KeywordRule}; +pub use pii::{builtin_rule, PiiAction, PiiGuardrail, PiiRule, BUILTIN_DETECTORS}; #[cfg(feature = "azure-content-safety")] pub use prompt_shield::PromptShieldGuardrail; #[cfg(feature = "azure-content-safety")] @@ -251,6 +254,31 @@ impl StreamOutputPolicy { /// Matches the Azure text-moderation buffer-mode default. pub const DEFAULT_STREAM_OUTPUT_BUFFER_BYTES: usize = 262_144; +/// One text-channel redaction outcome from +/// [`Guardrail::redact_input_text`] / [`Guardrail::redact_output_text`]: +/// the rewritten text plus per-detector match counts. Counts carry detector +/// NAMES only — the matched values are gone by construction, so this type +/// is safe to log and to attach to telemetry (#932 no-leak criterion). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Redaction { + pub text: String, + /// detector name → number of masked spans. + pub counts: std::collections::BTreeMap, +} + +impl Redaction { + /// Fold `other`'s counts into `self` (used by chains and by callers + /// merging per-field redactions into one per-request summary). + pub fn merge_counts( + into: &mut std::collections::BTreeMap, + other: &std::collections::BTreeMap, + ) { + for (k, v) in other { + *into.entry(k.clone()).or_insert(0) += v; + } + } +} + /// Pluggable content-policy hook. Production wires `Arc` /// in `ProxyState`; tests construct in-memory chains directly. #[async_trait] @@ -300,6 +328,38 @@ pub trait Guardrail: Send + Sync + 'static { fn runs_on_output(&self) -> bool { true } + + // --- redaction (#932) ------------------------------------------------- + // + // Redaction is a separate, synchronous, text→text capability rather + // than a mutation inside `check_input`/`check_output`: the check hooks + // scan ONE concatenated blob per request, while redaction must be + // applied per text FIELD (each message, each tool-call argument, each + // streamed channel) so the caller controls which wire fields are + // rewritten and structure is preserved. Callers run the check first + // (Block wins over Mask), then apply the redactor to each field. + + /// `true` when this guardrail can rewrite REQUEST text. Cheap probe so + /// call sites skip walking the body when nothing would change. + fn redacts_input(&self) -> bool { + false + } + + /// `true` when this guardrail can rewrite RESPONSE text. + fn redacts_output(&self) -> bool { + false + } + + /// Rewrite one request-side text field, masking sensitive spans. + /// `None` = no capability or no matches (caller keeps the original). + fn redact_input_text(&self, _text: &str) -> Option { + None + } + + /// Rewrite one response-side text field, masking sensitive spans. + fn redact_output_text(&self, _text: &str) -> Option { + None + } } #[cfg(test)] @@ -386,6 +446,7 @@ mod tests { supported_kinds(), &[ "keyword", + "pii", "azure_content_safety", "azure_content_safety_text_moderation", "aliyun_text_moderation", @@ -401,6 +462,10 @@ mod tests { "kind": "keyword", "patterns": [{"kind": "literal", "value": "x"}], }), + "pii" => serde_json::json!({ + "kind": "pii", + "detectors": [{"type": "email"}], + }), "azure_content_safety" => serde_json::json!({ "kind": "azure_content_safety", "endpoint": "https://x.cognitiveservices.azure.com", diff --git a/crates/aisix-guardrails/src/pii.rs b/crates/aisix-guardrails/src/pii.rs new file mode 100644 index 000000000..1240506e5 --- /dev/null +++ b/crates/aisix-guardrails/src/pii.rs @@ -0,0 +1,602 @@ +//! In-process sensitive-data detection + redaction guardrail (#932). +//! +//! Two detector sources mix in one rule list: +//! - built-in detectors (`BUILTIN_DETECTORS`): curated regexes for common +//! PII/secret shapes, several backed by a checksum validator (Luhn for +//! bank cards, ISO 7064 MOD 11-2 for Chinese national IDs) so a random +//! digit run doesn't false-positive; +//! - custom patterns: operator-supplied regexes. +//! +//! Each rule carries an action: +//! - `Block`: the request/response is rejected (422 content-filter) — +//! same enforcement path as the keyword blocklist. +//! - `Mask`: each matched span is rewritten to `[_REDACTED]` +//! and processing continues. Callers apply the rewrite through +//! [`crate::Guardrail::redact_input_text`] / +//! [`crate::Guardrail::redact_output_text`]. +//! +//! The detector NAME is the only thing that ever leaves this module — +//! block reasons, telemetry counts, and mask tokens all carry the name, +//! never the matched value (#153 anti-leak rule, and the #932 acceptance +//! criterion that redacted values must not appear in gateway logs). + +use std::borrow::Cow; +use std::collections::BTreeMap; + +use aisix_gateway::{ChatFormat, ChatResponse}; +use async_trait::async_trait; +use regex::Regex; + +use crate::{Guardrail, GuardrailVerdict, Redaction, StreamOutputPolicy}; + +/// What to do when a detector matches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PiiAction { + /// Rewrite the matched span to `[_REDACTED]` and continue. + Mask, + /// Reject the request/response (422 content-filter). + Block, +} + +impl PiiAction { + /// Parse the wire string. `None` for anything unrecognised — the + /// build layer decides whether to fall back or reject. + pub fn parse(s: &str) -> Option { + match s { + "mask" => Some(PiiAction::Mask), + "block" => Some(PiiAction::Block), + _ => None, + } + } +} + +/// Secondary validator applied to each regex match before it counts as a +/// detection. Cuts false positives on digit-run detectors where the regex +/// alone is too permissive. +type MatchValidator = fn(&str) -> bool; + +/// One compiled detector rule. +pub struct PiiRule { + /// Detector id (built-in) or operator-assigned custom-pattern name. + /// Surfaced in mask tokens / reasons / counts. + name: String, + regex: Regex, + action: PiiAction, + validate: Option, + /// Pre-computed `[_REDACTED]` token. + mask_token: String, +} + +impl PiiRule { + pub fn new( + name: impl Into, + pattern: &str, + action: PiiAction, + validate: Option, + ) -> Result { + let name = name.into(); + let mask_token = mask_token(&name); + Ok(Self { + regex: Regex::new(pattern)?, + name, + action, + validate, + mask_token, + }) + } + + /// `true` when `candidate` is a real detection (regex already matched; + /// this applies the optional checksum validator). + fn accepts(&self, candidate: &str) -> bool { + self.validate.is_none_or(|v| v(candidate)) + } + + /// First validated match in `text`, or `None`. + fn detects(&self, text: &str) -> bool { + self.regex.find_iter(text).any(|m| self.accepts(m.as_str())) + } + + /// Rewrite every validated match in `text` to the rule's mask token. + /// Returns the match count (0 = `text` returned unchanged). + fn mask_all<'t>(&self, text: &'t str) -> (Cow<'t, str>, u32) { + let mut count = 0u32; + let out = self.regex.replace_all(text, |caps: ®ex::Captures<'_>| { + let m = &caps[0]; + if self.accepts(m) { + count += 1; + self.mask_token.clone() + } else { + m.to_string() + } + }); + (out, count) + } +} + +/// `[_REDACTED]`, uppercased, non-alphanumerics folded to `_` so an +/// operator-assigned custom name can't inject brackets/spaces into the token. +fn mask_token(name: &str) -> String { + let cleaned: String = name + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() { + c.to_ascii_uppercase() + } else { + '_' + } + }) + .collect(); + format!("[{cleaned}_REDACTED]") +} + +/// One built-in detector: `(id, pattern, validator)`. +/// +/// Patterns lean conservative — `\b` boundaries and checksum validators — +/// because a mask false-positive corrupts user content (worse than a +/// keyword-blocklist false positive, which merely rejects). IDs are wire +/// contract: cp-api's validator and the dashboard's detector list carry +/// the same set. +pub const BUILTIN_DETECTORS: &[(&str, &str, Option)] = &[ + ( + "email", + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", + None, + ), + // Chinese mainland mobile: 1[3-9] + 9 digits, standalone digit run. + ("china_mobile", r"\b1[3-9]\d{9}\b", None), + // 18-char Chinese national ID (17 digits + check char), verified with + // the ISO 7064 MOD 11-2 checksum. + ( + "china_id_card", + r"\b\d{17}[0-9Xx]\b", + Some(china_id_checksum), + ), + // 13–19 digit PAN (bank/credit card), Luhn-verified. + ("bank_card", r"\b\d{13,19}\b", Some(luhn_checksum)), + ("us_ssn", r"\b\d{3}-\d{2}-\d{4}\b", None), + ( + "ip_address", + r"\b(?:(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\.){3}(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\b", + None, + ), + // Well-known credential signatures: OpenAI, AWS access key id, GitHub + // (classic + fine-grained), Slack, Google API key. + ( + "api_key", + r"\b(?:sk-[A-Za-z0-9_-]{20,}|AKIA[0-9A-Z]{16}|gh[pousr]_[A-Za-z0-9]{36}|github_pat_[A-Za-z0-9_]{22,}|xox[baprs]-[A-Za-z0-9-]{10,}|AIza[0-9A-Za-z_-]{35})\b", + None, + ), + ( + "jwt", + r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\b", + None, + ), + ( + "private_key", + r"(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----", + None, + ), +]; + +/// Look up a built-in detector by id and compile it with `action`. +pub fn builtin_rule(id: &str, action: PiiAction) -> Option { + BUILTIN_DETECTORS + .iter() + .find(|(name, _, _)| *name == id) + .map(|(name, pattern, validate)| { + PiiRule::new(*name, pattern, action, *validate) + .expect("built-in PII detector pattern must compile") + }) +} + +/// ISO 7064 MOD 11-2 check for the 18-char Chinese national ID. +fn china_id_checksum(s: &str) -> bool { + let bytes = s.as_bytes(); + if bytes.len() != 18 { + return false; + } + const WEIGHTS: [u32; 17] = [7, 9, 10, 5, 8, 4, 2, 1, 6, 3, 7, 9, 10, 5, 8, 4, 2]; + const CHECK: [u8; 11] = *b"10X98765432"; + let mut sum = 0u32; + for (i, b) in bytes[..17].iter().enumerate() { + if !b.is_ascii_digit() { + return false; + } + sum += u32::from(b - b'0') * WEIGHTS[i]; + } + let expected = CHECK[(sum % 11) as usize]; + bytes[17].to_ascii_uppercase() == expected +} + +/// Luhn check for payment-card numbers. +fn luhn_checksum(s: &str) -> bool { + let mut sum = 0u32; + let mut double = false; + for b in s.bytes().rev() { + if !b.is_ascii_digit() { + return false; + } + let mut d = u32::from(b - b'0'); + if double { + d *= 2; + if d > 9 { + d -= 9; + } + } + sum += d; + double = !double; + } + sum.is_multiple_of(10) +} + +/// The runtime guardrail for `kind: "pii"`. +pub struct PiiGuardrail { + rules: Vec, + check_input_enabled: bool, + check_output_enabled: bool, + max_buffer_bytes: usize, + on_buffer_exceeded_fail_open: bool, +} + +impl PiiGuardrail { + pub fn new( + rules: Vec, + hook_point: aisix_core::models::GuardrailHookPoint, + max_buffer_bytes: usize, + on_buffer_exceeded_fail_open: bool, + ) -> Self { + use aisix_core::models::GuardrailHookPoint as HP; + Self { + rules, + check_input_enabled: matches!(hook_point, HP::Input | HP::Both), + check_output_enabled: matches!(hook_point, HP::Output | HP::Both), + max_buffer_bytes, + on_buffer_exceeded_fail_open, + } + } + + /// First block-action rule that detects in `text`, for the verdict + /// reason. Mask-action rules never block. + fn first_block_match(&self, text: &str) -> Option<&PiiRule> { + self.rules + .iter() + .find(|r| r.action == PiiAction::Block && r.detects(text)) + } + + /// Apply every mask-action rule to `text`, in rule order. `None` when + /// nothing matched (caller keeps the original untouched). + fn mask_text(&self, text: &str) -> Option { + let mut counts: BTreeMap = BTreeMap::new(); + let mut current = Cow::Borrowed(text); + for rule in self.rules.iter().filter(|r| r.action == PiiAction::Mask) { + let (next, count) = rule.mask_all(¤t); + if count > 0 { + *counts.entry(rule.name.clone()).or_insert(0) += count; + current = Cow::Owned(next.into_owned()); + } + } + if counts.is_empty() { + None + } else { + Some(Redaction { + text: current.into_owned(), + counts, + }) + } + } +} + +#[async_trait] +impl Guardrail for PiiGuardrail { + fn name(&self) -> &'static str { + "pii" + } + + fn is_empty(&self) -> bool { + self.rules.is_empty() + } + + fn runs_on_output(&self) -> bool { + self.check_output_enabled + } + + /// Masking a streamed response requires the whole response held back — + /// a span can cross any chunk boundary, and a mask can't be applied + /// retroactively to bytes already on the wire. Cap + overflow policy + /// come from the row config. + fn stream_output_policy(&self) -> StreamOutputPolicy { + StreamOutputPolicy::BufferFull { + max_buffer_bytes: self.max_buffer_bytes, + on_exceeded_fail_open: self.on_buffer_exceeded_fail_open, + } + } + + async fn check_input(&self, req: &ChatFormat) -> GuardrailVerdict { + if !self.check_input_enabled { + return GuardrailVerdict::Allow; + } + let combined: String = req + .messages + .iter() + .map(crate::message_scan_text) + .filter(|s| !s.is_empty()) + .collect::>() + .join("\n"); + match self.first_block_match(&combined) { + // Reason carries the detector NAME only — never the matched + // value (#153 + #932 no-leak criterion). + Some(rule) => { + GuardrailVerdict::block(format!("input blocked by pii detector '{}'", rule.name)) + } + None => GuardrailVerdict::Allow, + } + } + + async fn check_output(&self, resp: &ChatResponse) -> GuardrailVerdict { + if !self.check_output_enabled { + return GuardrailVerdict::Allow; + } + let text = resp.guardrail_output_text(); + match self.first_block_match(&text) { + Some(rule) => { + GuardrailVerdict::block(format!("output blocked by pii detector '{}'", rule.name)) + } + None => GuardrailVerdict::Allow, + } + } + + fn redacts_input(&self) -> bool { + self.check_input_enabled && self.rules.iter().any(|r| r.action == PiiAction::Mask) + } + + fn redacts_output(&self) -> bool { + self.check_output_enabled && self.rules.iter().any(|r| r.action == PiiAction::Mask) + } + + fn redact_input_text(&self, text: &str) -> Option { + if !self.redacts_input() { + return None; + } + self.mask_text(text) + } + + fn redact_output_text(&self, text: &str) -> Option { + if !self.redacts_output() { + return None; + } + self.mask_text(text) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_core::models::GuardrailHookPoint; + use aisix_gateway::{ChatMessage, FinishReason, UsageStats}; + + fn guardrail(rules: Vec) -> PiiGuardrail { + PiiGuardrail::new(rules, GuardrailHookPoint::Both, 262_144, false) + } + + fn builtin(id: &str, action: PiiAction) -> PiiRule { + builtin_rule(id, action).expect("known builtin") + } + + fn req(msg: &str) -> ChatFormat { + ChatFormat::new("m", vec![ChatMessage::user(msg)]) + } + + fn resp(content: &str) -> ChatResponse { + ChatResponse { + id: "r".into(), + model: "m".into(), + message: ChatMessage::assistant(content), + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + } + } + + #[test] + fn email_masks_span_only() { + let g = guardrail(vec![builtin("email", PiiAction::Mask)]); + let r = g + .redact_input_text("contact me at alice@example.com please") + .unwrap(); + assert_eq!(r.text, "contact me at [EMAIL_REDACTED] please"); + assert_eq!(r.counts.get("email"), Some(&1)); + } + + #[test] + fn china_mobile_masks_standalone_number_only() { + let g = guardrail(vec![builtin("china_mobile", PiiAction::Mask)]); + let r = g + .redact_input_text("我的手机号是 13800138000，请回电") + .unwrap(); + assert_eq!(r.text, "我的手机号是 [CHINA_MOBILE_REDACTED]，请回电"); + // Embedded in a longer digit run → not a phone → untouched. + assert!(g.redact_input_text("订单号 913800138000123").is_none()); + } + + #[test] + fn china_id_card_requires_checksum() { + let g = guardrail(vec![builtin("china_id_card", PiiAction::Mask)]); + // Valid checksum (11010519491231002X is the canonical example). + let r = g + .redact_input_text("身份证 11010519491231002X 已登记") + .unwrap(); + assert_eq!(r.text, "身份证 [CHINA_ID_CARD_REDACTED] 已登记"); + // Same shape, broken check digit → untouched. + assert!(g + .redact_input_text("身份证 110105194912310021 已登记") + .is_none()); + } + + #[test] + fn bank_card_requires_luhn() { + let g = guardrail(vec![builtin("bank_card", PiiAction::Mask)]); + // 4111111111111111 passes Luhn. + let r = g.redact_input_text("card 4111111111111111 ok").unwrap(); + assert_eq!(r.text, "card [BANK_CARD_REDACTED] ok"); + assert!(g.redact_input_text("card 4111111111111112 ok").is_none()); + } + + #[test] + fn api_key_and_jwt_signatures_mask() { + let g = guardrail(vec![ + builtin("api_key", PiiAction::Mask), + builtin("jwt", PiiAction::Mask), + ]); + let r = g + .redact_input_text( + "key sk-abcdefghijklmnopqrstuv and token eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxIn0.dGVzdHNpZ25hdHVyZQ", + ) + .unwrap(); + assert_eq!(r.text, "key [API_KEY_REDACTED] and token [JWT_REDACTED]"); + assert_eq!(r.counts.len(), 2); + } + + #[test] + fn private_key_block_masks_across_lines() { + let g = guardrail(vec![builtin("private_key", PiiAction::Mask)]); + let text = "-----BEGIN RSA PRIVATE KEY-----\nMIIE...\n-----END RSA PRIVATE KEY-----"; + let r = g.redact_input_text(text).unwrap(); + assert_eq!(r.text, "[PRIVATE_KEY_REDACTED]"); + } + + #[test] + fn custom_pattern_masks_with_sanitised_token() { + let rule = PiiRule::new("employee id", r"\bEMP-\d{6}\b", PiiAction::Mask, None).unwrap(); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("badge EMP-123456").unwrap(); + assert_eq!(r.text, "badge [EMPLOYEE_ID_REDACTED]"); + assert_eq!(r.counts.get("employee id"), Some(&1)); + } + + #[tokio::test] + async fn block_action_blocks_and_names_detector_only() { + let g = guardrail(vec![builtin("china_id_card", PiiAction::Block)]); + let v = g.check_input(&req("id 11010519491231002X")).await; + match v { + GuardrailVerdict::Block { reason, .. } => { + assert!(reason.contains("china_id_card"), "reason: {reason}"); + assert!( + !reason.contains("11010519491231002X"), + "matched value must never appear in the reason: {reason}", + ); + } + other => panic!("expected Block, got {other:?}"), + } + } + + #[tokio::test] + async fn mask_action_does_not_block() { + let g = guardrail(vec![builtin("email", PiiAction::Mask)]); + assert_eq!( + g.check_input(&req("mail alice@example.com")).await, + GuardrailVerdict::Allow, + ); + assert_eq!( + g.check_output(&resp("mail alice@example.com")).await, + GuardrailVerdict::Allow, + ); + } + + #[tokio::test] + async fn block_action_checks_output_tool_calls() { + let g = guardrail(vec![builtin("email", PiiAction::Block)]); + let mut msg = ChatMessage::assistant(""); + msg.extra.insert( + "tool_calls".into(), + serde_json::json!([{ + "function": { "name": "send", "arguments": "{\"to\":\"bob@example.com\"}" } + }]), + ); + let r = ChatResponse { + id: "r".into(), + model: "m".into(), + message: msg, + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + }; + assert!(g.check_output(&r).await.is_block()); + } + + #[test] + fn hook_point_gates_redaction_sides() { + use aisix_core::models::GuardrailHookPoint as HP; + let input_only = PiiGuardrail::new( + vec![builtin("email", PiiAction::Mask)], + HP::Input, + 262_144, + false, + ); + assert!(input_only.redacts_input()); + assert!(!input_only.redacts_output()); + assert!(input_only.redact_output_text("a@b.co").is_none()); + assert!(!input_only.runs_on_output()); + + let output_only = PiiGuardrail::new( + vec![builtin("email", PiiAction::Mask)], + HP::Output, + 262_144, + false, + ); + assert!(!output_only.redacts_input()); + assert!(output_only.redacts_output()); + assert!(output_only.redact_input_text("a@b.co").is_none()); + } + + #[test] + fn stream_policy_uses_row_buffer_config() { + let g = PiiGuardrail::new( + vec![builtin("email", PiiAction::Mask)], + GuardrailHookPoint::Both, + 1_024, + true, + ); + assert_eq!( + g.stream_output_policy(), + StreamOutputPolicy::BufferFull { + max_buffer_bytes: 1_024, + on_exceeded_fail_open: true, + }, + ); + } + + #[test] + fn multiple_matches_count_per_detector() { + let g = guardrail(vec![builtin("email", PiiAction::Mask)]); + let r = g.redact_input_text("a@x.com then b@y.org").unwrap(); + assert_eq!(r.text, "[EMAIL_REDACTED] then [EMAIL_REDACTED]"); + assert_eq!(r.counts.get("email"), Some(&2)); + } + + #[test] + fn no_match_returns_none() { + let g = guardrail(vec![builtin("email", PiiAction::Mask)]); + assert!(g.redact_input_text("nothing sensitive here").is_none()); + } + + #[test] + fn ssn_and_ip_mask() { + let g = guardrail(vec![ + builtin("us_ssn", PiiAction::Mask), + builtin("ip_address", PiiAction::Mask), + ]); + let r = g + .redact_input_text("ssn 123-45-6789 from 192.168.1.100") + .unwrap(); + assert_eq!(r.text, "ssn [US_SSN_REDACTED] from [IP_ADDRESS_REDACTED]"); + } + + #[test] + fn every_builtin_detector_compiles() { + for (id, _, _) in BUILTIN_DETECTORS { + assert!( + builtin_rule(id, PiiAction::Mask).is_some(), + "builtin {id} must compile", + ); + } + assert!(builtin_rule("no_such_detector", PiiAction::Mask).is_none()); + } +} diff --git a/crates/aisix-obs/src/usage.rs b/crates/aisix-obs/src/usage.rs index 8e58f3824..96dae7a91 100644 --- a/crates/aisix-obs/src/usage.rs +++ b/crates/aisix-obs/src/usage.rs @@ -163,6 +163,15 @@ pub struct UsageEvent { #[serde(default, skip_serializing_if = "Vec::is_empty")] pub applied_guardrails: Vec, + /// Per-detector count of spans a `kind: "pii"` guardrail masked on this + /// request (input + output merged), e.g. `{"email": 2}`. Detector names + /// only — masked values are never captured (#932), so a compliance audit + /// sees THAT redaction happened without the event becoming a PII sink. + /// Empty (no redaction) is omitted from the wire; cp-api's `/dp/telemetry` + /// binds JSON leniently, so older CP images ignore the unknown field. + #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")] + pub redacted_entity_counts: std::collections::BTreeMap, + /// Cache outcome on this request. One of: /// /// - `"hit"` — cached response served the request without a diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 9ffd25ba9..53cc2301c 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -93,7 +93,7 @@ pub async fn chat_completions( let started = Instant::now(); let method = "POST"; let path = "/v1/chat/completions"; - let req = match body { + let mut req = match body { Ok(Json(r)) => r, Err(rej) => { use axum::extract::rejection::JsonRejection; @@ -129,14 +129,18 @@ pub async fn chat_completions( // read below to attach `applied_guardrails` to the telemetry event on // both the success and failure (guardrail-block) paths (#379). let mut applied_guardrails: Vec = Vec::new(); + // Filled by `dispatch` with per-detector PII mask counts (#932), same + // dual-path lifecycle as `applied_guardrails`. + let mut redaction_counts = crate::redact::RedactionCounts::new(); let outcome = dispatch( &state, &auth, - &req, + &mut req, &request_id, started, &client, &mut applied_guardrails, + &mut redaction_counts, ) .await; @@ -246,6 +250,7 @@ pub async fn chat_completions( error_message: String::new(), applied_guardrails: applied_guardrails.clone(), provider_key_id: success.provider_key_id.clone(), + redacted_entity_counts: redaction_counts.clone(), }, success.cost_usd, /* guardrail_blocked */ false, @@ -472,6 +477,9 @@ pub async fn chat_completions( // ultimately blocked on the output filter. applied_guardrails: applied_guardrails.clone(), provider_key_id: c.provider_key_id, + // Input-side masking happened before the output + // block — the audit trail keeps it. + redacted_entity_counts: redaction_counts.clone(), }, /* cost_usd */ 0.0, guardrail_blocked, @@ -499,6 +507,8 @@ pub async fn chat_completions( // the resolved guardrail chain (empty if it failed // before resolution). applied_guardrails: applied_guardrails.clone(), + // Input masking may have fired before the failure. + redacted_entity_counts: redaction_counts.clone(), ..UsageExtras::default() }, /* cost_usd */ 0.0, @@ -761,10 +771,14 @@ struct UpstreamCharge { routing: RoutingTelemetry, } +#[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, - req: &ChatFormat, + // `&mut` so mask-action PII guardrails (#932) can rewrite the request + // text in place before it reaches semantic routing, the cache key, or + // the upstream. + req: &mut ChatFormat, request_id: &str, started: Instant, client: &ClientContext, @@ -774,6 +788,12 @@ async fn dispatch( // without threading a field through every `Success`/`DispatchFailure` // construction site. Stays empty for requests rejected before resolution. applied_out: &mut Vec, + // Out-param: per-detector PII mask counts (#932), same lifecycle as + // `applied_out`. Input-side counts land here as soon as the request is + // rewritten; the non-streaming output side merges in later. Streaming + // output counts travel via `StreamCompletion` instead (the event is + // emitted at end-of-stream). + redactions_out: &mut crate::redact::RedactionCounts, ) -> Result { if req.messages.is_empty() { return Err(DispatchFailure::new( @@ -871,6 +891,16 @@ async fn dispatch( } } + // #932: mask-action PII rules rewrite the request in place AFTER the + // block check passes and BEFORE any downstream use — the semantic- + // routing embedding, the cache-key fingerprint, and the upstream + // dispatch all see the masked text, so the matched values never leave + // the gateway. + crate::redact::merge_counts( + redactions_out, + crate::redact::redact_chat_format(resolved_chain.as_ref(), req), + ); + // Budget pre-check via cp-api. The DP no longer owns budget state; // cp-api returns a cached/live decision per api_key. let decision = state.budgets.check(&auth.entry.id).await; @@ -980,6 +1010,7 @@ async fn dispatch( &model_id, &auth.entry.id, client, + redactions_out.clone(), ) .await; } @@ -1236,6 +1267,10 @@ async fn dispatch( // Applied guardrail set (#379), owned for the move into on_complete so // the streamed-response telemetry event records which guardrails ran. let applied_guardrails_for_telem = applied_guardrails.clone(); + // Input-side PII mask counts (#932), captured before the stream so + // the end-of-stream event merges them with the output-side counts + // accumulated in `comp.redacted_entity_counts`. + let input_redactions_for_telem = redactions_out.clone(); // Downstream client attribution (#492) moved into the on_complete // closure so streamed responses log the same IP/UA as non-streaming. let client_for_telem = client.clone(); @@ -1347,6 +1382,11 @@ async fn dispatch( error_message: String::new(), applied_guardrails: applied_guardrails_for_telem.clone(), provider_key_id: provider_key_id_for_telem.clone(), + redacted_entity_counts: { + let mut merged = input_redactions_for_telem.clone(); + crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); + merged + }, }, /* cost_usd */ 0.0, comp.guardrail_blocked, @@ -1523,7 +1563,7 @@ async fn dispatch( if let (Some(cache), Some(key)) = (policy_cache.as_ref(), cache_key.as_ref()) { match cache.get(key).await { - Ok(Some(cached)) => { + Ok(Some(mut cached)) => { reservation.commit_tokens(0).await; // #448: a cache hit is client-visible output just like a // fresh upstream response, so it must run output guardrails @@ -1553,6 +1593,14 @@ async fn dispatch( } _ => {} } + // #932: cache hits are client-visible output like a fresh + // upstream response — mask them too. Entries are stored + // masked (see the write below), but entries written before + // a rule was added/tightened must still be caught here. + crate::redact::merge_counts( + redactions_out, + crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut cached), + ); let prompt = cached.usage.prompt_tokens as u64; let completion = cached.usage.completion_tokens as u64; let total = cached.usage.total_tokens as u64; @@ -1832,7 +1880,7 @@ async fn dispatch( } } - let Some(upstream) = upstream else { + let Some(mut upstream) = upstream else { // Bubble the most recent BridgeError through ProxyError::Bridge. let err = last_err.unwrap_or_else(|| { BridgeError::Config("routing exhausted with no targets attempted".into()) @@ -1933,6 +1981,14 @@ async fn dispatch( } } + // #932: mask-action PII rules rewrite the response AFTER the block + // check passes and BEFORE the cache write below, so cached entries + // are stored masked — the matched values don't persist at rest. + crate::redact::merge_counts( + redactions_out, + crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut upstream), + ); + // Cache write is gated on the same policy as the lookup at the // top of dispatch — without a matching enabled cache_policy in // snapshot, we skip both read AND write so the cache backend @@ -2044,6 +2100,11 @@ async fn dispatch_ensemble( model_id: &str, api_key_id: &str, client: &ClientContext, + // Input-side PII mask counts (#932) captured by `dispatch` before the + // fan-out. Ensemble telemetry is emitted per sub-call inside this + // function (the handler's main emit is skipped), so the counts ride + // the judge event — the ensemble's terminal event. + input_redactions: crate::redact::RedactionCounts, ) -> Result { let started = Instant::now(); let with_model = |e: ProxyError| DispatchFailure::new(Some(model_id.to_string()), None, e); @@ -2359,6 +2420,9 @@ async fn dispatch_ensemble( let applied_guardrails_for_telem = applied_guardrails.to_vec(); let client_for_telem = client.clone(); let bypass_for_telem = bypass_reason.clone().unwrap_or_default(); + // Input-side PII mask counts (#932); merged with the streamed judge's + // output-side counts on the terminal (judge) event. + let input_redactions_for_telem = input_redactions.clone(); // Hold concurrency for the stream's full lifetime (#450). Snapshot the // keys BEFORE consuming the reservation into the owned guard. @@ -2461,6 +2525,11 @@ async fn dispatch_ensemble( attempt_model: judge_attempt_model.clone(), applied_guardrails: applied_guardrails_for_telem.clone(), provider_key_id: judge_provider_key_id.clone(), + redacted_entity_counts: { + let mut merged = input_redactions_for_telem.clone(); + crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); + merged + }, ..UsageExtras::default() }, /* cost_usd */ 0.0, @@ -2580,7 +2649,14 @@ async fn dispatch_ensemble( // usage to cp-api. The `emit_subcalls` closure is therefore defined // here and invoked from each branch of the guardrail match below — // never from the post-match continuation alone. - let emit_subcalls = |blocked: bool, bypass: &str| { + // Takes `outcome` as a parameter (not a capture) so the Allow branch + // below can mask `outcome.response` in place (#932) between the + // guardrail check and this emit. `redactions` lands on the judge + // event — the ensemble's terminal telemetry event. + let emit_subcalls = |outcome: &crate::ensemble::EnsembleOutcome, + blocked: bool, + bypass: &str, + redactions: &crate::redact::RedactionCounts| { for (index, member) in outcome.panel.iter().enumerate() { emit_panel_member(member, index, blocked, bypass); } @@ -2610,6 +2686,7 @@ async fn dispatch_ensemble( attempt_model: outcome.judge_model.clone(), applied_guardrails: applied_guardrails.to_vec(), provider_key_id: judge_provider_key_id, + redacted_entity_counts: redactions.clone(), ..UsageExtras::default() }, /* cost_usd */ 0.0, @@ -2639,7 +2716,12 @@ async fn dispatch_ensemble( ); // Block is not a bypass, so the output guardrail did not mutate // `bypass_reason`; carry the input-only bypass (if any). - emit_subcalls(true, &bypass_reason.clone().unwrap_or_default()); + emit_subcalls( + &outcome, + true, + &bypass_reason.clone().unwrap_or_default(), + &input_redactions, + ); return Err(DispatchFailure::new( Some(model_id.to_string()), None, @@ -2656,10 +2738,21 @@ async fn dispatch_ensemble( } } } - // Allow / Bypass continuation: emit the (non-blocked) sub-call events - // with the final bypass value (which an output bypass above may have - // set). - emit_subcalls(false, &bypass_reason.clone().unwrap_or_default()); + // Allow / Bypass continuation: mask the synthesized answer (#932) — + // the check above ran on the original text, the client gets the masked + // one — then emit the (non-blocked) sub-call events with the final + // bypass value (which an output bypass above may have set). + let mut ensemble_redactions = input_redactions.clone(); + crate::redact::merge_counts( + &mut ensemble_redactions, + crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut outcome.response), + ); + emit_subcalls( + &outcome, + false, + &bypass_reason.clone().unwrap_or_default(), + &ensemble_redactions, + ); // The synthesized answer is the client-facing response, rendered with // the requested (ensemble) model name. `telemetry_handled_by_stream` @@ -2884,6 +2977,7 @@ fn emit_usage_event( guardrail_blocked, guardrail_bypassed_reason: extras.bypass_reason, applied_guardrails: extras.applied_guardrails, + redacted_entity_counts: extras.redacted_entity_counts, cache_status: extras.cache_status, cache_hit_saved_input_tokens: extras.cache_hit_saved_input_tokens, cache_hit_saved_output_tokens: extras.cache_hit_saved_output_tokens, @@ -3026,6 +3120,10 @@ struct UsageExtras { /// emit events land in cp-api with the tag columns NULL. /// See AISIX-Cloud#436 / #302 M17. provider_key_id: String, + /// Per-detector PII mask counts for this request, input + output + /// merged (#932). Lands on `usage_events.redacted_entity_counts`. + /// Detector names only, never matched values. Empty = no redaction. + redacted_entity_counts: crate::redact::RedactionCounts, } /// Emit one zero-token `UsageEvent` per FAILED attempt of a request @@ -3211,6 +3309,10 @@ struct StreamCompletion { /// capture cap). Empty otherwise. Read by the on_complete telemetry /// closure; never reaches the CP sink. response_text: String, + /// Per-detector PII mask counts applied to the held stream at release + /// (#932). Merged with the input-side counts by the on_complete + /// telemetry closure. Detector names only, never matched values. + redacted_entity_counts: crate::redact::RedactionCounts, } /// Parameters needed to run output-guardrail evaluation at @@ -3378,9 +3480,11 @@ where .map(|ctx| ctx.chain.stream_output_policy()) .unwrap_or_default(); let hold_back = stream_policy.holds_back(); - // Rendered content events withheld from the wire until their - // window (or the whole response) scans clean. Hold-back path only. - let mut pending: Vec = Vec::new(); + // Content chunks withheld from the wire until their window (or the + // whole response) scans clean. Hold-back path only. Held PRE-render + // (#932): the BufferFull release rewrites masked spans across the + // held chunks before they are rendered + serialised at drain time. + let mut pending: Vec = Vec::new(); // Content accumulated since the last window flush (Window mode); // bounded to ~window_size, unlike content_buffer (whole response). let mut window_buf = String::new(); @@ -3403,8 +3507,25 @@ where // the truncated response as a successful one. let mut errored = false; let mut first_chunk_seen = false; + // Render + serialise one held/live chunk into an SSE Event. + // Serialisation of these plain structs can't realistically fail; + // the Err arm mirrors the pre-hold-back defensive error frame. + macro_rules! chunk_event { + ($chunk:expr) => {{ + let rendered = render_chunk(created, $chunk, &client_facing_model); + match serde_json::to_string(&rendered) { + Ok(json) => Event::default().data(json), + Err(err) => { + errored = true; + Event::default() + .event("error") + .data(error_frame_payload("internal_error", &err.to_string())) + } + } + }}; + } while let Some(item) = upstream.next().await { - let (ev, is_error_ev) = match item { + let maybe_chunk = match item { Ok(mut chunk) => { // Record TTFT on the first chunk carrying generated // output (content or tool calls). Skip role-only @@ -3524,42 +3645,34 @@ where if let Some(u) = chunk.usage.as_mut() { *u = u.saturating_add(&base_usage); } - let rendered = render_chunk(created, chunk, &client_facing_model); - match serde_json::to_string(&rendered) { - Ok(json) => (Event::default().data(json), false), - Err(err) => { - errored = true; - ( - Event::default() - .event("error") - .data(error_frame_payload("internal_error", &err.to_string())), - true, - ) - } - } + Some(chunk) } Err(err) => { errored = true; let etype = err.error_type(); - ( + yield Ok::<_, Infallible>( Event::default() .event("error") .data(error_frame_payload(etype, &err.to_string())), - true, - ) + ); + None } }; - if is_error_ev || !hold_back || cap_released { - // Error frames, the EndOfStreamCheck path, and a released - // BufferFull cap all forward straight to the wire. On the - // hold-back path an error drops the held (unscanned) - // content via the `errored` skip at end-of-stream (fail - // closed). + let Some(chunk) = maybe_chunk else { + continue; + }; + if !hold_back || cap_released { + // The EndOfStreamCheck path and a released BufferFull cap + // forward straight to the wire. (Error frames yielded in + // the Err arm above also bypass the hold; the `errored` + // skip at end-of-stream then drops any held unscanned + // content — fail closed.) + let ev = chunk_event!(chunk); yield Ok::<_, Infallible>(ev); } else { - // Hold-back: withhold this content event until its window - // (or the whole response) scans clean. - pending.push(ev); + // Hold-back: withhold this chunk until its window (or the + // whole response) scans clean. + pending.push(chunk); match &stream_policy { aisix_guardrails::StreamOutputPolicy::Window { size_chars, @@ -3629,11 +3742,16 @@ where _ => {} } // Clean (Allow / Bypass): release - // this window's events, then keep the + // this window's chunks, then keep the // trailing overlap as scan context for the - // next window (its events were already sent). - for e in pending.drain(..) { - yield Ok::<_, Infallible>(e); + // next window (its chunks were already sent). + // No mask rewrite here: a redacting (PII) + // guardrail always folds the chain to + // BufferFull, so Window mode implies no + // redactor in the chain. + for chunk in pending.drain(..) { + let ev = chunk_event!(chunk); + yield Ok::<_, Infallible>(ev); } // Clamp the retained overlap to cc-1 so a // misconfigured overlap >= window can't keep @@ -3656,9 +3774,14 @@ where let buffered = content_buffer.as_ref().map_or(0, |b| b.len()); if buffered > *max_buffer_bytes { if *on_exceeded_fail_open { + // Fail-open overflow releases the held + // chunks unscanned AND unmasked — the + // operator opted into that trade-off via + // `on_buffer_exceeded: fail_open`. cap_released = true; - for e in pending.drain(..) { - yield Ok::<_, Infallible>(e); + for chunk in pending.drain(..) { + let ev = chunk_event!(chunk); + yield Ok::<_, Infallible>(ev); } } else { tracing::warn!( @@ -3774,8 +3897,23 @@ where } }; if !blocked { - for e in pending.drain(..) { - yield Ok::<_, Infallible>(e); + // #932: the whole response is held here, so mask- + // action PII rules rewrite the held chunks before + // anything reaches the wire (channel reassembly — + // a masked span can cross chunk boundaries). + if let Some(ctx) = output_guardrail.as_ref() { + let counts = + crate::redact::redact_chat_chunks(ctx.chain.as_ref(), &mut pending); + if !counts.is_empty() { + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + counts, + ); + } + } + for chunk in pending.drain(..) { + let ev = chunk_event!(chunk); + yield Ok::<_, Infallible>(ev); } } } diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 816c39326..7c5dd0f5f 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -52,6 +52,9 @@ struct CompletionDispatchSuccess { /// or on a 200 with no `usage` block (rare edge). Handler /// gates UsageEvent emission on this being `Some`. usage: Option, + /// Per-detector PII mask counts (#932), input + output merged. + /// Attached to the emitted UsageEvent. Empty = no redaction. + redactions: crate::redact::RedactionCounts, /// True when the response leg was blocked by an OUTPUT guardrail /// AFTER the upstream billed for it (#911 [23]). The response body is /// the redacted 422, but `usage` still carries the billed counts so @@ -131,6 +134,7 @@ pub async fn completions( &usage, &client, success.guardrail_blocked, + success.redactions.clone(), ); } success.response @@ -196,14 +200,16 @@ fn completions_input_to_chat(model: &str, body: &Value) -> aisix_gateway::ChatFo async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, - body: Value, + mut body: Value, request_id: &str, source_ip: &str, ) -> Result { let model_name = body .get("model") .and_then(|v| v.as_str()) - .ok_or_else(|| ProxyError::InvalidRequest("missing `model` field".into()))?; + .ok_or_else(|| ProxyError::InvalidRequest("missing `model` field".into()))? + .to_string(); + let model_name = model_name.as_str(); let snapshot = state.snapshot.load(); @@ -249,6 +255,10 @@ async fn dispatch( } } + // #932: mask-action PII rules rewrite the prompt in place AFTER the + // block check passes, BEFORE the body is forwarded upstream. + let mut redactions = crate::redact::redact_completions_request(&resolved_chain, &mut body); + let model_rl = crate::quota::ModelRateLimit::from_model(model_name, &model_entry.id, &model_entry.value); let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; @@ -333,17 +343,27 @@ async fn dispatch( model_id: model_entry.id.to_string(), provider_key_id: pk_entry.id.to_string(), usage, + redactions, guardrail_blocked: true, }); } } + // #932: mask-action PII rules rewrite the reply text AFTER the + // block check passes. + let mut resp_json = resp_json; + crate::redact::merge_counts( + &mut redactions, + crate::redact::redact_completions_response(&resolved_chain, &mut resp_json), + ); + Ok(CompletionDispatchSuccess { response: Json(resp_json).into_response(), provider: provider_label, model_id: model_entry.id.to_string(), provider_key_id: pk_entry.id.to_string(), usage, + redactions, guardrail_blocked: false, }) } @@ -360,6 +380,7 @@ async fn dispatch( // gates emission on `usage.is_some()` so 501 stays // out of /logs noise (same convention as #402). usage: None, + redactions, guardrail_blocked: false, }) } @@ -442,6 +463,8 @@ fn emit_usage_event( usage: &CompletionUsage, client: &ClientContext, guardrail_blocked: bool, + // Per-detector PII mask counts (#932). Empty = no redaction. + redacted_entity_counts: crate::redact::RedactionCounts, ) { let snap = state.snapshot.load(); let mut event = UsageEvent { @@ -460,6 +483,7 @@ fn emit_usage_event( // #911 [23]: a billed-then-output-blocked completion surfaces on the // dashboard's Blocked tab while still carrying its billed token counts. guardrail_blocked, + redacted_entity_counts, ..Default::default() }; crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id); diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index d871c16eb..7c41ea3d9 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -172,6 +172,7 @@ pub async fn embeddings( elapsed, success.prompt_tokens, &client, + success.redactions.clone(), ); } success.response @@ -229,6 +230,8 @@ struct EmbedDispatchSuccess { /// The `{kind, hook}` set of guardrails that governed this request (#379 /// parity) — surfaced on the emitted UsageEvent. applied_guardrails: Vec, + /// Per-detector PII mask counts (#932) applied to the input. + redactions: crate::redact::RedactionCounts, prompt_tokens: u32, /// `true` when the dispatch produced a real 200 from the upstream /// (we have authoritative usage data to attribute). `false` for the @@ -245,7 +248,7 @@ struct EmbedDispatchSuccess { async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, - body: EmbeddingRequestBody, + mut body: EmbeddingRequestBody, request_id: &str, source_ip: &str, ) -> Result { @@ -310,6 +313,32 @@ async fn dispatch( } } + // #932: mask-action PII rules rewrite the embeddings input in place + // AFTER the block check passes — the masked text is what gets embedded, + // so the matched values never reach the provider. + let mut redactions = crate::redact::RedactionCounts::new(); + if aisix_guardrails::Guardrail::redacts_input(&resolved_chain) { + match &mut body.input { + InputField::Single(s) => { + if let Some(r) = aisix_guardrails::Guardrail::redact_input_text(&resolved_chain, s) + { + *s = r.text; + crate::redact::merge_counts(&mut redactions, r.counts); + } + } + InputField::Multi(items) => { + for s in items { + if let Some(r) = + aisix_guardrails::Guardrail::redact_input_text(&resolved_chain, s) + { + *s = r.text; + crate::redact::merge_counts(&mut redactions, r.counts); + } + } + } + } + } + let model_rl = crate::quota::ModelRateLimit::from_model(&body.model, &model_entry.id, &model_entry.value); let reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; @@ -360,6 +389,7 @@ async fn dispatch( model_id: model_entry.id.to_string(), provider_key_id: pk_entry.id.to_string(), applied_guardrails: applied_guardrails.clone(), + redactions: redactions.clone(), prompt_tokens, upstream_called: true, }) @@ -379,6 +409,7 @@ async fn dispatch( model_id: model_entry.id.to_string(), provider_key_id: pk_entry.id.to_string(), applied_guardrails: applied_guardrails.clone(), + redactions: redactions.clone(), prompt_tokens: 0, // No upstream call happened — the handler reads this // and skips UsageEvent emission. Distinguished from @@ -459,6 +490,8 @@ fn emit_usage_event( elapsed: Duration, prompt_tokens: u32, client: &ClientContext, + // Per-detector PII mask counts (#932) applied to the input. + redacted_entity_counts: crate::redact::RedactionCounts, ) { // Only populate fields meaningful to /v1/embeddings; rely on // UsageEvent's `#[derive(Default)]` for everything else. Wire-level @@ -498,6 +531,7 @@ fn emit_usage_event( status_code, inbound_protocol: "openai".to_string(), applied_guardrails: applied_guardrails.to_vec(), + redacted_entity_counts, client_source_ip: client.source_ip.clone(), client_user_agent: client.user_agent.clone(), ..Default::default() diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 429157e50..962c4c43b 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -48,6 +48,7 @@ mod model_resolve; mod models; mod passthrough; mod quota; +mod redact; mod render; mod request_id; mod rerank; diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 18f379bb5..f7fa8fb69 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -117,6 +117,10 @@ pub async fn messages( // read below to attach `applied_guardrails` to the telemetry event on both // the success and failure (input-block) paths (#379). let mut applied_guardrails: Vec = Vec::new(); + // Filled by `dispatch` with per-detector PII mask counts (#932), same + // dual-path lifecycle as `applied_guardrails`. Streaming output counts + // travel via the stream builders' end-of-stream emit instead. + let mut redaction_counts = crate::redact::RedactionCounts::new(); // #890 req-1: capture the client's streaming intent before dispatch // (which mutates the body). let stream_requested = body @@ -131,6 +135,7 @@ pub async fn messages( started, &client, &mut applied_guardrails, + &mut redaction_counts, ) .await { @@ -143,7 +148,11 @@ pub async fn messages( usage_handled_by_stream, routing, captured_content, + output_redactions, }) => { + // #932: fold the non-streaming response-side mask counts into + // the per-request total before the terminal emit below. + crate::redact::merge_counts(&mut redaction_counts, output_redactions); let elapsed = started.elapsed(); let status = response.status().as_u16(); emit_access_log( @@ -234,6 +243,7 @@ pub async fn messages( &client, attempt, applied_guardrails.clone(), + redaction_counts.clone(), captured_content, ); } @@ -325,6 +335,8 @@ pub async fn messages( ..Default::default() }, applied_guardrails.clone(), + // Input masking may have fired before the failure. + redaction_counts.clone(), /* content */ None, ); } @@ -377,11 +389,15 @@ fn emit_failed_attempts_anthropic( client, AttemptInfo::from_record(rec), applied_guardrails.to_vec(), + // Failed attempts carry no per-request redaction detail; the + // terminal (winner / pre-dispatch) event does. + crate::redact::RedactionCounts::new(), /* content */ None, ); } } +#[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, @@ -395,6 +411,10 @@ async fn dispatch( // rejected before resolution. The streaming paths capture the same set // directly from `resolved_chain` for their end-of-stream emit. applied_out: &mut Vec, + // Out-param: per-detector PII mask counts (#932), same lifecycle as + // `applied_out`. Streaming output counts travel via the stream + // builders' own end-of-stream emit instead. + redactions_out: &mut crate::redact::RedactionCounts, ) -> Result { let snapshot = state.snapshot.load(); @@ -459,6 +479,14 @@ async fn dispatch( ); } } + // #932: mask-action PII rules rewrite the Anthropic-native body in + // place AFTER the block check passes — both the passthrough and the + // cross-provider bridge forward from this body, so the masked text + // is what reaches the upstream. + crate::redact::merge_counts( + redactions_out, + crate::redact::redact_anthropic_request(resolved_chain.as_ref(), body), + ); } let model_rl = @@ -573,6 +601,7 @@ async fn dispatch( ..Default::default() }, &mut reservation, + redactions_out.clone(), ) .await { @@ -677,6 +706,10 @@ async fn dispatch_to_target( // post-stream token accounting into the end-of-stream guard. Left in place // on the non-streaming / error paths for the handler to commit or retry. reservation: &mut Option, + // Input-side PII mask counts (#932) — the streaming paths merge these + // into their end-of-stream telemetry emit (the non-streaming emit + // happens in `messages()`, which already holds them). + input_redactions: crate::redact::RedactionCounts, ) -> Result { let model = &target.model; let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; @@ -699,6 +732,7 @@ async fn dispatch_to_target( client, attempt, reservation, + input_redactions, ) .await; } @@ -721,6 +755,7 @@ async fn dispatch_to_target( client, attempt, reservation, + input_redactions, ) .await } @@ -748,6 +783,7 @@ async fn anthropic_passthrough_dispatch( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + input_redactions: crate::redact::RedactionCounts, ) -> Result { let mut body = body.clone(); let api_key = crate::dispatch::require_secret(pk_value, model)?; @@ -1056,6 +1092,13 @@ async fn anthropic_passthrough_dispatch( &client_ctx_c, attempt_c.clone(), applied_guardrails_c.clone(), + // #932: input-side mask counts captured before dispatch, + // merged with the hold-back release's output-side counts. + { + let mut merged = input_redactions.clone(); + crate::redact::merge_counts(&mut merged, usage.redacted_entity_counts); + merged + }, // Prompt captured up front; response assembled by the frame // parser into `usage.response_text`. Both gated on the cap. match (&captured_prompt_c, content_cap) { @@ -1108,6 +1151,8 @@ async fn anthropic_passthrough_dispatch( routing: RoutingTelemetry::default(), // Streaming content capture lands in C3b. captured_content: None, + // Streaming: the end-of-stream closure owns the counts. + output_redactions: crate::redact::RedactionCounts::new(), }) } else { // Non-streaming: deserialise and re-serialise as JSON. Decode @@ -1187,6 +1232,11 @@ async fn anthropic_passthrough_dispatch( } } + // #932: mask-action PII rules rewrite the passthrough response body + // (text blocks + tool_use input) AFTER the block check passes. + let output_redactions = + crate::redact::redact_anthropic_response(resolved_chain.as_ref(), &mut json_body); + // Capture the prompt (the outbound request body) + assembled assistant // text for content-capturing exporters (gated). Built here, before // `json_body` is rendered into the response; threaded to `fan_out` via @@ -1217,6 +1267,7 @@ async fn anthropic_passthrough_dispatch( usage_handled_by_stream: false, routing: RoutingTelemetry::default(), captured_content, + output_redactions, }) } } @@ -1308,6 +1359,7 @@ async fn cross_provider_dispatch( client: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + input_redactions: crate::redact::RedactionCounts, ) -> Result { use aisix_gateway::{Bridge, BridgeContext}; use aisix_provider_anthropic::{ @@ -1518,6 +1570,13 @@ async fn cross_provider_dispatch( &client_for_telem, attempt_for_telem.clone(), applied_guardrails_for_telem.clone(), + // #932: input-side mask counts captured before dispatch, + // merged with the hold-back release's output-side counts. + { + let mut merged = input_redactions.clone(); + crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); + merged + }, // Prompt captured up front, response assembled across the // stream into `comp.response_text`; both gated on the cap. match (&captured_prompt_for_telem, content_cap) { @@ -1556,11 +1615,13 @@ async fn cross_provider_dispatch( routing: RoutingTelemetry::default(), // Streaming content capture lands in C3b. captured_content: None, + // Streaming: the end-of-stream closure owns the counts. + output_redactions: crate::redact::RedactionCounts::new(), }); } // Non-streaming. - let resp = bridge.chat(&chat, &ctx).await.map_err(|err| { + let mut resp = bridge.chat(&chat, &ctx).await.map_err(|err| { if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) { state.runtime_status.mark_cooldown(model_id, ttl, reason); @@ -1591,6 +1652,10 @@ async fn cross_provider_dispatch( } } + // #932: mask-action PII rules rewrite the bridged response AFTER the + // block check passes, BEFORE it is rendered back as Anthropic JSON. + let output_redactions = crate::redact::redact_chat_response(resolved_chain.as_ref(), &mut resp); + let metrics = AnthropicUsageMetrics { prompt_tokens: resp.usage.prompt_tokens, completion_tokens: resp.usage.completion_tokens, @@ -1630,6 +1695,7 @@ async fn cross_provider_dispatch( usage_handled_by_stream: false, routing: RoutingTelemetry::default(), captured_content, + output_redactions, }) } @@ -1652,6 +1718,21 @@ fn build_anthropic_sse_stream( use futures::StreamExt; let mut encoder = encoder; + // #932 / #466-class: when the chain's streamed-output policy is the + // whole-response hold-back (BufferFull — keyword/pii/bedrock output + // guardrails), chunks are withheld from the encoder until the + // end-of-stream scan clears (and masks) them: a block keeps matched + // content off the wire entirely, and a mask can't rewrite bytes that + // already left. Window-policy guardrails (Azure/Aliyun) keep the + // pre-existing live-forward + end-of-stream check on this surface. + let hold_policy = output_guardrail.as_ref().and_then(|c| { + match aisix_guardrails::Guardrail::stream_output_policy(c.as_ref()) { + aisix_guardrails::StreamOutputPolicy::BufferFull { + max_buffer_bytes, .. + } => Some(max_buffer_bytes), + _ => None, + } + }); let stream = async_stream::stream! { let mut guard = CompleteAnthropicStreamOnDrop { slot: Some((on_complete, AnthropicStreamCompletion::default())), @@ -1659,15 +1740,19 @@ fn build_anthropic_sse_stream( let mut upstream = upstream; let mut first_chunk_seen = false; // Accumulate assistant text for the end-of-stream output guardrail - // (#448). Bytes are forwarded live (mirrors /v1/chat/completions and - // the common streaming-guardrail pattern), so a blocked response is - // signalled with a terminal `error` event rather than held back. + // (#448). Without a hold-back policy, bytes are forwarded live and + // a blocked response is signalled with a terminal `error` event. let mut content_text = String::new(); // Also collect streamed tool-call fragments so tool-call output is // scanned too (parity with the non-streaming path). Fragments are // kept raw — the guardrail scans their serialized text, no need to // reassemble by index. let mut tool_call_fragments: Vec = Vec::new(); + // Chunks withheld from the encoder until the end-of-stream scan + // clears them (hold-back policies only). Held PRE-encode so the + // mask rewrite can run on the normalised chunks. + let mut held_chunks: Vec = Vec::new(); + let mut held_bytes: usize = 0; while let Some(item) = upstream.next().await { match item { Ok(chunk) => { @@ -1712,6 +1797,23 @@ fn build_anthropic_sse_stream( } } } + if let Some(max_hold) = hold_policy { + // Hold-back: withhold the chunk until the end-of- + // stream scan clears it. Overflow fails closed — + // unscannable content must not be released. + held_bytes += chunk.delta.content.as_deref().map_or(0, str::len); + if held_bytes > max_hold { + tracing::warn!( + guardrail_hook = "output", + max_buffer_bytes = max_hold, + "streaming /v1/messages response exceeded hold-back cap; failing closed", + ); + yield Ok(bytes::Bytes::from(guardrail_block_frame(None))); + return; + } + held_chunks.push(chunk); + continue; + } for ev in encoder.next_events(&chunk) { yield Ok::<_, std::io::Error>(bytes::Bytes::from(ev.to_sse_string())); } @@ -1720,6 +1822,8 @@ fn build_anthropic_sse_stream( } } Err(e) => { + // Hold-back: the held (unscanned) chunks are dropped — + // fail closed; only the error frame reaches the client. let frame = format!( "event: error\ndata: {{\"type\":\"error\",\"error\":{{\"type\":\"{}\",\"message\":{}}}}}\n\n", e.error_type(), @@ -1763,12 +1867,37 @@ fn build_anthropic_sse_stream( reason = %reason, "guardrail blocked streaming /v1/messages response", ); + // Hold-back: the held chunks are dropped — the matched + // content never reached the wire. let frame = guardrail_block_frame(guardrail_name.as_deref()); yield Ok(bytes::Bytes::from(frame)); return; } } } + // Hold-back release (#932): the scan cleared — mask the held + // chunks (channel reassembly across chunk boundaries), then feed + // them through the encoder as if they had streamed live. + if !held_chunks.is_empty() { + if let Some(chain) = output_guardrail.as_ref() { + let counts = + crate::redact::redact_chat_chunks(chain.as_ref(), &mut held_chunks); + if !counts.is_empty() { + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + counts, + ); + } + } + for chunk in held_chunks.drain(..) { + for ev in encoder.next_events(&chunk) { + yield Ok::<_, std::io::Error>(bytes::Bytes::from(ev.to_sse_string())); + } + if encoder.is_finished() { + break; + } + } + } if !encoder.is_finished() { for ev in encoder.force_finish() { yield Ok(bytes::Bytes::from(ev.to_sse_string())); @@ -1821,6 +1950,9 @@ struct AnthropicStreamCompletion { /// capture cap). Empty otherwise. Read by the on_complete closure; never /// reaches the CP sink. response_text: String, + /// Per-detector PII mask counts applied to the held stream at release + /// (#932). Merged with the input-side counts by the on_complete emit. + redacted_entity_counts: crate::redact::RedactionCounts, } struct CompleteAnthropicStreamOnDrop { @@ -1864,6 +1996,11 @@ struct DispatchOutcome { /// or on the streaming path (filled at stream end). Forwarded only to /// `fan_out`, never to the CP telemetry sink. captured_content: Option, + /// Per-detector PII mask counts applied to the NON-streaming response + /// body (#932). Merged with the input-side counts by `messages()` + /// before the terminal emit. Empty on the streaming paths — their + /// end-of-stream closures own the output-side counts. + output_redactions: crate::redact::RedactionCounts, } /// Dispatch error carrying the per-attempt telemetry accumulated before @@ -1938,6 +2075,9 @@ fn emit_anthropic_usage_event( // The `{kind, hook}` set of guardrails that governed this request (#379). // Empty for the guardrail-free path and pre-resolution failures. applied_guardrails: Vec, + // Per-detector PII mask counts (#932), input + output merged. Detector + // names only, never matched values. Empty = no redaction. + redacted_entity_counts: crate::redact::RedactionCounts, content: Option, ) { // Per-PK telemetry attribution (#302 M17 / AISIX-Cloud#436). @@ -1989,6 +2129,7 @@ fn emit_anthropic_usage_event( client_source_ip: client.source_ip.clone(), client_user_agent: client.user_agent.clone(), applied_guardrails, + redacted_entity_counts, ..Default::default() }; // Handler label "messages" — Anthropic /v1/messages inbound @@ -2088,6 +2229,9 @@ struct AnthropicStreamUsage { /// Assistant text accumulated from `content_block_delta` frames, for /// the end-of-stream output guardrail (#448). response_text: String, + /// Per-detector PII mask counts applied to the held stream at release + /// (#932). Merged with the input-side counts by the on_complete emit. + redacted_entity_counts: crate::redact::RedactionCounts, } /// Update the accumulator from one parsed SSE `data:` JSON object. @@ -2335,6 +2479,21 @@ where { let delivered = Arc::new(AtomicU32::new(0)); let delivered_for_drop = Arc::clone(&delivered); + // #932 / #466-class: when the chain's streamed-output policy is the + // whole-response hold-back (BufferFull — keyword/pii/bedrock output + // guardrails), the passthrough must buffer the raw SSE bytes rather + // than forward them live: a block must keep matched content off the + // wire entirely, and a mask can't be applied to bytes already sent. + // Window-policy guardrails (Azure/Aliyun incremental release) keep the + // pre-existing live-forward + end-of-stream check on this surface. + let hold_policy = output_guardrail.as_ref().and_then(|c| { + match aisix_guardrails::Guardrail::stream_output_policy(c.as_ref()) { + aisix_guardrails::StreamOutputPolicy::BufferFull { + max_buffer_bytes, .. + } => Some(max_buffer_bytes), + _ => None, + } + }); let inner = async_stream::stream! { let mut guard = AnthropicStreamGuard { slot: Some((on_complete, AnthropicStreamUsage::default())), @@ -2343,6 +2502,8 @@ where futures::pin_mut!(upstream); let mut buf: Vec = Vec::new(); let mut first_token_seen = false; + // Whole-response hold-back buffer (BufferFull policies only). + let mut held: Vec = Vec::new(); while let Some(item) = upstream.next().await { if let Ok(bytes) = &item { // Side-channel parse: copy into the frame buffer (the @@ -2377,16 +2538,43 @@ where ); buf.clear(); } + if let Some(max_hold) = hold_policy { + // Hold-back: withhold the bytes until the end-of-stream + // scan clears (and masks) them. Overflow fails closed — + // content that can't be fully buffered to scan must not + // be released (mirrors /v1/responses). + if held.len() + bytes.len() > max_hold { + tracing::warn!( + guardrail_hook = "output", + max_buffer_bytes = max_hold, + "streaming /v1/messages passthrough exceeded hold-back cap; failing closed", + ); + yield Ok(Bytes::from(guardrail_block_frame(None))); + return; + } + held.extend_from_slice(bytes); + continue; + } } // Forward the original item verbatim (Ok bytes OR Err — an // upstream error mid-stream is passed through; the - // accumulator keeps whatever was captured before it). + // accumulator keeps whatever was captured before it). In + // hold-back mode an Err lands here too: it is forwarded and + // the held (unscanned) content is dropped — fail closed. + let errored = item.is_err(); yield item; + if errored && hold_policy.is_some() { + return; + } } // End-of-stream output guardrail (#448): scan the accumulated // assistant text. On a block, emit a terminal Anthropic `error` - // event (bytes were already forwarded verbatim, mirroring the - // cross-provider path and the common streaming-guardrail pattern). + // event. On the hold-back path (BufferFull) nothing has been + // forwarded yet, so a block keeps the matched content off the + // wire entirely; on the live-forward path (Window / + // EndOfStreamCheck) the bytes were already forwarded verbatim + // and the error frame is the trailing signal. + let mut blocked = false; if let Some(chain) = output_guardrail.as_ref() { // Clone (not take) when content capture is on, so the assembled // response survives for the on_complete content capture below; @@ -2415,11 +2603,31 @@ where reason = %reason, "guardrail blocked streaming /v1/messages passthrough response", ); + blocked = true; let frame = guardrail_block_frame(guardrail_name.as_deref()); yield Ok(Bytes::from(frame)); } } } + // Hold-back release (#932): the scan cleared — mask the held SSE + // bytes (channel reassembly across frames) and release them. + if hold_policy.is_some() && !blocked && !held.is_empty() { + match output_guardrail + .as_ref() + .and_then(|c| crate::redact::redact_anthropic_sse(c.as_ref(), &held)) + { + Some((rewritten, counts)) => { + crate::redact::merge_counts( + &mut guard.usage().redacted_entity_counts, + counts, + ); + yield Ok(Bytes::from(rewritten)); + } + None => { + yield Ok(Bytes::from(std::mem::take(&mut held))); + } + } + } // guard drops here → on_complete fires (delivery-gated). }; AnthropicDeliveryCounter { diff --git a/crates/aisix-proxy/src/redact.rs b/crates/aisix-proxy/src/redact.rs new file mode 100644 index 000000000..5aeb90a23 --- /dev/null +++ b/crates/aisix-proxy/src/redact.rs @@ -0,0 +1,1372 @@ +//! Application helpers for PII redaction (#932 / AISIX-Cloud#932). +//! +//! `aisix-guardrails` owns detection and the text→text rewrite +//! ([`Guardrail::redact_input_text`] / [`Guardrail::redact_output_text`]); +//! this module owns WHERE the rewrite is applied on each wire shape: +//! +//! - request side: the normalised [`ChatFormat`] (chat/completions), the +//! Anthropic-native `/v1/messages` body, the `/v1/responses` body, the +//! legacy completions `prompt`, and embeddings `input` — message text +//! only, mirroring the scan surface of `check_input`; +//! - response side: [`ChatResponse`] content + tool-call arguments, the +//! Anthropic-native response JSON, and buffered streamed chunks +//! (channel-reassembly: a masked span can cross chunk boundaries, so +//! each content channel is concatenated, rewritten once, and the full +//! rewritten text re-emitted on the channel's first chunk). +//! +//! Every helper returns per-detector match counts (detector names only, +//! never values) which callers merge into `usage_events +//! .redacted_entity_counts`. + +use std::collections::BTreeMap; + +use aisix_gateway::{ChatChunk, ChatFormat, ChatResponse}; +use aisix_guardrails::Guardrail; +use serde_json::Value; + +/// detector name → masked-span count. Mirrors +/// `UsageEvent::redacted_entity_counts`. +pub type RedactionCounts = BTreeMap; + +/// Merge `from` into `into` (repeated small helper; counts are tiny maps). +pub fn merge_counts(into: &mut RedactionCounts, from: RedactionCounts) { + for (k, v) in from { + *into.entry(k).or_insert(0) += v; + } +} + +/// Which side's redactor to run. The two sides can be configured +/// independently (`hook_point`), so every JSON-walking helper takes the +/// direction rather than hardcoding one. +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum Direction { + Input, + Output, +} + +fn redact_str( + chain: &dyn Guardrail, + dir: Direction, + text: &str, +) -> Option { + match dir { + Direction::Input => chain.redact_input_text(text), + Direction::Output => chain.redact_output_text(text), + } +} + +/// Rewrite one owned text field in place. No-op (and no allocation) when +/// nothing matches. +fn apply_to_string( + chain: &dyn Guardrail, + dir: Direction, + field: &mut String, + counts: &mut RedactionCounts, +) { + if field.is_empty() { + return; + } + if let Some(r) = redact_str(chain, dir, field) { + *field = r.text; + merge_counts(counts, r.counts); + } +} + +/// Rewrite a `Value::String` in place (helper for JSON-tree walking). +fn apply_to_value_string( + chain: &dyn Guardrail, + dir: Direction, + v: &mut Value, + counts: &mut RedactionCounts, +) { + if let Value::String(s) = v { + if !s.is_empty() { + if let Some(r) = redact_str(chain, dir, s) { + *s = r.text; + merge_counts(counts, r.counts); + } + } + } +} + +/// Recursively rewrite every string VALUE in a JSON tree (object values, +/// array elements). Keys and non-string scalars are untouched, so the +/// tree stays structurally valid — a phone number stored as a JSON number +/// is out of scope by design (rewriting it to a mask token would corrupt +/// the document). +pub fn redact_value_strings( + chain: &dyn Guardrail, + dir: Direction, + v: &mut Value, + counts: &mut RedactionCounts, +) { + match v { + Value::String(_) => apply_to_value_string(chain, dir, v, counts), + Value::Array(items) => { + for item in items { + redact_value_strings(chain, dir, item, counts); + } + } + Value::Object(map) => { + for (_, val) in map.iter_mut() { + redact_value_strings(chain, dir, val, counts); + } + } + _ => {} + } +} + +/// Rewrite a JSON-*encoded* string (OpenAI `function.arguments`): parse, +/// walk the string values, re-serialise — so a mask token can't corrupt +/// the embedded document (e.g. a phone number as a JSON number value +/// stays untouched rather than becoming invalid JSON). Falls back to a +/// raw text rewrite when the payload doesn't parse (a provider emitted +/// malformed/partial args — best effort beats leaking). +pub fn redact_json_encoded( + chain: &dyn Guardrail, + dir: Direction, + encoded: &mut String, + counts: &mut RedactionCounts, +) { + if encoded.is_empty() { + return; + } + match serde_json::from_str::(encoded) { + Ok(mut v) => { + let mut local = RedactionCounts::new(); + redact_value_strings(chain, dir, &mut v, &mut local); + if !local.is_empty() { + if let Ok(s) = serde_json::to_string(&v) { + *encoded = s; + merge_counts(counts, local); + } + } + } + Err(_) => apply_to_string(chain, dir, encoded, counts), + } +} + +// ─── Request side ──────────────────────────────────────────────────────────── + +/// Mask the request messages of a normalised [`ChatFormat`] in place: +/// the flat `content` string and the `text` field of typed content +/// blocks — the same surface `check_input` scans (`message_scan_text`). +/// Tool-call arguments replayed in history are covered too (they reach +/// the upstream verbatim). Returns the merged counts (empty = untouched). +pub fn redact_chat_format(chain: &dyn Guardrail, req: &mut ChatFormat) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_input() { + return counts; + } + for msg in &mut req.messages { + if let Some(content) = msg.content.as_mut() { + apply_to_string(chain, Direction::Input, content, &mut counts); + } + if let Some(blocks) = msg.content_blocks.as_mut() { + for block in blocks { + if block.get("type").and_then(Value::as_str) == Some("text") { + if let Some(text) = block.get_mut("text") { + apply_to_value_string(chain, Direction::Input, text, &mut counts); + } + } + } + } + // History-replay tool calls: arguments travel to the upstream + // verbatim through `extra`, so mask them like fresh content. + if let Some(tool_calls) = msg.extra.get_mut("tool_calls") { + redact_tool_call_arguments(chain, Direction::Input, tool_calls, &mut counts); + } + } + counts +} + +/// Mask `function.arguments` (JSON-encoded string) on each element of an +/// OpenAI-shaped `tool_calls` array. Names/ids are structural, not +/// content, and stay untouched. +fn redact_tool_call_arguments( + chain: &dyn Guardrail, + dir: Direction, + tool_calls: &mut Value, + counts: &mut RedactionCounts, +) { + let Some(items) = tool_calls.as_array_mut() else { + return; + }; + for tc in items { + if let Some(Value::String(s)) = tc.get_mut("function").and_then(|f| f.get_mut("arguments")) + { + let mut owned = std::mem::take(s); + redact_json_encoded(chain, dir, &mut owned, counts); + *s = owned; + } + } +} + +/// Mask an Anthropic-native `/v1/messages` request body in place: +/// `system` (string or text blocks) and `messages[].content` (string or +/// blocks — `text` blocks and nested `tool_result` content). `tool_use` +/// input objects in history are walked as JSON strings. +pub fn redact_anthropic_request(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_input() { + return counts; + } + if let Some(system) = body.get_mut("system") { + redact_anthropic_content(chain, Direction::Input, system, &mut counts); + } + if let Some(messages) = body.get_mut("messages").and_then(Value::as_array_mut) { + for msg in messages { + if let Some(content) = msg.get_mut("content") { + redact_anthropic_content(chain, Direction::Input, content, &mut counts); + } + } + } + counts +} + +/// Anthropic `content` is either a bare string or an array of typed +/// blocks. Rewrites `text` blocks, `tool_result` nested content, and +/// `tool_use` input objects; leaves image/document blocks alone. +fn redact_anthropic_content( + chain: &dyn Guardrail, + dir: Direction, + content: &mut Value, + counts: &mut RedactionCounts, +) { + match content { + Value::String(_) => apply_to_value_string(chain, dir, content, counts), + Value::Array(blocks) => { + for block in blocks { + match block.get("type").and_then(Value::as_str) { + Some("text") => { + if let Some(text) = block.get_mut("text") { + apply_to_value_string(chain, dir, text, counts); + } + } + Some("tool_result") => { + if let Some(inner) = block.get_mut("content") { + redact_anthropic_content(chain, dir, inner, counts); + } + } + Some("tool_use") => { + if let Some(input) = block.get_mut("input") { + redact_value_strings(chain, dir, input, counts); + } + } + _ => {} + } + } + } + _ => {} + } +} + +/// Mask an Anthropic-native `/v1/messages` RESPONSE body in place (the +/// non-streaming passthrough JSON): top-level `content` blocks (`text` + +/// `tool_use` input). +pub fn redact_anthropic_response(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_output() { + return counts; + } + if let Some(content) = body.get_mut("content") { + redact_anthropic_content(chain, Direction::Output, content, &mut counts); + } + counts +} + +/// Mask a `/v1/responses` request body in place: `instructions` and +/// `input` (bare string, or item list whose `message` items carry +/// `content` as a string or `input_text` parts). Function-call outputs +/// replayed as `function_call_output` items are walked too. +pub fn redact_responses_request(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_input() { + return counts; + } + if let Some(instructions) = body.get_mut("instructions") { + apply_to_value_string(chain, Direction::Input, instructions, &mut counts); + } + match body.get_mut("input") { + Some(v @ Value::String(_)) => { + apply_to_value_string(chain, Direction::Input, v, &mut counts) + } + Some(Value::Array(items)) => { + for item in items { + redact_responses_item(chain, Direction::Input, item, &mut counts); + } + } + _ => {} + } + counts +} + +/// One `/v1/responses` input/output item. `message` items carry +/// string-or-parts content (`input_text` / `output_text` / plain `text`); +/// `function_call` carries JSON-encoded `arguments`; +/// `function_call_output` carries a string `output`. +fn redact_responses_item( + chain: &dyn Guardrail, + dir: Direction, + item: &mut Value, + counts: &mut RedactionCounts, +) { + match item.get("type").and_then(Value::as_str) { + // An item without a `type` defaults to `message` on this API. + Some("message") | None => match item.get_mut("content") { + Some(v @ Value::String(_)) => apply_to_value_string(chain, dir, v, counts), + Some(Value::Array(parts)) => { + for part in parts { + if matches!( + part.get("type").and_then(Value::as_str), + Some("input_text") | Some("output_text") | Some("text") + ) { + if let Some(text) = part.get_mut("text") { + apply_to_value_string(chain, dir, text, counts); + } + } + } + } + _ => {} + }, + Some("function_call") => { + if let Some(Value::String(args)) = item.get_mut("arguments") { + let mut owned = std::mem::take(args); + redact_json_encoded(chain, dir, &mut owned, counts); + *args = owned; + } + } + Some("function_call_output") => { + if let Some(output) = item.get_mut("output") { + apply_to_value_string(chain, dir, output, counts); + } + } + _ => {} + } +} + +/// Mask a `/v1/responses` non-streaming RESPONSE body in place: every +/// item in `output` (message `output_text` parts, `function_call` +/// arguments) — the same surface the output check scans. +pub fn redact_responses_response(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_output() { + return counts; + } + if let Some(Value::Array(items)) = body.get_mut("output") { + for item in items { + redact_responses_item(chain, Direction::Output, item, &mut counts); + } + } + counts +} + +/// Mask a legacy `/v1/completions` request body in place: `prompt` as a +/// bare string or an array of strings (token-id arrays carry no text). +pub fn redact_completions_request(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_input() { + return counts; + } + match body.get_mut("prompt") { + Some(v @ Value::String(_)) => { + apply_to_value_string(chain, Direction::Input, v, &mut counts) + } + Some(Value::Array(items)) => { + for item in items { + if item.is_string() { + apply_to_value_string(chain, Direction::Input, item, &mut counts); + } + } + } + _ => {} + } + counts +} + +/// Mask a legacy `/v1/completions` RESPONSE body in place: `choices[].text`. +pub fn redact_completions_response(chain: &dyn Guardrail, body: &mut Value) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_output() { + return counts; + } + if let Some(Value::Array(choices)) = body.get_mut("choices") { + for choice in choices { + if let Some(text) = choice.get_mut("text") { + apply_to_value_string(chain, Direction::Output, text, &mut counts); + } + } + } + counts +} + +// ─── Response side (non-streaming) ─────────────────────────────────────────── + +/// Mask a normalised [`ChatResponse`] in place: assistant `content` plus +/// `tool_calls` function arguments (the same surface +/// `guardrail_output_text` scans). Reasoning content is excluded from +/// guardrail scope by design and stays untouched. +pub fn redact_chat_response(chain: &dyn Guardrail, resp: &mut ChatResponse) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_output() { + return counts; + } + if let Some(content) = resp.message.content.as_mut() { + apply_to_string(chain, Direction::Output, content, &mut counts); + } + if let Some(tool_calls) = resp.message.extra.get_mut("tool_calls") { + redact_tool_call_arguments(chain, Direction::Output, tool_calls, &mut counts); + } + counts +} + +// ─── Response side (streamed, buffered) ────────────────────────────────────── + +/// Mask a fully-buffered stream of normalised [`ChatChunk`]s in place — +/// the hold-back release path (BufferFull), where the whole response is +/// available before any byte reaches the wire. +/// +/// A masked span can cross chunk boundaries, so per-chunk rewriting would +/// miss it. Instead each content channel (delta content, and each +/// tool-call's streamed `arguments`) is concatenated across the buffered +/// chunks, rewritten once, and the FULL rewritten text re-emitted on the +/// channel's first carrying chunk; later chunks in that channel become +/// empty deltas. The stream is already released en bloc at this point, so +/// chunk-size distribution is not client-observable. Non-content fields +/// (ids, usage, finish_reason, reasoning) are untouched. +pub fn redact_chat_chunks(chain: &dyn Guardrail, chunks: &mut [ChatChunk]) -> RedactionCounts { + let mut counts = RedactionCounts::new(); + if !chain.redacts_output() { + return counts; + } + + // Content channel: all chunks stream one assistant message. + let content_sites: Vec = chunks + .iter() + .enumerate() + .filter(|(_, c)| c.delta.content.as_deref().is_some_and(|t| !t.is_empty())) + .map(|(i, _)| i) + .collect(); + if !content_sites.is_empty() { + let joined: String = content_sites + .iter() + .map(|&i| chunks[i].delta.content.as_deref().unwrap_or("")) + .collect(); + if let Some(r) = chain.redact_output_text(&joined) { + let mut first = true; + for &i in &content_sites { + chunks[i].delta.content = Some(if first { + first = false; + r.text.clone() + } else { + String::new() + }); + } + merge_counts(&mut counts, r.counts); + } + } + + // Tool-call channels: fragments carry an `index` discriminator; the + // concatenation of each channel's `function.arguments` strings is the + // complete JSON-encoded argument document. + let mut channels: BTreeMap> = BTreeMap::new(); + for (ci, chunk) in chunks.iter().enumerate() { + if let Some(tcs) = chunk.delta.tool_calls.as_ref() { + for (ti, tc) in tcs.iter().enumerate() { + let idx = tc.get("index").and_then(Value::as_u64).unwrap_or(0); + if tc + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(Value::as_str) + .is_some_and(|s| !s.is_empty()) + { + channels.entry(idx).or_default().push((ci, ti)); + } + } + } + } + for sites in channels.values() { + let joined: String = sites + .iter() + .map(|&(ci, ti)| { + chunks[ci].delta.tool_calls.as_ref().unwrap()[ti] + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(Value::as_str) + .unwrap_or("") + }) + .collect(); + let mut rewritten = joined.clone(); + let mut local = RedactionCounts::new(); + redact_json_encoded(chain, Direction::Output, &mut rewritten, &mut local); + if local.is_empty() { + continue; + } + let mut first = true; + for &(ci, ti) in sites { + let args = chunks[ci].delta.tool_calls.as_mut().unwrap()[ti] + .get_mut("function") + .and_then(|f| f.get_mut("arguments")) + .expect("site was selected for having arguments"); + *args = Value::String(if first { + first = false; + rewritten.clone() + } else { + String::new() + }); + } + merge_counts(&mut counts, local); + } + + counts +} + +// ─── Anthropic-native SSE (passthrough) rewrite ────────────────────────────── + +/// One parsed SSE frame from a buffered Anthropic-native byte stream. +struct SseFrame { + /// Original frame bytes (no trailing separator). Emitted verbatim + /// unless `data` was modified. + raw: Vec, + /// Parsed `data:` payload, when the frame carries one. + data: Option, + dirty: bool, +} + +impl SseFrame { + /// Re-render the frame: the first `data:` line is replaced with the + /// re-serialised payload (subsequent `data:` lines dropped); every + /// other line passes through untouched. + fn render(&self) -> Vec { + if !self.dirty { + return self.raw.clone(); + } + let Some(data) = self.data.as_ref() else { + return self.raw.clone(); + }; + let text = String::from_utf8_lossy(&self.raw); + let mut out = String::new(); + let mut data_written = false; + for line in text.split('\n') { + if line.starts_with("data:") { + if !data_written { + out.push_str("data: "); + out.push_str(&serde_json::to_string(data).unwrap_or_default()); + out.push('\n'); + data_written = true; + } + } else { + out.push_str(line); + out.push('\n'); + } + } + // Drop the final artificial newline added by the loop; the caller + // re-adds the frame separator. + if out.ends_with('\n') { + out.pop(); + } + out.into_bytes() + } +} + +/// Split a buffered SSE byte stream into frames on the blank-line +/// separator. Returns `(frames, trailing)` where `trailing` is a +/// partial frame with no terminator yet (forwarded verbatim). +fn split_sse_frames(raw: &[u8]) -> (Vec, &[u8]) { + let mut frames = Vec::new(); + let mut start = 0usize; + let mut i = 0usize; + while i + 1 < raw.len() { + if raw[i] == b'\n' && raw[i + 1] == b'\n' { + let frame_raw = &raw[start..i]; + let data = String::from_utf8_lossy(frame_raw) + .split('\n') + .find(|l| l.starts_with("data:")) + .and_then(|l| serde_json::from_str::(l["data:".len()..].trim()).ok()); + frames.push(SseFrame { + raw: frame_raw.to_vec(), + data, + dirty: false, + }); + start = i + 2; + i += 2; + } else { + i += 1; + } + } + (frames, &raw[start..]) +} + +/// Mask a fully-buffered Anthropic-native SSE response (the `/v1/messages` +/// passthrough hold-back). Text deltas are reassembled per content-block +/// `index` (a masked span can cross frame boundaries), masked once, and +/// the full masked text re-emitted on the channel's first frame; +/// `input_json_delta` (tool-use arguments) channels are masked as complete +/// JSON documents. `None` = nothing matched, forward the original bytes +/// byte-identical. +pub fn redact_anthropic_sse( + chain: &dyn Guardrail, + raw: &[u8], +) -> Option<(Vec, RedactionCounts)> { + if !chain.redacts_output() { + return None; + } + let (mut frames, trailing) = split_sse_frames(raw); + let mut counts = RedactionCounts::new(); + + // channel key → ordered (frame_idx, kind) sites. Kind distinguishes the + // JSON path to rewrite inside the frame payload. + #[derive(Clone, Copy)] + enum Site { + DeltaText, + DeltaPartialJson, + BlockStartText, + } + let mut text_channels: BTreeMap> = BTreeMap::new(); + let mut json_channels: BTreeMap> = BTreeMap::new(); + + for (fi, frame) in frames.iter().enumerate() { + let Some(data) = frame.data.as_ref() else { + continue; + }; + let index = data.get("index").and_then(Value::as_u64).unwrap_or(0); + match data.get("type").and_then(Value::as_str) { + Some("content_block_delta") => { + match data + .get("delta") + .and_then(|d| d.get("type")) + .and_then(Value::as_str) + { + Some("text_delta") => { + if data + .get("delta") + .and_then(|d| d.get("text")) + .and_then(Value::as_str) + .is_some_and(|t| !t.is_empty()) + { + text_channels + .entry(index) + .or_default() + .push((fi, Site::DeltaText)); + } + } + Some("input_json_delta") => { + if data + .get("delta") + .and_then(|d| d.get("partial_json")) + .and_then(Value::as_str) + .is_some_and(|t| !t.is_empty()) + { + json_channels + .entry(index) + .or_default() + .push((fi, Site::DeltaPartialJson)); + } + } + _ => {} + } + } + Some("content_block_start") => { + // A `text` block may open with non-empty initial text; it + // belongs at the head of the same channel as its deltas. + if data + .get("content_block") + .and_then(|b| b.get("text")) + .and_then(Value::as_str) + .is_some_and(|t| !t.is_empty()) + { + text_channels + .entry(index) + .or_default() + .push((fi, Site::BlockStartText)); + } + } + _ => {} + } + } + + fn site_text(data: &Value, site: Site) -> &str { + let path = match site { + Site::DeltaText => data.get("delta").and_then(|d| d.get("text")), + Site::DeltaPartialJson => data.get("delta").and_then(|d| d.get("partial_json")), + Site::BlockStartText => data.get("content_block").and_then(|b| b.get("text")), + }; + path.and_then(Value::as_str).unwrap_or("") + } + + fn site_slot(data: &mut Value, site: Site) -> Option<&mut Value> { + match site { + Site::DeltaText => data.get_mut("delta").and_then(|d| d.get_mut("text")), + Site::DeltaPartialJson => data + .get_mut("delta") + .and_then(|d| d.get_mut("partial_json")), + Site::BlockStartText => data + .get_mut("content_block") + .and_then(|b| b.get_mut("text")), + } + } + + let rewrite = |frames: &mut Vec, sites: &[(usize, Site)], new_text: String| { + let mut first = true; + for &(fi, site) in sites { + let frame = &mut frames[fi]; + if let Some(slot) = frame.data.as_mut().and_then(|d| site_slot(d, site)) { + *slot = Value::String(if first { + first = false; + new_text.clone() + } else { + String::new() + }); + frame.dirty = true; + } + } + }; + + for sites in text_channels.values() { + let joined: String = sites + .iter() + .map(|&(fi, site)| site_text(frames[fi].data.as_ref().unwrap(), site)) + .collect(); + if let Some(r) = chain.redact_output_text(&joined) { + rewrite(&mut frames, sites, r.text); + merge_counts(&mut counts, r.counts); + } + } + for sites in json_channels.values() { + let joined: String = sites + .iter() + .map(|&(fi, site)| site_text(frames[fi].data.as_ref().unwrap(), site)) + .collect(); + let mut rewritten = joined.clone(); + let mut local = RedactionCounts::new(); + redact_json_encoded(chain, Direction::Output, &mut rewritten, &mut local); + if !local.is_empty() { + rewrite(&mut frames, sites, rewritten); + merge_counts(&mut counts, local); + } + } + + if counts.is_empty() { + return None; + } + let mut out = Vec::with_capacity(raw.len()); + for frame in &frames { + out.extend_from_slice(&frame.render()); + out.extend_from_slice(b"\n\n"); + } + out.extend_from_slice(trailing); + Some((out, counts)) +} + +// ─── Responses-API SSE rewrite ─────────────────────────────────────────────── + +/// Mask a fully-buffered Responses-API SSE byte stream (the `/v1/responses` +/// verbatim hold-back and the cross-provider bridge release). Delta events +/// are reassembled per channel (`output_text.delta` by item, `function_call +/// _arguments.delta` by item), masked once, and re-emitted on the channel's +/// first frame; the aggregate events (`*.done`, `output_item.done`, +/// `response.completed`) carry complete texts and are masked directly — +/// deterministic masking keeps them consistent with the delta channels. +/// `None` = nothing matched, forward the original bytes byte-identical. +pub fn redact_responses_sse( + chain: &dyn Guardrail, + raw: &[u8], +) -> Option<(Vec, RedactionCounts)> { + if !chain.redacts_output() { + return None; + } + let (mut frames, trailing) = split_sse_frames(raw); + let mut counts = RedactionCounts::new(); + + // Delta channels: (event-type discriminant, channel key) → frame sites. + let mut text_channels: BTreeMap> = BTreeMap::new(); + let mut args_channels: BTreeMap> = BTreeMap::new(); + + fn channel_key(data: &Value) -> String { + // item_id is the stable discriminator; fall back to output_index + + // content_index for encoders that omit it. + match data.get("item_id").and_then(Value::as_str) { + Some(id) => format!( + "{id}/{}", + data.get("content_index") + .and_then(Value::as_u64) + .unwrap_or(0) + ), + None => format!( + "{}/{}", + data.get("output_index") + .and_then(Value::as_u64) + .unwrap_or(0), + data.get("content_index") + .and_then(Value::as_u64) + .unwrap_or(0) + ), + } + } + + for (fi, frame) in frames.iter().enumerate() { + let Some(data) = frame.data.as_ref() else { + continue; + }; + match data.get("type").and_then(Value::as_str) { + Some("response.output_text.delta") => { + if data + .get("delta") + .and_then(Value::as_str) + .is_some_and(|t| !t.is_empty()) + { + text_channels.entry(channel_key(data)).or_default().push(fi); + } + } + Some("response.function_call_arguments.delta") => { + if data + .get("delta") + .and_then(Value::as_str) + .is_some_and(|t| !t.is_empty()) + { + args_channels.entry(channel_key(data)).or_default().push(fi); + } + } + _ => {} + } + } + + // Rewrite the delta channels (first frame gets the full masked text). + let rewrite_channel = |frames: &mut Vec, sites: &[usize], new_text: String| { + let mut first = true; + for &fi in sites { + let frame = &mut frames[fi]; + if let Some(slot) = frame.data.as_mut().and_then(|d| d.get_mut("delta")) { + *slot = Value::String(if first { + first = false; + new_text.clone() + } else { + String::new() + }); + frame.dirty = true; + } + } + }; + for sites in text_channels.values() { + let joined: String = sites + .iter() + .map(|&fi| { + frames[fi] + .data + .as_ref() + .and_then(|d| d.get("delta")) + .and_then(Value::as_str) + .unwrap_or("") + }) + .collect(); + if let Some(r) = chain.redact_output_text(&joined) { + rewrite_channel(&mut frames, sites, r.text); + merge_counts(&mut counts, r.counts); + } + } + for sites in args_channels.values() { + let joined: String = sites + .iter() + .map(|&fi| { + frames[fi] + .data + .as_ref() + .and_then(|d| d.get("delta")) + .and_then(Value::as_str) + .unwrap_or("") + }) + .collect(); + let mut rewritten = joined.clone(); + let mut local = RedactionCounts::new(); + redact_json_encoded(chain, Direction::Output, &mut rewritten, &mut local); + if !local.is_empty() { + rewrite_channel(&mut frames, sites, rewritten); + merge_counts(&mut counts, local); + } + } + + // Aggregate events carry complete texts — mask them in place. Their + // counts are NOT merged into the totals: they duplicate the delta + // channels' matches (the audit count is per span served, not per + // wire occurrence). Only count them when the delta channel was absent + // (e.g. a `.done`-only encoder). + for frame in frames.iter_mut() { + let Some(data) = frame.data.as_mut() else { + continue; + }; + let ty = data + .get("type") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let mut local = RedactionCounts::new(); + match ty.as_str() { + "response.output_text.done" => { + if let Some(text) = data.get_mut("text") { + apply_to_value_string(chain, Direction::Output, text, &mut local); + } + } + "response.content_part.done" => { + if let Some(text) = data.get_mut("part").and_then(|p| p.get_mut("text")) { + apply_to_value_string(chain, Direction::Output, text, &mut local); + } + } + "response.function_call_arguments.done" => { + if let Some(Value::String(args)) = data.get_mut("arguments") { + let mut owned = std::mem::take(args); + redact_json_encoded(chain, Direction::Output, &mut owned, &mut local); + *args = owned; + } + } + "response.output_item.done" => { + if let Some(item) = data.get_mut("item") { + redact_responses_item(chain, Direction::Output, item, &mut local); + } + } + "response.completed" | "response.incomplete" | "response.failed" => { + if let Some(Value::Array(items)) = + data.get_mut("response").and_then(|r| r.get_mut("output")) + { + for item in items { + redact_responses_item(chain, Direction::Output, item, &mut local); + } + } + } + _ => {} + } + if !local.is_empty() { + frame.dirty = true; + if counts.is_empty() { + merge_counts(&mut counts, local); + } + } + } + + let any_dirty = frames.iter().any(|f| f.dirty); + if !any_dirty { + return None; + } + let mut out = Vec::with_capacity(raw.len()); + for frame in &frames { + out.extend_from_slice(&frame.render()); + out.extend_from_slice(b"\n\n"); + } + out.extend_from_slice(trailing); + Some((out, counts)) +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_gateway::{ChatDelta, ChatMessage}; + use aisix_guardrails::{builtin_rule, GuardrailChain, PiiAction, PiiGuardrail}; + use serde_json::json; + use std::sync::Arc; + + fn mask_chain(hook: aisix_core::models::GuardrailHookPoint) -> Arc { + let g = PiiGuardrail::new( + vec![ + builtin_rule("email", PiiAction::Mask).unwrap(), + builtin_rule("china_mobile", PiiAction::Mask).unwrap(), + ], + hook, + 262_144, + false, + ); + Arc::new(GuardrailChain::new(vec![Arc::new(g)])) + } + + fn both() -> Arc { + mask_chain(aisix_core::models::GuardrailHookPoint::Both) + } + + #[test] + fn chat_format_masks_content_blocks_and_history_tool_args() { + let chain = both(); + let mut req: ChatFormat = serde_json::from_value(json!({ + "model": "m", + "messages": [ + {"role": "user", "content": "mail a@x.com"}, + {"role": "user", "content": "", "content_blocks": [ + {"type": "text", "text": "call 13800138000"}, + {"type": "image_url", "image_url": {"url": "http://x"}} + ]}, + {"role": "assistant", "content": null, "tool_calls": [ + {"index": 0, "function": {"name": "send", "arguments": "{\"to\":\"b@y.org\"}"}} + ]} + ] + })) + .unwrap(); + let counts = redact_chat_format(chain.as_ref(), &mut req); + assert_eq!( + req.messages[0].content.as_deref(), + Some("mail [EMAIL_REDACTED]") + ); + let blocks = req.messages[1].content_blocks.as_ref().unwrap(); + assert_eq!( + blocks[0].get("text").unwrap().as_str().unwrap(), + "call [CHINA_MOBILE_REDACTED]", + ); + let args = req.messages[2].extra["tool_calls"][0]["function"]["arguments"] + .as_str() + .unwrap(); + assert_eq!(args, "{\"to\":\"[EMAIL_REDACTED]\"}"); + assert_eq!(counts.get("email"), Some(&2)); + assert_eq!(counts.get("china_mobile"), Some(&1)); + } + + #[test] + fn input_only_chain_skips_output_and_vice_versa() { + let input_only = mask_chain(aisix_core::models::GuardrailHookPoint::Input); + let mut resp = ChatResponse { + id: "r".into(), + model: "m".into(), + message: ChatMessage::assistant("mail a@x.com"), + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::new(0, 0), + }; + assert!(redact_chat_response(input_only.as_ref(), &mut resp).is_empty()); + assert_eq!(resp.message.content.as_deref(), Some("mail a@x.com")); + + let output_only = mask_chain(aisix_core::models::GuardrailHookPoint::Output); + let mut req: ChatFormat = serde_json::from_value(json!({ + "model": "m", + "messages": [{"role": "user", "content": "mail a@x.com"}] + })) + .unwrap(); + assert!(redact_chat_format(output_only.as_ref(), &mut req).is_empty()); + assert_eq!(req.messages[0].content.as_deref(), Some("mail a@x.com")); + } + + #[test] + fn chat_response_masks_content_and_tool_args_json_safely() { + let chain = both(); + let mut msg = ChatMessage::assistant("reach me at a@x.com"); + msg.extra.insert( + "tool_calls".into(), + json!([{ + "id": "call_1", "type": "function", + // A number-typed phone stays untouched (JSON preserved); + // the string email is masked. + "function": {"name": "f", "arguments": "{\"phone\":13800138000,\"mail\":\"b@y.org\"}"} + }]), + ); + let mut resp = ChatResponse { + id: "r".into(), + model: "m".into(), + message: msg, + finish_reason: aisix_gateway::FinishReason::Stop, + usage: aisix_gateway::UsageStats::new(0, 0), + }; + let counts = redact_chat_response(chain.as_ref(), &mut resp); + assert_eq!( + resp.message.content.as_deref(), + Some("reach me at [EMAIL_REDACTED]") + ); + let args = resp.message.extra["tool_calls"][0]["function"]["arguments"] + .as_str() + .unwrap(); + let parsed: Value = serde_json::from_str(args).expect("args stay valid JSON"); + assert_eq!(parsed["phone"], json!(13800138000u64)); + assert_eq!(parsed["mail"], json!("[EMAIL_REDACTED]")); + assert_eq!(counts.get("email"), Some(&2)); + } + + #[test] + fn anthropic_request_masks_system_text_blocks_and_tool_result() { + let chain = both(); + let mut body = json!({ + "model": "claude", + "system": [{"type": "text", "text": "user email a@x.com"}], + "messages": [ + {"role": "user", "content": "call 13800138000"}, + {"role": "user", "content": [ + {"type": "text", "text": "and b@y.org"}, + {"type": "tool_result", "tool_use_id": "t1", "content": [ + {"type": "text", "text": "result c@z.io"} + ]} + ]} + ] + }); + let counts = redact_anthropic_request(chain.as_ref(), &mut body); + assert_eq!(body["system"][0]["text"], "user email [EMAIL_REDACTED]"); + assert_eq!( + body["messages"][0]["content"], + "call [CHINA_MOBILE_REDACTED]" + ); + assert_eq!( + body["messages"][1]["content"][0]["text"], + "and [EMAIL_REDACTED]" + ); + assert_eq!( + body["messages"][1]["content"][1]["content"][0]["text"], + "result [EMAIL_REDACTED]", + ); + assert_eq!(counts.get("email"), Some(&3)); + } + + #[test] + fn anthropic_response_masks_text_and_tool_use_input() { + let chain = both(); + let mut body = json!({ + "content": [ + {"type": "text", "text": "email a@x.com"}, + {"type": "tool_use", "id": "t", "name": "send", + "input": {"to": "b@y.org", "count": 3}} + ] + }); + let counts = redact_anthropic_response(chain.as_ref(), &mut body); + assert_eq!(body["content"][0]["text"], "email [EMAIL_REDACTED]"); + assert_eq!(body["content"][1]["input"]["to"], "[EMAIL_REDACTED]"); + assert_eq!(body["content"][1]["input"]["count"], 3); + assert_eq!(counts.get("email"), Some(&2)); + } + + #[test] + fn responses_request_masks_string_and_item_forms() { + let chain = both(); + let mut body = json!({ + "model": "m", + "instructions": "never leak a@x.com", + "input": [ + {"type": "message", "role": "user", "content": "call 13800138000"}, + {"role": "user", "content": [ + {"type": "input_text", "text": "mail b@y.org"} + ]}, + {"type": "function_call_output", "call_id": "c", "output": "from c@z.io"} + ] + }); + let counts = redact_responses_request(chain.as_ref(), &mut body); + assert_eq!(body["instructions"], "never leak [EMAIL_REDACTED]"); + assert_eq!(body["input"][0]["content"], "call [CHINA_MOBILE_REDACTED]"); + assert_eq!( + body["input"][1]["content"][0]["text"], + "mail [EMAIL_REDACTED]" + ); + assert_eq!(body["input"][2]["output"], "from [EMAIL_REDACTED]"); + assert_eq!(counts.get("email"), Some(&3)); + + let mut simple = json!({"model": "m", "input": "mail a@x.com"}); + redact_responses_request(chain.as_ref(), &mut simple); + assert_eq!(simple["input"], "mail [EMAIL_REDACTED]"); + } + + fn content_chunk(text: &str) -> ChatChunk { + ChatChunk { + id: "c".into(), + model: "m".into(), + delta: ChatDelta { + content: Some(text.to_string()), + ..ChatDelta::default() + }, + finish_reason: None, + usage: None, + } + } + + #[test] + fn stream_chunks_mask_span_split_across_chunk_boundary() { + let chain = both(); + // "a@x.com" split across three chunks — per-chunk masking would miss it. + let mut chunks = vec![ + content_chunk("mail a@"), + content_chunk("x.c"), + content_chunk("om now"), + ]; + let counts = redact_chat_chunks(chain.as_ref(), &mut chunks); + assert_eq!(counts.get("email"), Some(&1)); + let reassembled: String = chunks + .iter() + .map(|c| c.delta.content.as_deref().unwrap_or("")) + .collect(); + assert_eq!(reassembled, "mail [EMAIL_REDACTED] now"); + // Full text lands on the first content chunk; the rest are empty. + assert_eq!( + chunks[0].delta.content.as_deref(), + Some("mail [EMAIL_REDACTED] now") + ); + assert_eq!(chunks[1].delta.content.as_deref(), Some("")); + } + + #[test] + fn stream_chunks_mask_tool_call_arguments_channel() { + let chain = both(); + let mut chunks = vec![ + ChatChunk { + id: "c".into(), + model: "m".into(), + delta: ChatDelta { + tool_calls: Some(vec![json!({ + "index": 0, "id": "call_1", "type": "function", + "function": {"name": "send", "arguments": "{\"to\":\"a@"} + })]), + ..ChatDelta::default() + }, + finish_reason: None, + usage: None, + }, + ChatChunk { + id: "c".into(), + model: "m".into(), + delta: ChatDelta { + tool_calls: Some(vec![json!({ + "index": 0, + "function": {"arguments": "x.com\"}"} + })]), + ..ChatDelta::default() + }, + finish_reason: None, + usage: None, + }, + ]; + let counts = redact_chat_chunks(chain.as_ref(), &mut chunks); + assert_eq!(counts.get("email"), Some(&1)); + let first_args = chunks[0].delta.tool_calls.as_ref().unwrap()[0]["function"]["arguments"] + .as_str() + .unwrap() + .to_string(); + let second_args = chunks[1].delta.tool_calls.as_ref().unwrap()[0]["function"]["arguments"] + .as_str() + .unwrap(); + assert_eq!(first_args, "{\"to\":\"[EMAIL_REDACTED]\"}"); + assert_eq!(second_args, ""); + } + + #[test] + fn stream_chunks_untouched_when_nothing_matches() { + let chain = both(); + let mut chunks = vec![content_chunk("hello "), content_chunk("world")]; + assert!(redact_chat_chunks(chain.as_ref(), &mut chunks).is_empty()); + assert_eq!(chunks[0].delta.content.as_deref(), Some("hello ")); + assert_eq!(chunks[1].delta.content.as_deref(), Some("world")); + } + + #[test] + fn anthropic_sse_masks_text_delta_across_frames() { + let chain = both(); + let raw = concat!( + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":3}}}\n\n", + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"mail a@\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"x.com ok\"}}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ); + let (out, counts) = redact_anthropic_sse(chain.as_ref(), raw.as_bytes()).unwrap(); + let out = String::from_utf8(out).unwrap(); + assert!(out.contains("mail [EMAIL_REDACTED] ok"), "out: {out}"); + assert!(!out.contains("a@x.com")); + // Second delta emptied; frame structure + unrelated frames intact. + assert!( + out.contains("{\"type\":\"text_delta\",\"text\":\"\"}") + || out.contains("\"text\":\"\"") + ); + assert!(out.contains("message_start")); + assert!(out.contains("message_stop")); + assert_eq!(counts.get("email"), Some(&1)); + } + + #[test] + fn anthropic_sse_masks_tool_use_input_json_channel() { + let chain = both(); + let raw = concat!( + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"id\":\"t\",\"name\":\"send\",\"input\":{}}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"to\\\":\\\"a@\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"x.com\\\"}\"}}\n\n", + ); + let (out, counts) = redact_anthropic_sse(chain.as_ref(), raw.as_bytes()).unwrap(); + let out = String::from_utf8(out).unwrap(); + assert!(out.contains("[EMAIL_REDACTED]"), "out: {out}"); + assert!(!out.contains("a@"), "no split original fragments: {out}"); + assert_eq!(counts.get("email"), Some(&1)); + } + + #[test] + fn responses_sse_masks_delta_channel_and_aggregate_events() { + let chain = both(); + let raw = concat!( + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"delta\":\"mail a@\"}\n\n", + "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"output_index\":0,\"content_index\":0,\"delta\":\"x.com ok\"}\n\n", + "event: response.output_text.done\ndata: {\"type\":\"response.output_text.done\",\"item_id\":\"msg_1\",\"text\":\"mail a@x.com ok\"}\n\n", + "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"mail a@x.com ok\"}]}]}}\n\n", + ); + let (out, counts) = redact_responses_sse(chain.as_ref(), raw.as_bytes()).unwrap(); + let out = String::from_utf8(out).unwrap(); + assert!(!out.contains("a@x.com"), "original must be gone: {out}"); + // Delta channel: full masked text on the first delta; done + + // completed events masked consistently. + assert!( + out.contains("\"delta\":\"mail [EMAIL_REDACTED] ok\""), + "out: {out}" + ); + assert!( + out.contains("\"text\":\"mail [EMAIL_REDACTED] ok\""), + "out: {out}" + ); + // Aggregate events don't double-count the same span. + assert_eq!(counts.get("email"), Some(&1)); + } + + #[test] + fn responses_sse_masks_function_call_args_channel() { + let chain = both(); + let raw = concat!( + "event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"{\\\"to\\\":\\\"a@\"}\n\n", + "event: response.function_call_arguments.delta\ndata: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"x.com\\\"}\"}\n\n", + "event: response.function_call_arguments.done\ndata: {\"type\":\"response.function_call_arguments.done\",\"item_id\":\"fc_1\",\"arguments\":\"{\\\"to\\\":\\\"a@x.com\\\"}\"}\n\n", + ); + let (out, counts) = redact_responses_sse(chain.as_ref(), raw.as_bytes()).unwrap(); + let out = String::from_utf8(out).unwrap(); + assert!(!out.contains("a@"), "original fragments gone: {out}"); + assert!(out.contains("[EMAIL_REDACTED]"), "out: {out}"); + assert_eq!(counts.get("email"), Some(&1)); + } + + #[test] + fn responses_sse_returns_none_when_clean() { + let chain = both(); + let raw = "event: response.output_text.delta\ndata: {\"type\":\"response.output_text.delta\",\"item_id\":\"m\",\"delta\":\"hello\"}\n\n"; + assert!(redact_responses_sse(chain.as_ref(), raw.as_bytes()).is_none()); + } + + #[test] + fn responses_response_json_masks_output_items() { + let chain = both(); + let mut body = json!({ + "id": "resp_1", + "output": [ + {"type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": "mail a@x.com"} + ]}, + {"type": "function_call", "call_id": "c", "name": "send", + "arguments": "{\"to\":\"b@y.org\"}"} + ] + }); + let counts = redact_responses_response(chain.as_ref(), &mut body); + assert_eq!( + body["output"][0]["content"][0]["text"], + "mail [EMAIL_REDACTED]" + ); + assert_eq!( + body["output"][1]["arguments"], + "{\"to\":\"[EMAIL_REDACTED]\"}" + ); + assert_eq!(counts.get("email"), Some(&2)); + } + + #[test] + fn anthropic_sse_returns_none_when_clean() { + let chain = both(); + let raw = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n"; + assert!(redact_anthropic_sse(chain.as_ref(), raw.as_bytes()).is_none()); + } + + #[test] + fn malformed_tool_args_fall_back_to_raw_text_masking() { + let chain = both(); + let mut encoded = String::from("not json but has a@x.com inside"); + let mut counts = RedactionCounts::new(); + redact_json_encoded(chain.as_ref(), Direction::Output, &mut encoded, &mut counts); + assert_eq!(encoded, "not json but has [EMAIL_REDACTED] inside"); + assert_eq!(counts.get("email"), Some(&1)); + } +} diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 480eabfe9..06adeeef9 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -71,6 +71,11 @@ struct ResponseDispatchSuccess { /// guard owns the UsageEvent emit (parsed from the terminal SSE event), /// so the top-level handler must NOT emit the winner event again (#808). usage_handled_by_stream: bool, + /// Per-detector PII mask counts applied to the response body (#932), + /// non-streaming + buffered paths. Merged with the input-side counts by + /// the handler before the terminal emit. Empty on the live streaming + /// paths — their end-of-stream closures own the output-side counts. + output_redactions: crate::redact::RedactionCounts, } /// Dispatch error carrying the per-attempt telemetry accumulated before @@ -124,7 +129,7 @@ pub async fn responses( State(state): State, auth: AuthenticatedKey, client: ClientContext, - Json(body): Json, + Json(mut body): Json, ) -> Response { let started = Instant::now(); let request_id = new_request_id(); @@ -136,8 +141,24 @@ pub async fn responses( .unwrap_or("") .to_string(); - match dispatch(&state, &auth, &body, &request_id, started, &client).await { + // Filled by `dispatch` with per-detector PII mask counts (#932); attached + // to the terminal usage event on both the success and failure paths. + let mut redaction_counts = crate::redact::RedactionCounts::new(); + match dispatch( + &state, + &auth, + &mut body, + &request_id, + started, + &client, + &mut redaction_counts, + ) + .await + { Ok(success) => { + // #932: fold the non-streaming response-side mask counts into + // the per-request total before the terminal emit below. + crate::redact::merge_counts(&mut redaction_counts, success.output_redactions.clone()); let elapsed = started.elapsed(); let status = success.response.status().as_u16(); emit_access_log( @@ -201,6 +222,7 @@ pub async fn responses( &client, attempt, success.guardrail_blocked, + redaction_counts.clone(), ); } } @@ -258,6 +280,8 @@ pub async fn responses( error_class: err.kind().to_string(), ..Default::default() }, + // Input masking may have fired before the failure. + redaction_counts.clone(), ); } err.into_response() @@ -268,13 +292,20 @@ pub async fn responses( async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, - body: &Value, + // `&mut` so mask-action PII guardrails (#932) can rewrite the request + // text in place before it reaches the upstream. + body: &mut Value, request_id: &str, // Request-scoped clock + downstream client attribution, threaded so the // streaming path's Drop guard can stamp latency + client IP/UA on the // end-of-stream UsageEvent it emits (#808). started: Instant, client: &ClientContext, + // Out-param: per-detector PII mask counts (#932). Input-side counts land + // here as soon as the request is rewritten; the non-streaming output side + // arrives via `ResponseDispatchSuccess::output_redactions`; streaming + // output counts travel via the stream completion instead. + redactions_out: &mut crate::redact::RedactionCounts, ) -> Result { let snapshot = state.snapshot.load(); @@ -340,6 +371,13 @@ async fn dispatch( .into(), ); } + // #932: mask-action PII rules rewrite the Responses body in place + // AFTER the block check passes — both the verbatim passthrough and + // the cross-provider bridge forward from this body. + crate::redact::merge_counts( + redactions_out, + crate::redact::redact_responses_request(resolved_chain.as_ref(), body), + ); } let model_rl = @@ -438,6 +476,7 @@ async fn dispatch( client, attempt, &mut reservation, + redactions_out.clone(), ) .await } else { @@ -455,6 +494,7 @@ async fn dispatch( client, attempt, &mut reservation, + redactions_out.clone(), ) .await }; @@ -655,6 +695,10 @@ async fn responses_to_target( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // Input-side PII mask counts (#932) for the verbatim streaming path's + // end-of-stream emit; the non-streaming/buffered emits happen in the + // handler, which already holds them. + input_redactions: crate::redact::RedactionCounts, ) -> Result { let mut body = body.clone(); let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; @@ -898,6 +942,16 @@ async fn responses_to_target( crate::error::guardrail_block_message("response", guardrail_name.as_deref()), )); } + // #932: the whole SSE response is held here — mask the frames + // (channel reassembly) before anything reaches the wire. + let mut output_redactions = crate::redact::RedactionCounts::new(); + let buf = match crate::redact::redact_responses_sse(chain, &buf) { + Some((rewritten, counts)) => { + output_redactions = counts; + rewritten + } + None => buf, + }; // #808: the whole SSE response is buffered here, so parse its // terminal event for usage and let the handler emit (the body is // a single complete chunk now, not a live stream). @@ -913,6 +967,7 @@ async fn responses_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: false, usage_handled_by_stream: false, + output_redactions, }); } @@ -1011,6 +1066,10 @@ async fn responses_to_target( &client_c, attempt, /* guardrail_blocked */ false, + // Live-forward path: no output masking possible (an + // output-masking guardrail forces the buffered branch), + // so only the input-side counts apply. + input_redactions.clone(), ); }); let mut response = @@ -1027,6 +1086,8 @@ async fn responses_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: false, usage_handled_by_stream: true, + // The Drop guard's emit carries the counts. + output_redactions: crate::redact::RedactionCounts::new(), }) } else { let json_body: Value = upstream_resp @@ -1084,10 +1145,16 @@ async fn responses_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: true, usage_handled_by_stream: false, + output_redactions: crate::redact::RedactionCounts::new(), }); } } + // #932: mask-action PII rules rewrite the response body AFTER the + // block check passes. + let mut json_body = json_body; + let output_redactions = crate::redact::redact_responses_response(chain, &mut json_body); + Ok(ResponseDispatchSuccess { response: Json(json_body).into_response(), provider: provider_label, @@ -1097,6 +1164,7 @@ async fn responses_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: false, usage_handled_by_stream: false, + output_redactions, }) } } @@ -1122,6 +1190,9 @@ async fn responses_cross_provider_to_target( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // Input-side PII mask counts (#932), merged into the streamed judge + // path's end-of-stream emit; non-streaming emits happen in the handler. + input_redactions: crate::redact::RedactionCounts, ) -> Result { use aisix_gateway::{Bridge, BridgeContext}; @@ -1289,6 +1360,13 @@ async fn responses_cross_provider_to_target( &client_c, attempt_c, comp.guardrail_blocked, + // #932: input-side counts merged with the hold-back + // release's output-side counts. + { + let mut merged = input_redactions.clone(); + crate::redact::merge_counts(&mut merged, comp.redacted_entity_counts); + merged + }, ); }, ); @@ -1315,11 +1393,13 @@ async fn responses_cross_provider_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: false, usage_handled_by_stream: true, + // The stream's end-of-stream emit carries the counts. + output_redactions: crate::redact::RedactionCounts::new(), }); } // Non-streaming. - let resp = bridge.chat(&chat, &ctx).await.map_err(|err| { + let mut resp = bridge.chat(&chat, &ctx).await.map_err(|err| { if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) { state.runtime_status.mark_cooldown(model_id, ttl, reason); @@ -1369,10 +1449,15 @@ async fn responses_cross_provider_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: true, usage_handled_by_stream: false, + output_redactions: crate::redact::RedactionCounts::new(), }); } } + // #932: mask-action PII rules rewrite the bridged response AFTER the + // block check passes, BEFORE it is re-encoded as Responses JSON. + let output_redactions = crate::redact::redact_chat_response(chain.as_ref(), &mut resp); + let created_at = chrono::Utc::now().timestamp(); let json_body = crate::responses_bridge::chat_response_to_responses_json( &resp, @@ -1394,6 +1479,7 @@ async fn responses_cross_provider_to_target( routing: RoutingTelemetry::default(), guardrail_blocked: false, usage_handled_by_stream: false, + output_redactions, }) } @@ -1758,6 +1844,9 @@ fn emit_usage_event( client: &ClientContext, attempt: AttemptInfo, guardrail_blocked: bool, + // Per-detector PII mask counts (#932), input + output merged. Detector + // names only, never matched values. Empty = no redaction. + redacted_entity_counts: crate::redact::RedactionCounts, ) { let snap = state.snapshot.load(); let tags = provider_telemetry_tags(&snap, provider_key_id); @@ -1791,6 +1880,7 @@ fn emit_usage_event( client_source_ip: client.source_ip.clone(), client_user_agent: client.user_agent.clone(), guardrail_blocked, + redacted_entity_counts, ..Default::default() }; state.usage_sink.try_emit("responses", event.clone()); @@ -1814,6 +1904,9 @@ fn emit_zero_token_event( elapsed: Duration, client: &ClientContext, attempt: AttemptInfo, + // Per-detector PII mask counts (#932): input masking may have fired + // before the failure. Empty for most failure classes. + redacted_entity_counts: crate::redact::RedactionCounts, ) { let snap = state.snapshot.load(); let tags = provider_telemetry_tags(&snap, provider_key_id); @@ -1823,6 +1916,7 @@ fn emit_zero_token_event( model_id: model_id.to_string(), api_key_id: api_key_id.to_string(), requested_model: requested_model.to_string(), + redacted_entity_counts, latency_ms: elapsed.as_millis().min(u32::MAX as u128) as u32, status_code, inbound_protocol: "openai".to_string(), @@ -1873,6 +1967,9 @@ fn emit_failed_attempts( Duration::from_millis(u64::from(rec.latency_ms)), client, AttemptInfo::from_record(rec), + // Failed attempts carry no per-request redaction detail; the + // terminal event does. + crate::redact::RedactionCounts::new(), ); } } diff --git a/crates/aisix-proxy/src/responses_bridge.rs b/crates/aisix-proxy/src/responses_bridge.rs index ce7d8b173..a24c0c773 100644 --- a/crates/aisix-proxy/src/responses_bridge.rs +++ b/crates/aisix-proxy/src/responses_bridge.rs @@ -898,6 +898,9 @@ pub struct ResponsesStreamCompletion { /// the non-streaming path so the dashboard's Blocked tab + budget ledger /// see it. pub guardrail_blocked: bool, + /// Per-detector PII mask counts applied to the held stream at release + /// (#932). Merged with the input-side counts by the on_complete emit. + pub redacted_entity_counts: crate::redact::RedactionCounts, } struct CompleteOnDrop { @@ -1072,7 +1075,25 @@ pub fn build_responses_bridge_stream( return; } } - // Passed: release the held events verbatim. + // Passed (#932): mask the held SSE frames (channel reassembly) + // before release, then hand them to the client. + if !held.is_empty() && aisix_guardrails::Guardrail::redacts_output(chain.as_ref()) { + let mut joined: Vec = Vec::with_capacity(held_bytes); + for b in &held { + joined.extend_from_slice(b); + } + if let Some((rewritten, counts)) = + crate::redact::redact_responses_sse(chain.as_ref(), &joined) + { + crate::redact::merge_counts( + &mut guard.comp().redacted_entity_counts, + counts, + ); + yield Ok(bytes::Bytes::from(rewritten)); + return; + } + } + // Release the held events verbatim. for b in held { yield Ok(b); } diff --git a/crates/aisix-server/src/heartbeat.rs b/crates/aisix-server/src/heartbeat.rs index aa00739b5..e5d657067 100644 --- a/crates/aisix-server/src/heartbeat.rs +++ b/crates/aisix-server/src/heartbeat.rs @@ -641,6 +641,7 @@ mod tests { body["supported_guardrail_kinds"], serde_json::json!([ "keyword", + "pii", "azure_content_safety", "azure_content_safety_text_moderation", "aliyun_text_moderation", diff --git a/schemas/resources/guardrail.schema.json b/schemas/resources/guardrail.schema.json index 3266978c7..e55899d75 100644 --- a/schemas/resources/guardrail.schema.json +++ b/schemas/resources/guardrail.schema.json @@ -148,6 +148,51 @@ "type": "object" } ] + }, + "PiiCustomPattern": { + "additionalProperties": false, + "description": "One operator-supplied regex detector for `kind: \"pii\"`.", + "properties": { + "action": { + "description": "Per-pattern action override: `mask` or `block`. Falls back to the guardrail's `default_action` when omitted.", + "type": "string" + }, + "name": { + "description": "Detector name surfaced in the mask token (`[_REDACTED]`), telemetry counts, and block reasons. Never the matched value.", + "maxLength": 64, + "minLength": 1, + "type": "string" + }, + "regex": { + "description": "Regular expression the DP compiles at chain build. Invalid patterns are logged and the row is skipped (same contract as keyword rules).", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "name", + "regex" + ], + "type": "object" + }, + "PiiDetectorConfig": { + "additionalProperties": false, + "description": "One built-in detector selection for `kind: \"pii\"`. The `type` names a detector compiled into the DP (`aisix-guardrails::pii::BUILTIN_DETECTORS`); `action` optionally overrides the guardrail-level `default_action` for this detector only.", + "properties": { + "action": { + "description": "Per-detector action override: `mask` or `block`. Falls back to the guardrail's `default_action` when omitted.", + "type": "string" + }, + "type": { + "description": "Built-in detector id: `email`, `china_mobile`, `china_id_card`, `bank_card`, `us_ssn`, `ip_address`, `api_key`, `jwt`, `private_key`.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "type" + ], + "type": "object" } }, "description": "Content policy evaluated before or after upstream calls.", @@ -522,6 +567,54 @@ "region" ], "type": "object" + }, + { + "description": "In-process sensitive-data detection + redaction (#932): built-in detectors + custom regex, per-detector `mask`/`block` actions, on input and/or output, including streaming output. Always available.", + "properties": { + "custom_patterns": { + "default": [], + "description": "Operator-supplied regex detectors, evaluated after the built-ins.", + "items": { + "$ref": "#/definitions/PiiCustomPattern" + }, + "type": "array" + }, + "default_action": { + "default": "mask", + "description": "Action for detectors that don't set their own: `mask` (default) or `block`.", + "type": "string" + }, + "detectors": { + "default": [], + "description": "Built-in detectors to enable. Unknown detector ids are rejected at build time (row skipped + warn), so a typo can't silently disarm the policy.", + "items": { + "$ref": "#/definitions/PiiDetectorConfig" + }, + "type": "array" + }, + "kind": { + "enum": [ + "pii" + ], + "type": "string" + }, + "max_buffer_bytes": { + "default": 262144, + "description": "Max bytes buffered for a streamed response before `on_buffer_exceeded` applies.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + }, + "on_buffer_exceeded": { + "default": "fail_closed", + "description": "Buffer-overflow policy. Use `fail_open` to release output unscanned (and unmasked) when the buffer cap is hit; the default `fail_closed` blocks the response instead.", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" } ], "properties": { diff --git a/tests/e2e/src/cases/guardrail-pii-redaction-e2e.test.ts b/tests/e2e/src/cases/guardrail-pii-redaction-e2e.test.ts new file mode 100644 index 000000000..2f0fa8a88 --- /dev/null +++ b/tests/e2e/src/cases/guardrail-pii-redaction-e2e.test.ts @@ -0,0 +1,455 @@ +import { createHash, randomUUID } from "node:crypto"; +import OpenAI, { APIError } from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: `kind: "pii"` guardrail (#932 / AISIX-Cloud#932) — in-process +// sensitive-data detection with per-detector `mask` / `block` actions. +// +// - mask on the REQUEST: the caller's prompt PII is rewritten to +// [_REDACTED] before it reaches the upstream (verified via the +// mock upstream's received body). +// - mask on the RESPONSE (non-streaming + streaming): the model's reply is +// rewritten before it reaches the caller; the streaming case splits the +// value across chunk boundaries to pin the channel-reassembly path. +// - block: a block-action detector rejects with the standard 422 +// content_filter envelope, and the matched value never appears in it. +// +// Detector values below are synthetic: the china_id_card sample is the +// canonical ISO 7064 documentation example, the bank card is the classic +// Luhn test number. + +const CALLER = "sk-pii-e2e-caller"; +const hash = (s: string) => createHash("sha256").update(s).digest("hex"); + +const EMAIL = "alice@example.com"; +const CN_ID = "11010519491231002X"; // valid ISO 7064 MOD 11-2 check digit +const CARD = "4111111111111111"; // passes Luhn + +describe("pii guardrail e2e: mask + block on request and response", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let streamUpstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcd: EtcdClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + // Non-streaming upstream: echoes a reply CONTAINING an email, so the + // output mask has something to rewrite. + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-pii", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: `you can reach the customer at ${EMAIL} today`, + }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + + // Streaming upstream: the SAME email split across two delta chunks — + // per-chunk masking would miss it; only the hold-back channel + // reassembly catches the span. + streamUpstream = await startOpenAiUpstream({ + streamEvents: [ + '{"id":"strm-pii","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + '{"id":"strm-pii","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"mail alice@exam"},"finish_reason":null}]}', + '{"id":"strm-pii","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"ple.com now"},"finish_reason":null}]}', + '{"id":"strm-pii","object":"chat.completion.chunk","model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "[DONE]", + ], + eventDelayMs: 20, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "pii-e2e-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "pii-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: hash(CALLER), + allowed_models: ["pii-e2e"], + }); + + const streamPk = await admin.createProviderKey({ + display_name: "pii-stream-e2e-pk", + secret: "sk-mock", + api_base: `${streamUpstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "pii-stream-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: streamPk.id, + }); + await admin.createApiKey({ + key_hash: hash(`${CALLER}-stream`), + allowed_models: ["pii-stream-e2e"], + }); + + // One env-wide pii guardrail: email masks (redact-and-continue), + // china_id_card blocks (reject). + await admin.json("POST", "/admin/v1/guardrails", { + name: "pii-e2e-guard", + enabled: true, + hook_point: "both", + kind: "pii", + detectors: [ + { type: "email", action: "mask" }, + { type: "china_id_card", action: "block" }, + ], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + await streamUpstream?.close(); + }); + + const client = () => + new OpenAI({ + apiKey: CALLER, + baseURL: `${app!.proxyUrl}/v1`, + maxRetries: 0, + }); + + test("mask: request PII is rewritten before the upstream, response PII before the caller", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + // Propagation probe: once the guardrail is live, the (email-bearing) + // mock reply comes back masked. + await waitConfigPropagation(async () => { + const r = await client().chat.completions.create({ + model: "pii-e2e", + messages: [{ role: "user", content: "probe" }], + }); + return (r.choices[0]?.message?.content ?? "").includes("[EMAIL_REDACTED]"); + }); + + const res = await client().chat.completions.create({ + model: "pii-e2e", + messages: [ + { role: "user", content: `contact me at ${EMAIL} about the order` }, + ], + }); + + // Response side: the model's reply had the email; the caller sees the + // mask token and never the value. + const reply = res.choices[0]?.message?.content ?? ""; + expect(reply).toContain("[EMAIL_REDACTED]"); + expect(reply).not.toContain(EMAIL); + + // Request side: the upstream received the MASKED prompt — the value + // never left the gateway. (Structure preserved: only the span is + // replaced.) + const lastReq = upstream.receivedRequests.at(-1); + expect(lastReq).toBeDefined(); + const upstreamBody = lastReq!.body; + expect(upstreamBody).toContain("[EMAIL_REDACTED]"); + expect(upstreamBody).toContain("about the order"); + expect(upstreamBody).not.toContain(EMAIL); + }); + + test("block: a block-action detector rejects with 422 content_filter, value not echoed", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + const upstreamHitsBefore = upstream.receivedRequests.length; + let caught: unknown; + try { + await client().chat.completions.create({ + model: "pii-e2e", + messages: [{ role: "user", content: `my id number is ${CN_ID}` }], + }); + } catch (e) { + caught = e; + } + + expect(caught).toBeInstanceOf(APIError); + if (!(caught instanceof APIError)) throw new Error("unreachable"); + expect(caught.status).toBe(422); + expect((caught.error as { type?: unknown })?.type).toBe("content_filter"); + // #153 / #932 no-leak: the matched value never appears in the envelope; + // the guardrail name does (#519 B.4b). + const blob = JSON.stringify(caught.error ?? {}) + (caught.message ?? ""); + expect(blob).not.toContain(CN_ID); + expect(blob).toContain("guardrail 'pii-e2e-guard'"); + // Input block fires pre-dispatch: the upstream is never hit. + expect(upstream.receivedRequests.length).toBe(upstreamHitsBefore); + }); + + test("mask does NOT block: a mask-only match still gets a 200", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + // A bank card would only matter to a block detector — none is + // configured for it, and email is mask-action, so the request goes + // through (masked) rather than 422ing. + const res = await client().chat.completions.create({ + model: "pii-e2e", + messages: [{ role: "user", content: `card ${CARD} email ${EMAIL}` }], + }); + expect(res.choices[0]?.message?.content ?? "").toContain("[EMAIL_REDACTED]"); + }); + + test("streaming mask: a span split across delta chunks is reassembled and masked (#932)", async (ctx) => { + if (!etcdReachable || !app || !streamUpstream) { + ctx.skip(); + return; + } + + const streamCaller = `${CALLER}-stream`; + const doStream = () => + fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${streamCaller}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "pii-stream-e2e", + messages: [{ role: "user", content: "innocent prompt" }], + stream: true, + }), + }); + + await waitConfigPropagation(async () => { + const probe = await doStream(); + if (probe.status !== 200) { + await probe.text(); + return false; + } + return (await probe.text()).includes("[EMAIL_REDACTED]"); + }); + + const res = await doStream(); + expect(res.status).toBe(200); + const wire = await res.text(); + + // The email was split "alice@exam" + "ple.com" across two chunks — + // channel reassembly at the hold-back release must still catch it. + expect(wire).toContain("[EMAIL_REDACTED]"); + expect(wire).not.toContain(EMAIL); + expect(wire).not.toContain("alice@exam"); + // Clean stream contract: [DONE] present, no error event. + expect(wire).toContain("data: [DONE]"); + expect(wire).not.toContain("event: error"); + // Non-content fields survive the rewrite (finish_reason intact). + expect(wire).toContain('"finish_reason":"stop"'); + }); + + test("monitor mode: enforcement_mode=monitor observes but does not mask", async (ctx) => { + if (!etcdReachable) { + ctx.skip(); + return; + } + + // Dedicated app instance (its own etcd prefix / env), so the env-wide + // masking guardrail from the main suite cannot interfere — this env + // carries ONLY the monitor-mode pii guardrail. + const monUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-mon", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: `reply with ${EMAIL}` }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + const monApp = await spawnApp(); + const monAdmin = new AdminClient(monApp.adminUrl, monApp.adminKey); + const monPk = await monAdmin.createProviderKey({ + display_name: "pii-mon-pk", + secret: "sk-mock", + api_base: `${monUpstream.baseUrl}/v1`, + }); + await monAdmin.createModel({ + display_name: "pii-mon-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: monPk.id, + }); + const monCaller = `${CALLER}-mon`; + await monAdmin.createApiKey({ + key_hash: hash(monCaller), + allowed_models: ["pii-mon-e2e"], + }); + await monAdmin.json("POST", "/admin/v1/guardrails", { + name: "pii-mon-guard", + enabled: true, + hook_point: "both", + enforcement_mode: "monitor", + kind: "pii", + detectors: [{ type: "email", action: "mask" }], + }); + + const monClient = new OpenAI({ + apiKey: monCaller, + baseURL: `${monApp.proxyUrl}/v1`, + maxRetries: 0, + }); + await waitConfigPropagation(async () => { + try { + const r = await monClient.chat.completions.create({ + model: "pii-mon-e2e", + messages: [{ role: "user", content: "probe" }], + }); + return (r.choices[0]?.message?.content ?? "").length > 0; + } catch { + return false; + } + }); + + const res = await monClient.chat.completions.create({ + model: "pii-mon-e2e", + messages: [{ role: "user", content: `mask me maybe: ${EMAIL}` }], + }); + // Monitor mode: content flows UNCHANGED in both directions; the + // would-be mask counts land in ops logs only. + expect(res.choices[0]?.message?.content ?? "").toContain(EMAIL); + const lastReq = monUpstream.receivedRequests.at(-1); + expect(lastReq!.body).toContain(EMAIL); + + await monApp.exit(); + await monUpstream.close(); + }); + + test("/v1/messages passthrough: request masked before the Anthropic upstream", async (ctx) => { + if (!etcdReachable || !app || !admin) { + ctx.skip(); + return; + } + + // Anthropic-native mock upstream: capture the received body, return a + // minimal non-streaming message whose reply carries an email so the + // output mask has something to rewrite too. + const received: unknown[] = []; + const { createServer } = await import("node:http"); + const anthUpstream = createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (c: Buffer) => chunks.push(c)); + req.on("end", () => { + received.push(JSON.parse(Buffer.concat(chunks).toString())); + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + id: "msg_pii", + type: "message", + role: "assistant", + model: "claude-3-5-haiku-20241022", + content: [{ type: "text", text: `reach them at ${EMAIL} ok` }], + stop_reason: "end_turn", + usage: { input_tokens: 4, output_tokens: 6 }, + }), + ); + }); + }); + await new Promise((resolve) => anthUpstream.listen(0, resolve)); + const anthPort = (anthUpstream.address() as { port: number }).port; + + const anthPk = await admin.createProviderKey({ + display_name: "pii-anth-pk", + secret: "sk-ant-mock", + api_base: `http://127.0.0.1:${anthPort}`, + provider: "anthropic", + adapter: "anthropic", + }); + await admin.createModel({ + display_name: "pii-anth-e2e", + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: anthPk.id, + }); + const anthCaller = `${CALLER}-anth`; + await admin.createApiKey({ + key_hash: hash(anthCaller), + allowed_models: ["pii-anth-e2e"], + }); + + const call = () => + fetch(`${app!.proxyUrl}/v1/messages`, { + method: "POST", + headers: { "content-type": "application/json", "x-api-key": anthCaller }, + body: JSON.stringify({ + model: "pii-anth-e2e", + max_tokens: 32, + messages: [ + { role: "user", content: `write to ${EMAIL} please` }, + ], + }), + }); + + await waitConfigPropagation(async () => { + const r = await call(); + if (r.status !== 200) { + await r.text(); + return false; + } + const b = (await r.json()) as { + content?: Array<{ text?: string }>; + }; + return (b.content?.[0]?.text ?? "").includes("[EMAIL_REDACTED]"); + }); + + const res = await call(); + expect(res.status).toBe(200); + const body = (await res.json()) as { content?: Array<{ text?: string }> }; + // Response side masked for the caller… + expect(body.content?.[0]?.text ?? "").toContain("[EMAIL_REDACTED]"); + expect(JSON.stringify(body)).not.toContain(EMAIL); + // …and the request side was masked before the upstream. + const lastReq = received.at(-1); + const upstreamBlob = JSON.stringify(lastReq); + expect(upstreamBlob).toContain("[EMAIL_REDACTED]"); + expect(upstreamBlob).not.toContain(EMAIL); + + await new Promise((resolve, reject) => + anthUpstream.close((e) => (e ? reject(e) : resolve())), + ); + }); +}); diff --git a/tests/e2e/src/cases/messages-streaming-output-guardrail-e2e.test.ts b/tests/e2e/src/cases/messages-streaming-output-guardrail-e2e.test.ts index 4b8c36f80..5cc7faf55 100644 --- a/tests/e2e/src/cases/messages-streaming-output-guardrail-e2e.test.ts +++ b/tests/e2e/src/cases/messages-streaming-output-guardrail-e2e.test.ts @@ -96,7 +96,11 @@ describe("streaming /v1/messages output guardrail (#448)", () => { const res = await stream(); expect(res.status).toBe(200); // stream starts 200; the block is in-band const body = await res.text(); - expect(body, "streamed content is forwarded verbatim").toContain(FORBIDDEN); + // #932 / #466-class: keyword output guardrails carry the BufferFull + // hold-back policy, so /v1/messages streaming now withholds the whole + // response until it scans clean — the matched content must NOT reach + // the wire (pre-fix it was forwarded verbatim before the error frame). + expect(body, "hold-back keeps the matched content off the wire").not.toContain(FORBIDDEN); expect(body, "stream must end with a content_filter error event").toContain("content_filter"); }); }); diff --git a/tests/e2e/src/cases/messages-streaming-toolcall-guardrail-e2e.test.ts b/tests/e2e/src/cases/messages-streaming-toolcall-guardrail-e2e.test.ts index 2f41b5bbc..c94871d5a 100644 --- a/tests/e2e/src/cases/messages-streaming-toolcall-guardrail-e2e.test.ts +++ b/tests/e2e/src/cases/messages-streaming-toolcall-guardrail-e2e.test.ts @@ -85,7 +85,10 @@ describe("streaming /v1/messages tool_use output guardrail (#448)", () => { await waitConfigPropagation(async () => (await stream().then((r) => r.text())).includes("content_filter")); const body = await stream().then((r) => r.text()); - expect(body, "tool_use arguments are forwarded verbatim").toContain(FORBIDDEN); + // #932 / #466-class: keyword output guardrails carry the BufferFull + // hold-back policy, so the tool_use arguments are withheld until the + // end-of-stream scan — the matched content must NOT reach the wire. + expect(body, "hold-back keeps the matched arguments off the wire").not.toContain(FORBIDDEN); expect(body, "stream must end with a content_filter error event").toContain("content_filter"); }); });