i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-12 12:09:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(ws): restore redacted PII in provider frames before client delivery

Browse Source 
Responses WebSocket 只实现了脱敏的一半：请求侧 mask 之后，provider 事件帧在推给
客户端之前没有还原，于是 session 映射内的占位符以 <AETHER:EMAIL:...> 的形式直接
透给客户端。这里补齐响应侧，语义与 HTTP 路径对齐。

- 还原点是 relay loop 的最后一跳（send_client_message 之前、capture_client_frame
  之前），对应 HTTP 的 restore_sync_response_body / StreamingResponseRestorer 所在
  位置。审计与终态观测继续消费脱敏态事件，只有发往客户端的那一份拷贝被还原。
- 复用 privacy::restore_json_strings（改为 pub(crate)）与
  RedactionSession::restore_text，不复制任何还原逻辑：只还原本 session mask 过的
  映射，未映射的占位符原样保留；type / model / id 等协议字段不可能命中 sentinel，
  因此不受影响。批量 {"chunks":[...]} 帧一并递归还原。
- session 生命周期：mask 仍然是 per-turn（slot 依旧每轮新建），但 session 改由连接
  持有，按有界 FIFO 留最近 8 轮。理由是 WS 的会话历史留在上游，continuation 只发
  增量输入，per-turn 释放会漏还原后续响应里回显的更早轮次占位符；HTTP 不会漏，是
  因为它每次重发整段历史、重新 mask 会派生出同一个 sentinel。被挤出窗口的轮次退回
  「占位符原样透传」，不会错误还原成别的值。
- 未命中还原时不改写字节；连接上没有任何 mask session 时（未启用脱敏）连事件 clone
  都不做。

测试：redaction.rs 新增 8 条单测（还原命中/批量帧/未映射占位符原样/未命中不改写/
无 session 不介入/空 session 不留存/审计侧入参不被改写/跨轮还原/窗口有界）；
responses_websocket_e2e 新增一条用例，mock 上游回显收到的 input，断言上游只看到
占位符而客户端拿到真实邮箱。
This commit is contained in:
AAEE86
2026-08-17 14:53:46 +08:00
committed by ZheFox
parent dbf2809bd6
commit 6fc02dad3e
8 changed files with 603 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/websocket/responses/redaction.rs
+388 -14
View File
@@ -1,4 +1,4 @@
//! Responses WebSocket 请求侧 PII 脱敏。
//! Responses WebSocket 两侧的 PII 脱敏：请求侧 mask + 响应侧 restore。
//!
//! HTTP 路径在前门建 `RedactionSessionSlot` 并塞进 `parts.extensions`，planner
//! 只有拿到这个 slot 才会脱敏。WS 的 planning Parts 是合成的：四个规划入口
@@ -8,6 +8,37 @@
//!
//! 因此约定：**进入任何下游用途之前，客户端 `response.create` 只在这里脱敏一次**，
//! 之后所有路径都只看脱敏后的事件。
//!
//! # 响应侧
//!
//! 只 mask 不 restore 是半个实现：HTTP 在把响应交给客户端之前会把占位符换回真实值
//! （`privacy::restore_sync_response_body` / `privacy::StreamingResponseRestorer`），
//! WS 少了这一步，客户端就会直接看到 `<AETHER:EMAIL:...>`。
//! [`ResponsesWebSocketRedactionRestorer`] 补上这一跳，语义与 HTTP 完全一致：
//! 复用 `privacy::restore_json_strings`，只还原本连接自己 mask 出来的映射，
//! 未映射的占位符原样透传。
//!
//! ## session 为什么活在连接上而不是活在这一轮里
//!
//! mask session 由 planner 写进 per-turn 的 slot，而 slot 随 planning Parts 在
//! 规划结束时就被丢弃，响应帧到达时已经无处可取。可选的存活范围有两个：
//!
//! * 挂在 `LogicalTurn` 上：这一轮结束即释放，是 HTTP「一个请求一个 session」的
//! 直译。但 WS 的会话历史留在上游：continuation 只发增量输入，第 1 轮的
//! `input` 不会在第 3 轮重发。于是第 3 轮的响应里若回显了第 1 轮的占位符
//! （"你刚才给我的邮箱是……"），本轮 session 里没有这条映射，占位符就漏给客户端。
//! HTTP 不会漏，是因为它每次都重发整段历史，重新 mask 同一个值会派生出同一个
//! sentinel（HMAC over 规则 + bucket + 值），所以映射天然齐备。
//! * 挂在连接上（当前实现）：每轮仍然各自 mask、各自持有独立 session
//! （per-turn 语义不变），连接只是把最近若干轮的 session 留下来一起参与还原，
//! 凑出的映射集合正好等于「等价 HTTP 请求会拥有的那一份」。
//!
//! 选后者。代价是每帧最多对 [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 个 session
//! 各扫一遍，以及这些 session 的映射会驻留到连接结束；用有界 FIFO 兜住上限。
//! 窗口不够用或每帧成本变高时，正确的下一步是在 `privacy` 侧提供跨 session 的
//! 合并匹配器，而不是把这个窗口调大。
use std::collections::VecDeque;
use serde_json::Value;
@@ -15,6 +46,7 @@ use crate::ai_serving::{
resolve_local_decision_execution_runtime_auth_context, resolve_provider_chat_pii_redaction,
};
use crate::control::GatewayControlDecision;
use crate::privacy::{restore_json_strings, RedactionSession, RedactionSessionSlot};
use crate::{AppState, GatewayError};
/// Responses WebSocket 只承载 `openai:responses`，脱敏规则按这个客户端格式选取。
@@ -26,9 +58,24 @@ const RESPONSES_WEBSOCKET_CLIENT_API_FORMAT: &str = "openai:responses";
/// planner 后续用真实 candidate_id 再取一次配置时，body 已是脱敏态、不会重复写入。
const WEBSOCKET_TURN_REDACTION_CANDIDATE_ID: &str = "responses_websocket_turn";
/// 一条连接最多留几轮的 mask session 用于响应侧还原。
///
/// 取值权衡见模块文档：调大会线性增加每帧还原成本和常驻映射量，调小则更容易漏还原
/// 上游历史里更早那几轮的占位符。8 覆盖的是「上游最可能回显的最近窗口」。
const MAX_RETAINED_TURN_REDACTION_SESSIONS: usize = 8;
/// 一轮客户端 `response.create` 的请求侧脱敏结果。
#[derive(Debug)]
pub(super) struct ResponsesWebSocketTurnRedaction {
/// 脱敏后的客户端事件；这一轮之后所有下游路径都只看它。
pub(super) client_event: Value,
/// 这一轮 mask 出来的映射表，响应侧还原只能靠它。
pub(super) session: RedactionSession,
}
/// 对一条客户端 `response.create` 做请求侧脱敏。
///
/// 返回 `Some(脱敏后的事件)` 仅当脱敏真正命中；`None` 表示未启用或没有命中，调用方
/// 返回 `Some(..)` 仅当脱敏真正命中；`None` 表示未启用或没有命中，调用方
/// 继续用原事件即可（避免未开启脱敏时多一次整包 clone）。
///
/// 脱敏只改写 `instructions` / `input`（见 `privacy::mask_openai_responses_request_value`），
@@ -42,7 +89,7 @@ pub(super) async fn redact_responses_websocket_client_event(
parts: &http::request::Parts,
control_decision: &GatewayControlDecision,
client_event: &Value,
) -> Result<Option<Value>, GatewayError> {
) -> Result<Option<ResponsesWebSocketTurnRedaction>, GatewayError> {
let Some(auth_context) =
resolve_local_decision_execution_runtime_auth_context(control_decision)
else {
@@ -60,7 +107,73 @@ pub(super) async fn redact_responses_websocket_client_event(
if !redaction.redacted {
return Ok(None);
}
Ok(Some(redaction.body_json.into_owned()))
// mask 命中时 `resolve_provider_chat_pii_redaction` 必定把 session 写进 slot。
// 取不到就是内部契约被破坏了，此时继续下发意味着这一轮的响应无法还原、占位符
// 会漏给客户端；按本模块既有的「脱敏链路出错就让这一轮失败」处理，不做降级。
let Some(session) = parts
.extensions
.get::<RedactionSessionSlot>()
.and_then(|slot| slot.take_for_candidate(Some(WEBSOCKET_TURN_REDACTION_CANDIDATE_ID)))
else {
return Err(GatewayError::Internal(
"chat pii redaction masked a Responses WebSocket turn without retaining its session"
.to_string(),
));
};
Ok(Some(ResponsesWebSocketTurnRedaction {
client_event: redaction.body_json.into_owned(),
session,
}))
}
/// 一条连接上「我们 mask 过哪些映射」的留存集合，供响应侧还原使用。
///
/// 每轮一个独立 session（per-turn mask 语义不变），连接按 FIFO 留最近
/// [`MAX_RETAINED_TURN_REDACTION_SESSIONS`] 轮。上游重绑不清空：客户端仍在同一段
/// 对话里，旧占位符可能随重发的输入再次出现。
#[derive(Default)]
pub(super) struct ResponsesWebSocketRedactionRestorer {
sessions: VecDeque<RedactionSession>,
}
impl ResponsesWebSocketRedactionRestorer {
/// 登记这一轮的 mask session。
pub(super) fn register(&mut self, session: RedactionSession) {
if session.mapping_count() == 0 {
return;
}
self.sessions.push_back(session);
while self.sessions.len() > MAX_RETAINED_TURN_REDACTION_SESSIONS {
self.sessions.pop_front();
}
}
/// 把一帧 provider 事件里的占位符换回真实值，返回要发给客户端的帧文本。
///
/// `None` 表示这一帧没有任何东西要还原，调用方必须原样转发上游字节：未启用
/// 脱敏（没有任何 session）时连 clone 都不做。
///
/// 入参只读：审计与终态观测继续消费脱敏态的事件，还原只作用于发往客户端的
/// 那一份拷贝，和 HTTP 侧「审计存脱敏体、线上还原」保持一致。
pub(super) fn restore_provider_frame_text(&self, event: &Value) -> Option<String> {
if self.sessions.is_empty() {
return None;
}
let mut restored_event = event.clone();
let mut restored = false;
for session in &self.sessions {
// 逐 session 还原而不是合并映射：每个 session 只认自己 mask 过的
// sentinel（`RedactionSession::restore_text`），跨 session 合并会绕开
// 这条边界。同一个值在不同轮派生出的 sentinel 相同，所以顺序无关。
restored |= restore_json_strings(&mut restored_event, session);
}
if !restored {
return None;
}
// 刚从 JSON 解析出来的 Value 再序列化不会失败；真失败时宁可让客户端看到
// 占位符，也不能丢掉这一帧——丢帧会让客户端的协议状态机卡死。
serde_json::to_string(&restored_event).ok()
}
}
#[cfg(test)]
@@ -80,7 +193,10 @@ mod tests {
};
use super::super::turn_state::LogicalTurn;
use super::super::turn::prepare_responses_websocket_turn_decision;
use super::redact_responses_websocket_client_event;
use super::{
redact_responses_websocket_client_event, ResponsesWebSocketRedactionRestorer,
ResponsesWebSocketTurnRedaction, MAX_RETAINED_TURN_REDACTION_SESSIONS,
};
use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
use crate::handlers::proxy::websocket::ingress::WebSocketRequestContext;
@@ -89,6 +205,11 @@ mod tests {
const TEST_USER_ID: &str = "user-responses-ws-redaction";
const TEST_API_KEY_ID: &str = "api-key-responses-ws-redaction";
const TEST_EMAIL: &str = "[email protected]";
/// 另一轮用的 PII，用来证明连接级还原覆盖到更早的轮次。
const OTHER_TEST_EMAIL: &str = "[email protected]";
/// 不是本连接 mask 出来的占位符：格式合法（符合 sentinel 正则），但没有任何
/// session 记过它，必须原样透传。
const FOREIGN_SENTINEL: &str = "<AETHER:EMAIL:AAAAAAAAAAAAAAAAAAAA>";
fn auth_export_record() -> StoredAuthApiKeyExportRecord {
StoredAuthApiKeyExportRecord::new(
@@ -191,6 +312,10 @@ mod tests {
}
fn client_event() -> Value {
client_event_with_email(TEST_EMAIL)
}
fn client_event_with_email(email: &str) -> Value {
json!({
"type": "response.create",
"model": "public-model",
@@ -198,11 +323,46 @@ mod tests {
"generate": false,
"input": [{
"role": "user",
"content": [{"type": "input_text", "text": format!("mail {TEST_EMAIL}")}]
"content": [{"type": "input_text", "text": format!("mail {email}")}]
}]
})
}
/// 真跑一遍请求侧脱敏，拿到这一轮的生效事件和 mask session。
async fn turn_redaction(
state: &AppState,
decision: &GatewayControlDecision,
email: &str,
) -> ResponsesWebSocketTurnRedaction {
let context = websocket_context(decision.clone());
let parts = build_planning_parts(&context);
let event = client_event_with_email(email);
redact_responses_websocket_client_event(state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted")
}
/// 这一轮为 `email` 派生出的占位符。
fn sentinel_for(redaction: &ResponsesWebSocketTurnRedaction, email: &str) -> String {
redaction
.session
.sentinel_for_original(email)
.expect("a masked email must have a sentinel")
.to_string()
}
/// 上游回显占位符的一帧 provider 事件。
fn provider_delta_frame(text: &str) -> Value {
json!({
"type": "response.output_text.delta",
"item_id": "msg_ws",
"output_index": 0,
"content_index": 0,
"delta": text,
})
}
#[tokio::test]
async fn websocket_client_event_is_redacted_without_losing_protocol_fields() {
let state = redaction_enabled_state();
@@ -214,7 +374,8 @@ mod tests {
redact_responses_websocket_client_event(&state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted");
.expect("an email in the request should be redacted")
.client_event;
let serialized = serde_json::to_string(&redacted).expect("event should serialize");
assert!(!serialized.contains(TEST_EMAIL), "{serialized}");
@@ -240,7 +401,8 @@ mod tests {
redact_responses_websocket_client_event(&state, &parts, &context.decision, &event)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted");
.expect("an email in the request should be redacted")
.client_event;
// 复用同一个 parts/slot，和 re-plan 在同一 turn 内二次脱敏的情形一致。
let second_pass =
@@ -273,13 +435,9 @@ mod tests {
/// 真跑一遍脱敏，拿到这一轮的「生效事件」。
async fn redacted_client_event(state: &AppState, decision: &GatewayControlDecision) -> Value {
let context = websocket_context(decision.clone());
let parts = build_planning_parts(&context);
let event = client_event();
redact_responses_websocket_client_event(state, &parts, &context.decision, &event)
turn_redaction(state, decision, TEST_EMAIL)
.await
.expect("redaction should resolve")
.expect("an email in the request should be redacted")
.client_event
}
/// 只有 `action` 没有 serde 默认值，其余字段都能省略。
@@ -477,4 +635,220 @@ mod tests {
);
assert_redacted_json(report_context, "quota retry report context");
}
// -----------------------------------------------------------------------
// 响应侧还原
// -----------------------------------------------------------------------
/// 本次修复的核心：上游把占位符回显在事件里，客户端必须拿到真实值。
#[tokio::test]
async fn provider_frame_placeholders_are_restored_before_client_delivery() {
let state = redaction_enabled_state();
let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await;
let sentinel = sentinel_for(&redaction, TEST_EMAIL);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(redaction.session);
let frame = provider_delta_frame(&format!("your mail is {sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("a frame echoing this turn's sentinel must be restored");
assert!(
restored.contains(TEST_EMAIL),
"the client must receive the real value: {restored}"
);
assert!(
!restored.contains(&sentinel),
"no sentinel may survive to the client: {restored}"
);
// 协议字段不受影响，客户端的状态机照旧。
let restored: Value = serde_json::from_str(&restored).expect("restored frame is JSON");
assert_eq!(restored["type"], "response.output_text.delta");
assert_eq!(restored["item_id"], "msg_ws");
assert_eq!(restored["output_index"], 0);
}
/// Codex 把多个事件批量塞进 `{"chunks":[...]}`，还原必须走进批量里。
#[tokio::test]
async fn placeholders_batched_inside_a_chunks_envelope_are_restored() {
let state = redaction_enabled_state();
let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await;
let sentinel = sentinel_for(&redaction, TEST_EMAIL);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(redaction.session);
let frame = json!({
"chunks": [
provider_delta_frame("plain delta"),
provider_delta_frame(&format!("mail {sentinel}")),
]
});
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("a batched sentinel must be restored");
assert!(restored.contains(TEST_EMAIL), "{restored}");
assert!(!restored.contains(&sentinel), "{restored}");
}
/// 只还原本连接 mask 过的映射，和 `RedactionSession::restore_text` 一致：
/// 别处来的占位符（比如客户端自己发的、或上一条连接的）保持原样。
#[tokio::test]
async fn an_unmapped_placeholder_is_left_untouched() {
let state = redaction_enabled_state();
let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await;
let sentinel = sentinel_for(&redaction, TEST_EMAIL);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(redaction.session);
let frame = provider_delta_frame(&format!("{FOREIGN_SENTINEL} and {sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("the mapped sentinel is still restored");
assert!(restored.contains(TEST_EMAIL), "{restored}");
assert!(
restored.contains(FOREIGN_SENTINEL),
"an unmapped placeholder must survive verbatim: {restored}"
);
}
/// 没有命中还原时必须让调用方原样转发上游字节。
#[tokio::test]
async fn a_frame_without_known_placeholders_is_not_rewritten() {
let state = redaction_enabled_state();
let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await;
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(redaction.session);
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame("nothing to restore"))
.is_none(),
"a frame with no mapped sentinel must be relayed byte-for-byte"
);
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL))
.is_none(),
"a frame that only carries unmapped placeholders must not be rewritten"
);
}
/// 未启用脱敏（或这条连接从没 mask 到东西）时，还原器必须完全不介入：
/// 连 clone 都不做，输出就是上游原字节。
#[tokio::test]
async fn a_restorer_without_sessions_never_rewrites_a_frame() {
let restorer = ResponsesWebSocketRedactionRestorer::default();
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL))
.is_none());
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(TEST_EMAIL))
.is_none());
}
/// 空 session（启用了脱敏但这一轮没命中任何规则）不该被留下来白扫每一帧。
#[tokio::test]
async fn a_session_without_mappings_is_not_retained() {
let state = redaction_enabled_state();
let hmac_key = state
.encryption_key()
.expect("the test state carries an encryption key")
.as_bytes()
.to_vec();
let empty_session = crate::privacy::RedactionSession::new(
crate::privacy::RedactionSessionConfig::default_ttl(hmac_key, 0),
);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(empty_session);
assert!(restorer
.restore_provider_frame_text(&provider_delta_frame(FOREIGN_SENTINEL))
.is_none());
}
/// 还原只作用于发给客户端的那一份拷贝：审计和终态观测消费的事件必须保持脱敏态。
#[tokio::test]
async fn restoring_does_not_mutate_the_event_the_audit_path_keeps() {
let state = redaction_enabled_state();
let redaction = turn_redaction(&state, &control_decision(), TEST_EMAIL).await;
let sentinel = sentinel_for(&redaction, TEST_EMAIL);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(redaction.session);
let frame = provider_delta_frame(&format!("mail {sentinel}"));
let before = frame.clone();
let _ = restorer
.restore_provider_frame_text(&frame)
.expect("the frame is restored for the client");
assert_eq!(
frame, before,
"capture_client_frame / 终态观测拿到的事件必须仍是脱敏态"
);
}
/// 连接级持有的意义：WS 的会话历史留在上游，continuation 只发增量输入，
/// 所以第 2 轮的响应可能回显第 1 轮的占位符。per-turn 持有会漏掉这一条。
#[tokio::test]
async fn a_later_turn_restores_a_placeholder_first_masked_by_an_earlier_turn() {
let state = redaction_enabled_state();
let decision = control_decision();
let first = turn_redaction(&state, &decision, TEST_EMAIL).await;
let second = turn_redaction(&state, &decision, OTHER_TEST_EMAIL).await;
let first_sentinel = sentinel_for(&first, TEST_EMAIL);
let second_sentinel = sentinel_for(&second, OTHER_TEST_EMAIL);
assert_ne!(first_sentinel, second_sentinel);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(first.session);
restorer.register(second.session);
let frame = provider_delta_frame(&format!("{first_sentinel} then {second_sentinel}"));
let restored = restorer
.restore_provider_frame_text(&frame)
.expect("both turns' sentinels are restorable on this connection");
assert!(restored.contains(TEST_EMAIL), "{restored}");
assert!(restored.contains(OTHER_TEST_EMAIL), "{restored}");
assert!(!restored.contains(&first_sentinel), "{restored}");
assert!(!restored.contains(&second_sentinel), "{restored}");
}
/// 留存窗口是有界的：长连接不能无限累积映射，代价是更早的轮次会退回
/// 「占位符原样透传」而不是被错误还原成别的值。
#[tokio::test]
async fn the_retained_session_window_is_bounded() {
let state = redaction_enabled_state();
let decision = control_decision();
let oldest = turn_redaction(&state, &decision, TEST_EMAIL).await;
let oldest_sentinel = sentinel_for(&oldest, TEST_EMAIL);
let mut restorer = ResponsesWebSocketRedactionRestorer::default();
restorer.register(oldest.session);
// 再灌满整个窗口，最老的那一轮必须被挤出去。
let mut newest_sentinel = String::new();
for index in 0..MAX_RETAINED_TURN_REDACTION_SESSIONS {
let email = format!("ws.turn{index}@example.com");
let redaction = turn_redaction(&state, &decision, &email).await;
newest_sentinel = sentinel_for(&redaction, &email);
restorer.register(redaction.session);
}
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(&oldest_sentinel))
.is_none(),
"the evicted turn's sentinel is relayed verbatim, never mis-restored"
);
assert!(
restorer
.restore_provider_frame_text(&provider_delta_frame(&newest_sentinel))
.is_some(),
"the most recent turns stay restorable"
);
}
}