diff --git a/crates/aisix-proxy/src/json_splice.rs b/crates/aisix-proxy/src/json_splice.rs new file mode 100644 index 00000000..2d946451 --- /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 ae44e969..84f35cec 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 9130e4bf..d983c7d2 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,51 @@ 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, 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 + // `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 +789,69 @@ 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,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); + } + 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[3].is_key("description") + || path[3].is_key("title"))) + || (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 +889,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 +906,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 +925,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 +2001,173 @@ 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" + ); + } + + /// 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 new file mode 100644 index 00000000..b8b57c64 --- /dev/null +++ b/tests/e2e/src/cases/guardrail-mcp-mask-writeback-e2e.test.ts @@ -0,0 +1,304 @@ +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. 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: ***"); + 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 2378a3fc..8ccd908f 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