From 01a88350875b612323177d3a6b0f1a7c961abd41 Mon Sep 17 00:00:00 2001 From: Coffee <95295094+hanakannzashi@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:50:48 +0800 Subject: [PATCH] fix: gateway-sign auto-redacted responses --- crates/api/src/routes/completions.rs | 99 ++++++++++++++++++++----- crates/api/tests/e2e_all/auto_redact.rs | 98 ++++++++++++++++++++---- crates/services/src/completions/mod.rs | 6 +- 3 files changed, 168 insertions(+), 35 deletions(-) diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index 3b4813162..1a87122f5 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -465,6 +465,13 @@ fn chat_stream_usage_mode( } } +fn auto_redact_requires_gateway_signature( + auto_redact_enabled: bool, + model_attestation_supported: Option, +) -> bool { + auto_redact_enabled && model_attestation_supported.unwrap_or(false) +} + #[cfg(test)] fn prepare_chat_stream_chunk_for_client( chunk: &mut inference_providers::models::ChatCompletionChunk, @@ -1447,29 +1454,24 @@ async fn chat_completions_inner( .resolve_alias_cached(&request.model) .await; let resolved_model_name = alias_canonical.as_deref().unwrap_or(&request.model); - let model_attestation_supported = if request.stream == Some(true) { - match app_state.models_service.get_models_with_pricing().await { - Ok(models) => models - .iter() - .find(|model| model.model_name == resolved_model_name) - .map(|model| model.attestation_supported), - Err(error) => { - tracing::warn!( - model = %request.model, - error = %error, - "Failed to read cached model metadata for stream usage shaping; preserving raw passthrough" - ); - None - } + let model_attestation_supported = match app_state.models_service.get_models_with_pricing().await + { + Ok(models) => models + .iter() + .find(|model| model.model_name == resolved_model_name) + .map(|model| model.attestation_supported), + Err(error) => { + tracing::warn!( + model = %request.model, + error = %error, + "Failed to read cached model metadata for attestation signing decisions" + ); + None } - } else { - None }; let usage_mode = chat_stream_usage_mode(&request, model_attestation_supported, e2ee_active); let rewrite_public_stream_usage = usage_mode.rewrite_public_stream_usage; - let gateway_signature_enabled = usage_mode.gateway_signature_enabled; let strip_intermediate_usage = usage_mode.strip_intermediate_usage; - service_request.skip_provider_chat_signature = gateway_signature_enabled; // Auto-redact (opt-in via x-auto-redact header or auto_redact body field). // On success this may rewrite service_request.messages to substitute @@ -1516,6 +1518,19 @@ async fn chat_completions_inner( // Anthropic wire adapter bypass that mutation with the original body. service_request.original_request = None; } + + // Auto-redact re-serializes response bytes after restoring the original + // values. A provider signature covers the pre-redaction bytes, so an + // attested response must use a Gateway signature over the bytes returned + // to the client instead. + let gateway_signature_enabled = usage_mode.gateway_signature_enabled + || auto_redact_requires_gateway_signature(auto_redact_enabled, model_attestation_supported); + // Never publish a provider signature over bytes that auto-redact changes. + // If metadata is unavailable, we cannot safely create a Gateway signature + // either, but omitting a signature is still better than returning one that + // cannot verify. + service_request.skip_provider_chat_signature = + usage_mode.gateway_signature_enabled || auto_redact_enabled; let redaction_map = Arc::new(redaction_map); // Check if streaming is requested @@ -1733,7 +1748,7 @@ async fn chat_completions_inner( )); } } - if gateway_signature_enabled { + if gateway_signature_enabled || auto_redact_enabled { let mut chat_id = public_signature_chat_id.lock().await; if chat_id.is_none() { *chat_id = Some(chat.id.clone()); @@ -2004,6 +2019,14 @@ async fn chat_completions_inner( .release_chat_signature_pin(&chat_id) .await; } + } else if auto_redact_enabled { + if let Some(chat_id) = + public_signature_chat_id_for_chain.lock().await.clone() + { + attestation_service_for_chain + .release_chat_signature_pin(&chat_id) + .await; + } } if combined.is_empty() { @@ -2114,6 +2137,10 @@ async fn chat_completions_inner( Ok(b) => b, Err(e) => { tracing::error!(error = %e, "failed to re-serialize unredacted chat response"); + app_state + .attestation_service + .release_chat_signature_pin(&response_with_bytes.response.id) + .await; return ( StatusCode::INTERNAL_SERVER_ERROR, ResponseJson(ErrorResponse::new( @@ -2148,6 +2175,32 @@ async fn chat_completions_inner( _ => body_bytes, }; + if auto_redact_enabled { + if gateway_signature_enabled { + let response_hash = hex::encode(Sha256::digest(&body_bytes)); + if let Err(error) = app_state + .attestation_service + .store_chat_signature_and_unpin( + &response_with_bytes.response.id, + request_hash.clone(), + response_hash, + ) + .await + { + tracing::error!( + chat_id = %response_with_bytes.response.id, + error = %error, + "Failed to store auto-redacted chat completion signature" + ); + } + } else { + app_state + .attestation_service + .release_chat_signature_pin(&response_with_bytes.response.id) + .await; + } + } + let mut response_builder = Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "application/json"); @@ -3227,6 +3280,14 @@ mod tests { assert!(chat_stream_continuous_usage_requested(&request)); } + #[test] + fn auto_redact_requires_gateway_signature_for_attested_models() { + assert!(auto_redact_requires_gateway_signature(true, Some(true))); + assert!(!auto_redact_requires_gateway_signature(false, Some(true))); + assert!(!auto_redact_requires_gateway_signature(true, Some(false))); + assert!(!auto_redact_requires_gateway_signature(true, None)); + } + #[test] fn include_usage_rewrites_and_signs_attested_streams() { let request = chat_request_with_include_usage(Some(true)); diff --git a/crates/api/tests/e2e_all/auto_redact.rs b/crates/api/tests/e2e_all/auto_redact.rs index 810dcf2f6..83ebf2132 100644 --- a/crates/api/tests/e2e_all/auto_redact.rs +++ b/crates/api/tests/e2e_all/auto_redact.rs @@ -6,6 +6,7 @@ use crate::common::*; use api::models::BatchUpdateModelApiRequest; +use bytes::Bytes; /// Pull `choices[0].message.content` out of a chat completion response as /// a `&str`. Returns empty string if the path isn't there — tests assert @@ -97,18 +98,22 @@ async fn auto_redact_header_redacts_prompt_and_restores_response() { )) .await; + let request_body = serde_json::json!({ + "model": E2E_QWEN_MODEL_NAME, + "messages": [{ + "role": "user", + "content": "Please reach out to alice@example.com" + }], + }); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); + let resp = server .post("/v1/chat/completions") .add_header("Authorization", format!("Bearer {api_key}")) .add_header("User-Agent", MOCK_USER_AGENT) .add_header("x-auto-redact", "on") - .json(&serde_json::json!({ - "model": E2E_QWEN_MODEL_NAME, - "messages": [{ - "role": "user", - "content": "Please reach out to alice@example.com" - }], - })) + .content_type("application/json") + .bytes(Bytes::from(request_json.clone())) .await; assert_eq!(resp.status_code(), 200); @@ -128,7 +133,9 @@ async fn auto_redact_header_redacts_prompt_and_restores_response() { ); // Client must see the original PII in the response. - let body: serde_json::Value = resp.json(); + let response_text = resp.text(); + let body: serde_json::Value = + serde_json::from_str(&response_text).expect("chat response should be JSON"); let content = extract_choice_text(&body); assert!( content.contains("alice@example.com"), @@ -138,6 +145,31 @@ async fn auto_redact_header_redacts_prompt_and_restores_response() { !content.contains("redacted1@example.com"), "placeholder should have been swapped back; got {content}" ); + + let chat_id = body["id"] + .as_str() + .expect("chat response should have an id"); + let signature_response = server + .get(format!("/v1/signature/{chat_id}?signing_algo=ecdsa").as_str()) + .add_header("Authorization", format!("Bearer {api_key}")) + .await; + assert_eq!( + signature_response.status_code(), + 200, + "gateway signature should be available: {}", + signature_response.text() + ); + let signature = signature_response.json::(); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&response_text) + ), + "gateway signature must bind the exact request and response bytes returned to the client" + ); } #[tokio::test] @@ -189,7 +221,7 @@ async fn auto_redact_body_field_equivalent_to_header() { } #[tokio::test] -async fn auto_redact_streaming_splits_placeholder_across_chunks() { +async fn auto_redact_continuous_stream_uses_gateway_signature() { let (server, _pool, mock_provider, _db) = setup_test_server_with_pool().await; setup_qwen_model(&server).await; setup_privacy_filter_model(&server).await; @@ -206,21 +238,30 @@ async fn auto_redact_streaming_splits_placeholder_across_chunks() { )) .await; + // Continuous usage preserves the provider's stream bytes unless another + // route feature rewrites them. Auto-redact does rewrite them, which is the + // regression covered by #892. + let request_body = serde_json::json!({ + "model": E2E_QWEN_MODEL_NAME, + "messages": [{ "role": "user", "content": "email alice@example.com" }], + "stream": true, + "stream_options": { "continuous_usage_stats": true }, + }); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); + let resp = server .post("/v1/chat/completions") .add_header("Authorization", format!("Bearer {api_key}")) .add_header("User-Agent", MOCK_USER_AGENT) .add_header("x-auto-redact", "on") - .json(&serde_json::json!({ - "model": E2E_QWEN_MODEL_NAME, - "messages": [{ "role": "user", "content": "email alice@example.com" }], - "stream": true, - })) + .content_type("application/json") + .bytes(Bytes::from(request_json.clone())) .await; assert_eq!(resp.status_code(), 200); // Concatenate all `delta.content` from the SSE stream. let body_text = resp.text(); + let mut chat_id = None::; let mut assembled = String::new(); for line in body_text.lines() { let payload = match line.strip_prefix("data: ") { @@ -233,6 +274,12 @@ async fn auto_redact_streaming_splits_placeholder_across_chunks() { let Ok(chunk) = serde_json::from_str::(payload) else { continue; }; + if chat_id.is_none() { + chat_id = chunk + .get("id") + .and_then(|id| id.as_str()) + .map(ToOwned::to_owned); + } if let Some(choices) = chunk.get("choices").and_then(|c| c.as_array()) { for ch in choices { if let Some(content) = ch @@ -254,6 +301,29 @@ async fn auto_redact_streaming_splits_placeholder_across_chunks() { !assembled.contains("redacted1@example.com"), "no placeholder should leak to client; got: {assembled:?}" ); + + let chat_id = chat_id.expect("stream should include a chat id"); + let signature_response = server + .get(format!("/v1/signature/{chat_id}?signing_algo=ecdsa").as_str()) + .add_header("Authorization", format!("Bearer {api_key}")) + .await; + assert_eq!( + signature_response.status_code(), + 200, + "gateway signature should be available: {}", + signature_response.text() + ); + let signature = signature_response.json::(); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&body_text) + ), + "gateway signature must bind the exact request and SSE response bytes returned to the client" + ); } #[tokio::test] diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index 71d1f988d..f7708a77e 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -1838,8 +1838,10 @@ impl ports::CompletionServiceTrait for CompletionServiceImpl { let backend_latency = provider_start_time.elapsed(); let queue_time = provider_start_time.duration_since(service_start_time); - // Store attestation signature (only for models that support TEE attestation) - if model.attestation_supported { + // Store a model-side signature only when the API route returns the + // provider's exact bytes. Routes that rewrite the public response + // store a Gateway signature after the final bytes are available. + if model.attestation_supported && !request.skip_provider_chat_signature { let attestation_service = self.attestation_service.clone(); let chat_id = response_with_bytes.response.id.clone(); let model_name = model.model_name.clone();