feat(security): harden gateway boundaries and usage policies

Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
@@ -433,7 +433,8 @@ impl AttemptBodyCapture {
if bytes.is_empty() || self.truncated {
return;
}
let max_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES;
let max_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES
.min(crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES);
if self.buffer.len() >= max_bytes {
self.truncated = true;
return;
@@ -447,11 +448,22 @@ impl AttemptBodyCapture {
}
pub(crate) fn encode(&self) -> (Option<String>, Option<UsageBodyCaptureState>) {
let body = (!self.buffer.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&self.buffer));
let state = if self.truncated {
self.encode_with_limit(
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES
.min(crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES),
)
}
fn encode_with_limit(
&self,
max_bytes: usize,
) -> (Option<String>, Option<UsageBodyCaptureState>) {
let captured = &self.buffer[..self.buffer.len().min(max_bytes)];
let body = (!captured.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(captured));
let state = if self.truncated || captured.len() < self.buffer.len() {
UsageBodyCaptureState::Truncated
} else if self.buffer.is_empty() {
} else if captured.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
@@ -1418,15 +1430,14 @@ mod stage_tests {
);
}
/// body capture 的编码状态。截断分支这里到不了:共享的
/// `DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES` 是 `usize::MAX`,
/// 也就是默认不限长;截断只在 usage 侧把上限调低后才可能发生。
/// body capture 的编码状态。Full 记录级别仍受 gateway 的硬上限约束,
/// 这样长连接不会把审计副本无限累积。
#[test]
fn body_capture_encodes_inline_and_empty_states() {
assert_eq!(
super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
usize::MAX,
"the default capture limit is unbounded; truncation is not reachable here"
super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES
.min(crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES),
crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES
);
let mut capture = AttemptBodyCapture::default();
@@ -1454,6 +1465,20 @@ mod stage_tests {
state,
Some(aether_data_contracts::repository::usage::UsageBodyCaptureState::None)
);
let defensive = AttemptBodyCapture {
buffer: b"abcdef".to_vec(),
truncated: false,
};
let (body, state) = defensive.encode_with_limit(3);
let decoded = base64::engine::general_purpose::STANDARD
.decode(body.expect("bounded capture should be encoded"))
.expect("capture is valid base64");
assert_eq!(decoded, b"abc");
assert_eq!(
state,
Some(aether_data_contracts::repository::usage::UsageBodyCaptureState::Truncated)
);
}
/// candidate 行的 error_type 映射:投递失败与供应商侧失败必须各有名字。
File diff suppressed because it is too large Load Diff
@@ -3,3 +3,13 @@
pub(crate) use aether_gateway_execution::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
};
// Usage/audit captures are secondary copies of the stream. Keep a hard
// ceiling even when the configurable "full" record level is otherwise
// unbounded; this does not limit bytes forwarded to the client.
pub(crate) const MAX_STREAM_BODY_CAPTURE_BYTES: usize = 64 * 1024 * 1024;
// Stream frames are newline-delimited JSON. Binary response chunks are base64
// encoded before framing, so this must be larger than the normal 64 MiB raw
// response limit while still bounding an attacker-controlled unterminated line.
pub(crate) const MAX_EXECUTION_STREAM_FRAME_LINE_BYTES: usize = 128 * 1024 * 1024;
@@ -355,6 +355,7 @@ pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -
#[cfg(test)]
mod tests {
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use std::collections::BTreeSet;
use aether_contracts::{ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionResult};
@@ -435,6 +436,15 @@ mod tests {
}
fn sample_key() -> StoredProviderCatalogKey {
let credential_state = AppState::new()
.expect("credential state should build")
.with_data_state_for_tests(
GatewayDataState::disabled()
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let encrypted_api_key = credential_state
.seal_provider_catalog_key_api_key("provider-1", "key-1", "plain-upstream-key")
.expect("api key should encrypt");
StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
@@ -446,7 +456,7 @@ mod tests {
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
"plain-upstream-key".to_string(),
encrypted_api_key,
None,
None,
Some(serde_json::json!({"openai:chat": 1})),
@@ -466,7 +476,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
AppState::new()
.expect("state should build")
File diff suppressed because it is too large Load Diff
@@ -21,11 +21,13 @@ use crate::execution_runtime::kiro_cache::{
};
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::{
DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
apply_upstream_response_body_limit, decode_base64_body_with_limit,
json_value_fits_serialized_limit, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
};
use crate::AppState;
const WEB_SEARCH_TOOL_NAME: &str = "web_search";
const MAX_KIRO_MCP_JSON_BYTES: usize = 8 * 1024 * 1024;
const WEB_SEARCH_TOOL_TYPE_PREFIX: &str = "web_search";
const WEB_SEARCH_QUERY_PREFIX: &str = "Perform a web search for the query: ";
@@ -73,8 +75,6 @@ struct McpResponse {
struct McpError {
#[serde(default)]
code: Option<i64>,
#[serde(default)]
message: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -146,14 +146,14 @@ pub(crate) async fn maybe_execute_kiro_web_search_stream(
request_id = %plan.request_id,
candidate_id = ?plan.candidate_id,
status_code = mcp_execution.result.status_code,
mcp_url = %mcp_execution.url,
mcp_origin = %crate::handlers::shared::security_log_url_origin(&mcp_execution.url),
profile_arn_present = mcp_execution.profile_arn_present,
"gateway executed Kiro web_search through MCP endpoint"
);
if !(200..300).contains(&mcp_execution.result.status_code) {
return Ok(Some(KiroWebSearchStream {
frame_stream: execution_result_frame_stream(&mcp_execution.result),
frame_stream: execution_result_failure_frame_stream(&mcp_execution.result),
report_context: report_context.cloned(),
}));
}
@@ -183,7 +183,7 @@ pub(crate) async fn maybe_execute_kiro_web_search_stream(
cache_usage,
)
.map_err(ExecutionRuntimeTransportError::BodyEncode)?;
let mut synthetic_context = synthetic_report_context(report_context, mcp_execution.url);
let mut synthetic_context = synthetic_report_context(report_context);
if let Some(context) = synthetic_context.as_mut().and_then(Value::as_object_mut) {
context.insert("kiro_web_search_mcp".to_string(), Value::Bool(true));
}
@@ -214,13 +214,13 @@ async fn kiro_simulated_cache_enabled(state: &AppState, plan: &ExecutionPlan) ->
.is_some_and(|provider| {
kiro_simulated_cache_enabled_from_provider_config(provider.config.as_ref())
}),
Err(err) => {
Err(_) => {
warn!(
event_name = "kiro_simulated_cache_config_read_failed",
log_type = "event",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
error = ?err,
error_category = "provider_catalog_read_failed",
"failed to read Kiro simulated cache provider config; defaulting disabled"
);
false
@@ -228,30 +228,28 @@ async fn kiro_simulated_cache_enabled(state: &AppState, plan: &ExecutionPlan) ->
}
}
fn execute_result_body_bytes(result: &ExecutionResult) -> Vec<u8> {
let Some(body) = result.body.as_ref() else {
return Vec::new();
};
if let Some(json_body) = body.json_body.as_ref() {
return serde_json::to_vec(json_body).unwrap_or_default();
}
body.body_bytes_b64
.as_deref()
.and_then(|body| base64::engine::general_purpose::STANDARD.decode(body).ok())
.unwrap_or_default()
}
fn execution_result_frame_stream(
fn execution_result_failure_frame_stream(
result: &ExecutionResult,
) -> BoxStream<'static, Result<Bytes, IoError>> {
let body = serde_json::to_vec(&kiro_mcp_failure_body(result.status_code)).unwrap_or_default();
raw_response_frame_stream(
result.status_code,
result.headers.clone(),
Bytes::from(execute_result_body_bytes(result)),
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
Bytes::from(body),
result.telemetry.clone(),
)
}
fn kiro_mcp_failure_body(status_code: u16) -> Value {
json!({
"error": {
"type": "kiro_web_search_error",
"message": "Kiro web search request failed",
"code": status_code,
}
})
}
fn sse_frame_stream(body: Bytes) -> BoxStream<'static, Result<Bytes, IoError>> {
raw_response_frame_stream(
200,
@@ -328,12 +326,7 @@ async fn execute_mcp_request(
request: &McpRequest,
) -> Result<KiroMcpExecution, ExecutionRuntimeTransportError> {
let mcp_url = aether_provider_transport::kiro::build_kiro_mcp_url_from_resolved_url(&plan.url)
.ok_or_else(|| {
ExecutionRuntimeTransportError::UpstreamRequest(format!(
"failed to build Kiro MCP url from {}",
plan.url
))
})?;
.ok_or_else(kiro_mcp_url_build_error)?;
let mut request_context = build_mcp_request_context(state, plan).await;
if !request_context.profile_arn_present {
if let Some(profile_arn) = discover_kiro_profile_arn(state, plan, &request_context).await? {
@@ -346,7 +339,7 @@ async fn execute_mcp_request(
}
let body_json =
serde_json::to_value(request).map_err(ExecutionRuntimeTransportError::BodyEncode)?;
let mcp_plan = ExecutionPlan {
let mut mcp_plan = ExecutionPlan {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
provider_name: plan.provider_name.clone(),
@@ -367,6 +360,7 @@ async fn execute_mcp_request(
transport_profile: plan.transport_profile.clone(),
timeouts: plan.timeouts.clone(),
};
apply_upstream_response_body_limit(&mut mcp_plan, MAX_KIRO_MCP_JSON_BYTES);
let result = DirectSyncExecutionRuntime::new()
.execute_sync(&mcp_plan)
.await?;
@@ -390,7 +384,7 @@ async fn build_mcp_request_context(
{
Ok(Some(transport)) => transport,
Ok(None) => return fallback(),
Err(err) => {
Err(_) => {
warn!(
event_name = "kiro_web_search_transport_snapshot_unavailable",
log_type = "ops",
@@ -399,7 +393,7 @@ async fn build_mcp_request_context(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
error = ?err,
error_category = "transport_snapshot_read_failed",
"gateway could not read Kiro transport snapshot for web_search MCP"
);
return fallback();
@@ -574,7 +568,7 @@ async fn discover_kiro_profile_arn_in_region(
if let Some(token) = next_token.as_deref() {
body.insert("nextToken".to_string(), Value::String(token.to_string()));
}
let list_plan = ExecutionPlan {
let mut list_plan = ExecutionPlan {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
provider_name: plan.provider_name.clone(),
@@ -598,6 +592,7 @@ async fn discover_kiro_profile_arn_in_region(
transport_profile: plan.transport_profile.clone(),
timeouts: plan.timeouts.clone(),
};
apply_upstream_response_body_limit(&mut list_plan, MAX_KIRO_MCP_JSON_BYTES);
let result = DirectSyncExecutionRuntime::new()
.execute_sync(&list_plan)
.await?;
@@ -676,6 +671,10 @@ fn kiro_runtime_base_url_for_region(region: &str) -> String {
}
fn kiro_runtime_host_for_region(region: &str) -> String {
// Region values can originate in persisted OAuth metadata. Normalize
// before interpolation so a crafted value cannot turn the host header or
// URL into an attacker-controlled origin.
let region = aether_provider_transport::kiro::normalize_kiro_region(region);
match region {
"us-gov-east-1" | "us-gov-west-1" => format!("q-fips.{region}.amazonaws.com"),
"us-iso-east-1" => "q.us-iso-east-1.c2s.ic.gov".to_string(),
@@ -959,7 +958,6 @@ fn parse_mcp_search_results(result: &ExecutionResult) -> Option<WebSearchResults
event_name = "kiro_web_search_mcp_error",
log_type = "event",
code = error.code.unwrap_or_default(),
message = error.message.as_deref().unwrap_or("unknown"),
"Kiro MCP web_search returned JSON-RPC error"
);
return None;
@@ -983,12 +981,18 @@ fn parse_mcp_search_results(result: &ExecutionResult) -> Option<WebSearchResults
fn execution_result_body_json(result: &ExecutionResult) -> Option<Value> {
let body = result.body.as_ref()?;
if let Some(json_body) = body.json_body.as_ref() {
return Some(json_body.clone());
return json_value_fits_serialized_limit(json_body, MAX_KIRO_MCP_JSON_BYTES)
.then(|| json_body.clone());
}
let body = body.body_bytes_b64.as_deref()?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(body)
.ok()?;
// MCP responses are parsed into an owned JSON tree. Keep a fixed ceiling
// here even if the operator disables the general internal body cap, so a
// forged execution result cannot trigger an unbounded base64 allocation.
let bytes = decode_base64_body_with_limit(
body,
crate::headers::max_internal_buffered_body_bytes().min(MAX_KIRO_MCP_JSON_BYTES),
)
.ok()?;
serde_json::from_slice(&bytes).ok()
}
@@ -1198,13 +1202,14 @@ fn generate_search_summary(query: &str, results: Option<&WebSearchResults>) -> S
}
summary.push_str(&format!(" Source: {}\n\n", result.url));
}
if let Some(error) = results
if results
.error
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
{
summary.push_str(&format!("Search warning: {error}\n\n"));
summary.push_str("Search warning: the provider reported an incomplete result.\n\n");
}
}
_ => summary.push_str("No results found.\n"),
@@ -1237,12 +1242,15 @@ fn estimate_text_tokens(text: &str) -> u64 {
((text.len() as u64 + 3) / 4).max(1)
}
fn synthetic_report_context(report_context: Option<&Value>, mcp_url: String) -> Option<Value> {
fn kiro_mcp_url_build_error() -> ExecutionRuntimeTransportError {
ExecutionRuntimeTransportError::UpstreamRequest("failed to build Kiro MCP URL".to_string())
}
fn synthetic_report_context(report_context: Option<&Value>) -> Option<Value> {
let mut context = report_context.cloned()?;
if let Some(object) = context.as_object_mut() {
object.insert("has_envelope".to_string(), Value::Bool(false));
object.insert("needs_conversion".to_string(), Value::Bool(false));
object.insert("upstream_url".to_string(), Value::String(mcp_url));
object.remove("envelope_name");
}
Some(context)
@@ -1257,7 +1265,9 @@ mod tests {
use super::{
build_mcp_headers_from_plan, build_web_search_sse_body, detect_kiro_web_search_request,
parse_mcp_search_results, strip_search_query_prefix, KiroPromptCacheUsage,
execution_result_body_json, generate_search_summary, kiro_mcp_failure_body,
kiro_mcp_url_build_error, parse_mcp_search_results, strip_search_query_prefix,
KiroPromptCacheUsage, WebSearchResult, WebSearchResults,
};
fn sample_plan(body: serde_json::Value) -> ExecutionPlan {
@@ -1443,6 +1453,43 @@ mod tests {
);
}
#[test]
fn kiro_mcp_url_build_error_omits_resolved_url() {
let sensitive_url = "https://token:[email protected]/path";
let message = kiro_mcp_url_build_error().to_string();
assert_eq!(
message,
"failed to execute upstream request: failed to build Kiro MCP URL"
);
assert!(!message.contains(sensitive_url));
}
#[test]
fn kiro_mcp_errors_and_search_warnings_do_not_copy_upstream_secrets() {
let secret = "authorization=Bearer upstream-secret";
let failure = kiro_mcp_failure_body(502).to_string();
let summary = generate_search_summary(
"test",
Some(&WebSearchResults {
results: vec![WebSearchResult {
title: "Safe result".to_string(),
url: "https://example.test/result".to_string(),
snippet: None,
published_date: None,
}],
total_results: Some(1),
query: None,
error: Some(secret.to_string()),
}),
);
assert!(failure.contains("Kiro web search request failed"));
assert!(summary.contains("provider reported an incomplete result"));
assert!(!failure.contains("upstream-secret"));
assert!(!summary.contains("upstream-secret"));
}
#[test]
fn parses_mcp_search_result_text_payload() {
let result = aether_contracts::ExecutionResult {
@@ -1474,6 +1521,29 @@ mod tests {
assert_eq!(parsed.results[0].title, "Example");
}
#[test]
fn kiro_mcp_result_rejects_oversized_base64_before_decode() {
let encoded_limit =
crate::execution_runtime::transport::maximum_base64_len_for_decoded_limit(
super::MAX_KIRO_MCP_JSON_BYTES,
);
let result = aether_contracts::ExecutionResult {
request_id: "req-oversized".to_string(),
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some("A".repeat(encoded_limit + 1)),
}),
telemetry: None,
error: None,
};
assert!(execution_result_body_json(&result).is_none());
}
#[test]
fn builds_anthropic_web_search_sse() {
let sse = build_web_search_sse_body(
@@ -30,7 +30,8 @@ pub(crate) use self::admission::{
};
pub(crate) use self::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
pub(crate) use self::constants::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
MAX_ERROR_BODY_BYTES, MAX_EXECUTION_STREAM_FRAME_LINE_BYTES, MAX_STREAM_BODY_CAPTURE_BYTES,
MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
};
pub(crate) use self::fallback::{
analyze_local_candidate_failover_sync, local_failover_response_text,
@@ -3,21 +3,56 @@ use aether_contracts::{ExecutionPlan, ExecutionResult};
use crate::constants::TRACE_ID_HEADER;
use crate::{AppState, GatewayError};
fn remote_runtime_request_error_kind(error: &reqwest::Error) -> &'static str {
if error.is_timeout() {
"timeout"
} else if error.is_connect() {
"connect"
} else if error.is_request() {
"request"
} else if error.is_body() {
"body"
} else if error.is_decode() {
"decode"
} else {
"transport"
}
}
fn build_remote_execution_runtime_request(
state: &AppState,
remote_execution_runtime_base_url: &str,
path: &str,
trace_id: Option<&str>,
plan: &ExecutionPlan,
) -> reqwest::RequestBuilder {
) -> Result<reqwest::RequestBuilder, GatewayError> {
let envelope_limit = crate::execution_runtime::transport::execution_result_envelope_limit_bytes(
crate::headers::max_internal_buffered_body_bytes(),
);
let body = crate::execution_runtime::transport::serialize_serializable_with_limit(
plan,
envelope_limit,
)
.map_err(|error| {
let kind = match error {
crate::execution_runtime::transport::ExecutionRuntimeTransportError::BodyTooLarge {
..
} => "too_large",
_ => "encode",
};
GatewayError::Internal(format!(
"remote execution runtime request body failed ({kind})"
))
})?;
let mut request = state
.client
.post(format!("{remote_execution_runtime_base_url}{path}"))
.json(plan);
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body(body);
if let Some(trace_id) = trace_id.map(str::trim).filter(|value| !value.is_empty()) {
request = request.header(TRACE_ID_HEADER, trace_id);
}
request
Ok(request)
}
pub(crate) async fn post_sync_plan_to_remote_execution_runtime(
@@ -32,10 +67,15 @@ pub(crate) async fn post_sync_plan_to_remote_execution_runtime(
"/v1/execute/sync",
trace_id,
plan,
)
)?
.send()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| {
GatewayError::Internal(format!(
"remote execution runtime request failed ({})",
remote_runtime_request_error_kind(&err)
))
})
}
pub(crate) async fn post_stream_plan_to_remote_execution_runtime(
@@ -50,10 +90,15 @@ pub(crate) async fn post_stream_plan_to_remote_execution_runtime(
"/v1/execute/stream",
trace_id,
plan,
)
)?
.send()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| {
GatewayError::Internal(format!(
"remote execution runtime request failed ({})",
remote_runtime_request_error_kind(&err)
))
})
}
pub(crate) async fn execute_sync_plan_via_remote_execution_runtime(
@@ -76,8 +121,27 @@ pub(crate) async fn execute_sync_plan_via_remote_execution_runtime(
)));
}
response
.json::<ExecutionResult>()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
let body = aether_http::read_response_bytes_with_limit(
response,
crate::execution_runtime::transport::execution_result_envelope_limit_bytes(
crate::headers::max_internal_buffered_body_bytes(),
),
)
.await
.map_err(|err| {
GatewayError::Internal(format!(
"remote execution runtime response body failed ({})",
match err {
aether_http::ResponseBodyReadError::TooLarge { .. } => "too_large",
aether_http::ResponseBodyReadError::Read(error) => {
remote_runtime_request_error_kind(&error)
}
}
))
})?;
serde_json::from_slice::<ExecutionResult>(&body).map_err(|_| {
GatewayError::Internal(
"remote execution runtime returned invalid execution JSON".to_string(),
)
})
}
File diff suppressed because it is too large Load Diff
@@ -191,6 +191,46 @@ pub(super) fn inspect_prefetched_stream_body(
}
}
fn append_error_frame_payload(
body: &mut Vec<u8>,
chunk_b64: Option<&str>,
text: Option<&str>,
) -> Result<bool, GatewayError> {
let remaining = MAX_ERROR_BODY_BYTES.saturating_sub(body.len());
if remaining == 0 {
return Ok(false);
}
if let Some(chunk_b64) = chunk_b64 {
// Do not decode an attacker-controlled megabyte-scale base64 value
// merely to retain the first 16 KiB of an error body. A standard
// base64 value representing at most `remaining` bytes cannot exceed
// this length.
let max_encoded_len = remaining
.saturating_add(2)
.checked_div(3)
.unwrap_or(usize::MAX)
.saturating_mul(4);
if chunk_b64.len() > max_encoded_len {
warn!(
encoded_bytes = chunk_b64.len(),
max_encoded_bytes = max_encoded_len,
"execution runtime error frame body exceeded capture limit"
);
return Ok(false);
}
let chunk = base64::engine::general_purpose::STANDARD
.decode(chunk_b64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
body.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
} else if let Some(text) = text {
let text_bytes = text.as_bytes();
body.extend_from_slice(&text_bytes[..text_bytes.len().min(remaining)]);
}
Ok(body.len() < MAX_ERROR_BODY_BYTES)
}
pub(super) async fn collect_error_body<R>(
lines: &mut FramedRead<R, LinesCodec>,
) -> Result<Vec<u8>, GatewayError>
@@ -201,23 +241,19 @@ where
while let Some(frame) = read_next_frame(lines).await? {
match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => {
let chunk = if let Some(chunk_b64) = chunk_b64 {
base64::engine::general_purpose::STANDARD
.decode(chunk_b64)
.map_err(|err| GatewayError::Internal(err.to_string()))?
} else {
text.unwrap_or_default().into_bytes()
};
body.extend_from_slice(&chunk);
if body.len() >= MAX_ERROR_BODY_BYTES {
body.truncate(MAX_ERROR_BODY_BYTES);
if !append_error_frame_payload(&mut body, chunk_b64.as_deref(), text.as_deref())? {
break;
}
}
StreamFramePayload::Telemetry { .. } => {}
StreamFramePayload::Eof { .. } => break,
StreamFramePayload::Error { error } => {
warn!(error = %error.message, "execution runtime stream emitted error frame while collecting error body");
warn!(
error_kind = ?error.kind,
error_phase = ?error.phase,
upstream_status = ?error.upstream_status,
"execution runtime stream emitted error frame while collecting error body"
);
break;
}
StreamFramePayload::Headers { .. } => {}
@@ -242,3 +278,32 @@ where
}
Ok(None)
}
#[cfg(test)]
mod tests {
use super::{append_error_frame_payload, MAX_ERROR_BODY_BYTES};
#[test]
fn oversized_base64_error_frame_is_rejected_before_decode() {
let mut body = Vec::new();
let encoded = "x".repeat((MAX_ERROR_BODY_BYTES + 2) / 3 * 4 + 1);
let keep_reading = append_error_frame_payload(&mut body, Some(&encoded), None)
.expect("oversized frame should be handled without a decode error");
assert!(!keep_reading);
assert!(body.is_empty());
}
#[test]
fn error_frame_payload_is_capped_to_remaining_capture_budget() {
let mut body = vec![b'a'; MAX_ERROR_BODY_BYTES - 2];
let keep_reading = append_error_frame_payload(&mut body, None, Some("hello"))
.expect("text payload should append");
assert!(!keep_reading);
assert_eq!(body.len(), MAX_ERROR_BODY_BYTES);
assert_eq!(&body[MAX_ERROR_BODY_BYTES - 2..], b"he");
}
}
File diff suppressed because it is too large Load Diff
@@ -68,6 +68,60 @@ enum StreamFailureHandling {
HonorLocalFailover,
}
const UPSTREAM_STREAM_FAILURE_MESSAGE: &str = "Upstream response stream failed";
const EXECUTION_STREAM_PROTOCOL_FAILURE_MESSAGE: &str = "Execution runtime stream protocol failed";
const EXECUTION_STREAM_PROCESSING_FAILURE_MESSAGE: &str =
"Execution runtime stream processing failed";
fn encode_bounded_stream_capture(body: &[u8]) -> Option<String> {
encode_stream_capture_with_limit(
body,
crate::execution_runtime::MAX_STREAM_BODY_CAPTURE_BYTES,
)
}
fn encode_stream_capture_with_limit(body: &[u8], max_bytes: usize) -> Option<String> {
let captured = &body[..body.len().min(max_bytes)];
(!captured.is_empty()).then(|| base64::engine::general_purpose::STANDARD.encode(captured))
}
fn stable_stream_failure_message(error_type: &str) -> &'static str {
match error_type {
"first_byte_timeout" => "Upstream response timed out before the first byte",
"read_timeout" => "Upstream response stream timed out",
"execution_runtime_stream_read_error" => UPSTREAM_STREAM_FAILURE_MESSAGE,
"execution_runtime_stream_frame_decode_error"
| "execution_runtime_stream_chunk_decode_error" => {
EXECUTION_STREAM_PROTOCOL_FAILURE_MESSAGE
}
"execution_runtime_sync_json_stream_bridge_error"
| "execution_runtime_stream_rewrite_error"
| "execution_runtime_stream_rewrite_flush_error" => {
EXECUTION_STREAM_PROCESSING_FAILURE_MESSAGE
}
_ => "Execution runtime stream failed",
}
}
fn public_execution_error_message(error: &ExecutionError) -> String {
match &error.kind {
ExecutionErrorKind::ConnectTimeout => "Upstream connection timed out".to_string(),
ExecutionErrorKind::FirstByteTimeout => {
"Upstream response timed out before the first byte".to_string()
}
ExecutionErrorKind::ReadTimeout => "Upstream response stream timed out".to_string(),
ExecutionErrorKind::Upstream4xx | ExecutionErrorKind::Upstream5xx => error
.upstream_status
.map(|status| format!("Upstream request returned HTTP {status}"))
.unwrap_or_else(|| "Upstream request failed".to_string()),
ExecutionErrorKind::TlsError => "Upstream TLS connection failed".to_string(),
ExecutionErrorKind::ProxyError => "Upstream proxy request failed".to_string(),
ExecutionErrorKind::ProtocolError => UPSTREAM_STREAM_FAILURE_MESSAGE.to_string(),
ExecutionErrorKind::Cancelled => "Request was cancelled".to_string(),
ExecutionErrorKind::Internal => "Execution runtime stream failed".to_string(),
}
}
impl StreamFailureReport {
fn into_body_jsons(self) -> (Value, Option<Value>) {
let Self {
@@ -110,11 +164,11 @@ impl StreamFailureReport {
pub(super) fn build_stream_failure_report(
error_type: impl Into<String>,
error_message: impl Into<String>,
_error_message: impl Into<String>,
status_code: u16,
) -> StreamFailureReport {
let error_type = error_type.into();
let error_message = error_message.into();
let error_message = stable_stream_failure_message(error_type.as_str()).to_string();
StreamFailureReport {
status_code,
error_type,
@@ -129,13 +183,14 @@ pub(super) fn build_stream_failure_report(
pub(super) fn build_stream_transport_failure_report(
error_type: impl Into<String>,
error_message: impl Into<String>,
_error_message: impl Into<String>,
status_code: u16,
) -> StreamFailureReport {
let error_type = error_type.into();
StreamFailureReport {
status_code,
error_type: error_type.into(),
error_message: error_message.into(),
error_message: stable_stream_failure_message(error_type.as_str()).to_string(),
error_type,
upstream_status_code: None,
transport_error: true,
honor_http_failover: false,
@@ -163,7 +218,7 @@ pub(super) fn build_stream_failure_from_execution_error(
.ok()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "internal".to_string());
let error_message = error.message.trim().to_string();
let error_message = public_execution_error_message(error);
let phase = serde_json::to_value(&error.phase).unwrap_or(Value::Null);
let mut error_object = Map::from_iter([
("phase".to_string(), phase),
@@ -323,8 +378,7 @@ fn build_stream_failure_sync_payload(
headers,
body_json: Some(body),
client_body_json: client_body,
body_base64: (!provider_buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(provider_buffered_body)),
body_base64: encode_bounded_stream_capture(provider_buffered_body),
telemetry,
}
}
@@ -524,8 +578,7 @@ pub(super) async fn handle_prefetch_provider_private_stream_error(
headers,
body_json: Some(body_json),
client_body_json: None,
body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(buffered_body)),
body_base64: encode_bounded_stream_capture(buffered_body),
telemetry,
};
let failure_analysis = record_stream_sync_failure(
@@ -813,11 +866,11 @@ pub(super) async fn submit_midstream_stream_failure(
started_at_unix_ms: u64,
failure: StreamFailureReport,
) {
let Some(report_kind) =
direct_stream_finalize_kind.and_then(resolve_core_error_background_report_kind)
else {
return;
};
let background_report_kind =
direct_stream_finalize_kind.and_then(resolve_core_error_background_report_kind);
let submit_background_report = background_report_kind.is_some();
let report_kind =
background_report_kind.unwrap_or_else(|| "execution_runtime_stream_error".to_string());
let candidate_status_code = failure.upstream_status_code;
let payload = build_stream_failure_sync_payload(
@@ -839,7 +892,10 @@ pub(super) async fn submit_midstream_stream_failure(
StreamFailureHandling::Terminal,
)
.await;
if let Err(err) = submit_sync_report(state, payload).await {
if !submit_background_report {
return;
}
if let Err(_err) = submit_sync_report(state, payload).await {
let request_id = short_request_id(plan.request_id.as_str());
warn!(
event_name = "execution_report_submit_failed",
@@ -848,7 +904,7 @@ pub(super) async fn submit_midstream_stream_failure(
request_id = %request_id,
candidate_id = ?plan.candidate_id,
report_scope = "stream_failure",
error = ?err,
error_category = "stream_report_submit_failed",
"gateway failed to submit sync execution report for terminal stream failure"
);
}
@@ -864,9 +920,22 @@ mod tests {
use super::{
build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body,
build_stream_failure_sync_payload, build_stream_transport_failure_report,
build_stream_failure_report, build_stream_failure_sync_payload,
build_stream_transport_failure_report, encode_stream_capture_with_limit,
};
#[test]
fn failure_capture_encoding_defensively_caps_an_oversized_slice() {
let encoded = encode_stream_capture_with_limit(b"abcdef", 3)
.expect("bounded failure capture should be encoded");
let decoded = base64::engine::general_purpose::STANDARD
.decode(encoded)
.expect("capture should be valid base64");
assert_eq!(decoded, b"abc");
assert!(encode_stream_capture_with_limit(b"abc", 0).is_none());
}
#[test]
fn committed_transport_failure_has_no_upstream_status() {
for status_code in [502, 504] {
@@ -913,6 +982,61 @@ mod tests {
assert!(!failure.transport_error);
}
#[test]
fn execution_error_details_are_not_projected_to_stream_clients() {
let secret = "Bearer stream-secret https://user:[email protected]/private";
let failure = build_stream_failure_from_execution_error(&ExecutionError {
kind: ExecutionErrorKind::ProtocolError,
phase: ExecutionPhase::StreamRead,
message: secret.to_string(),
upstream_status: None,
retryable: true,
failover_recommended: true,
});
assert_eq!(failure.error_message, "Upstream response stream failed");
let client_body = failure
.to_json_string()
.expect("stream failure should serialize");
assert!(!client_body.contains(secret));
assert!(!client_body.contains("stream-secret"));
}
#[test]
fn upstream_execution_error_keeps_only_http_status_diagnostics() {
let secret = "authorization=Bearer upstream-secret";
let failure = build_stream_failure_from_execution_error(&ExecutionError {
kind: ExecutionErrorKind::Upstream4xx,
phase: ExecutionPhase::FirstByte,
message: secret.to_string(),
upstream_status: Some(429),
retryable: true,
failover_recommended: true,
});
assert_eq!(failure.error_message, "Upstream request returned HTTP 429");
assert!(!failure
.to_json_string()
.expect("stream failure should serialize")
.contains(secret));
}
#[test]
fn internal_stream_failure_details_are_replaced_with_stable_text() {
let secret = "failed near /Users/admin/.config with token=stream-secret";
let failure =
build_stream_failure_report("execution_runtime_stream_rewrite_error", secret, 502);
assert_eq!(
failure.error_message,
"Execution runtime stream processing failed"
);
assert!(!failure
.to_json_string()
.expect("stream failure should serialize")
.contains(secret));
}
#[test]
fn midstream_failure_trace_uses_terminal_error_instead_of_buffered_sse() {
let provider_buffered_body = concat!(
@@ -22,13 +22,14 @@ use crate::ai_serving::api::{
};
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::{
append_upstream_response_body_chunk, decode_response_body_bytes, format_hyper_error_chain,
format_wreq_upstream_request_error, stream_first_byte_timeout_message, DirectUpstreamResponse,
append_upstream_response_body_chunk, decode_response_body_bytes,
stream_first_byte_timeout_message, DirectUpstreamResponse,
};
use crate::execution_runtime::DirectUpstreamStreamExecution;
use crate::GatewayError;
const STREAM_USAGE_OBSERVER_MAX_LINE_BYTES: usize = 1024 * 1024;
const UPSTREAM_STREAM_READ_ERROR_MESSAGE: &str = "Upstream response stream failed";
pub(crate) fn build_direct_execution_frame_stream(
execution: DirectUpstreamStreamExecution,
@@ -39,6 +40,7 @@ pub(crate) fn build_direct_execution_frame_stream(
candidate_id: _,
status_code,
headers,
upstream_content_length,
provider_api_format,
stream_summary_report_context,
prefetched_body,
@@ -74,7 +76,11 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
let mut observer_buffered = Vec::new();
if should_buffer_non_stream_response(&headers, &observer_context) {
if should_buffer_non_stream_response(
&headers,
upstream_content_length,
&observer_context,
) {
let original_headers = headers.clone();
match buffer_non_sse_upstream_body(
prefetched_body,
@@ -104,8 +110,10 @@ pub(crate) fn build_direct_execution_frame_stream(
summary = outcome.terminal_summary;
}
Ok(None) => {}
Err(err) => {
yield Err(IoError::other(format!("{err:?}")));
Err(_err) => {
yield Err(IoError::other(
"Execution runtime stream conversion failed",
));
return;
}
}
@@ -250,16 +258,16 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
Err(message) => {
Err(_message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
error_category = "prefetched_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
match encode_error_frame(UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string()) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -333,14 +341,14 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
Err(err) => {
let message = format_error_chain(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
error_category = "reqwest_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
@@ -415,14 +423,14 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
Err(err) => {
let message = format_hyper_error_chain(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
error_category = "hyper_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
@@ -497,14 +505,14 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
Err(err) => {
let message = format_wreq_upstream_request_error(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
error_category = "browser_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
@@ -575,16 +583,16 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
Ok(None) => break,
Err(message) => {
Err(_message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error = %message,
error_category = "tunnel_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
match encode_error_frame(UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string()) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -664,14 +672,14 @@ fn encode_data_frame(chunk: &Bytes) -> Result<Bytes, IoError> {
})
}
fn encode_error_frame(message: String) -> Result<Bytes, IoError> {
fn encode_error_frame(_message: String) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::ProtocolError,
phase: ExecutionPhase::StreamRead,
message,
message: UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string(),
upstream_status: None,
retryable: true,
failover_recommended: true,
@@ -770,6 +778,7 @@ fn should_treat_upstream_response_as_stream(
fn should_buffer_non_stream_response(
headers: &BTreeMap<String, String>,
upstream_content_length: Option<u64>,
report_context: &Value,
) -> bool {
if should_treat_upstream_response_as_stream(headers, report_context) {
@@ -784,6 +793,22 @@ fn should_buffer_non_stream_response(
return true;
}
// `content-length` is intentionally removed from the response header map
// before it reaches the execution stream. Retain its parsed value as
// internal metadata so only a declared fixed-length JSON response is
// converted to the client's SSE contract.
if report_context
.get("upstream_is_stream")
.and_then(Value::as_bool)
== Some(true)
&& upstream_content_length.is_some()
&& headers
.get("content-type")
.is_some_and(|value| value.to_ascii_lowercase().contains("json"))
{
return true;
}
headers
.get("content-length")
.and_then(|value| value.trim().parse::<u64>().ok())
@@ -813,9 +838,9 @@ async fn buffer_non_sse_upstream_body(
&mut upstream_bytes,
)?;
}
Err(message) => {
Err(_message) => {
return Err(BufferedUpstreamBodyError {
message,
message: UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string(),
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
@@ -864,13 +889,13 @@ async fn buffer_non_sse_upstream_body(
&mut upstream_bytes,
)?;
}
Err(err) => {
let message = format_error_chain(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error = %message,
error_category = "reqwest_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
@@ -922,13 +947,13 @@ async fn buffer_non_sse_upstream_body(
&mut upstream_bytes,
)?;
}
Err(err) => {
let message = format_hyper_error_chain(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error = %message,
error_category = "hyper_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
@@ -980,13 +1005,13 @@ async fn buffer_non_sse_upstream_body(
&mut upstream_bytes,
)?;
}
Err(err) => {
let message = format_wreq_upstream_request_error(&err);
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error = %message,
error_category = "browser_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
@@ -1034,16 +1059,16 @@ async fn buffer_non_sse_upstream_body(
)?;
}
Ok(None) => break,
Err(message) => {
Err(_message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error = %message,
error_category = "tunnel_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message,
message: UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string(),
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
@@ -1071,14 +1096,16 @@ fn maybe_bridge_non_sse_sync_json_to_stream(
return Ok(None);
}
let decoded_body_bytes = decode_response_body_bytes(headers, body_bytes)
.map_err(|error| GatewayError::Internal(error.to_string()))?;
let decoded_body_bytes = decode_response_body_bytes(headers, body_bytes).map_err(|_error| {
GatewayError::Internal("execution runtime response decode failed".to_string())
})?;
if !response_body_is_json(headers, decoded_body_bytes.as_ref()) {
return Ok(None);
}
let body_json: Value = serde_json::from_slice(decoded_body_bytes.as_ref())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_json: Value = serde_json::from_slice(decoded_body_bytes.as_ref()).map_err(|_err| {
GatewayError::Internal("execution runtime response JSON decode failed".to_string())
})?;
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
@@ -1114,17 +1141,6 @@ fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8])
serde_json::from_slice::<Value>(body_bytes).is_ok()
}
fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
let mut message = err.to_string();
let mut source = err.source();
while let Some(cause) = source {
message.push_str(": ");
message.push_str(&cause.to_string());
source = cause.source();
}
message
}
fn observe_stream_chunk(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
@@ -1135,10 +1151,8 @@ fn observe_stream_chunk(
let normalized = if let Some(normalizer) = private_stream_normalizer {
match normalizer.push_chunk(chunk) {
Ok(normalized) => normalized,
Err(err) => {
observer.disable_with_error(format!(
"failed to normalize provider private stream chunk: {err:?}"
));
Err(_err) => {
observer.disable_with_error("provider stream normalization failed");
return;
}
}
@@ -1160,23 +1174,21 @@ fn finalize_stream_terminal_summary(
Ok(flushed) => {
observe_normalized_bytes(observer, report_context, observer_buffered, &flushed)
}
Err(err) => observer.disable_with_error(format!(
"failed to flush provider private stream normalization: {err:?}"
)),
Err(_err) => observer.disable_with_error("provider stream normalization failed"),
}
}
if !observer_buffered.is_empty() {
let line = std::mem::take(observer_buffered);
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
if let Err(_err) = observer.push_line(report_context, line) {
observer.disable_with_error("stream usage parsing failed");
}
}
match observer.finish(report_context) {
Ok(summary) => summary,
Err(err) => {
observer.disable_with_error(err.to_string());
Err(_err) => {
observer.disable_with_error("stream usage parsing failed");
observer.latest_summary().cloned()
}
}
@@ -1216,8 +1228,8 @@ fn observe_normalized_bytes(
remaining = &remaining[line_part_len..];
if observer_buffered.last() == Some(&b'\n') {
let line = std::mem::take(observer_buffered);
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
if let Err(_err) = observer.push_line(report_context, line) {
observer.disable_with_error("stream usage parsing failed");
observer_buffered.clear();
return;
}
@@ -1232,7 +1244,10 @@ mod tests {
use std::sync::Arc;
use std::time::Duration;
use aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNode};
use async_stream::stream;
use axum::body::{Body, Bytes};
use axum::extract::ws::Message;
@@ -1245,9 +1260,9 @@ mod tests {
use tokio::sync::watch;
use super::{
build_direct_execution_frame_stream, observe_normalized_bytes,
build_direct_execution_frame_stream, encode_error_frame, observe_normalized_bytes,
should_buffer_non_stream_response, should_treat_upstream_response_as_stream,
STREAM_USAGE_OBSERVER_MAX_LINE_BYTES,
STREAM_USAGE_OBSERVER_MAX_LINE_BYTES, UPSTREAM_STREAM_READ_ERROR_MESSAGE,
};
use crate::ai_serving::api::StreamingStandardTerminalObserver;
use crate::execution_runtime::transport::{
@@ -1267,6 +1282,66 @@ mod tests {
}
}
const LOCAL_TUNNEL_TEST_PSK: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
const LOCAL_TUNNEL_TEST_GENERATION: &str = "stream-pump-test-generation-1";
fn authenticated_local_tunnel_test_state() -> AppState {
let node = StoredProxyNode::new(
"node-1".to_string(),
"Node 1".to_string(),
"127.0.0.1".to_string(),
0,
false,
"online".to_string(),
30,
1,
0,
0,
0,
0,
true,
true,
1,
)
.expect("tunnel node should build")
.with_runtime_fields(
None,
None,
None,
None,
Some(serde_json::json!({
"tunnel_security": {
"mode": TUNNEL_SECURITY_NON_TLS_REQUIRED,
"encryption_key": LOCAL_TUNNEL_TEST_PSK,
}
})),
None,
None,
None,
None,
None,
None,
)
.with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string());
let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::new(
InMemoryProxyNodeRepository::seed([node]),
))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data)
}
async fn recv_tunnel_test_frame(
proxy_rx: &mut aether_runtime::BoundedQueueReceiver<Message>,
description: &str,
) -> Message {
tokio::time::timeout(Duration::from_secs(5), proxy_rx.recv())
.await
.unwrap_or_else(|_| panic!("timed out waiting for {description}"))
.unwrap_or_else(|| panic!("proxy channel closed before {description}"))
}
#[test]
fn treats_kiro_eventstream_envelope_as_stream_even_when_content_type_is_json() {
let headers = BTreeMap::from([("content-type".into(), "application/json".into())]);
@@ -1295,10 +1370,12 @@ mod tests {
assert!(!should_buffer_non_stream_response(
&BTreeMap::from([("content-type".into(), "application/json".into())]),
None,
&streaming_context
));
assert!(should_buffer_non_stream_response(
&BTreeMap::from([("content-type".into(), "application/json".into())]),
None,
&non_stream_context
));
assert!(should_buffer_non_stream_response(
@@ -1306,14 +1383,27 @@ mod tests {
("content-type".into(), "application/json".into()),
("content-length".into(), "128".into()),
]),
Some(128),
&streaming_context
));
assert!(!should_buffer_non_stream_response(
&BTreeMap::from([("content-type".into(), "text/event-stream".into())]),
None,
&non_stream_context
));
}
#[test]
fn error_frames_do_not_include_transport_details() {
let secret = "Bearer stream-secret https://user:[email protected]/private";
let frame = encode_error_frame(secret.to_string()).expect("error frame should encode");
let frame = String::from_utf8(frame.to_vec()).expect("error frame should be utf8");
assert!(frame.contains(UPSTREAM_STREAM_READ_ERROR_MESSAGE));
assert!(!frame.contains(secret));
assert!(!frame.contains("stream-secret"));
}
#[test]
fn oversized_usage_line_disables_observation_without_retaining_the_line() {
let mut observer = StreamingStandardTerminalObserver::default();
@@ -1920,20 +2010,24 @@ mod tests {
}
#[tokio::test]
async fn direct_execution_frame_stream_preserves_local_tunnel_stream_error_message() {
let state = AppState::new().expect("app state should build");
async fn direct_execution_frame_stream_sanitizes_local_tunnel_stream_error_message() {
let state = authenticated_local_tunnel_test_state();
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
801,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
2,
)));
tunnel_app.hub.register_proxy(Arc::new(
TunnelProxyConn::new(
801,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
2,
)
.with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string())
.with_authenticated_key(LOCAL_TUNNEL_TEST_PSK.to_string()),
));
let plan = ExecutionPlan {
request_id: "req-local-stream-error-1".into(),
@@ -1967,7 +2061,7 @@ mod tests {
execute_stream_plan_via_local_tunnel(&state_for_task, &plan_for_task).await
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
let request_headers = match recv_tunnel_test_frame(&mut proxy_rx, "headers frame").await {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
@@ -1975,7 +2069,7 @@ mod tests {
.expect("request header frame should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
let request_body = match recv_tunnel_test_frame(&mut proxy_rx, "body frame").await {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
@@ -2057,28 +2151,29 @@ mod tests {
.and_then(Value::as_str)
.expect("error frame should include a message");
assert_eq!(error_message, original_error);
assert!(
!error_message.contains("unexpected EOF during chunk size line"),
"local tunnel path should preserve the original proxy error text"
);
assert_eq!(error_message, UPSTREAM_STREAM_READ_ERROR_MESSAGE);
assert!(!error_message.contains(original_error));
}
#[tokio::test]
async fn second_local_tunnel_request_works_after_first_completes() {
let state = AppState::new().expect("app state should build");
let state = authenticated_local_tunnel_test_state();
let tunnel_app = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
900,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
2,
)));
tunnel_app.hub.register_proxy(Arc::new(
TunnelProxyConn::new(
900,
"node-1".to_string(),
"Node 1".to_string(),
proxy_tx,
proxy_close_tx,
16,
2,
)
.with_tunnel_generation(LOCAL_TUNNEL_TEST_GENERATION.to_string())
.with_authenticated_key(LOCAL_TUNNEL_TEST_PSK.to_string()),
));
let plan = ExecutionPlan {
request_id: "req-reuse-1".into(),
@@ -2115,13 +2210,13 @@ mod tests {
);
// Read request frames from proxy side
let req1_headers = match proxy_rx.recv().await.expect("req1 headers") {
let req1_headers = match recv_tunnel_test_frame(&mut proxy_rx, "req1 headers").await {
Message::Binary(data) => data,
other => panic!("unexpected: {other:?}"),
};
let req1_header =
tunnel_protocol::FrameHeader::parse(&req1_headers).expect("req1 header parse");
let _req1_body = proxy_rx.recv().await.expect("req1 body");
let _req1_body = recv_tunnel_test_frame(&mut proxy_rx, "req1 body").await;
// Simulate proxy response
let resp_meta = serde_json::to_vec(&tunnel_protocol::ResponseMeta {
@@ -2188,10 +2283,7 @@ mod tests {
);
// Read second request's frames
let req2_headers = tokio::time::timeout(Duration::from_secs(2), proxy_rx.recv())
.await
.expect("second request should arrive within 2s")
.expect("req2 headers");
let req2_headers = recv_tunnel_test_frame(&mut proxy_rx, "req2 headers").await;
let req2_data = match req2_headers {
Message::Binary(data) => data,
other => panic!("unexpected: {other:?}"),
@@ -9,9 +9,9 @@ use crate::api::response::build_client_response_from_parts;
use crate::control::GatewayControlDecision;
use crate::usage::spawn_sync_report;
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
use aether_usage_runtime::decode_internal_report_body_base64;
use axum::body::Body;
use axum::http::{Response, StatusCode};
use base64::Engine as _;
use tracing::warn;
#[derive(Clone, Debug)]
@@ -144,9 +144,8 @@ fn build_local_core_sync_finalize_fallback_response(
}
if let Some(body_base64) = payload.body_base64.as_ref() {
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_bytes =
decode_internal_report_body_base64(body_base64).map_err(GatewayError::Internal)?;
return build_local_sync_response_from_bytes(trace_id, decision, payload, body_bytes);
}
@@ -299,9 +298,8 @@ fn resolve_local_sync_source_body_json(
let body_json = if let Some(body_json) = payload.body_json.clone() {
body_json
} else if let Some(body_base64) = payload.body_base64.as_deref() {
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_bytes =
decode_internal_report_body_base64(body_base64).map_err(GatewayError::Internal)?;
let stripped = strip_utf8_bom_and_ws(&body_bytes);
let Ok(body_json) = serde_json::from_slice::<serde_json::Value>(stripped) else {
return Ok(None);
@@ -328,9 +326,8 @@ fn decode_local_sync_body_text(
let Some(body_base64) = payload.body_base64.as_deref() else {
return Ok(None);
};
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_bytes =
decode_internal_report_body_base64(body_base64).map_err(GatewayError::Internal)?;
let stripped = strip_utf8_bom_and_ws(&body_bytes);
let body_text = String::from_utf8_lossy(stripped).trim().to_string();
if body_text.is_empty() {
@@ -60,8 +60,8 @@ use crate::execution_runtime::transport::{
build_request_body, collect_response_headers, decode_response_body_bytes_with_limit,
execution_plan_response_body_limit_bytes, execution_response_body_mode,
format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error,
response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError,
response_body_is_json, safe_transport_error_message, send_request, DirectHttpResponse,
DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
use crate::execution_runtime::{
@@ -101,7 +101,7 @@ mod policy;
#[path = "execution/response.rs"]
mod response;
use policy::decode_execution_result_body;
use policy::{decode_execution_result_body, decode_execution_result_body_with_limit};
pub(crate) use response::{
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
@@ -115,6 +115,11 @@ const SYNC_EXECUTION_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES: &[u8] = b"\n";
const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL: Duration = Duration::from_secs(5);
// Progress parsing must retain only the incomplete SSE record. This is not a
// response-body limit: the full upstream body is still handled by the normal
// execution body policy, while malformed streams cannot grow telemetry state
// forever by withholding a record separator.
const OPENAI_IMAGE_SYNC_PROGRESS_MAX_BUFFER_BYTES: usize = 16 * 1024 * 1024;
const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE: &str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.";
fn elapsed_ms_since(started_at: Instant) -> u64 {
@@ -174,6 +179,39 @@ struct SyncAttemptTerminalGuard {
armed: bool,
}
/// Keep forced-terminal records useful for operations without copying an
/// arbitrary `GatewayError` into the candidate/usage stores. Gateway errors
/// can wrap provider URLs, credentials, query strings, or database details;
/// those values belong in the internal logging path only.
fn persisted_sync_abort_message(error: &GatewayError) -> &'static str {
match error {
GatewayError::UpstreamUnavailable { .. } => {
"local sync attempt aborted before terminal finalization: upstream unavailable"
}
GatewayError::ControlUnavailable { .. } => {
"local sync attempt aborted before terminal finalization: control unavailable"
}
GatewayError::LocalExecutionPlanningTimeout { .. } => {
"local sync attempt aborted before terminal finalization: planning timeout"
}
GatewayError::AdmissionTimeout { .. } => {
"local sync attempt aborted before terminal finalization: admission timeout"
}
GatewayError::Client { .. } => {
"local sync attempt aborted before terminal finalization: client error"
}
GatewayError::PlanUsageLimited(_) => {
"local sync attempt aborted before terminal finalization: usage limit"
}
GatewayError::LastActiveAdminUpdateDenied | GatewayError::LastActiveAdminDeleteDenied => {
"local sync attempt aborted before terminal finalization: policy denied"
}
GatewayError::Internal(_) => {
"local sync attempt aborted before terminal finalization: internal error"
}
}
}
impl SyncAttemptTerminalGuard {
fn new(
state: &AppState,
@@ -213,7 +251,7 @@ impl SyncAttemptTerminalGuard {
RequestCandidateStatus::Failed,
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
"local_sync_attempt_aborted",
format!("Local sync attempt failed before terminal finalization: {error:?}"),
persisted_sync_abort_message(error),
)
.await;
}
@@ -346,7 +384,7 @@ impl SyncExecutionFailure {
error_type: fallback_kind
.map(SyncExecutionFailureFallbackKind::error_type)
.unwrap_or("execution_runtime_unavailable"),
message: err.to_string(),
message: safe_transport_error_message(&err),
status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()),
latency_ms: None,
fallback_kind,
@@ -1132,10 +1170,42 @@ impl<'a> OpenAiImageSyncProgressRecorder<'a> {
if chunk.is_empty() {
return;
}
self.buffer.extend_from_slice(chunk);
let mut blocks = Vec::new();
let mut remaining = chunk;
let mut parser_overflowed = false;
loop {
while let Some(block_end) = find_sse_block_end(&self.buffer) {
blocks.push(self.buffer.drain(..block_end).collect::<Vec<_>>());
}
if remaining.is_empty() {
break;
}
let capacity =
OPENAI_IMAGE_SYNC_PROGRESS_MAX_BUFFER_BYTES.saturating_sub(self.buffer.len());
if capacity == 0 {
// The progress recorder is observational. Drop an incomplete
// oversized record and keep the client-facing response alive.
self.buffer.clear();
parser_overflowed = true;
break;
}
let take = find_sse_block_end(remaining)
.map_or(remaining.len(), |block_end| block_end)
.min(capacity);
if take == 0 {
self.buffer.clear();
parser_overflowed = true;
break;
}
self.buffer.extend_from_slice(&remaining[..take]);
remaining = &remaining[take..];
}
if parser_overflowed {
debug!("openai image sync progress parser dropped an oversized incomplete SSE record");
}
let mut force_persist = false;
while let Some(block_end) = find_sse_block_end(&self.buffer) {
let block = self.buffer.drain(..block_end).collect::<Vec<_>>();
for block in blocks {
let Some(frame) = parse_openai_image_sync_sse_frame(&block) else {
continue;
};
@@ -1813,12 +1883,16 @@ async fn openai_image_sync_json_heartbeat_final_bytes(
{
Ok(bytes) if !bytes.is_empty() => bytes.to_vec(),
Ok(_) => openai_image_sync_json_heartbeat_error_body("empty sync image response"),
Err(err) => openai_image_sync_json_heartbeat_error_body(&err.to_string()),
Err(_err) => {
openai_image_sync_json_heartbeat_error_body("sync image response read failed")
}
},
Ok(None) => openai_image_sync_json_heartbeat_error_body(
"sync image execution ended without a local response",
),
Err(err) => openai_image_sync_json_heartbeat_error_body(&format!("{err:?}")),
// Do not serialize the internal error: Debug output can contain
// upstream URLs, credentials, or other request-specific details.
Err(_err) => openai_image_sync_json_heartbeat_error_body("sync image execution failed"),
}
}
@@ -1840,7 +1914,7 @@ async fn apply_sync_success_effects(
) {
if let Some(report_context) = report_context {
crate::ai_serving::persist_converted_response_history(
state.runtime_state(),
state,
report_context,
payload
.client_body_json
@@ -2155,7 +2229,7 @@ async fn execute_execution_runtime_sync_impl(
}
},
Err(err) => {
let transport_error_message = err.to_string();
let transport_error_message = safe_transport_error_message(&err);
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
@@ -2167,7 +2241,7 @@ async fn execute_execution_runtime_sync_impl(
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
error = %transport_error_message,
"gateway ChatGPT-Web image execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
@@ -2207,7 +2281,7 @@ async fn execute_execution_runtime_sync_impl(
}
}
Err(err) => {
let transport_error_message = err.to_string();
let transport_error_message = safe_transport_error_message(&err);
warn!(
event_name = "grok_execution_unavailable",
log_type = "ops",
@@ -2219,7 +2293,7 @@ async fn execute_execution_runtime_sync_impl(
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
error = %transport_error_message,
"gateway Grok execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
@@ -2262,7 +2336,7 @@ async fn execute_execution_runtime_sync_impl(
match (override_fn.0)(&plan) {
Ok(result) => result,
Err(err) => {
let transport_error_message = format!("{err:?}");
let transport_error_message = persisted_sync_abort_message(&err).to_string();
warn!(
event_name = "sync_execution_runtime_test_override_failed",
log_type = "ops",
@@ -2408,7 +2482,7 @@ async fn execute_execution_runtime_sync_impl(
}
},
Err(err) => {
let transport_error_message = err.to_string();
let transport_error_message = safe_transport_error_message(&err);
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
@@ -2420,7 +2494,7 @@ async fn execute_execution_runtime_sync_impl(
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
error = %transport_error_message,
"gateway ChatGPT-Web image execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
@@ -2459,7 +2533,7 @@ async fn execute_execution_runtime_sync_impl(
}
},
Err(err) => {
let transport_error_message = err.to_string();
let transport_error_message = safe_transport_error_message(&err);
warn!(
event_name = "grok_execution_unavailable",
log_type = "ops",
@@ -2471,7 +2545,7 @@ async fn execute_execution_runtime_sync_impl(
key_id,
model_name,
candidate_index = candidate_index.as_str(),
error = %err,
error = %transport_error_message,
"gateway Grok execution unavailable"
);
let terminal_unix_secs = current_request_candidate_unix_ms();
@@ -2572,8 +2646,26 @@ async fn execute_execution_runtime_sync_impl(
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let mut headers = std::mem::take(&mut result.headers);
let (body_bytes, mut body_json, body_base64) =
decode_execution_result_body(result.body.take(), &mut headers)?;
let result_body = result.body.take();
let chatgpt_web_image_result = plan
.provider_api_format
.eq_ignore_ascii_case("openai:image")
&& (plan.headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("x-aether-chatgpt-web-image") && value == "1"
}) || report_context
.as_ref()
.and_then(|value| value.get("chatgpt_web_image"))
.and_then(Value::as_bool)
.unwrap_or(false));
let (body_bytes, mut body_json, body_base64) = if chatgpt_web_image_result {
decode_execution_result_body_with_limit(
result_body,
&mut headers,
crate::execution_runtime::chatgpt_web_image::chatgpt_web_image_sse_envelope_limit_bytes(),
)?
} else {
decode_execution_result_body(result_body, &mut headers)?
};
if let Some(message) = invalid_gemini_provider_success_message(
&plan,
report_context.as_ref(),
@@ -3382,7 +3474,7 @@ async fn execute_sync_via_remote_execution_runtime(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
error_message: Some(persisted_sync_abort_message(&err).to_string()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
@@ -3423,9 +3515,15 @@ async fn execute_sync_via_remote_execution_runtime(
}
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
let mut result = response
.json::<ExecutionResult>()
.await
let response_body = aether_http::read_response_bytes_with_limit(
response,
crate::execution_runtime::transport::execution_result_envelope_limit_bytes(
crate::headers::max_internal_buffered_body_bytes(),
),
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let mut result = serde_json::from_slice::<ExecutionResult>(&response_body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
result
.response_observation
@@ -3893,6 +3991,20 @@ mod tests {
assert!(message.is_none());
}
#[test]
fn forced_sync_abort_message_does_not_copy_internal_error_details() {
let secret = "https://user:[email protected]/v1?api_key=should-not-persist";
let error = GatewayError::Internal(secret.to_string());
let message = persisted_sync_abort_message(&error);
assert_eq!(
message,
"local sync attempt aborted before terminal finalization: internal error"
);
assert!(!message.contains("password"));
assert!(!message.contains("api_key"));
}
#[tokio::test]
async fn sync_attempt_terminal_guard_marks_dropped_pending_attempt_cancelled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -4482,4 +4594,22 @@ mod tests {
Some(&report_context),
));
}
#[tokio::test]
async fn openai_image_sync_json_heartbeat_hides_internal_error_details() {
let secret = "https://user:[email protected]/private?api_key=top-secret";
let bytes = openai_image_sync_json_heartbeat_final_bytes(Err(GatewayError::Internal(
secret.to_string(),
)))
.await;
let body: Value = serde_json::from_slice(&bytes).expect("heartbeat error body is JSON");
assert_eq!(
body.pointer("/error/message").and_then(Value::as_str),
Some("sync image execution failed")
);
let body_text = String::from_utf8(bytes).expect("heartbeat body is UTF-8");
assert!(!body_text.contains(secret));
assert!(!body_text.contains("top-secret"));
}
}
@@ -1,15 +1,35 @@
use aether_contracts::ResponseBody;
use serde_json::Value;
use std::collections::BTreeMap;
use aether_contracts::ResponseBody;
use base64::Engine as _;
use crate::execution_runtime::transport::{
decode_base64_body_with_limit,
serialize_json_body_with_limit as serialize_transport_json_body_with_limit,
};
use crate::GatewayError;
type DecodedBody = (Vec<u8>, Option<serde_json::Value>, Option<String>);
fn serialize_json_body_with_limit(body: &Value, limit: usize) -> Result<Vec<u8>, GatewayError> {
serialize_transport_json_body_with_limit(body, limit)
.map_err(|error| GatewayError::Internal(error.to_string()))
}
pub(super) fn decode_execution_result_body(
body: Option<ResponseBody>,
headers: &mut BTreeMap<String, String>,
) -> Result<DecodedBody, GatewayError> {
decode_execution_result_body_with_limit(
body,
headers,
crate::headers::max_internal_buffered_body_bytes(),
)
}
pub(super) fn decode_execution_result_body_with_limit(
body: Option<ResponseBody>,
headers: &mut BTreeMap<String, String>,
body_limit: usize,
) -> Result<DecodedBody, GatewayError> {
let Some(body) = body else {
return Ok((Vec::new(), None, None));
@@ -18,10 +38,8 @@ pub(super) fn decode_execution_result_body(
json_body,
body_bytes_b64,
} = body;
if let Some(body_bytes_b64) = body_bytes_b64 {
let bytes = base64::engine::general_purpose::STANDARD
.decode(&body_bytes_b64)
let bytes = decode_base64_body_with_limit(&body_bytes_b64, body_limit)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok((bytes, json_body, Some(body_bytes_b64)));
}
@@ -32,8 +50,7 @@ pub(super) fn decode_execution_result_body(
headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
let bytes = serde_json::to_vec(&json_body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let bytes = serialize_json_body_with_limit(&json_body, body_limit)?;
headers.insert("content-length".to_string(), bytes.len().to_string());
return Ok((bytes, Some(json_body), None));
}
@@ -118,4 +135,34 @@ mod tests {
Some(raw_len.as_str())
);
}
#[test]
fn scoped_decode_limit_can_cover_a_bounded_synthetic_envelope() {
let raw = vec![b'x'; 65 * 1024];
let encoded = base64::engine::general_purpose::STANDARD.encode(&raw);
let mut headers = BTreeMap::new();
let (decoded, json, retained) = super::decode_execution_result_body_with_limit(
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(encoded.clone()),
}),
&mut headers,
raw.len(),
)
.expect("body at the scoped limit should decode");
assert_eq!(decoded, raw);
assert_eq!(json, None);
assert_eq!(retained.as_deref(), Some(encoded.as_str()));
assert!(super::decode_execution_result_body_with_limit(
Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(encoded),
}),
&mut headers,
raw.len() - 1,
)
.is_err());
}
}
@@ -2,13 +2,10 @@ use std::collections::BTreeMap;
use aether_contracts::ExecutionPlan;
use axum::body::Body;
use axum::http::header::HeaderValue;
use axum::http::Response;
use serde_json::json;
use crate::api::response::{
build_client_response_from_parts, build_client_response_from_parts_with_mutator,
};
use crate::api::response::build_client_response_from_parts;
use crate::async_task::VideoTaskService;
use crate::control::GatewayControlDecision;
use crate::video_tasks::{
@@ -156,35 +153,55 @@ pub(crate) fn maybe_build_local_video_error_response(
return Ok(None);
}
let empty_body = json!({});
let response_body = payload.body_json.as_ref().unwrap_or(&empty_body);
let body_bytes =
serde_json::to_vec(response_body).map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_len = body_bytes.len().to_string();
let response_body =
local_video_error_response_body(payload.report_kind.as_str(), payload.status_code);
let body_bytes = serde_json::to_vec(&response_body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let headers = BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("content-length".to_string(), body_bytes.len().to_string()),
]);
Ok(Some(build_client_response_from_parts_with_mutator(
Ok(Some(build_client_response_from_parts(
payload.status_code,
&payload.headers,
&headers,
Body::from(body_bytes),
trace_id,
Some(decision),
|headers| {
headers.remove(http::header::CONTENT_ENCODING);
headers.remove(http::header::CONTENT_LENGTH);
headers.insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_str(body_len.as_str())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
);
Ok(())
},
)?))
}
fn local_video_error_response_body(report_kind: &str, status_code: u16) -> serde_json::Value {
let (code, gemini_status) = match status_code {
400 => ("invalid_request", "INVALID_ARGUMENT"),
401 => ("authentication_error", "UNAUTHENTICATED"),
403 => ("permission_denied", "PERMISSION_DENIED"),
404 => ("not_found", "NOT_FOUND"),
429 => ("rate_limit_exceeded", "RESOURCE_EXHAUSTED"),
503 => ("server_error", "UNAVAILABLE"),
500..=599 => ("server_error", "INTERNAL"),
_ => ("provider_error", "UNKNOWN"),
};
if report_kind.starts_with("gemini_video_") {
json!({
"error": {
"code": status_code,
"message": "Video generation failed",
"status": gemini_status,
}
})
} else {
json!({
"error": {
"message": "Video generation failed",
"type": code,
"code": code,
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -193,7 +210,7 @@ mod tests {
use serde_json::json;
#[tokio::test]
async fn local_video_error_response_rewrites_headers_without_mutating_payload() {
async fn local_video_error_response_does_not_expose_upstream_payload_or_headers() {
let decision = GatewayControlDecision::synthetic(
"/v1/videos",
Some("ai_public".to_string()),
@@ -212,12 +229,15 @@ mod tests {
headers: BTreeMap::from([
("content-encoding".to_string(), "gzip".to_string()),
("content-length".to_string(), "999".to_string()),
("x-upstream-id".to_string(), "video-123".to_string()),
(
"x-upstream-debug".to_string(),
"Authorization: Bearer header-secret".to_string(),
),
]),
body_json: Some(json!({
"error": {
"type": "video_backend_error",
"message": "backend failed",
"message": "Authorization: Bearer body-secret at https://internal.test/?key=secret",
}
})),
client_body_json: None,
@@ -239,13 +259,7 @@ mod tests {
Some("application/json")
);
assert_eq!(response.headers().get(http::header::CONTENT_ENCODING), None);
assert_eq!(
response
.headers()
.get("x-upstream-id")
.and_then(|value| value.to_str().ok()),
Some("video-123")
);
assert_eq!(response.headers().get("x-upstream-debug"), None);
assert_eq!(
payload.headers.get("content-encoding").map(String::as_str),
Some("gzip")
@@ -258,12 +272,36 @@ mod tests {
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let response_body =
serde_json::from_slice::<serde_json::Value>(&body).expect("response body should parse");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("response body should parse"),
payload
.body_json
.clone()
.expect("payload body should exist")
response_body,
json!({
"error": {
"message": "Video generation failed",
"type": "server_error",
"code": "server_error",
}
})
);
let encoded = response_body.to_string();
for sensitive in ["Bearer", "body-secret", "header-secret", "internal.test"] {
assert!(!encoded.contains(sensitive));
}
assert!(payload.body_json.is_some());
}
#[test]
fn gemini_video_error_response_uses_fixed_schema() {
assert_eq!(
local_video_error_response_body("gemini_video_create_sync_finalize", 429),
json!({
"error": {
"code": 429,
"message": "Video generation failed",
"status": "RESOURCE_EXHAUSTED",
}
})
);
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff