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
@@ -2325,16 +2325,11 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 2);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
assert!(stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.is_none());
}
#[test]
@@ -2624,16 +2619,11 @@ mod tests {
);
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[1].candidate_index, 1);
assert_eq!(
stored[1]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
assert!(stored[1]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.is_none());
}
#[test]
@@ -2703,7 +2693,7 @@ mod tests {
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("ranking metadata should persist as object extra data");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert!(extra_data.get("existing").is_none());
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
@@ -2716,14 +2706,7 @@ mod tests {
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
assert_eq!(
extra_data
.get("dispatch_ref")
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("ranked-key"))
);
assert!(extra_data.get("dispatch_ref").is_none());
}
#[tokio::test]
@@ -3061,7 +3044,7 @@ mod tests {
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("skipped ranking metadata should persist");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert!(extra_data.get("existing").is_none());
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
@@ -278,13 +278,21 @@ mod tests {
assert_eq!(metadata["transport_diagnostics"]["provider_type"], "codex");
assert_eq!(
metadata["transport_diagnostics"]["fingerprint"]["transport_profile"]["profile_id"],
"chrome_136"
metadata["transport_diagnostics"]["key_fingerprint_configured"],
Value::Bool(true)
);
assert_eq!(
metadata["transport_diagnostics"]["key_transport_profile_configured"],
Value::Bool(true)
);
assert_eq!(
metadata["transport_diagnostics"]["resolved_transport_profile_id"],
"chrome_136"
);
assert_eq!(
metadata["transport_diagnostics"]["resolved_transport_profile"]["profile_id"],
"chrome_136"
);
assert_eq!(
metadata["transport_diagnostics"]["request_pair"]["conversion_enabled"],
Value::Bool(true)
@@ -251,6 +251,7 @@ mod tests {
use aether_ai_serving::{
ai_ranking_context, build_ai_rankable_candidate, AiRankableCandidateParts,
};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::{
provider_catalog::InMemoryProviderCatalogReadRepository,
routing_profiles::InMemoryRoutingGroupRepository,
@@ -566,6 +567,15 @@ mod tests {
api_formats: Option<serde_json::Value>,
allowed_models: Option<serde_json::Value>,
) -> 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_id, id, "plain-upstream-key")
.expect("api key should encrypt");
StoredProviderCatalogKey::new(
id.to_string(),
provider_id.to_string(),
@@ -577,7 +587,7 @@ mod tests {
.expect("key should build")
.with_transport_fields(
api_formats,
"plain-upstream-key".to_string(),
encrypted_api_key,
None,
None,
Some(json!({"openai:chat": 1})),
@@ -691,7 +701,7 @@ mod tests {
let observed_at_unix_secs = current_unix_secs();
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![
("provider_priority_mode".to_string(), json!("provider")),
@@ -700,6 +710,7 @@ mod tests {
serde_json::to_value(TunnelAttachmentRecord {
gateway_instance_id: "gateway-b".to_string(),
relay_base_url: "http://gateway-b:8080".to_string(),
tunnel_generation: "test-generation-remote".to_string(),
conn_count: 1,
observed_at_unix_secs,
})
@@ -710,6 +721,7 @@ mod tests {
serde_json::to_value(TunnelAttachmentRecord {
gateway_instance_id: "gateway-a".to_string(),
relay_base_url: "http://gateway-a:8080".to_string(),
tunnel_generation: "test-generation-local".to_string(),
conn_count: 1,
observed_at_unix_secs,
})
@@ -768,7 +780,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -821,7 +833,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
@@ -878,7 +890,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -953,7 +965,7 @@ mod tests {
.expect("routing strategy should be created");
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_routing_group_repository_for_tests(routing_repository);
let state = AppState::new()
@@ -1015,7 +1027,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
@@ -1081,7 +1093,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1134,7 +1146,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1208,7 +1220,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1288,7 +1300,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1364,7 +1376,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1431,7 +1443,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1514,7 +1526,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1579,7 +1591,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1668,7 +1680,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1754,7 +1766,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1851,7 +1863,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1956,7 +1968,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -2050,7 +2062,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1473,6 +1473,7 @@ mod tests {
use super::*;
use crate::data::GatewayDataState;
use crate::AppState;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::DataLayerError;
@@ -2174,6 +2175,19 @@ mod tests {
None,
)
.expect("endpoint transport should build");
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(
row.provider_id.as_str(),
row.key_id.as_str(),
"plain-upstream-key",
)
.expect("api key should encrypt");
let key = StoredProviderCatalogKey::new(
row.key_id.clone(),
row.provider_id.clone(),
@@ -2185,7 +2199,7 @@ mod tests {
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!([row.endpoint_api_format.clone()])),
"plain-upstream-key".to_string(),
encrypted_api_key,
None,
None,
None,
@@ -2540,7 +2554,7 @@ mod tests {
provider_repository,
candidate_repository,
)
.with_encryption_key_for_tests("development-key");
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
@@ -2660,7 +2674,7 @@ mod tests {
provider_repository,
candidate_repository,
)
.with_encryption_key_for_tests("development-key")
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY)
// Legacy keys deliberately disagree with the routing policy: the
// resolved policy must be the only source of scheduler ordering.
.with_system_config_values_for_tests([
@@ -11,6 +11,7 @@ use crate::ai_serving::planner::route::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
};
use crate::ai_serving::{resolve_decision_execution_runtime_auth_context, GatewayControlDecision};
use crate::state::VideoTaskRouteAccess;
use crate::{AiExecutionDecision, AppState, GatewayError};
pub(crate) async fn maybe_build_stream_decision_payload(
@@ -155,16 +156,37 @@ async fn maybe_build_local_video_task_content_stream_decision_payload(
return Ok(None);
}
let _ = state
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
.await?;
let Some(user_id) = decision
.auth_context
.as_ref()
.filter(|auth_context| auth_context.access_allowed)
.map(|auth_context| auth_context.user_id.trim())
.filter(|value| !value.is_empty())
else {
return Err(crate::video_tasks::not_found_error());
};
if state
.hydrate_video_task_for_route_for_user(
decision.route_family.as_deref(),
parts.uri.path(),
user_id,
)
.await?
!= VideoTaskRouteAccess::Allowed
{
return Err(crate::video_tasks::not_found_error());
}
let Some(action) = state.video_tasks.prepare_openai_content_stream_action(
parts.uri.path(),
parts.uri.query(),
trace_id,
) else {
return Ok(None);
let Some(action) = state
.video_tasks
.prepare_openai_content_stream_action_for_user(
parts.uri.path(),
parts.uri.query(),
trace_id,
user_id,
)
else {
return Err(crate::video_tasks::not_found_error());
};
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
@@ -16,6 +16,7 @@ use crate::ai_serving::{
build_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
GatewayControlDecision,
};
use crate::state::VideoTaskRouteAccess;
use crate::{AiExecutionDecision, AppState, GatewayError};
pub(crate) async fn maybe_build_sync_decision_payload(
@@ -191,10 +192,6 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
return Ok(None);
}
let _ = state
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
.await?;
let auth_context = resolve_execution_runtime_auth_context(
state,
decision,
@@ -204,16 +201,30 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
)
.await?;
let Some(auth_context) = auth_context else {
return Ok(None);
return Err(crate::video_tasks::not_found_error());
};
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan(
if !auth_context.access_allowed || auth_context.user_id.trim().is_empty() {
return Err(crate::video_tasks::not_found_error());
}
if state
.hydrate_video_task_for_route_for_user(
decision.route_family.as_deref(),
parts.uri.path(),
&auth_context.user_id,
)
.await?
!= VideoTaskRouteAccess::Allowed
{
return Err(crate::video_tasks::not_found_error());
}
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan_for_user(
plan_kind,
parts.uri.path(),
Some(body_json),
Some(&auth_context),
trace_id,
) else {
return Ok(None);
return Err(crate::video_tasks::not_found_error());
};
let aether_video_tasks_core::LocalVideoTaskFollowUpPlan {
@@ -236,8 +247,7 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
downstream_path = %parts.uri.path(),
provider_api_format = %plan.provider_api_format,
client_api_format = %plan.client_api_format,
upstream_base_url = ?upstream_base_url,
upstream_url = %plan.url,
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
"gateway built local video follow-up sync decision payload"
);
@@ -37,6 +37,10 @@ const ROUTING_GROUP_SELECTION_CACHE_TTL: Duration = Duration::from_secs(30);
const ROUTING_GROUP_SELECTION_CACHE_STALE_TTL: Duration = Duration::from_secs(120);
const CODEX_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
const CODEX_FEDRAMP_HEADER: &str = "x-openai-fedramp";
const INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE: &str =
"routing provider request violates provider contract";
const INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE: &str =
"invalid provider request headers in routing mutation";
#[derive(Debug, Clone)]
pub(crate) struct ResolvedLocalDecisionAuthInput {
@@ -312,10 +316,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
)
}
}
.map_err(|violation| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("routing provider_request violates provider contract: {violation:?}"),
})?;
.map_err(|_| invalid_routing_provider_contract())?;
}
let provider_model = provider_request_body
.get("model")
@@ -863,10 +864,32 @@ fn routing_selection_error(error: GatewayRoutingSelectionError) -> GatewayError
status: StatusCode::SERVICE_UNAVAILABLE,
message: "no enabled routing strategy is configured for this request".to_string(),
},
error => GatewayError::Client {
GatewayRoutingSelectionError::NotFound(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: error.to_string(),
message: "requested routing group was not found".to_string(),
},
GatewayRoutingSelectionError::Disabled(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: "requested routing group is not enabled".to_string(),
},
GatewayRoutingSelectionError::Forbidden(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: "requested routing group is not allowed for this principal".to_string(),
},
}
}
fn invalid_routing_provider_contract() -> GatewayError {
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE.to_string(),
}
}
fn invalid_routing_provider_headers() -> GatewayError {
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE.to_string(),
}
}
@@ -921,14 +944,9 @@ fn btree_headers_to_header_map(
) -> Result<HeaderMap, GatewayError> {
let mut output = HeaderMap::new();
for (name, value) in headers {
let name = HeaderName::from_bytes(name.as_bytes()).map_err(|err| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("invalid provider request header name in routing mutation: {err}"),
})?;
let value = HeaderValue::from_str(value).map_err(|err| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("invalid provider request header value in routing mutation: {err}"),
})?;
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(|_| invalid_routing_provider_headers())?;
let value = HeaderValue::from_str(value).map_err(|_| invalid_routing_provider_headers())?;
output.insert(name, value);
}
Ok(output)
@@ -1151,6 +1169,50 @@ mod tests {
}
}
#[test]
fn routing_selection_errors_do_not_echo_explicit_group() {
let secret = "private-group?token=Bearer-secret";
for error in [
GatewayRoutingSelectionError::NotFound(secret.to_string()),
GatewayRoutingSelectionError::Disabled(secret.to_string()),
GatewayRoutingSelectionError::Forbidden(secret.to_string()),
] {
let error = routing_selection_error(error);
assert!(matches!(
error,
GatewayError::Client {
status: StatusCode::FORBIDDEN,
ref message,
} if !message.contains(secret)
));
}
}
#[test]
fn routing_provider_errors_do_not_echo_dynamic_details() {
let secret = "https://internal.example/?token=Bearer-secret";
let contract_error = invalid_routing_provider_contract();
let header_error = btree_headers_to_header_map(&BTreeMap::from([(
format!("Authorization: {secret}"),
secret.to_string(),
)]))
.expect_err("invalid header should fail");
for (error, expected_message) in [
(contract_error, INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE),
(header_error, INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE),
] {
assert!(matches!(
error,
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
ref message,
} if message == expected_message && !message.contains(secret)
));
}
}
#[tokio::test]
async fn explicit_auth_snapshot_override_does_not_fall_back_to_the_planner_cache() {
// AppState::new has no auth snapshot repository. Without the explicit
@@ -76,10 +76,12 @@ pub(crate) fn build_local_execution_report_context(
let RequestOrigin {
client_ip,
user_agent,
forwarded_headers_trusted,
} = parts
.request_origin
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
let original_headers = crate::ai_serving::collect_control_headers(parts.original_headers);
let original_headers =
collect_report_context_original_headers(parts.original_headers, forwarded_headers_trusted);
let original_request_body = crate::ai_serving::build_report_context_original_request_echo(
parts.original_request_body_json,
parts.original_request_body_base64,
@@ -106,10 +108,12 @@ pub(crate) fn build_local_execution_report_context(
value,
);
}
if let Some(incoming_tls) =
crate::ai_serving::tls_fingerprint_from_headers(parts.original_headers)
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
if forwarded_headers_trusted {
if let Some(incoming_tls) =
crate::ai_serving::tls_fingerprint_from_headers(parts.original_headers)
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
}
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
insert_scheduler_affinity_policy_report_context_field(&mut extra_fields, parts.routing_policy);
@@ -189,6 +193,17 @@ pub(crate) fn build_local_execution_report_context(
})
}
fn collect_report_context_original_headers(
headers: &http::HeaderMap,
forwarded_headers_trusted: bool,
) -> BTreeMap<String, String> {
let mut collected = crate::ai_serving::collect_control_headers(headers);
if !forwarded_headers_trusted {
collected.retain(|name, _| !name.starts_with("x-aether-tls-"));
}
collected
}
fn insert_request_path_fields(
extra_fields: &mut Map<String, Value>,
request_path: Option<&str>,
@@ -258,8 +273,8 @@ mod tests {
use serde_json::{json, Map, Value};
use super::{
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
LocalExecutionReportContextParts,
build_local_execution_report_context, collect_report_context_original_headers,
provider_stream_event_api_format_for_provider_type, LocalExecutionReportContextParts,
};
use crate::ai_serving::ExecutionRuntimeAuthContext;
use crate::ai_serving::RequestOrigin;
@@ -289,6 +304,26 @@ mod tests {
);
}
#[test]
fn untrusted_tls_forwarding_headers_are_excluded_from_report_context() {
let mut headers = http::HeaderMap::new();
headers.insert("x-aether-tls-ja3", "spoofed-ja3".parse().unwrap());
headers.insert(http::header::USER_AGENT, "test-client".parse().unwrap());
let untrusted = collect_report_context_original_headers(&headers, false);
assert!(!untrusted.contains_key("x-aether-tls-ja3"));
assert_eq!(
untrusted.get("user-agent").map(String::as_str),
Some("test-client")
);
let trusted = collect_report_context_original_headers(&headers, true);
assert_eq!(
trusted.get("x-aether-tls-ja3").map(String::as_str),
Some("spoofed-ja3")
);
}
#[test]
fn local_execution_report_context_records_request_origin_and_session_affinity() {
let auth_context = ExecutionRuntimeAuthContext {
@@ -339,6 +374,7 @@ mod tests {
request_origin: Some(RequestOrigin {
client_ip: Some("203.0.113.8".to_string()),
user_agent: Some("Claude-Code/1.0".to_string()),
forwarded_headers_trusted: false,
}),
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
@@ -491,7 +527,11 @@ mod tests {
original_headers: &original_headers,
request_path: None,
request_query_string: None,
request_origin: None,
request_origin: Some(RequestOrigin {
client_ip: None,
user_agent: None,
forwarded_headers_trusted: true,
}),
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: None,
@@ -8,7 +8,10 @@ use crate::ai_serving::transport::{
GeminiFilesRequestBodyError,
};
use crate::ai_serving::GEMINI_FILES_UPLOAD_PLAN_KIND;
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::ai_serving::{
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
};
use crate::AppState;
use super::support::{
@@ -47,6 +50,26 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
let transport = &attempt.eligible.transport;
let effective_headers = input.effective_headers(&parts.headers);
if matches!(
spec_metadata.decision_kind,
GEMINI_FILES_GET_PLAN_KIND
| GEMINI_FILES_DELETE_PLAN_KIND
| GEMINI_FILES_DOWNLOAD_PLAN_KIND
) && !candidate_matches_owned_gemini_file_mapping(state, parts, input, attempt).await
{
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"gemini_file_mapping_mismatch",
)
.await;
return None;
}
if let Some(skip_reason) =
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
{
@@ -191,3 +214,64 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
file_name,
})
}
async fn candidate_matches_owned_gemini_file_mapping(
state: &AppState,
parts: &http::request::Parts,
input: &LocalGeminiFilesDecisionInput,
attempt: &LocalGeminiFilesCandidateAttempt,
) -> bool {
let Some(file_name) = normalize_gemini_file_name_from_path(parts.uri.path()) else {
return false;
};
let user_id = input.auth_context.user_id.trim();
if user_id.is_empty() || !state.has_gemini_file_mapping_data_reader() {
return false;
}
let Ok(Some(mapping)) = state
.find_active_gemini_file_mapping_for_owner(
file_name.as_str(),
&attempt.eligible.transport.key.id,
user_id,
crate::clock::current_unix_secs(),
)
.await
else {
return false;
};
mapping.user_id.as_deref().map(str::trim) == Some(user_id)
&& mapping.key_id == attempt.eligible.transport.key.id
}
pub(crate) fn normalize_gemini_file_name_from_path(path: &str) -> Option<String> {
let suffix = path.strip_prefix("/v1beta/files/")?.trim_matches('/');
let suffix = suffix.strip_suffix(":download").unwrap_or(suffix).trim();
let suffix = suffix.strip_prefix("files/").unwrap_or(suffix).trim();
if suffix.is_empty() || suffix.contains('/') {
return None;
}
Some(format!("files/{suffix}"))
}
#[cfg(test)]
mod tests {
use super::normalize_gemini_file_name_from_path;
#[test]
fn normalizes_supported_gemini_file_object_paths() {
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/file-123"),
Some("files/file-123".to_string())
);
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/file-123:download"),
Some("files/file-123".to_string())
);
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/files/abc-123"),
Some("files/abc-123".to_string())
);
assert_eq!(normalize_gemini_file_name_from_path("/v1beta/files"), None);
}
}
@@ -587,7 +587,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
state,
body_json,
spec_metadata.api_format,
provider_api_format,
@@ -1,7 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use tracing::debug;
use super::super::{
augment_sync_report_context, build_ai_execution_plan_from_decision,
@@ -10,7 +9,6 @@ use super::super::{
AiStreamAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
@@ -157,21 +155,16 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
return Ok(None);
};
let (url, url_source) = if let Some(upstream_url) =
take_non_empty_string(&mut payload.upstream_url)
{
(upstream_url, "upstream_url")
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
upstream_url
} else {
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
return Ok(None);
};
(
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
),
"upstream_base_url",
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
)
};
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
@@ -238,16 +231,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
tracing::debug!(
event_name = "local_openai_responses_stream_plan_built",
log_type = "debug",
request_id = %plan.request_id,
@@ -255,12 +239,13 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
downstream_path = %parts.uri.path(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = effective_upstream_is_stream,
@@ -1,7 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use tracing::debug;
use super::super::{
augment_sync_report_context, build_ai_execution_plan_from_decision,
@@ -10,7 +9,6 @@ use super::super::{
AiSyncAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -142,21 +140,16 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
return Ok(None);
};
let (url, url_source) = if let Some(upstream_url) =
take_non_empty_string(&mut payload.upstream_url)
{
(upstream_url, "upstream_url")
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
upstream_url
} else {
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
return Ok(None);
};
(
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
),
"upstream_base_url",
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
)
};
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
@@ -205,16 +198,7 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
tracing::debug!(
event_name = "local_openai_responses_sync_plan_built",
log_type = "debug",
request_id = %plan.request_id,
@@ -222,12 +206,13 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
downstream_path = %parts.uri.path(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = payload.upstream_is_stream,
@@ -3,7 +3,6 @@ use tracing::debug;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision_with_websocket_mode;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_native_client_envelope_name,
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
@@ -205,12 +204,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
&resolved.transport,
);
let log_base_url = sanitize_upstream_url_for_log(resolved.transport.endpoint.base_url.as_str());
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(resolved.upstream_url.as_str());
debug!(
event_name = "local_openai_responses_decision_payload_built",
log_type = "debug",
@@ -227,9 +224,12 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_api_format = spec_metadata.api_format,
provider_api_format = %resolved.provider_api_format,
request_path = %parts.uri.path(),
request_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&resolved.upstream_url),
request_query = ?log_request_query,
upstream_base_url = %log_base_url,
upstream_url = %log_upstream_url,
upstream_is_stream = resolved.upstream_is_stream,
has_envelope = resolved.envelope_name.is_some(),
"gateway built local openai responses decision payload"
@@ -24,7 +24,6 @@ use crate::ai_serving::planner::gemini_cli::{
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
sanitize_upstream_url_for_log,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
@@ -428,7 +427,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
state,
body_json,
spec_metadata.api_format,
provider_api_format,
@@ -867,17 +866,10 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
let log_base_url = sanitize_upstream_url_for_log(transport.endpoint.base_url.as_str());
let log_custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_upstream_url_resolved",
@@ -893,12 +885,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
base_url = %log_base_url,
custom_path = ?log_custom_path,
request_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&upstream_url),
request_path = %parts.uri.path(),
request_query = ?log_request_query,
mapped_model = %mapped_model,
upstream_url = %log_upstream_url,
upstream_is_stream,
"gateway resolved local openai responses upstream url"
);
@@ -2011,8 +2005,6 @@ async fn build_kiro_openai_responses_payload_parts(
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_kiro_upstream_url_resolved",
log_type = "debug",
@@ -2027,7 +2019,7 @@ async fn build_kiro_openai_responses_payload_parts(
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
upstream_url = %log_upstream_url,
upstream_origin = %crate::handlers::shared::security_log_url_origin(&upstream_url),
upstream_is_stream,
"gateway resolved local openai responses kiro upstream url"
);