From fb90ace46241d42358b1fa7c8403292336bd8aa3 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 20 Aug 2026 14:43:07 +0800 Subject: [PATCH 1/5] feat(guardrails): capture-group scoping + configurable replacement for pii custom patterns (AISIX-Cloud#1334) - A custom pattern regex with at least one capture group now rewrites only group 1 of each match, keeping the rest of the match verbatim, so a rule can replace a value while preserving its key/label ("version": "12.1" -> "version": "***" stays parseable JSON). Patterns without capture groups keep the whole-match semantics. - Checksum validators (Luhn / ISO 7064) now run on the replaced span (group 1 when present), so a prefixed pattern cannot silently disable its validator. - PiiCustomPattern gains an optional replacement field overriding the default [_REDACTED] token; empty string deletes the span; the text is literal (no group expansion). - replacement on a pattern whose effective action is block rejects the row at build time (a knob is enforced as written or rejected, never accepted-but-unread). - e2e: capture-group rules drive a real DP end to end - request and response rewritten in place (zh + en), hard negatives byte-identical, embedded JSON still parses; regenerated guardrail.schema.json. --- crates/aisix-core/src/models/guardrail.rs | 27 +- crates/aisix-guardrails/src/build.rs | 87 ++++++- crates/aisix-guardrails/src/pii.rs | 246 ++++++++++++++++-- schemas/resources/guardrail.schema.json | 7 +- .../guardrail-pii-capture-group-e2e.test.ts | 196 ++++++++++++++ 5 files changed, 540 insertions(+), 23 deletions(-) create mode 100644 tests/e2e/src/cases/guardrail-pii-capture-group-e2e.test.ts diff --git a/crates/aisix-core/src/models/guardrail.rs b/crates/aisix-core/src/models/guardrail.rs index 8a680b167..c15218e5f 100644 --- a/crates/aisix-core/src/models/guardrail.rs +++ b/crates/aisix-core/src/models/guardrail.rs @@ -445,12 +445,26 @@ pub struct PiiCustomPattern { pub name: String, /// Regular expression AISIX compiles when building the guardrail chain. /// An invalid pattern makes AISIX log and skip the guardrail. + /// + /// When the expression declares at least one capture group, a `mask` + /// action rewrites only the first capture group of each match and keeps + /// the rest of the match unchanged. Use this to replace a value while + /// preserving its surrounding key or label, for example + /// `"version"\s*:\s*"([^"]*)"`. Without capture groups, the whole match + /// is rewritten. #[schemars(length(min = 1))] pub regex: String, /// Per-pattern action override. Falls back to the guardrail's /// `default_action` when omitted. #[serde(default, skip_serializing_if = "Option::is_none")] pub action: Option, + /// Literal text that replaces the masked span, such as `***`. When + /// omitted, the span is rewritten to `[_REDACTED]`. An empty + /// string removes the span. Only valid when the pattern's effective + /// action is `mask`. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schemars(length(max = 256))] + pub replacement: Option, } /// Config block for `kind: "pii"`. Built-in sensitive-data detection and @@ -1464,7 +1478,8 @@ mod tests { "kind": "pii", "default_action": "block", "custom_patterns": [ - { "name": "employee_id", "regex": "\\bEMP-\\d{6}\\b", "action": "mask" } + { "name": "employee_id", "regex": "\\bEMP-\\d{6}\\b", "action": "mask" }, + { "name": "eda_version", "regex": "version: (\\S+)", "action": "mask", "replacement": "***" } ], "max_buffer_bytes": 1024, "on_buffer_exceeded": "fail_open" @@ -1473,15 +1488,23 @@ mod tests { match g.config { GuardrailKind::Pii(ref c) => { assert!(c.detectors.is_empty()); - assert_eq!(c.custom_patterns.len(), 1); + assert_eq!(c.custom_patterns.len(), 2); assert_eq!(c.custom_patterns[0].name, "employee_id"); assert_eq!(c.custom_patterns[0].action.as_deref(), Some("mask")); + assert_eq!(c.custom_patterns[0].replacement, None); + assert_eq!(c.custom_patterns[1].replacement.as_deref(), Some("***")); 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"), } + // A row without `replacement` serialises without the key + // (skip_serializing_if), so older documents round-trip untouched. + let ser = serde_json::to_value(&g).unwrap(); + let pats = ser["custom_patterns"].as_array().unwrap(); + assert!(pats[0].get("replacement").is_none()); + assert_eq!(pats[1]["replacement"], "***"); } #[test] diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index fd4a424c6..a3abc36d3 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -249,12 +249,22 @@ fn build_one_inner( value: s.to_owned(), })?, }; - let rule = PiiRule::new(p.name.clone(), &p.regex, action, None).map_err(|e| { - BuildError::InvalidRegex { + // `replacement` only means something to a mask rewrite; on a + // block-action pattern it would be accepted-but-unread config + // (the #962 class), so the row is rejected instead. cp-api + // validates the same combination at write time; this covers + // the declarative-file source and defends the invariant. + if p.replacement.is_some() && action == PiiAction::Block { + return Err(BuildError::ReplacementOnBlock { + name: p.name.clone(), + }); + } + let rule = PiiRule::new(p.name.clone(), &p.regex, action, None) + .map_err(|e| BuildError::InvalidRegex { pattern: p.regex.clone(), source: e, - } - })?; + })? + .with_replacement(p.replacement.clone()); rules.push(rule); } let on_exceeded_fail_open = cfg.on_buffer_exceeded == "fail_open"; @@ -459,6 +469,14 @@ enum BuildError { /// skipped + warned rather than silently running a weaker policy. #[error("invalid {field} value {value:?}")] InvalidValue { field: &'static str, value: String }, + /// A `custom_patterns[].replacement` on a pattern whose effective + /// action is `block` — the replacement would never be read (never + /// half-honor a knob, #963). Carries the pattern NAME only, never + /// the replacement text or a matched value. + #[error( + "custom_patterns[].replacement requires action=mask (pattern {name:?} resolves to block)" + )] + ReplacementOnBlock { name: 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 @@ -1459,6 +1477,67 @@ mod tests { ); } + /// AISIX-Cloud#1334: a capture-group custom pattern with a + /// `replacement` builds and rewrites only group 1 through the chain. + #[test] + fn pii_custom_pattern_replacement_and_group_build_from_row() { + let table: ResourceTable = ResourceTable::default(); + table.insert(entry( + "eda", + "g-1", + parse( + r#"{ + "name": "eda", + "kind": "pii", + "custom_patterns": [{ + "name": "eda_version", + "regex": "version\\s*:\\s*(\\d+(?:\\.\\d+)+)", + "action": "mask", + "replacement": "***" + }] + }"#, + ), + )); + let chain = build_chain_from_snapshot(&table, None); + assert_eq!(chain.len(), 1); + let r = chain.redact_input_text("tool version: 12.1 done").unwrap(); + assert_eq!(r.text, "tool version: *** done"); + } + + /// A `replacement` on a pattern whose effective action is `block` + /// (explicit or via `default_action`) rejects the row — the knob + /// would otherwise be accepted but never read (#963). + #[test] + fn pii_replacement_on_block_action_rejects_row() { + for row in [ + // Explicit per-pattern block. + r#"{ + "name": "bad-explicit", + "kind": "pii", + "custom_patterns": [{ + "name": "p", "regex": "x(y)", "action": "block", "replacement": "*" + }] + }"#, + // Inherited block via default_action. + r#"{ + "name": "bad-inherited", + "kind": "pii", + "default_action": "block", + "custom_patterns": [{ + "name": "p", "regex": "x(y)", "replacement": "*" + }] + }"#, + ] { + let table: ResourceTable = ResourceTable::default(); + table.insert(entry("bad", "g-1", parse(row))); + let chain = build_chain_from_snapshot(&table, None); + assert!( + chain.is_empty(), + "replacement+block row must be skipped, not half-honored", + ); + } + } + /// A stub remote guardrail that always fails open (returns `Bypass`), /// standing in for a Bedrock/Azure guardrail whose upstream is down. struct AlwaysBypass; diff --git a/crates/aisix-guardrails/src/pii.rs b/crates/aisix-guardrails/src/pii.rs index 1240506e5..a49782ce6 100644 --- a/crates/aisix-guardrails/src/pii.rs +++ b/crates/aisix-guardrails/src/pii.rs @@ -10,15 +10,27 @@ //! 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`] / +//! - `Mask`: each matched span is rewritten to the rule's replacement +//! text (`[_REDACTED]` by default, or an operator-configured +//! `replacement`) and processing continues. Callers apply the rewrite +//! through [`crate::Guardrail::redact_input_text`] / //! [`crate::Guardrail::redact_output_text`]. //! +//! Capture-group scoping (AISIX-Cloud#1334): a rule whose regex declares +//! at least one capture group replaces (and checksum-validates) ONLY +//! group 1 of each match; the rest of the match is kept verbatim. This is +//! how a rule expresses "replace the value, keep the key" for shapes like +//! `"version": "12.1"` — the `regex` crate has no lookaround, so context +//! can only be consumed by the match and preserved via the group split. +//! A regex without capture groups keeps the original whole-match +//! semantics. +//! //! 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). +//! criterion that redacted values must not appear in gateway logs). An +//! operator-configured `replacement` is config, not matched content, so +//! it may appear in rewritten payloads by design. use std::borrow::Cow; use std::collections::BTreeMap; @@ -32,7 +44,8 @@ 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. + /// Rewrite the matched span to the rule's replacement text + /// (`[_REDACTED]` by default) and continue. Mask, /// Reject the request/response (422 content-filter). Block, @@ -63,8 +76,14 @@ pub struct PiiRule { regex: Regex, action: PiiAction, validate: Option, - /// Pre-computed `[_REDACTED]` token. + /// The literal text a masked span is rewritten to: the pre-computed + /// `[_REDACTED]` token, or the operator-configured + /// `replacement` override (see [`PiiRule::with_replacement`]). mask_token: String, + /// `true` when the regex declares at least one capture group — the + /// rule then replaces and validates ONLY group 1 of each match, + /// keeping the rest of the match verbatim (module docs). + group_scoped: bool, } impl PiiRule { @@ -76,37 +95,84 @@ impl PiiRule { ) -> Result { let name = name.into(); let mask_token = mask_token(&name); + let regex = Regex::new(pattern)?; + // captures_len() counts group 0 (the whole match), so > 1 means + // the pattern declares explicit groups. `(?:…)` doesn't count. + let group_scoped = regex.captures_len() > 1; Ok(Self { - regex: Regex::new(pattern)?, + regex, name, action, validate, mask_token, + group_scoped, }) } + /// Override the default `[_REDACTED]` mask token with an + /// operator-configured literal replacement (AISIX-Cloud#1334). + /// `None` keeps the default; an empty string deletes the span. + /// The text is used verbatim — no `$n` group expansion. + pub fn with_replacement(mut self, replacement: Option) -> Self { + if let Some(r) = replacement { + self.mask_token = r; + } + self + } + /// `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)) } + /// The span of `caps` this rule replaces and validates: group 1 for a + /// group-scoped rule, the whole match otherwise. `None` when group 1 + /// did not participate in this match (an alternation branch without + /// it) — the match is then not a detection and stays untouched. + fn target_span<'t>(&self, caps: ®ex::Captures<'t>) -> Option> { + caps.get(if self.group_scoped { 1 } else { 0 }) + } + /// 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())) + if self.group_scoped { + // Group-scoped rules validate group 1, so the (slower) + // captures iterator is required; the built-in detectors have + // no groups and stay on the find_iter fast path below. + self.regex.captures_iter(text).any(|caps| { + self.target_span(&caps) + .is_some_and(|g| self.accepts(g.as_str())) + }) + } else { + 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). + /// Rewrite every validated match in `text` to the rule's mask token — + /// only group 1 of the match for a group-scoped rule. 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() + let whole = caps.get(0).expect("group 0 always participates"); + match self.target_span(caps) { + Some(target) if self.accepts(target.as_str()) => { + count += 1; + if self.group_scoped { + // Keep the match's bytes outside group 1 verbatim. + let m = whole.as_str(); + let mut out = + String::with_capacity(m.len() - target.len() + self.mask_token.len()); + out.push_str(&m[..target.start() - whole.start()]); + out.push_str(&self.mask_token); + out.push_str(&m[target.end() - whole.start()..]); + out + } else { + self.mask_token.clone() + } + } + _ => whole.as_str().to_string(), } }); (out, count) @@ -589,6 +655,154 @@ mod tests { assert_eq!(r.text, "ssn [US_SSN_REDACTED] from [IP_ADDRESS_REDACTED]"); } + // ---- capture-group scoping + replacement (AISIX-Cloud#1334) ---- + + #[test] + fn group_scoped_rule_replaces_group_1_only() { + let rule = PiiRule::new( + "eda_version", + r"version\s*:\s*(\d+(?:\.\d+)+)", + PiiAction::Mask, + None, + ) + .unwrap() + .with_replacement(Some("***".into())); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("tool version: 12.1 loaded").unwrap(); + assert_eq!(r.text, "tool version: *** loaded"); + assert_eq!(r.counts.get("eda_version"), Some(&1)); + } + + #[test] + fn group_scoped_rule_keeps_json_parseable() { + let rule = PiiRule::new( + "eda_version", + r#""version"\s*:\s*"([^"]*)""#, + PiiAction::Mask, + None, + ) + .unwrap() + .with_replacement(Some("***".into())); + let g = guardrail(vec![rule]); + let doc = r#"{"tool":"eda","version": "12.1","cells":42}"#; + let r = g.redact_input_text(doc).unwrap(); + assert_eq!(r.text, r#"{"tool":"eda","version": "***","cells":42}"#); + // The acceptance criterion is structural, not textual: the + // rewritten document must still parse. + let v: serde_json::Value = serde_json::from_str(&r.text).unwrap(); + assert_eq!(v["version"], "***"); + assert_eq!(v["cells"], 42); + } + + #[test] + fn group_scoped_rule_skips_hard_negatives_byte_for_byte() { + let rule = PiiRule::new( + "eda_version", + r"(?:version|版本)\s*[:：]\s*(\d+(?:\.\d+)+)", + PiiAction::Mask, + None, + ) + .unwrap() + .with_replacement(Some("***".into())); + let g = guardrail(vec![rule]); + // Dot-separated numbers everywhere, none anchored by the keyword: + // zero hits, and None means the caller keeps the original bytes. + let log = "Elapsed: 12.345s Memory: 4.2 GB node 0.13um top.v:12:1 at 10.2.255.1"; + assert!(g.redact_input_text(log).is_none()); + // Mixed hit + negatives: only the anchored value is rewritten, + // every other byte survives verbatim (exact-string assertion). + let mixed = "Elapsed: 12.345s 版本：12.1 at 10.2.255.1"; + let r = g.redact_input_text(mixed).unwrap(); + assert_eq!(r.text, "Elapsed: 12.345s 版本：*** at 10.2.255.1"); + } + + #[test] + fn group_scoped_validator_checks_group_not_whole_match() { + // A prefixed card-number rule: the whole match includes "card " + // which can never pass Luhn. The validator must run on group 1 — + // otherwise a prefixed pattern silently disables its checksum. + let rule = PiiRule::new( + "prefixed_card", + r"card (\d{13,19})", + PiiAction::Mask, + Some(luhn_checksum), + ) + .unwrap() + .with_replacement(Some("####".into())); + let g = guardrail(vec![rule]); + // 4111111111111111 passes Luhn → masked, prefix kept. + let r = g.redact_input_text("card 4111111111111111 ok").unwrap(); + assert_eq!(r.text, "card #### ok"); + // Same shape, broken check digit → untouched. + assert!(g.redact_input_text("card 4111111111111112 ok").is_none()); + } + + #[test] + fn group_scoped_block_rule_validates_group_too() { + let rule = PiiRule::new( + "prefixed_card", + r"card (\d{13,19})", + PiiAction::Block, + Some(luhn_checksum), + ) + .unwrap(); + let g = guardrail(vec![rule]); + assert!(g.first_block_match("card 4111111111111111").is_some()); + assert!(g.first_block_match("card 4111111111111112").is_none()); + } + + #[test] + fn group_not_participating_leaves_match_untouched() { + // Alternation where only one branch carries group 1: the other + // branch has no sensitive segment to replace, so it stays as-is. + let rule = PiiRule::new("opt_group", r"ver=(\d+)|verless", PiiAction::Mask, None) + .unwrap() + .with_replacement(Some("***".into())); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("ver=42 and verless mode").unwrap(); + assert_eq!(r.text, "ver=*** and verless mode"); + assert_eq!(r.counts.get("opt_group"), Some(&1)); + assert!(g.redact_input_text("verless only").is_none()); + } + + #[test] + fn replacement_defaults_to_name_token_and_supports_empty() { + // No replacement → the [_REDACTED] default, group-scoped. + let rule = PiiRule::new("ver", r"version: (\S+)", PiiAction::Mask, None).unwrap(); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("version: 12.1").unwrap(); + assert_eq!(r.text, "version: [VER_REDACTED]"); + // Empty replacement deletes the span. + let rule = PiiRule::new("ver", r"version: (\S+)", PiiAction::Mask, None) + .unwrap() + .with_replacement(Some(String::new())); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("version: 12.1 end").unwrap(); + assert_eq!(r.text, "version: end"); + } + + #[test] + fn replacement_without_groups_replaces_whole_match() { + // Back-compat: no capture group + custom replacement still swaps + // the entire match. + let rule = PiiRule::new("ver", r"\bv\d+\.\d+\b", PiiAction::Mask, None) + .unwrap() + .with_replacement(Some("vX.Y".into())); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("running v12.1 now").unwrap(); + assert_eq!(r.text, "running vX.Y now"); + } + + #[test] + fn replacement_dollar_is_literal_not_group_expansion() { + let rule = PiiRule::new("ver", r"version: (\S+)", PiiAction::Mask, None) + .unwrap() + .with_replacement(Some("$1".into())); + let g = guardrail(vec![rule]); + let r = g.redact_input_text("version: 12.1").unwrap(); + assert_eq!(r.text, "version: $1"); + } + #[test] fn every_builtin_detector_compiles() { for (id, _, _) in BUILTIN_DETECTORS { diff --git a/schemas/resources/guardrail.schema.json b/schemas/resources/guardrail.schema.json index 5bcd79d5c..2ef693637 100644 --- a/schemas/resources/guardrail.schema.json +++ b/schemas/resources/guardrail.schema.json @@ -175,9 +175,14 @@ "type": "string" }, "regex": { - "description": "Regular expression AISIX compiles when building the guardrail chain. An invalid pattern makes AISIX log and skip the guardrail.", + "description": "Regular expression AISIX compiles when building the guardrail chain. An invalid pattern makes AISIX log and skip the guardrail.\n\nWhen the expression declares at least one capture group, a `mask` action rewrites only the first capture group of each match and keeps the rest of the match unchanged. Use this to replace a value while preserving its surrounding key or label, for example `\"version\"\\s*:\\s*\"([^\"]*)\"`. Without capture groups, the whole match is rewritten.", "minLength": 1, "type": "string" + }, + "replacement": { + "description": "Literal text that replaces the masked span, such as `***`. When omitted, the span is rewritten to `[_REDACTED]`. An empty string removes the span. Only valid when the pattern's effective action is `mask`.", + "maxLength": 256, + "type": "string" } }, "required": [ diff --git a/tests/e2e/src/cases/guardrail-pii-capture-group-e2e.test.ts b/tests/e2e/src/cases/guardrail-pii-capture-group-e2e.test.ts new file mode 100644 index 000000000..b0e0430d9 --- /dev/null +++ b/tests/e2e/src/cases/guardrail-pii-capture-group-e2e.test.ts @@ -0,0 +1,196 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + ProxyClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: capture-group custom patterns + configurable `replacement` +// (AISIX-Cloud#1334) — the "replace the value, keep the key" semantics. +// +// Rules under test (seeded as kind=pii custom_patterns): +// - plain form `(?:version|版本)\s*[:：]\s*(\d+(?:\.\d+)+)` → group 1 +// becomes `***`, the keyword and separator stay verbatim; anchoring on +// the keyword is what keeps dot-separated negatives (durations, sizes, +// IPs, file:line) untouched. +// - JSON form `"version"\s*:\s*"([^"]*)"` → only the value inside the +// quotes is rewritten, so the document still parses. +// +// Acceptance (issue #1334): both forms rewrite on request AND response, +// Chinese and English labels covered, JSON stays parseable (asserted by +// parsing, not string comparison), hard negatives hit zero, non-matched +// content survives byte-for-byte (exact-equality assertions), and the +// whole flow stays 200 — mask never blocks. + +const CALLER = "sk-pii-capture-e2e"; +const hash = (s: string) => createHash("sha256").update(s).digest("hex"); + +// Hard negatives from the issue: dot-separated numbers that are NOT +// version values (duration, size, process node, file:line, IPv4). +const NEGATIVES = + "Elapsed: 12.345s Memory: 4.2 GB node 0.13um top.v:12:1 ip 10.2.255.1"; + +// Mixed prompt: an English hit, the negatives, and a Chinese hit. +const PROMPT = `EDA version: 12.1 ${NEGATIVES} 工具版本：2022.4 结束`; +const PROMPT_MASKED = `EDA version: *** ${NEGATIVES} 工具版本：*** 结束`; + +// The model reply embeds the JSON key-value form. +const REPLY_JSON = `{"tool":"vcs","version": "12.1","cells":42}`; +const REPLY = `config ${REPLY_JSON} ok ${NEGATIVES}`; + +describe("pii capture-group + replacement e2e", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-capture", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: REPLY }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 8, total_tokens: 13 }, + }, + }); + + app = await spawnApp(); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: "pii-capture-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "pii-capture-e2e", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await seed.createGuardrail({ + name: "pii-capture-guard", + enabled: true, + hook_point: "both", + kind: "pii", + custom_patterns: [ + { + name: "eda_version", + regex: "(?:version|版本)\\s*[:：]\\s*(\\d+(?:\\.\\d+)+)", + action: "mask", + replacement: "***", + }, + { + name: "eda_version_json", + regex: '"version"\\s*:\\s*"([^"]*)"', + action: "mask", + replacement: "***", + }, + ], + }); + + // Caller key LAST: once it authenticates, every resource seeded + // before it (guardrail included) is in the snapshot (AGENTS.md gate + // rule — the gate must not exercise the behavior under test). + await seed.createApiKey({ + key_hash: hash(CALLER), + allowed_models: ["pii-capture-e2e"], + }); + await waitConfigPropagation(async () => { + const r = await new ProxyClient(app!.proxyUrl, CALLER).listModels(); + return r.status === 200; + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + const client = () => + new OpenAI({ + apiKey: CALLER, + baseURL: `${app!.proxyUrl}/v1`, + maxRetries: 0, + }); + + const chat = (content: string) => + client().chat.completions.create({ + model: "pii-capture-e2e", + messages: [{ role: "user", content }], + }); + + test("request: value replaced in place, key + negatives byte-identical, zh+en", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + const res = await chat(PROMPT); + // Mask never blocks: the call succeeded (a 422 would have thrown). + expect(res.choices[0]?.message?.content ?? "").not.toBe(""); + + // The upstream saw the ENTIRE prompt with only the two anchored + // values rewritten — exact equality is the byte-for-byte assertion + // for everything outside the hits (negatives included). + const sent = JSON.parse(upstream.receivedRequests.at(-1)!.body) as { + messages: Array<{ content: string }>; + }; + expect(sent.messages[0].content).toBe(PROMPT_MASKED); + }); + + test("request: a no-hit prompt passes through byte-identical", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + await chat(NEGATIVES); + const sent = JSON.parse(upstream.receivedRequests.at(-1)!.body) as { + messages: Array<{ content: string }>; + }; + expect(sent.messages[0].content).toBe(NEGATIVES); + }); + + test("response: JSON form masked in place and still parseable", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const res = await chat("show me the config"); + const reply = res.choices[0]?.message?.content ?? ""; + // Only the value inside the quotes changed; prefix/suffix and the + // negatives after the JSON survive byte-for-byte. + expect(reply).toBe( + `config {"tool":"vcs","version": "***","cells":42} ok ${NEGATIVES}`, + ); + // Structural acceptance: parse the embedded JSON, don't just diff it. + const embedded = reply.slice(reply.indexOf("{"), reply.lastIndexOf("}") + 1); + const parsed = JSON.parse(embedded) as { + tool: string; + version: string; + cells: number; + }; + expect(parsed.version).toBe("***"); + expect(parsed.tool).toBe("vcs"); + expect(parsed.cells).toBe(42); + expect(reply).not.toContain("12.1"); + }); +}); From 9291b34729422996e3ecf2fa36e5e02cd98b4122 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 20 Aug 2026 15:43:17 +0800 Subject: [PATCH 2/5] test(guardrails): pin that built-in pii detectors declare no capture groups group_scoped is auto-detected from the pattern, so an accidental capturing group in a future builtin would silently narrow its replacement to group 1. Assert captures_len() == 1 for every builtin (audit rider on #1007). --- crates/aisix-guardrails/src/pii.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/crates/aisix-guardrails/src/pii.rs b/crates/aisix-guardrails/src/pii.rs index a49782ce6..06138b7e0 100644 --- a/crates/aisix-guardrails/src/pii.rs +++ b/crates/aisix-guardrails/src/pii.rs @@ -806,10 +806,18 @@ mod tests { #[test] fn every_builtin_detector_compiles() { for (id, _, _) in BUILTIN_DETECTORS { - assert!( - builtin_rule(id, PiiAction::Mask).is_some(), - "builtin {id} must compile", + let rule = builtin_rule(id, PiiAction::Mask) + .unwrap_or_else(|| panic!("builtin {id} must compile")); + // group_scoped is auto-detected from the pattern, so an + // accidental capturing `(...)` in a builtin would silently + // narrow its replacement to group 1. Builtins must stay + // whole-match: use `(?:…)` for grouping. + assert_eq!( + rule.regex.captures_len(), + 1, + "builtin {id} must not declare capture groups", ); + assert!(!rule.group_scoped); } assert!(builtin_rule("no_such_detector", PiiAction::Mask).is_none()); } From dfdc4e50b74adf34f6ef18779c8b9310a0f65293 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 20 Aug 2026 16:25:24 +0800 Subject: [PATCH 3/5] feat(mcp): mask write-back channel for tool calls (AISIX-Cloud#1330) - json_splice: byte-splicing rewrite of JSON string values selected by a path predicate. Decodes only the selected leaves, re-encodes the replacements, and splices them into the original buffer - key order, whitespace, number spellings, and escape choices outside the masked spans survive byte-for-byte (a Value round-trip cannot promise that: BTreeMap re-sorts keys and numbers re-serialise canonically). - /mcp input hook: string leaves under params.arguments are rewritten through the chain's sync redactor after the block check; the inner gateway receives the masked body (Content-Length refreshed). A splice failure fails closed. - /mcp output hook: output_guardrail_block becomes verdict + write-back (apply_output_guardrails). Masked spans are spliced in place across result.content[].text, result.content[].resource.text, and every string leaf under result.structuredContent; the client receives the original bytes everywhere else. - scan surface: embedded-resource text (type=resource, resource.text) now enters the output scan set - previously a sibling text block kept the set non-empty and the whole log body went unread. base64 blob resources are deliberately not decoded yet (design point pending). - usage: MCP events now carry redacted_entity_counts, and full-content exporters receive the POST-MASK tool args/result via CapturedContent (capture cloned after the write-backs, same order as the LLM path). --- crates/aisix-proxy/src/json_splice.rs | 412 ++++++++++++++++++++++ crates/aisix-proxy/src/lib.rs | 1 + crates/aisix-proxy/src/mcp.rs | 488 ++++++++++++++++++++++---- 3 files changed, 836 insertions(+), 65 deletions(-) create mode 100644 crates/aisix-proxy/src/json_splice.rs diff --git a/crates/aisix-proxy/src/json_splice.rs b/crates/aisix-proxy/src/json_splice.rs new file mode 100644 index 000000000..2d9464515 --- /dev/null +++ b/crates/aisix-proxy/src/json_splice.rs @@ -0,0 +1,412 @@ +//! Byte-splicing rewrite of JSON string VALUES (AISIX-Cloud#1330). +//! +//! The MCP write-back channel must return every byte outside a masked +//! span verbatim. A `serde_json::Value` round-trip cannot promise that: +//! this workspace's `Map` is a BTreeMap (keys re-sort), and numbers +//! re-serialise canonically (`1e3` → `1000.0`). So this module never +//! re-serialises the document — it scans the raw bytes once, decodes +//! only the string values a path predicate selects, and splices the +//! re-encoded replacements back into the original buffer. Everything +//! else — key order, whitespace, number spellings, escape choices — +//! survives byte-for-byte. +//! +//! Object KEYS are never offered for rewrite (they are schema, not +//! data — same rule as `collect_string_leaves` in the MCP scan path), +//! but they ARE decoded to build the path handed to the predicate. +//! +//! The scanner assumes syntactically valid JSON (callers run it on +//! bytes `serde_json` has already parsed) and still fails safe: any +//! unexpected byte, overrun, or depth blow-up returns an error rather +//! than a partially rewritten document. Callers decide the failure +//! policy (the MCP output hook fails closed). + +use std::ops::Range; + +/// One step of the path from the document root to a value. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PathSeg { + /// Object member, key decoded (escapes resolved). + Key(String), + /// Array element index. + Index(usize), +} + +impl PathSeg { + /// `true` when this segment is `Key(name)`. + pub fn is_key(&self, name: &str) -> bool { + matches!(self, PathSeg::Key(k) if k == name) + } +} + +/// Scanner failure. Carries no document content (the byte offset only), +/// so an error can be logged without leaking the payload. +#[derive(Debug, thiserror::Error)] +#[error("json splice scan failed at byte {at}")] +pub struct SpliceError { + at: usize, +} + +/// Depth cap. `serde_json` refuses documents deeper than 128, so bytes +/// that reached a splice call can never hit this; it bounds the scanner +/// on its own anyway. +const MAX_DEPTH: usize = 256; + +/// Rewrite the string values of `input` selected by `should_rewrite`, +/// leaving every other byte untouched. +/// +/// For each string VALUE (never a key) whose path satisfies the +/// predicate, the decoded text is offered to `rewrite`; `Some(new)` +/// replaces that value's bytes with the JSON encoding of `new`. +/// +/// Returns `Ok(None)` when nothing changed (callers keep the original +/// buffer — the no-hit case allocates nothing), `Ok(Some(bytes))` with +/// the spliced document otherwise. +pub fn rewrite_string_values( + input: &[u8], + mut should_rewrite: impl FnMut(&[PathSeg]) -> bool, + mut rewrite: impl FnMut(&str) -> Option, +) -> Result>, SpliceError> { + enum Frame { + Object, + Array, + } + + let err = |at: usize| SpliceError { at }; + let mut splices: Vec<(Range, String)> = Vec::new(); + let mut path: Vec = Vec::new(); + let mut frames: Vec = Vec::new(); + let mut pos = 0usize; + + let skip_ws = |pos: &mut usize| { + while *pos < input.len() && matches!(input[*pos], b' ' | b'\t' | b'\n' | b'\r') { + *pos += 1; + } + }; + // Span of the string token starting at `start` (must be `"`), + // inclusive of both quotes. + let scan_string = |start: usize| -> Result { + let mut i = start + 1; + while i < input.len() { + match input[i] { + b'\\' => i += 2, // skips the escaped byte; `\uXXXX` needs no care (hex only) + b'"' => return Ok(i + 1), + _ => i += 1, + } + } + Err(SpliceError { at: start }) + }; + let decode_str = |range: Range| -> Result { + let at = range.start; + serde_json::from_slice::(&input[range]).map_err(|_| SpliceError { at }) + }; + + // `true` → the loop continues at a VALUE position; `false` → the + // value just ended and the closer/comma logic below runs. + 'value: loop { + skip_ws(&mut pos); + let b = *input.get(pos).ok_or_else(|| err(pos))?; + match b { + b'{' => { + frames.push(Frame::Object); + if frames.len() > MAX_DEPTH { + return Err(err(pos)); + } + pos += 1; + skip_ws(&mut pos); + match input.get(pos) { + Some(b'}') => { + pos += 1; + frames.pop(); + // fall through to after-value + } + Some(b'"') => { + let end = scan_string(pos)?; + path.push(PathSeg::Key(decode_str(pos..end)?)); + pos = end; + skip_ws(&mut pos); + if input.get(pos) != Some(&b':') { + return Err(err(pos)); + } + pos += 1; + continue 'value; + } + _ => return Err(err(pos)), + } + } + b'[' => { + frames.push(Frame::Array); + if frames.len() > MAX_DEPTH { + return Err(err(pos)); + } + pos += 1; + skip_ws(&mut pos); + if input.get(pos) == Some(&b']') { + pos += 1; + frames.pop(); + // fall through to after-value + } else { + path.push(PathSeg::Index(0)); + continue 'value; + } + } + b'"' => { + let end = scan_string(pos)?; + if should_rewrite(&path) { + let decoded = decode_str(pos..end)?; + if let Some(new) = rewrite(&decoded) { + // to_string of a String is infallible. + let encoded = serde_json::to_string(&new).map_err(|_| err(pos))?; + splices.push((pos..end, encoded)); + } + } + pos = end; + } + // Number / true / false / null. The scanner does not + // re-validate the token — the bytes already parsed upstream — + // it only needs the token's extent. + b'-' | b'0'..=b'9' | b't' | b'f' | b'n' => { + while pos < input.len() + && matches!(input[pos], + b'-' | b'+' | b'.' | b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z') + { + pos += 1; + } + } + _ => return Err(err(pos)), + } + + // A value just ended: unwind closers, then either continue with + // the next member/element or finish. + loop { + skip_ws(&mut pos); + let Some(frame) = frames.last() else { + // Root value complete: only trailing whitespace may follow. + if pos != input.len() { + return Err(err(pos)); + } + break 'value; + }; + match (frame, input.get(pos)) { + (Frame::Object, Some(b',')) => { + pos += 1; + path.pop(); + skip_ws(&mut pos); + if input.get(pos) != Some(&b'"') { + return Err(err(pos)); + } + let end = scan_string(pos)?; + path.push(PathSeg::Key(decode_str(pos..end)?)); + pos = end; + skip_ws(&mut pos); + if input.get(pos) != Some(&b':') { + return Err(err(pos)); + } + pos += 1; + continue 'value; + } + (Frame::Object, Some(b'}')) => { + pos += 1; + path.pop(); + frames.pop(); + } + (Frame::Array, Some(b',')) => { + pos += 1; + match path.last_mut() { + Some(PathSeg::Index(i)) => *i += 1, + _ => return Err(err(pos)), + } + continue 'value; + } + (Frame::Array, Some(b']')) => { + pos += 1; + path.pop(); + frames.pop(); + } + _ => return Err(err(pos)), + } + } + } + + if splices.is_empty() { + return Ok(None); + } + // Splices were recorded in scan order (strictly ascending, disjoint). + let mut out = Vec::with_capacity(input.len()); + let mut copied = 0usize; + for (range, replacement) in splices { + out.extend_from_slice(&input[copied..range.start]); + out.extend_from_slice(replacement.as_bytes()); + copied = range.end; + } + out.extend_from_slice(&input[copied..]); + Ok(Some(out)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rewrite_all(input: &str, f: impl FnMut(&str) -> Option) -> Option { + rewrite_string_values(input.as_bytes(), |_| true, f) + .unwrap() + .map(|b| String::from_utf8(b).unwrap()) + } + + #[test] + fn rewrites_only_the_selected_leaf_bytes() { + // Deliberately hostile formatting: odd whitespace, exotic number + // spellings, escape choices — none of it may change. + let doc = "{ \"a\" : 1e3,\"b\":[ true, \"secret\" ,null] , \"c\": 0.1000 }"; + let out = rewrite_all(doc, |s| (s == "secret").then(|| "MASK".to_string())).unwrap(); + assert_eq!( + out, + "{ \"a\" : 1e3,\"b\":[ true, \"MASK\" ,null] , \"c\": 0.1000 }" + ); + } + + #[test] + fn no_change_returns_none() { + let doc = r#"{"a": "x", "b": 2}"#; + assert!(rewrite_string_values(doc.as_bytes(), |_| true, |_| None) + .unwrap() + .is_none()); + } + + #[test] + fn keys_are_never_offered_but_shape_the_path() { + let doc = r#"{"secret": {"inner": "value"}}"#; + let mut offered = Vec::new(); + let mut paths = Vec::new(); + rewrite_string_values( + doc.as_bytes(), + |p| { + paths.push(p.to_vec()); + true + }, + |s| { + offered.push(s.to_owned()); + None + }, + ) + .unwrap(); + // Only the value is offered — the keys "secret"/"inner" are not. + assert_eq!(offered, vec!["value"]); + assert_eq!( + paths, + vec![vec![ + PathSeg::Key("secret".into()), + PathSeg::Key("inner".into()) + ]], + ); + } + + #[test] + fn array_indices_and_nesting_track_correctly() { + let doc = r#"{"params":{"arguments":{"xs":["a",{"y":"b"},[],"c"],"n":7}},"id":"z"}"#; + let mut seen = Vec::new(); + rewrite_string_values( + doc.as_bytes(), + |p| { + p.first().is_some_and(|s| s.is_key("params")) + && p.get(1).is_some_and(|s| s.is_key("arguments")) + }, + |s| { + seen.push(s.to_owned()); + None + }, + ) + .unwrap(); + // "z" (outside params.arguments) is filtered by the predicate. + assert_eq!(seen, vec!["a", "b", "c"]); + } + + #[test] + fn escaped_key_decodes_for_the_predicate() { + // `param\u0073` decodes to "params" — the predicate must see the + // decoded spelling or a smuggled escape would bypass the scope. + let doc = r#"{"param\u0073":{"arguments":{"t":"hit"}}}"#; + let out = rewrite_string_values( + doc.as_bytes(), + |p| p.first().is_some_and(|s| s.is_key("params")), + |s| (s == "hit").then(|| "X".to_string()), + ) + .unwrap() + .unwrap(); + // The key's original escape spelling is untouched; only the value changed. + assert_eq!( + String::from_utf8(out).unwrap(), + r#"{"param\u0073":{"arguments":{"t":"X"}}}"# + ); + } + + #[test] + fn escaped_and_multibyte_values_reencode_correctly() { + let doc = r#"{"a":"line\nbreak \"q\" 版本","b":"清 洁"}"#; + let out = rewrite_all(doc, |s| { + (s == "line\nbreak \"q\" 版本").then(|| "打码\"了\"".to_string()) + }) + .unwrap(); + // serde_json re-encodes the replacement; the untouched leaf keeps + // its original bytes. + assert_eq!(out, r#"{"a":"打码\"了\"","b":"清 洁"}"#); + let v: serde_json::Value = serde_json::from_str(&out).unwrap(); + assert_eq!(v["a"], "打码\"了\""); + } + + #[test] + fn multiple_rewrites_splice_in_order() { + let doc = r#"["one","keep","two"]"#; + let out = rewrite_all(doc, |s| match s { + "one" => Some("1".into()), + "two" => Some("2".into()), + _ => None, + }) + .unwrap(); + assert_eq!(out, r#"["1","keep","2"]"#); + } + + #[test] + fn empty_containers_and_scalars_pass_through() { + for doc in [ + r#"{}"#, + r#"[]"#, + r#"{"a":[],"b":{}}"#, + "42", + "null", + r#""s""#, + ] { + let got = rewrite_string_values(doc.as_bytes(), |_| true, |_| None).unwrap(); + assert!(got.is_none(), "{doc}"); + } + // A bare root string IS a value and can be rewritten. + let out = rewrite_all(r#""s""#, |_| Some("t".into())).unwrap(); + assert_eq!(out, r#""t""#); + } + + #[test] + fn malformed_input_errors_instead_of_partial_output() { + for doc in [ + r#"{"a": }"#, + r#"{"a":"x""#, + r#"{"a":"x"} trailing"#, + r#"{'a':1}"#, + ] { + assert!( + rewrite_string_values(doc.as_bytes(), |_| true, |_| Some("m".into())).is_err(), + "{doc}", + ); + } + } + + #[test] + fn depth_cap_errors() { + let mut doc = String::new(); + for _ in 0..300 { + doc.push('['); + } + for _ in 0..300 { + doc.push(']'); + } + assert!(rewrite_string_values(doc.as_bytes(), |_| true, |_| None).is_err()); + } +} diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index ae44e9692..84f35cec3 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -47,6 +47,7 @@ pub mod health; mod http_client; mod images; mod jobs; +mod json_splice; mod jwt; mod mcp; mod mcp_auth; diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index 9130e4bfe..f411bd74b 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -267,7 +267,9 @@ async fn dispatch( // Buffer the body so the JSON-RPC method can be inspected, then rebuilt for // the gateway. The global body-limit layer has already capped the size. - let (parts, body) = request.into_parts(); + // `parts` is mutable so the input mask write-back below can refresh + // Content-Length when it changes the body. + let (mut parts, body) = request.into_parts(); let bytes = match to_bytes( body, crate::error::body_read_cap(state.request_body_limit_bytes), @@ -354,6 +356,8 @@ async fn dispatch( Duration::ZERO, false, Vec::new(), + crate::redact::RedactionCounts::new(), + None, trace, /* dispatched */ false, ); @@ -427,6 +431,8 @@ async fn dispatch( Duration::ZERO, true, monitor_hits, + crate::redact::RedactionCounts::new(), + None, trace, /* dispatched */ false, ); @@ -434,6 +440,78 @@ async fn dispatch( } } + // Content capture opt-in (AISIX-Cloud#1330): the largest cap any + // full-content exporter requests, same gate as the LLM handlers. + // Resolved before the write-backs so the response is buffered when + // capture needs it even without a guardrail chain. + let capture_cap = if is_tool_call { + aisix_obs::content_capture_cap( + snapshot + .observability_exporters + .entries() + .iter() + .map(|e| &e.value), + ) + } else { + None + }; + + // Input mask write-back (AISIX-Cloud#1330): rewrite the string leaves + // under `params.arguments`, splicing on the raw bytes so every byte + // outside a masked span reaches the gateway verbatim (`json_splice`). + // Runs AFTER the block check — block rules judge the original text — + // and only when the chain has a sync redactor (kind=pii mask rules). + let mut redaction_counts = crate::redact::RedactionCounts::new(); + let bytes = match &guardrail_chain { + Some(chain) if aisix_guardrails::Guardrail::redacts_input(chain) => { + match rewrite_tool_arguments(chain, &bytes) { + Ok(None) => bytes, + Ok(Some((rewritten, counts))) => { + crate::redact::merge_counts(&mut redaction_counts, counts); + // The body length changed; the inner service must not + // trust a stale Content-Length. + parts.headers.insert( + axum::http::header::CONTENT_LENGTH, + axum::http::HeaderValue::from(rewritten.len()), + ); + axum::body::Bytes::from(rewritten) + } + Err(err) => { + // Structurally impossible (the body already parsed as + // JSON for the peek) — fail closed rather than forward + // content the operator's mask policy should have hidden. + tracing::warn!( + tool = %mcp_tool, + error = %err, + "mcp input mask splice failed; blocking tool call", + ); + emit_tool_call_usage( + state, + &snapshot, + &auth, + request_id, + &mcp_server, + &mcp_tool, + StatusCode::OK.as_u16(), + Duration::ZERO, + true, + monitor_hits, + crate::redact::RedactionCounts::new(), + None, + trace, + /* dispatched */ false, + ); + return jsonrpc_guardrail_block(rpc_id, "tool call", None); + } + } + } + _ => bytes, + }; + // Post-mask by construction: cloned AFTER the write-back above, so a + // capturing exporter can never archive a value the mask removed + // (same ordering rule as the LLM path — mask first, capture after). + let captured_args = capture_cap.map(|_| bytes.clone()); + // Scope the gateway to the tools this caller's key permits — resolved // from the key together with the environment/team MCP access policies — // so MCP tool access is governed by the same key object as LLM access. @@ -479,11 +557,14 @@ async fn dispatch( }; let latency = started.elapsed(); - // Output guardrails: scan the tool result before returning it. The response - // body is only buffered when a guardrail chain is attached. - let response = if let Some(chain) = &guardrail_chain { - let (resp_parts, resp_body) = response.into_parts(); - let resp_bytes = match to_bytes( + // Output guardrails + mask write-back: scan the tool result before + // returning it, rewriting masked spans in place. The response body is + // buffered when a guardrail chain is attached OR a full-content + // exporter wants the result captured. + let mut captured_result: Option = None; + let response = if guardrail_chain.is_some() || capture_cap.is_some() { + let (mut resp_parts, resp_body) = response.into_parts(); + let mut resp_bytes = match to_bytes( resp_body, crate::error::body_read_cap(state.request_body_limit_bytes), ) @@ -494,31 +575,53 @@ async fn dispatch( return (StatusCode::BAD_GATEWAY, "invalid upstream response").into_response() } }; - if let Some(guardrail_name) = - output_guardrail_block(chain, &resp_bytes, &mcp_tool, &mut monitor_hits).await - { - emit_tool_call_usage( - state, - &snapshot, - &auth, - request_id, - &mcp_server, - &mcp_tool, - StatusCode::OK.as_u16(), - latency, - true, - monitor_hits, - trace, - /* dispatched */ true, - ); - return jsonrpc_guardrail_block(rpc_id, "tool result", guardrail_name.as_deref()); + if let Some(chain) = &guardrail_chain { + match apply_output_guardrails(chain, &resp_bytes, &mcp_tool, &mut monitor_hits).await { + ToolResultOutcome::Block(guardrail_name) => { + emit_tool_call_usage( + state, + &snapshot, + &auth, + request_id, + &mcp_server, + &mcp_tool, + StatusCode::OK.as_u16(), + latency, + true, + monitor_hits, + redaction_counts, + None, + trace, + /* dispatched */ true, + ); + return jsonrpc_guardrail_block( + rpc_id, + "tool result", + guardrail_name.as_deref(), + ); + } + ToolResultOutcome::Allow(Some((rewritten, counts))) => { + crate::redact::merge_counts(&mut redaction_counts, counts); + resp_parts.headers.insert( + axum::http::header::CONTENT_LENGTH, + axum::http::HeaderValue::from(rewritten.len()), + ); + resp_bytes = axum::body::Bytes::from(rewritten); + } + ToolResultOutcome::Allow(None) => {} + } } + // Post-mask by construction — cloned after the write-back above. + captured_result = capture_cap.map(|_| resp_bytes.clone()); Response::from_parts(resp_parts, Body::from(resp_bytes)) } else { response }; if is_tool_call { + let capture = capture_cap.map(|cap| { + tool_call_capture(cap, captured_args.as_deref(), captured_result.as_deref()) + }); emit_tool_call_usage( state, &snapshot, @@ -530,6 +633,8 @@ async fn dispatch( latency, false, monitor_hits, + redaction_counts, + capture.as_ref(), trace, /* dispatched */ true, ); @@ -537,18 +642,77 @@ async fn dispatch( response } -/// Run the output guardrail chain over an MCP tool result. Returns `Some(_)` to -/// block — the inner value is the firing guardrail's name, or `None` for a -/// fail-closed block on a body that cannot be parsed — and `None` to allow. The -/// tool result's text is fed to `check_output` as assistant text, the same hook -/// the LLM response path uses; a protocol-level error envelope (no `result`) has -/// nothing to scan and is allowed. -async fn output_guardrail_block( +/// Splice-rewrite the string leaves under `params.arguments` of a +/// `tools/call` body through the chain's input redactor (AISIX-Cloud#1330). +/// `Ok(None)` = nothing masked, keep the original bytes. +fn rewrite_tool_arguments( + chain: &aisix_guardrails::GuardrailChain, + body: &[u8], +) -> Result, crate::redact::RedactionCounts)>, crate::json_splice::SpliceError> { + let mut counts = crate::redact::RedactionCounts::new(); + let out = crate::json_splice::rewrite_string_values( + body, + |path| { + path.first().is_some_and(|s| s.is_key("params")) + && path.get(1).is_some_and(|s| s.is_key("arguments")) + }, + |text| { + aisix_guardrails::Guardrail::redact_input_text(chain, text).map(|r| { + crate::redact::merge_counts(&mut counts, r.counts); + r.text + }) + }, + )?; + Ok(out.map(|bytes| (bytes, counts))) +} + +/// The post-mask content-capture pair for a tool call: `prompt` is the +/// (rewritten) `params.arguments`, `response` the (rewritten) `result`. +/// Both re-serialise through `Value` — capture is telemetry, not wire +/// bytes, so canonical key order is fine here. +fn tool_call_capture( + cap: u32, + request_bytes: Option<&[u8]>, + result_bytes: Option<&[u8]>, +) -> aisix_obs::CapturedContent { + let prompt = request_bytes + .and_then(|b| serde_json::from_slice::(b).ok()) + .and_then(|p| p.params) + .and_then(|p| p.arguments) + .map(|v| v.to_string()) + .unwrap_or_default(); + let response = result_bytes + .and_then(|b| serde_json::from_slice::(b).ok()) + .and_then(|mut v| v.get_mut("result").map(serde_json::Value::take)) + .map(|v| v.to_string()) + .unwrap_or_default(); + aisix_obs::CapturedContent::new(&prompt, &response, cap as usize) +} + +/// Outcome of the output-hook guardrail pass over an MCP tool result. +enum ToolResultOutcome { + /// Reject the tool result. The inner value is the firing guardrail's + /// name, or `None` for a fail-closed block (unparseable body / splice + /// failure). + Block(Option), + /// Release the tool result; `Some` carries the mask-rewritten body + /// bytes and the per-detector counts. + Allow(Option<(Vec, crate::redact::RedactionCounts)>), +} + +/// Run the output guardrail chain over an MCP tool result: verdict AND +/// mask write-back (AISIX-Cloud#1330). The tool result's text is fed to +/// `check_output` as assistant text, the same hook the LLM response path +/// uses; a protocol-level error envelope (no `result`) has nothing to +/// scan and is allowed. When the chain carries a sync redactor, masked +/// spans are spliced back into the raw body bytes — every byte outside a +/// masked span reaches the client verbatim. +async fn apply_output_guardrails( chain: &aisix_guardrails::GuardrailChain, response_bytes: &[u8], tool: &str, monitor_hits: &mut Vec, -) -> Option> { +) -> ToolResultOutcome { // Fail closed on an unparseable body. The `/mcp` gateway is configured // `json_response = true`, so a `tools/call` returns a single // `application/json` object; a body that does not parse (e.g. if that ever @@ -556,28 +720,46 @@ async fn output_guardrail_block( // guardrail — block rather than allow. let value: serde_json::Value = match serde_json::from_slice(response_bytes) { Ok(value) => value, - Err(_) => return Some(None), + Err(_) => return ToolResultOutcome::Block(None), }; // A protocol-level error envelope (no `result`) has no tool output to scan. - let result = value.get("result")?; - // Scan the client-visible tool text — the `text`-type content blocks the - // result carries — not the serialized JSON envelope. This keeps MCP output - // and LLM output on the same representation: a keyword guardrail sees the - // decoded prose, so envelope field names (`content`, `type`, `text`) can't - // trip a false positive, and escaped characters can't hide blocked content. - let mut scanned: Vec = result - .get("content") - .and_then(|c| c.as_array()) - .map(|blocks| { - blocks - .iter() - .filter(|b| b.get("type").and_then(|t| t.as_str()) == Some("text")) - .filter_map(|b| b.get("text").and_then(|t| t.as_str())) - .filter(|text| !text.is_empty()) - .map(str::to_owned) - .collect() - }) - .unwrap_or_default(); + let Some(result) = value.get("result") else { + return ToolResultOutcome::Allow(None); + }; + // Scan the client-visible tool text — the decoded content-block strings, + // not the serialized JSON envelope. This keeps MCP output and LLM output + // on the same representation: a keyword guardrail sees the decoded prose, + // so envelope field names (`content`, `type`, `text`) can't trip a false + // positive, and escaped characters can't hide blocked content. + let mut scanned: Vec = Vec::new(); + if let Some(blocks) = result.get("content").and_then(|c| c.as_array()) { + for block in blocks { + // The `text` string of any block. (Previously filtered to + // `type == "text"`; any block-level `text` is still a string + // VALUE, never an envelope key, so scanning it is safe and + // strictly wider.) + if let Some(text) = block.get("text").and_then(|t| t.as_str()) { + if !text.is_empty() { + scanned.push(text.to_owned()); + } + } + // Embedded resource (`type: "resource"`): the payload rides + // `resource.text` — the natural shape for a log file — and + // previously escaped the scan entirely whenever a sibling + // text block kept `scanned` non-empty (AISIX-Cloud#1330). + // base64 `blob` resources are NOT decoded here (deliberate: + // mimeType allowlist + size caps are an open design point). + if let Some(text) = block + .get("resource") + .and_then(|r| r.get("text")) + .and_then(|t| t.as_str()) + { + if !text.is_empty() { + scanned.push(text.to_owned()); + } + } + } + } // `structuredContent` is serialized to the client ALONGSIDE `content`, and // the spec only RECOMMENDS mirroring it into a text block — so a tool can // return clean prose and carry the sensitive value here. Scan its string @@ -602,20 +784,64 @@ async fn output_guardrail_block( }; let (verdict, hits) = aisix_guardrails::Guardrail::check_output_observed(chain, &resp).await; monitor_hits.extend(hits); - match verdict { - aisix_guardrails::GuardrailVerdict::Block { - reason, - guardrail_name, - } => { + if let aisix_guardrails::GuardrailVerdict::Block { + reason, + guardrail_name, + } = verdict + { + tracing::warn!( + guardrail_hook = "output", + tool = %tool, + reason = %reason, + "guardrail blocked MCP tool result" + ); + return ToolResultOutcome::Block(guardrail_name); + } + // Mask write-back over the same surface the scan covers: + // `result.content[i].text`, `result.content[i].resource.text`, and + // every string leaf under `result.structuredContent`. + if !aisix_guardrails::Guardrail::redacts_output(chain) { + return ToolResultOutcome::Allow(None); + } + let mut counts = crate::redact::RedactionCounts::new(); + let rewritten = crate::json_splice::rewrite_string_values( + response_bytes, + |path| { + if !path.first().is_some_and(|s| s.is_key("result")) { + return false; + } + match path.get(1) { + Some(seg) if seg.is_key("structuredContent") => true, + Some(seg) if seg.is_key("content") => { + matches!(path.get(2), Some(crate::json_splice::PathSeg::Index(_))) + && ((path.len() == 4 && path[3].is_key("text")) + || (path.len() == 5 + && path[3].is_key("resource") + && path[4].is_key("text"))) + } + _ => false, + } + }, + |text| { + aisix_guardrails::Guardrail::redact_output_text(chain, text).map(|r| { + crate::redact::merge_counts(&mut counts, r.counts); + r.text + }) + }, + ); + match rewritten { + Ok(None) => ToolResultOutcome::Allow(None), + Ok(Some(bytes)) => ToolResultOutcome::Allow(Some((bytes, counts))), + Err(err) => { + // Structurally impossible (the body parsed above); fail closed + // rather than release a result the mask policy should rewrite. tracing::warn!( - guardrail_hook = "output", tool = %tool, - reason = %reason, - "guardrail blocked MCP tool result" + error = %err, + "mcp output mask splice failed; blocking tool result", ); - Some(guardrail_name) + ToolResultOutcome::Block(None) } - _ => None, } } @@ -653,6 +879,12 @@ fn emit_tool_call_usage( latency: Duration, guardrail_blocked: bool, guardrail_monitor_hits: Vec, + // Per-detector mask counts from the write-back passes (names only, + // never values — #932 no-leak). + redacted_entity_counts: crate::redact::RedactionCounts, + // Post-mask captured args/result for full-content exporters + // (AISIX-Cloud#1330); `None` when no exporter captures content. + content: Option<&aisix_obs::CapturedContent>, trace: Option<&std::sync::Arc>, // Whether the tool call reached its upstream MCP server — false for a // quota rejection or an input-guardrail block, which refuse before any @@ -664,6 +896,7 @@ fn emit_tool_call_usage( occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), api_key_id: auth.entry.id.clone(), status_code, + redacted_entity_counts, // Single-attempt endpoint: the attempt spans the whole request, so // the upstream figure and what the caller waited for coincide. upstream_latency_ms: latency.as_millis().min(u32::MAX as u128) as u32, @@ -682,15 +915,15 @@ fn emit_tool_call_usage( // ONE label set across every handler (AISIX-Cloud#1317). // #698: both emit legs (CP sink + per-env exporter fan-out) go through // the shared chokepoint — pre-fix MCP usage reached only the CP sink, so - // exporters never saw /mcp traffic. No content capture (tool args/results - // are a separate surface from prompt/response). + // exporters never saw /mcp traffic. `content` carries the POST-MASK + // tool args/result for full-content exporters (AISIX-Cloud#1330). crate::usage_attr::emit_usage( state, snap, "mcp", event, aisix_obs::UsageEventLabels::default(), - None, + content, trace, /* terminal */ true, dispatched, @@ -1758,6 +1991,131 @@ mod tests { const INPUT_GUARD: &str = r#"{"name":"mcp-input-guard","kind":"keyword","patterns":[{"kind":"literal","value":"forbidden-token"}]}"#; const OUTPUT_GUARD: &str = r#"{"name":"mcp-output-guard","kind":"keyword","hook_point":"output","patterns":[{"kind":"literal","value":"forbidden-token"}]}"#; + /// Verdict-only view over [`apply_output_guardrails`] for the block + /// tests: `Some(name)` = blocked, `None` = allowed (rewritten or not). + async fn output_guardrail_block( + chain: &aisix_guardrails::GuardrailChain, + response_bytes: &[u8], + tool: &str, + monitor_hits: &mut Vec, + ) -> Option> { + match apply_output_guardrails(chain, response_bytes, tool, monitor_hits).await { + ToolResultOutcome::Block(name) => Some(name), + ToolResultOutcome::Allow(_) => None, + } + } + + /// Mask-action pii guardrail with a capture-group custom pattern + + /// literal replacement (AISIX-Cloud#1334), for the write-back tests. + fn pii_mask_guard(hook: &str) -> String { + format!( + r#"{{"name":"pii-mask","kind":"pii","hook_point":"{hook}","custom_patterns":[{{"name":"eda_version","regex":"version\\s*:\\s*(\\d+(?:\\.\\d+)+)","action":"mask","replacement":"***"}}]}}"#, + ) + } + + fn env_chain_with(guardrail_json: &str) -> aisix_guardrails::GuardrailChain { + use aisix_guardrails::{LiveGuardrailIndex, RequestContext}; + let handle = SnapshotHandle::new(snapshot_with_key()); + seed_guardrail(&handle, guardrail_json); + LiveGuardrailIndex::new(handle, None).resolve(&RequestContext { + passthrough_route_id: "", + model_id: "", + mcp_server_id: "", + api_key_id: "ak-1", + team_id: None, + }) + } + + /// AISIX-Cloud#1330: the output hook rewrites masked spans IN PLACE — + /// content text, embedded resource text, structuredContent leaves — + /// and every byte outside the hits (key order, whitespace, number + /// spellings, the envelope) survives verbatim. + #[tokio::test] + async fn output_mask_rewrites_in_place_byte_identical_elsewhere() { + let chain = env_chain_with(&pii_mask_guard("output")); + let body = concat!( + r#"{"jsonrpc":"2.0","id":1,"result":{"content":["#, + r#"{"type":"text","text":"tool version: 12.1 ok"},"#, + r#"{"type":"resource","resource":{"uri":"file:///run.log","mimeType":"text/plain","text":"Compile version: 2022.4 Elapsed: 12.345s"}}"#, + r#"], "structuredContent":{"log":"cfg version: 9.0 end","cells": 1e3}}}"#, + ); + let out = match apply_output_guardrails(&chain, body.as_bytes(), "report", &mut Vec::new()) + .await + { + ToolResultOutcome::Allow(Some((bytes, counts))) => { + assert_eq!(counts.get("eda_version"), Some(&3)); + String::from_utf8(bytes).unwrap() + } + other => panic!( + "expected a rewritten Allow, got {}", + match other { + ToolResultOutcome::Block(_) => "Block", + ToolResultOutcome::Allow(None) => "Allow(None)", + ToolResultOutcome::Allow(_) => unreachable!(), + } + ), + }; + assert_eq!( + out, + concat!( + r#"{"jsonrpc":"2.0","id":1,"result":{"content":["#, + r#"{"type":"text","text":"tool version: *** ok"},"#, + r#"{"type":"resource","resource":{"uri":"file:///run.log","mimeType":"text/plain","text":"Compile version: *** Elapsed: 12.345s"}}"#, + r#"], "structuredContent":{"log":"cfg version: *** end","cells": 1e3}}}"#, + ), + ); + // Structural acceptance: the rewritten body still parses. + let v: serde_json::Value = serde_json::from_str(&out).unwrap(); + assert_eq!( + v["result"]["content"][1]["resource"]["mimeType"], + "text/plain" + ); + + // A no-hit result is returned with NO rewrite at all. + let clean = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"Elapsed: 12.345s"}]}}"#; + assert!(matches!( + apply_output_guardrails(&chain, clean, "report", &mut Vec::new()).await, + ToolResultOutcome::Allow(None), + )); + } + + /// The input hook rewrites only the string leaves under + /// `params.arguments`; the method, tool name, id, and every other + /// byte are untouched. + #[test] + fn input_mask_rewrites_only_arguments() { + let chain = env_chain_with(&pii_mask_guard("input")); + let body = r#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"eda__echo","arguments":{"text":"build version: 12.1 done","note":"version untouched"}}}"#; + let (bytes, counts) = rewrite_tool_arguments(&chain, body.as_bytes()) + .unwrap() + .expect("a hit rewrites"); + assert_eq!(counts.get("eda_version"), Some(&1)); + assert_eq!( + String::from_utf8(bytes).unwrap(), + r#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"eda__echo","arguments":{"text":"build version: *** done","note":"version untouched"}}}"#, + ); + // No hit → no allocation, original bytes forwarded. + let clean = br#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"eda__echo","arguments":{"text":"Elapsed: 12.345s"}}}"#; + assert!(rewrite_tool_arguments(&chain, clean).unwrap().is_none()); + } + + /// AISIX-Cloud#1330 scan-surface fix: a forbidden token that appears + /// ONLY in an embedded resource's text — next to a clean text block — + /// must block. Pre-fix the `type == "text"` filter dropped the + /// resource and the non-empty scan set suppressed the fallback, so + /// exactly this shape (text summary + resource log) went unread. + #[tokio::test] + async fn output_guardrail_scans_embedded_resource_text() { + let chain = env_chain_with(OUTPUT_GUARD); + let body = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"summary ok"},{"type":"resource","resource":{"uri":"file:///run.log","text":"log carries forbidden-token here"}}]}}"#; + assert!( + output_guardrail_block(&chain, body, "report", &mut Vec::new()) + .await + .is_some(), + "resource.text must be scanned even when a text block exists" + ); + } + fn tools_call_with_args(arguments: serde_json::Value) -> HttpRequest { mcp_request( "tools/call", From 96a1000f063f547aba408a47121b603b8fff7562 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 20 Aug 2026 16:29:54 +0800 Subject: [PATCH 4/5] test(mcp): e2e for the mask write-back channel (AISIX-Cloud#1330) - harness: the mock MCP upstream records raw request bodies and gains a report tool returning fixed rich content - a text summary block, an embedded resource (resource.text log), and structuredContent with a string leaf plus a numeric field. - e2e: two DP instances share one upstream; the unguarded instance is the byte-for-byte baseline (kept exporter-less so raw values never reach the SOC target legitimately). Pins: upstream receives masked arguments (byte-diff, ids normalised); the client body is a full-body byte-diff against the baseline with only the masked spans changed and still parses; text block, resource.text, and structuredContent leaves all rewrite (zh + en); rewrite never blocks (200, no error, no isError); the SLS export carries post-mask content and detector counts, never the raw values. --- .../guardrail-mcp-mask-writeback-e2e.test.ts | 300 ++++++++++++++++++ tests/e2e/src/harness/upstream-mcp.ts | 52 ++- 2 files changed, 350 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts diff --git a/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts b/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts new file mode 100644 index 000000000..c3625fbeb --- /dev/null +++ b/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts @@ -0,0 +1,300 @@ +import { createHash, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + decodedTextFor, + EtcdClient, + ProxyClient, + SeedClient, + spawnApp, + startMcpUpstream, + startMockSls, + waitConfigPropagation, + waitForToken, + type McpUpstream, + type MockSls, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: the MCP mask write-back channel (AISIX-Cloud#1330), against a real +// DP + etcd + a real MCP upstream (official TypeScript SDK server) + the +// SLS mock as the SOC export target. +// +// Two DP instances share ONE upstream: +// - appG: pii mask guardrail (hook both) + a full-content SLS exporter; +// - appP: no guardrail, no exporter — the byte-for-byte baseline. +// The baseline lives in its own app so raw sensitive values never reach +// the SOC export legitimately; anything raw in SLS is therefore a leak. +// +// Pinned contract: +// - request direction: the upstream receives the tool arguments with ONLY +// the masked spans rewritten (byte-diff against the baseline app's +// upstream request); +// - response direction: the client receives the tool result with ONLY the +// masked spans rewritten — full-body byte-diff — covering a text block, +// an embedded resource's `resource.text` (the compile-log shape that +// previously escaped scanning entirely), and `structuredContent` leaves; +// the body still parses as JSON; +// - rewrite never blocks: HTTP 200, no JSON-RPC error, no `isError`; +// - the SOC export carries the POST-MASK content and the detector counts, +// never the raw values; +// - Chinese and English label forms both rewrite. + +const KEY = "sk-mcp-writeback-e2e"; +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); + +const CREDENTIAL_REF = "mock"; +const SLS_PROJECT = "aisix-e2e-obs"; +const FULL_LOGSTORE = "mcp-writeback-full"; + +/** SOC-searchable marker: rides the guarded request arguments only. */ +const MARKER = "mcp-soc-probe-7f3a"; + +// Hard negatives: dot-separated numbers that are NOT version values. +const NEG = "Elapsed: 12.345s Memory: 4.2 GB 0.13um top.v:12:1 10.2.255.1"; + +// Request-side text (en + zh hits + negatives + the SOC marker). +const ARG_TEXT = `${MARKER} build version: 12.1 ${NEG} 工具版本：2022.4 完成`; +const ARG_MASKED = `${MARKER} build version: *** ${NEG} 工具版本：*** 完成`; + +// Fixed `report` tool content (response side). +const SUMMARY = `阶段汇总 版本：2022.4 用时 12.345s`; +const LOG = `Compile OK version: 12.1\n${NEG}`; +const STRUCT_LOG = `config {"version": "12.1"} ok`; + +interface RpcReply { + status: number; + body: string; + json?: { + result?: { + content?: Array<{ + type: string; + text?: string; + resource?: { uri?: string; mimeType?: string; text?: string }; + }>; + structuredContent?: Record; + isError?: boolean; + }; + error?: { code: number; message: string }; + }; +} + +describe("mcp mask write-back e2e: /mcp", () => { + let appG: SpawnedApp | undefined; + let appP: SpawnedApp | undefined; + let upstream: McpUpstream | undefined; + let sls: MockSls | undefined; + let etcdReachable = false; + + const post = async (app: SpawnedApp, body: unknown): Promise => { + const res = await fetch(`${app.proxyUrl}/mcp`, { + method: "POST", + headers: { + authorization: `Bearer ${KEY}`, + "content-type": "application/json", + accept: "application/json, text/event-stream", + }, + body: JSON.stringify(body), + }); + const text = await res.text(); + let json: RpcReply["json"]; + try { + json = text ? (JSON.parse(text) as RpcReply["json"]) : undefined; + } catch { + json = undefined; + } + return { status: res.status, body: text, json }; + }; + + /** Per-operation handshake; both apps serve the stateless endpoint. */ + const callTool = async ( + app: SpawnedApp, + name: string, + args: Record, + ): Promise => { + await post(app, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "mcp-writeback-e2e", version: "0.1" }, + }, + }); + return post(app, { + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name, arguments: args }, + }); + }; + + /** Upstream JSON-RPC ids are minted per ephemeral client; strip them so + * the request byte-diff compares everything else exactly. */ + const stripIds = (s: string) => s.replace(/"id":\s*(?:"[^"]*"|\d+)/g, '"id":0'); + + const seedEnv = async ( + app: SpawnedApp, + opts: { guarded: boolean }, + ): Promise => { + const seed = new SeedClient(new EtcdClient(), app.etcdPrefix); + await seed.update("mcp_servers", randomUUID(), { + display_name: "eda", + url: upstream!.url, + enabled: true, + }); + if (opts.guarded) { + await seed.createObservabilityExporter({ + name: "sls-mcp-writeback", + enabled: true, + kind: "aliyun_sls", + endpoint: sls!.url, + project: SLS_PROJECT, + logstore: FULL_LOGSTORE, + credential_ref: CREDENTIAL_REF, + content_mode: "full", + }); + await seed.createGuardrail({ + name: "mcp-writeback-guard", + enabled: true, + hook_point: "both", + kind: "pii", + custom_patterns: [ + { + name: "eda_version", + regex: "(?:version|版本)\\s*[:：]\\s*(\\d+(?:\\.\\d+)+)", + action: "mask", + replacement: "***", + }, + { + name: "eda_version_json", + regex: '"version"\\s*:\\s*"([^"]*)"', + action: "mask", + replacement: "***", + }, + ], + }); + } + // Caller key LAST (AGENTS.md gate rule): the key authenticating + // implies every row above is in the snapshot. + await seed.createApiKey({ + key_hash: sha256(KEY), + allowed_models: [], + mcp_access: { allow: ["*"] }, + }); + const proxy = new ProxyClient(app.proxyUrl, KEY); + await waitConfigPropagation(async () => (await proxy.listModels()).status === 200); + }; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + upstream = await startMcpUpstream("eda", { + reportContent: { summary: SUMMARY, log: LOG, structuredLog: STRUCT_LOG }, + }); + sls = await startMockSls(); + appG = await spawnApp({ + extraEnv: { + [`SLS_CRED_${CREDENTIAL_REF.toUpperCase()}_AK_ID`]: "mock-akid", + [`SLS_CRED_${CREDENTIAL_REF.toUpperCase()}_AK_SECRET`]: "mock-secret", + }, + }); + appP = await spawnApp(); + await seedEnv(appG, { guarded: true }); + await seedEnv(appP, { guarded: false }); + }, 90_000); + + afterAll(async () => { + await appG?.exit(); + await appP?.exit(); + await upstream?.close(); + await sls?.close(); + }); + + test("request: upstream receives the masked arguments, byte-identical elsewhere", async (ctx) => { + if (!etcdReachable || !appG || !appP || !upstream) return ctx.skip(); + + const before = upstream.received.length; + const guarded = await callTool(appG, "eda__echo", { text: ARG_TEXT }); + const baseline = await callTool(appP, "eda__echo", { text: ARG_TEXT }); + expect(guarded.status).toBe(200); + expect(baseline.status).toBe(200); + + // The two upstream-received tools/call bodies differ ONLY in the + // masked spans (and per-connection rpc ids, normalised out). + const calls = upstream.received + .slice(before) + .filter((b) => b.includes('"tools/call"')); + expect(calls).toHaveLength(2); + const [guardedRaw, baselineRaw] = calls; + expect(stripIds(guardedRaw)).toBe( + stripIds(baselineRaw) + .replace("version: 12.1", "version: ***") + .replace("版本：2022.4", "版本：***"), + ); + // The raw values never reached the upstream on the guarded path. + expect(guardedRaw).not.toContain("version: 12.1"); + expect(guardedRaw).not.toContain("2022.4"); + + // The echo reply reflects what the upstream actually saw: the masked + // text — and the negatives byte-identical inside it. + expect(guarded.json?.result?.isError).toBeFalsy(); + expect(guarded.json?.result?.content?.[0]?.text).toBe(`eda:${ARG_MASKED}`); + expect(baseline.json?.result?.content?.[0]?.text).toBe(`eda:${ARG_TEXT}`); + }); + + test("response: text + embedded resource + structuredContent masked in place, full-body byte-diff, still JSON", async (ctx) => { + if (!etcdReachable || !appG || !appP) return ctx.skip(); + + const guarded = await callTool(appG, "eda__report", { text: "go" }); + const baseline = await callTool(appP, "eda__report", { text: "go" }); + expect(guarded.status).toBe(200); + expect(baseline.status).toBe(200); + + // Rewrite never blocks: 200, no protocol error, no tool error. + expect(guarded.json?.error).toBeUndefined(); + expect(guarded.json?.result?.isError).toBeFalsy(); + + // Full-body byte-diff: both apps return the same client-facing bytes + // (same request id, same upstream content) except the masked spans. + // In the raw body the structuredContent hit appears JSON-escaped. + expect(guarded.body).toBe( + baseline.body + .replace("version: 12.1", "version: ***") // resource.text log line + .replace("版本：2022.4", "版本：***") // summary text block + .replace('{\\"version\\": \\"12.1\\"}', '{\\"version\\": \\"***\\"}'), // structured leaf + ); + expect(guarded.body).not.toContain("12.1"); + expect(guarded.body).not.toContain("2022.4"); + + // Structural acceptance: parse the rewritten body, don't just diff it. + const result = guarded.json?.result; + expect(result?.content?.[0]?.text).toBe("阶段汇总 版本：*** 用时 12.345s"); + const resource = result?.content?.[1]?.resource; + expect(resource?.mimeType).toBe("text/plain"); + expect(resource?.text).toBe(`Compile OK version: ***\n${NEG}`); + expect(result?.structuredContent).toEqual({ + log: 'config {"version": "***"} ok', + cells: 42, + }); + }); + + test("SOC export: captured content is the post-mask text with detector counts, never the raw values", async (ctx) => { + if (!etcdReachable || !appG || !sls) return ctx.skip(); + + // The guarded echo call from the request test carried the MARKER; its + // usage event (with captured content) lands on the full logstore. + await waitForToken(sls, FULL_LOGSTORE, MARKER); + const decoded = decodedTextFor(sls, FULL_LOGSTORE); + // Post-mask capture on both directions... + expect(decoded).toContain("version: ***"); + expect(decoded).toContain("版本：***"); + // ...the detector name rides the event (counts, names only)... + expect(decoded).toContain("eda_version"); + // ...and the raw values never reach the SOC target. + expect(decoded).not.toContain("version: 12.1"); + expect(decoded).not.toContain("2022.4"); + }); +}); diff --git a/tests/e2e/src/harness/upstream-mcp.ts b/tests/e2e/src/harness/upstream-mcp.ts index 2378a3fc7..8ccd908fe 100644 --- a/tests/e2e/src/harness/upstream-mcp.ts +++ b/tests/e2e/src/harness/upstream-mcp.ts @@ -15,6 +15,12 @@ import { export interface McpUpstream { /** Streamable HTTP endpoint of this upstream (`http://127.0.0.1:/mcp`). */ url: string; + /** + * Raw request bodies this upstream received, in arrival order — exactly + * the bytes the gateway's MCP client sent, so a masking suite can + * byte-diff what reached the upstream. + */ + received: string[]; close(): Promise; } @@ -26,6 +32,14 @@ export interface McpUpstreamOptions { * every other suite asserts on stays `echo` + `reverse`. */ structuredTool?: boolean; + /** + * Also expose a `report` tool returning fixed rich content regardless of + * its arguments: a text summary block, an embedded resource carrying + * `log` as `resource.text` (the natural shape for a compile/sim log), + * and a `structuredContent` object with a `log` string leaf plus a + * numeric field. The mask write-back suite owns the strings. + */ + reportContent?: { summary: string; log: string; structuredLog: string }; } /** @@ -45,8 +59,9 @@ export async function startMcpUpstream( label: string, options: McpUpstreamOptions = {}, ): Promise { + const received: string[] = []; const httpServer: HttpServer = createServer((req, res) => { - void handle(label, req, res, options); + void handle(label, req, res, options, received); }); await new Promise((resolve) => httpServer.listen(0, "127.0.0.1", resolve), @@ -57,6 +72,7 @@ export async function startMcpUpstream( } return { url: `http://127.0.0.1:${address.port}/mcp`, + received, close: () => new Promise((resolve) => httpServer.close(() => resolve())), }; @@ -67,6 +83,7 @@ async function handle( req: IncomingMessage, res: ServerResponse, options: McpUpstreamOptions, + received: string[], ): Promise { try { if (req.method !== "POST") { @@ -75,7 +92,9 @@ async function handle( } const chunks: Buffer[] = []; for await (const chunk of req) chunks.push(chunk as Buffer); - const body: unknown = JSON.parse(Buffer.concat(chunks).toString("utf8")); + const raw = Buffer.concat(chunks).toString("utf8"); + received.push(raw); + const body: unknown = JSON.parse(raw); const server = new Server( { name: `${label}-upstream`, version: "0.1.0" }, @@ -111,10 +130,39 @@ async function handle( }, ] : []), + ...(options.reportContent + ? [ + { + name: "report", + description: "return a fixed rich report (text + resource + structured)", + inputSchema: { + type: "object" as const, + properties: { text: { type: "string" } }, + }, + }, + ] + : []), ], })); server.setRequestHandler(CallToolRequestSchema, async (request) => { const text = String(request.params.arguments?.text ?? ""); + if (request.params.name === "report" && options.reportContent) { + const { summary, log, structuredLog } = options.reportContent; + return { + content: [ + { type: "text", text: summary }, + { + type: "resource", + resource: { + uri: "file:///run.log", + mimeType: "text/plain", + text: log, + }, + }, + ], + structuredContent: { log: structuredLog, cells: 42 }, + }; + } if (request.params.name === "lookup") { // The text block is deliberately constant and clean: only // `structuredContent` carries the caller's value, which is exactly From e449703741a22a814796328defe47e738831a541 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 20 Aug 2026 16:47:35 +0800 Subject: [PATCH 5/5] fix(mcp): cover resource_link description/title in scan + mask; deflake the SOC e2e assertion Audit riders on #1008: - A resource_link block carries its data in block-level description and title; with any sibling text block the non-empty scan set suppressed the fallback, so a rule anchored only there never fired and PII was never masked - the same silent class the PR fixes for resource.text, one sibling over. Both fields now scan and rewrite; name/uri stay untouched by design (they address the resource, and rewriting an identifier breaks the client's follow-up fetch). The unit test pins the exclusion with a name that WOULD match the mask rule. - The SLS leak assertion raced the report event: it waited only for the echo marker, so the response-direction negatives could false-pass before the report record flushed. The summary prefix is now a second wait token. --- crates/aisix-proxy/src/mcp.rs | 72 ++++++++++++++++--- .../guardrail-mcp-mask-writeback-e2e.test.ts | 6 +- 2 files changed, 67 insertions(+), 11 deletions(-) diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index f411bd74b..d983c7d2a 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -734,13 +734,18 @@ async fn apply_output_guardrails( let mut scanned: Vec = Vec::new(); if let Some(blocks) = result.get("content").and_then(|c| c.as_array()) { for block in blocks { - // The `text` string of any block. (Previously filtered to - // `type == "text"`; any block-level `text` is still a string - // VALUE, never an envelope key, so scanning it is safe and - // strictly wider.) - if let Some(text) = block.get("text").and_then(|t| t.as_str()) { - if !text.is_empty() { - scanned.push(text.to_owned()); + // The `text` string of any block, plus the block-level + // `description`/`title` a `resource_link` carries — all string + // VALUES, never envelope keys, so scanning them is safe and + // strictly wider than the old `type == "text"` filter. A + // link's `name`/`uri` are deliberately NOT scanned or masked: + // they address the resource, and rewriting an identifier + // breaks the client's follow-up fetch. + for key in ["text", "description", "title"] { + if let Some(text) = block.get(key).and_then(|t| t.as_str()) { + if !text.is_empty() { + scanned.push(text.to_owned()); + } } } // Embedded resource (`type: "resource"`): the payload rides @@ -798,8 +803,10 @@ async fn apply_output_guardrails( return ToolResultOutcome::Block(guardrail_name); } // Mask write-back over the same surface the scan covers: - // `result.content[i].text`, `result.content[i].resource.text`, and - // every string leaf under `result.structuredContent`. + // `result.content[i].{text,description,title}`, + // `result.content[i].resource.text`, and every string leaf under + // `result.structuredContent`. (`name`/`uri` stay untouched — they + // address a resource; see the scan-loop comment.) if !aisix_guardrails::Guardrail::redacts_output(chain) { return ToolResultOutcome::Allow(None); } @@ -814,7 +821,10 @@ async fn apply_output_guardrails( Some(seg) if seg.is_key("structuredContent") => true, Some(seg) if seg.is_key("content") => { matches!(path.get(2), Some(crate::json_splice::PathSeg::Index(_))) - && ((path.len() == 4 && path[3].is_key("text")) + && ((path.len() == 4 + && (path[3].is_key("text") + || path[3].is_key("description") + || path[3].is_key("title"))) || (path.len() == 5 && path[3].is_key("resource") && path[4].is_key("text"))) @@ -2116,6 +2126,48 @@ mod tests { ); } + /// Same silent class one sibling over (#1008 audit): a + /// `resource_link` block carries its data in block-level + /// `description`/`title`. Both must scan (block) and mask (rewrite); + /// `name`/`uri` address the resource and stay untouched. + #[tokio::test] + async fn output_guardrail_covers_resource_link_description_and_title() { + // Block rule anchored only in the link description, with a clean + // sibling text block suppressing the fallback. + let chain = env_chain_with(OUTPUT_GUARD); + let body = br#"{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"summary ok"},{"type":"resource_link","uri":"file:///a.log","name":"a.log","description":"holds forbidden-token data"}]}}"#; + assert!( + output_guardrail_block(&chain, body, "list", &mut Vec::new()) + .await + .is_some(), + "resource_link description must be scanned" + ); + + // Mask rule: description and title rewrite in place; name/uri and + // every other byte survive verbatim. + let chain = env_chain_with(&pii_mask_guard("output")); + let body = concat!( + r#"{"jsonrpc":"2.0","id":1,"result":{"content":["#, + r#"{"type":"resource_link","uri":"file:///v.log","name":"run version: 9.9.log","title":"run version: 3.4","description":"log for version: 12.1"}"#, + r#"]}}"#, + ); + match apply_output_guardrails(&chain, body.as_bytes(), "list", &mut Vec::new()).await { + ToolResultOutcome::Allow(Some((bytes, counts))) => { + assert_eq!(counts.get("eda_version"), Some(&2)); + assert_eq!( + String::from_utf8(bytes).unwrap(), + concat!( + r#"{"jsonrpc":"2.0","id":1,"result":{"content":["#, + r#"{"type":"resource_link","uri":"file:///v.log","name":"run version: 9.9.log","title":"run version: ***","description":"log for version: ***"}"#, + r#"]}}"#, + ), + "description/title masked; uri/name untouched", + ); + } + _ => panic!("expected a rewritten Allow"), + } + } + fn tools_call_with_args(arguments: serde_json::Value) -> HttpRequest { mcp_request( "tools/call", diff --git a/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts b/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts index c3625fbeb..b8b57c64e 100644 --- a/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts +++ b/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts @@ -285,8 +285,12 @@ describe("mcp mask write-back e2e: /mcp", () => { if (!etcdReachable || !appG || !sls) return ctx.skip(); // The guarded echo call from the request test carried the MARKER; its - // usage event (with captured content) lands on the full logstore. + // usage event (with captured content) lands on the full logstore. The + // report event must have flushed too before the raw-value negatives + // below can prove anything about the RESPONSE direction — its summary + // prefix is the wait token (#1008 audit MEDIUM-2). await waitForToken(sls, FULL_LOGSTORE, MARKER); + await waitForToken(sls, FULL_LOGSTORE, "阶段汇总"); const decoded = decodedTextFor(sls, FULL_LOGSTORE); // Post-mask capture on both directions... expect(decoded).toContain("version: ***");