mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
fix(ws): restore redacted PII in provider frames before client delivery
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
@@ -189,7 +189,11 @@ pub(super) async fn forward_client_message(
|
||||
)
|
||||
.await;
|
||||
let client_event = match redacted_client_event {
|
||||
Ok(Some(redacted)) => redacted,
|
||||
Ok(Some(redaction)) => {
|
||||
// 这一轮的映射登记到连接上，响应帧才能在最后一跳还原回真实值。
|
||||
bound.redaction_restorer.register(redaction.session);
|
||||
redaction.client_event
|
||||
}
|
||||
Ok(None) => client_event,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::extract::ws::WebSocket;
|
||||
use axum::extract::ws::{Message as AxumWsMessage, WebSocket};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde_json::Value;
|
||||
use tokio::time::sleep;
|
||||
@@ -486,10 +486,22 @@ pub(super) async fn relay_bound_connection(
|
||||
detach_exhausted_upstream(bound, directive, &context.trace_id).await;
|
||||
continue;
|
||||
}
|
||||
if let Err(error) = send_client_message(
|
||||
client_socket,
|
||||
upstream_message_to_client(upstream_message.clone()),
|
||||
).await {
|
||||
// 响应侧还原：HTTP 在把响应体交给客户端之前会把占位符换回真实值
|
||||
// （`privacy::restore_sync_response_body` /
|
||||
// `privacy::StreamingResponseRestorer`），这里是 WS 的同一个位置
|
||||
// ——最后一跳之前，并且在 `capture_client_frame` 之前，所以审计与
|
||||
// 终态观测继续消费脱敏态的事件。没有命中还原时保持上游原字节。
|
||||
let restored_client_frame = parsed_upstream_frame
|
||||
.as_ref()
|
||||
.map(ParsedResponsesWebSocketFrame::event)
|
||||
.and_then(|event| {
|
||||
bound.redaction_restorer.restore_provider_frame_text(event)
|
||||
});
|
||||
let client_frame = match restored_client_frame {
|
||||
Some(restored) => AxumWsMessage::Text(restored.into()),
|
||||
None => upstream_message_to_client(upstream_message.clone()),
|
||||
};
|
||||
if let Err(error) = send_client_message(client_socket, client_frame).await {
|
||||
warn!(
|
||||
event_name = "responses_websocket_client_send_failed",
|
||||
log_type = "ops",
|
||||
|
||||
@@ -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 = "ws.user@example.com";
|
||||
/// 另一轮用的 PII，用来证明连接级还原覆盖到更早的轮次。
|
||||
const OTHER_TEST_EMAIL: &str = "ws.other@example.com";
|
||||
/// 不是本连接 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -227,9 +227,11 @@ pub(super) async fn run_responses_websocket(
|
||||
&first_event,
|
||||
)
|
||||
.await;
|
||||
let first_event = match redacted_first_event {
|
||||
Ok(Some(redacted)) => redacted,
|
||||
Ok(None) => first_event,
|
||||
// 首轮的 mask session 要活到响应帧还原，但连接此刻还没绑定，只能先接住，
|
||||
// 等 `bind_responses_upstream` 之后登记到连接上。
|
||||
let (first_event, first_turn_redaction_session) = match redacted_first_event {
|
||||
Ok(Some(redaction)) => (redaction.client_event, Some(redaction.session)),
|
||||
Ok(None) => (first_event, None),
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "responses_websocket_redaction_failed",
|
||||
@@ -413,6 +415,9 @@ pub(super) async fn run_responses_websocket(
|
||||
};
|
||||
first_turn.mark_upstream_request_sent();
|
||||
first_turn.set_provider_response_headers(bound.upstream_response_headers.clone());
|
||||
if let Some(session) = first_turn_redaction_session {
|
||||
bound.redaction_restorer.register(session);
|
||||
}
|
||||
bound.turn_state.begin(
|
||||
LogicalTurn::new(first_event, 1, first_logical_turn_id),
|
||||
ActiveProviderAttempt::new(&state, first_turn),
|
||||
@@ -524,6 +529,7 @@ mod tests {
|
||||
observe_active_response_rebind_safety, record_exhausted_bound_key,
|
||||
should_request_full_continuation_retry,
|
||||
};
|
||||
use super::super::redaction::ResponsesWebSocketRedactionRestorer;
|
||||
use super::super::request::{
|
||||
changed_followup_response_create_model, continuation_requires_same_upstream,
|
||||
normalize_followup_response_create, planned_response_create_event,
|
||||
@@ -1212,6 +1218,7 @@ mod tests {
|
||||
pending_adapter_observation: None,
|
||||
exhausted_exclusions: ExhaustedResponsesWebSocketExclusions::default(),
|
||||
pending_turn_finalization: None,
|
||||
redaction_restorer: ResponsesWebSocketRedactionRestorer::default(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ use tokio::task::JoinHandle;
|
||||
|
||||
use super::adapter::{ResponsesWebSocketDrainDirective, ResponsesWebSocketProtocolAdapter};
|
||||
use super::binding::UpstreamBindingIdentity;
|
||||
use super::redaction::ResponsesWebSocketRedactionRestorer;
|
||||
use super::turn_state::ResponsesTurnState;
|
||||
use crate::ai_serving::{AiExecutionDecision, ResponsesWebSocketBodyNormalization};
|
||||
|
||||
@@ -28,6 +29,12 @@ pub(super) struct BoundResponsesConnection {
|
||||
pub(super) binding_identity: UpstreamBindingIdentity,
|
||||
/// 这条连接上「有没有正在进行的 logical turn」的唯一事实来源。
|
||||
pub(super) turn_state: ResponsesTurnState,
|
||||
/// 这条连接迄今 mask 出来的映射，用于把 provider 事件里的占位符换回真实值。
|
||||
///
|
||||
/// 刻意按连接持有而不是按 turn 持有：WS 的会话历史留在上游，continuation 只发
|
||||
/// 增量输入，所以后面几轮的响应可能回显更早那几轮的占位符（理由详见
|
||||
/// [`super::redaction`]）。上游重绑时不重置。
|
||||
pub(super) redaction_restorer: ResponsesWebSocketRedactionRestorer,
|
||||
pub(super) next_turn_index: u64,
|
||||
pub(super) upstream_response_headers: BTreeMap<String, String>,
|
||||
pub(super) pending_adapter_drain: Option<ResponsesWebSocketDrainDirective>,
|
||||
|
||||
@@ -7,6 +7,7 @@ use wreq::ws::message::Message as WreqWsMessage;
|
||||
|
||||
use super::adapter::ResponsesWebSocketProtocolAdapter;
|
||||
use super::binding::{UpstreamBindingIdentity, UpstreamBindingIdentityError};
|
||||
use super::redaction::ResponsesWebSocketRedactionRestorer;
|
||||
use super::request::planned_response_create_event;
|
||||
use super::state::{BoundResponsesConnection, ExhaustedResponsesWebSocketExclusions};
|
||||
use super::turn_state::ResponsesTurnState;
|
||||
@@ -122,6 +123,8 @@ async fn bind_responses_upstream_inner(
|
||||
// 首条 response.create 已经发出，但这一轮的 logical turn 和 attempt 由调用方
|
||||
// 通过 `ResponsesTurnState::begin` 装上：绑定本身不持有记账状态。
|
||||
turn_state: ResponsesTurnState::Idle,
|
||||
// 同理，这一轮的 mask session 也由调用方登记：绑定看不到脱敏链路。
|
||||
redaction_restorer: ResponsesWebSocketRedactionRestorer::default(),
|
||||
next_turn_index: 2,
|
||||
upstream_response_headers: upstream.response_headers,
|
||||
pending_adapter_drain: None,
|
||||
|
||||
@@ -2523,7 +2523,14 @@ fn restore_json_response_body(
|
||||
})
|
||||
}
|
||||
|
||||
fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool {
|
||||
/// 递归把 JSON 里的占位符换回真实值，只认本 `session` 记录过的映射。
|
||||
///
|
||||
/// 同步响应体（[`restore_sync_response_body`]）和 Responses WebSocket 的
|
||||
/// provider 事件帧（`handlers::proxy::websocket::responses::redaction`）共用它，
|
||||
/// 两边因此保持同一套还原语义：未映射的占位符原样保留，`type` / `model` / `id`
|
||||
/// 这类协议字段虽然也被遍历，但它们不可能包含本 session 派生出的 sentinel，
|
||||
/// 所以不会被改写。
|
||||
pub(crate) fn restore_json_strings(value: &mut Value, session: &RedactionSession) -> bool {
|
||||
match value {
|
||||
Value::String(text) => {
|
||||
let restored = session.restore_text(text);
|
||||
|
||||
Reference in New Issue
Block a user