mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 08:27:46 +08:00
fix: harden frontdoor and usage ingestion
This commit is contained in:
@@ -317,12 +317,7 @@ pub(super) async fn execute_provider_quota_plan(
|
||||
match state.execute_execution_runtime_sync_plan(None, &plan).await {
|
||||
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
|
||||
Err(err) => {
|
||||
let error = match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
};
|
||||
let error = err.into_message();
|
||||
let proxy_node_id = plan
|
||||
.proxy
|
||||
.as_ref()
|
||||
|
||||
@@ -301,12 +301,7 @@ fn admin_provider_ops_decode_response_bytes(
|
||||
}
|
||||
|
||||
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
|
||||
match error {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
}
|
||||
error.into_message()
|
||||
}
|
||||
|
||||
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {
|
||||
|
||||
@@ -16,7 +16,12 @@ use aether_runtime_state::{DataLayerError, RuntimeState};
|
||||
use futures_util::future::join_all;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tracing::warn;
|
||||
use tracing::{info, warn};
|
||||
|
||||
const DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 512;
|
||||
const MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 10_000;
|
||||
const POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV: &str =
|
||||
"AETHER_GATEWAY_ADMIN_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT";
|
||||
|
||||
fn current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
@@ -29,6 +34,20 @@ fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bo
|
||||
pool_config.probing_enabled
|
||||
}
|
||||
|
||||
fn pool_runtime_window_metric_key_limit() -> usize {
|
||||
std::env::var(POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV)
|
||||
.ok()
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
|
||||
.clamp(1, MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
|
||||
}
|
||||
|
||||
fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[String] {
|
||||
let end = key_ids.len().min(limit.max(1));
|
||||
&key_ids[..end]
|
||||
}
|
||||
|
||||
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
|
||||
runtime: &RuntimeState,
|
||||
provider_ids: &[String],
|
||||
@@ -54,8 +73,21 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
||||
) -> AdminProviderPoolRuntimeState {
|
||||
let mut state = AdminProviderPoolRuntimeState::default();
|
||||
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
|
||||
let cost_keys = pool_cost_keys(provider_id, key_ids);
|
||||
let latency_keys = pool_latency_keys(provider_id, key_ids);
|
||||
let metric_key_limit = pool_runtime_window_metric_key_limit();
|
||||
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
|
||||
if metric_key_ids.len() < key_ids.len() {
|
||||
info!(
|
||||
event_name = "admin_pool_runtime_window_metrics_truncated",
|
||||
log_type = "event",
|
||||
provider_id,
|
||||
total_key_count = key_ids.len(),
|
||||
scanned_key_count = metric_key_ids.len(),
|
||||
metric_key_limit,
|
||||
"gateway limited admin pool runtime cost/latency window reads"
|
||||
);
|
||||
}
|
||||
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
|
||||
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
|
||||
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
|
||||
&& admin_provider_pool_cache_affinity_enabled(pool_config);
|
||||
|
||||
@@ -179,7 +211,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
||||
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
|
||||
)
|
||||
.await;
|
||||
for (key_id, members) in key_ids.iter().zip(cost_results) {
|
||||
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
|
||||
let total = members
|
||||
.unwrap_or_default()
|
||||
.iter()
|
||||
@@ -197,7 +229,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
||||
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
|
||||
)
|
||||
.await;
|
||||
for (key_id, members) in key_ids.iter().zip(latency_results) {
|
||||
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
|
||||
let samples = members
|
||||
.unwrap_or_default()
|
||||
.iter()
|
||||
@@ -265,3 +297,30 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
|
||||
.kv_get(&pool_cooldown_key(provider_id, key_id))
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::bounded_runtime_window_metric_key_ids;
|
||||
|
||||
#[test]
|
||||
fn runtime_window_metric_key_ids_are_bounded() {
|
||||
let key_ids = vec![
|
||||
"key-1".to_string(),
|
||||
"key-2".to_string(),
|
||||
"key-3".to_string(),
|
||||
];
|
||||
|
||||
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 2);
|
||||
|
||||
assert_eq!(bounded, &key_ids[..2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_window_metric_key_ids_keep_at_least_one_key() {
|
||||
let key_ids = vec!["key-1".to_string(), "key-2".to_string()];
|
||||
|
||||
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 0);
|
||||
|
||||
assert_eq!(bounded, &key_ids[..1]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -662,10 +662,5 @@ fn admin_provider_oauth_decode_response_bytes(
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
|
||||
match error {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
}
|
||||
error.into_message()
|
||||
}
|
||||
|
||||
@@ -354,12 +354,7 @@ pub(crate) async fn maybe_build_internal_finalize_video_response(
|
||||
}
|
||||
|
||||
pub(crate) fn gateway_error_message(error: GatewayError) -> String {
|
||||
match error {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
}
|
||||
error.into_message()
|
||||
}
|
||||
|
||||
pub(crate) fn build_internal_tunnel_heartbeat_ack(
|
||||
|
||||
@@ -21,13 +21,13 @@ use crate::constants::{
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
|
||||
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_INVALID_REQUEST,
|
||||
EXECUTION_PATH_LOCAL_OVERLOADED, EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED,
|
||||
EXECUTION_PATH_LOCAL_RATE_LIMITED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND,
|
||||
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH, EXECUTION_RUNTIME_LOOP_GUARD_HEADER,
|
||||
FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER, FORWARDED_PROTO_HEADER, GATEWAY_HEADER,
|
||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
||||
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||
EXECUTION_PATH_LOCAL_INVALID_REQUEST, EXECUTION_PATH_LOCAL_OVERLOADED,
|
||||
EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED, EXECUTION_PATH_LOCAL_RATE_LIMITED,
|
||||
EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH,
|
||||
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER,
|
||||
FORWARDED_PROTO_HEADER, GATEWAY_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||
TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
|
||||
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
|
||||
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
|
||||
};
|
||||
@@ -65,7 +65,11 @@ use axum::extract::{ConnectInfo, Request, State};
|
||||
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
|
||||
use futures_util::StreamExt;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::{collections::BTreeMap, time::Instant};
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
error::Error as StdError,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
|
||||
@@ -89,23 +93,192 @@ const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
|
||||
"Gateway detected an execution runtime request loop back into the local frontdoor";
|
||||
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
|
||||
"当前 API Key 并发请求数已达上限,请稍后重试";
|
||||
const REQUEST_BODY_READ_TIMEOUT_DETAIL: &str =
|
||||
"Request body read timed out before the gateway could route the request";
|
||||
const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
|
||||
const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str =
|
||||
"当前 AI 请求在本地执行规划阶段超时,请稍后重试";
|
||||
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
|
||||
const MANAGEMENT_TOKEN_PREFIX: &str = "ae-";
|
||||
const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_";
|
||||
|
||||
fn build_request_body_normalization_error_response(
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct RequestBodyBufferPolicy {
|
||||
max_bytes: u64,
|
||||
read_timeout: Duration,
|
||||
}
|
||||
|
||||
impl RequestBodyBufferPolicy {
|
||||
fn from_state(state: &AppState) -> Self {
|
||||
Self {
|
||||
max_bytes: crate::headers::max_request_body_bytes(),
|
||||
read_timeout: state.frontdoor_runtime_guards.request_body_read_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn for_tests(max_bytes: u64, read_timeout: Duration) -> Self {
|
||||
Self {
|
||||
max_bytes,
|
||||
read_timeout,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum RequestBodyBufferError {
|
||||
Normalization(RequestBodyNormalizationError),
|
||||
TooLarge { limit_bytes: u64 },
|
||||
Timeout { timeout_ms: u64 },
|
||||
ReadFailed { message: String },
|
||||
}
|
||||
|
||||
impl RequestBodyBufferError {
|
||||
fn http_status(&self) -> http::StatusCode {
|
||||
match self {
|
||||
Self::Normalization(error) => error.http_status(),
|
||||
Self::TooLarge { .. } => http::StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Self::Timeout { .. } => http::StatusCode::REQUEST_TIMEOUT,
|
||||
Self::ReadFailed { .. } => http::StatusCode::BAD_REQUEST,
|
||||
}
|
||||
}
|
||||
|
||||
fn client_message(&self) -> String {
|
||||
match self {
|
||||
Self::Normalization(error) => error.client_message(),
|
||||
Self::TooLarge { limit_bytes } => format!("Request body exceeds {limit_bytes} bytes"),
|
||||
Self::Timeout { .. } => REQUEST_BODY_READ_TIMEOUT_DETAIL.to_string(),
|
||||
Self::ReadFailed { .. } => REQUEST_BODY_READ_FAILED_DETAIL.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn reason(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Normalization(error) => match error {
|
||||
RequestBodyNormalizationError::UnsupportedContentEncoding(_) => {
|
||||
"unsupported_content_encoding"
|
||||
}
|
||||
RequestBodyNormalizationError::DecodeFailed { .. } => "decode_failed",
|
||||
RequestBodyNormalizationError::DecompressedBodyTooLarge { .. } => {
|
||||
"decompressed_body_too_large"
|
||||
}
|
||||
RequestBodyNormalizationError::RequestBodyTooLarge { .. } => {
|
||||
"request_body_too_large"
|
||||
}
|
||||
},
|
||||
Self::TooLarge { .. } => "request_body_too_large",
|
||||
Self::Timeout { .. } => "request_body_read_timeout",
|
||||
Self::ReadFailed { .. } => "request_body_read_failed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn buffer_and_normalize_request_body(
|
||||
request_body: &mut Option<Body>,
|
||||
headers: &mut http::HeaderMap,
|
||||
body_owner_expectation: &'static str,
|
||||
trace_id: &str,
|
||||
method: &http::Method,
|
||||
path_and_query: &str,
|
||||
phase: &'static str,
|
||||
policy: RequestBodyBufferPolicy,
|
||||
) -> Result<Bytes, RequestBodyBufferError> {
|
||||
if let Err(err) =
|
||||
crate::headers::check_request_content_length_with_limit(headers, policy.max_bytes)
|
||||
{
|
||||
return Err(RequestBodyBufferError::Normalization(err));
|
||||
}
|
||||
|
||||
let read_started_at = Instant::now();
|
||||
let timeout_ms = policy.read_timeout.as_millis() as u64;
|
||||
info!(
|
||||
event_name = "frontdoor_request_body_buffer_started",
|
||||
log_type = "event",
|
||||
trace_id,
|
||||
method = %method,
|
||||
path = %path_and_query,
|
||||
phase,
|
||||
max_body_bytes = policy.max_bytes,
|
||||
timeout_ms,
|
||||
"gateway started buffering request body"
|
||||
);
|
||||
|
||||
let body_limit = usize::try_from(policy.max_bytes).unwrap_or(usize::MAX);
|
||||
let body = match tokio::time::timeout(
|
||||
policy.read_timeout,
|
||||
to_bytes(
|
||||
request_body.take().expect(body_owner_expectation),
|
||||
body_limit,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(body)) => body,
|
||||
Ok(Err(err)) if request_body_collection_exceeded_limit(&err) => {
|
||||
return Err(RequestBodyBufferError::TooLarge {
|
||||
limit_bytes: policy.max_bytes,
|
||||
});
|
||||
}
|
||||
Ok(Err(err)) => {
|
||||
return Err(RequestBodyBufferError::ReadFailed {
|
||||
message: err.to_string(),
|
||||
});
|
||||
}
|
||||
Err(_) => {
|
||||
return Err(RequestBodyBufferError::Timeout { timeout_ms });
|
||||
}
|
||||
};
|
||||
|
||||
let normalized = crate::headers::normalize_request_body_headers_and_bytes_with_limit(
|
||||
headers,
|
||||
body,
|
||||
policy.max_bytes,
|
||||
)
|
||||
.map_err(RequestBodyBufferError::Normalization)?;
|
||||
info!(
|
||||
event_name = "frontdoor_request_body_buffer_completed",
|
||||
log_type = "event",
|
||||
trace_id,
|
||||
method = %method,
|
||||
path = %path_and_query,
|
||||
phase,
|
||||
body_bytes = normalized.len(),
|
||||
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
|
||||
"gateway completed request body buffering"
|
||||
);
|
||||
Ok(normalized)
|
||||
}
|
||||
|
||||
fn request_body_collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
|
||||
let mut current = Some(error);
|
||||
while let Some(error) = current {
|
||||
if error.to_string().contains("length limit exceeded") {
|
||||
return true;
|
||||
}
|
||||
current = error.source();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn build_request_body_buffer_error_response(
|
||||
trace_id: &str,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
error: &RequestBodyNormalizationError,
|
||||
error: &RequestBodyBufferError,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
warn!(
|
||||
event_name = "frontdoor_request_body_normalization_failed",
|
||||
event_name = "frontdoor_request_body_buffer_failed",
|
||||
log_type = "ops",
|
||||
trace_id,
|
||||
method = %request_context.request_method,
|
||||
path = %request_context.request_path_and_query(),
|
||||
error = %error,
|
||||
"gateway rejected request with invalid encoded body"
|
||||
status_code = error.http_status().as_u16(),
|
||||
reason = error.reason(),
|
||||
detail = %error.client_message(),
|
||||
read_error = match error {
|
||||
RequestBodyBufferError::ReadFailed { message } => message.as_str(),
|
||||
_ => "",
|
||||
},
|
||||
"gateway rejected request body before local execution planning"
|
||||
);
|
||||
build_local_http_error_response(
|
||||
trace_id,
|
||||
@@ -115,36 +288,16 @@ fn build_request_body_normalization_error_response(
|
||||
)
|
||||
}
|
||||
|
||||
async fn buffer_and_normalize_request_body(
|
||||
request_body: &mut Option<Body>,
|
||||
headers: &mut http::HeaderMap,
|
||||
body_owner_expectation: &'static str,
|
||||
) -> Result<Result<Bytes, RequestBodyNormalizationError>, GatewayError> {
|
||||
if let Err(err) = crate::headers::check_request_content_length(headers) {
|
||||
return Ok(Err(err));
|
||||
}
|
||||
let body = to_bytes(
|
||||
request_body.take().expect(body_owner_expectation),
|
||||
usize::MAX,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
Ok(crate::headers::normalize_request_body_headers_and_bytes(
|
||||
headers, body,
|
||||
))
|
||||
}
|
||||
|
||||
fn finalize_request_body_normalization_rejection(
|
||||
fn finalize_request_body_buffer_rejection(
|
||||
state: &AppState,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
remote_addr: &std::net::SocketAddr,
|
||||
started_at: &std::time::Instant,
|
||||
trace_id: &str,
|
||||
request_permit: Option<aether_runtime::AdmissionPermit>,
|
||||
error: &RequestBodyNormalizationError,
|
||||
error: &RequestBodyBufferError,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
let response =
|
||||
build_request_body_normalization_error_response(trace_id, request_context, error)?;
|
||||
let response = build_request_body_buffer_error_response(trace_id, request_context, error)?;
|
||||
Ok(finalize_gateway_response_with_context(
|
||||
state,
|
||||
response,
|
||||
@@ -156,6 +309,59 @@ fn finalize_request_body_normalization_rejection(
|
||||
))
|
||||
}
|
||||
|
||||
fn local_execution_planning_timeout_parts(error: &GatewayError) -> Option<(&'static str, u64)> {
|
||||
match error {
|
||||
GatewayError::LocalExecutionPlanningTimeout {
|
||||
phase, timeout_ms, ..
|
||||
} => Some((*phase, *timeout_ms)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn finalize_local_execution_planning_timeout(
|
||||
state: &AppState,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
remote_addr: &std::net::SocketAddr,
|
||||
started_at: &std::time::Instant,
|
||||
trace_id: &str,
|
||||
request_permit: Option<aether_runtime::AdmissionPermit>,
|
||||
control_decision: Option<&GatewayControlDecision>,
|
||||
phase: &'static str,
|
||||
timeout_ms: u64,
|
||||
) -> Result<Response<Body>, GatewayError> {
|
||||
warn!(
|
||||
event_name = "frontdoor_local_execution_planning_timeout",
|
||||
log_type = "ops",
|
||||
trace_id,
|
||||
method = %request_context.request_method,
|
||||
path = %request_context.request_path_and_query(),
|
||||
route_family = control_decision
|
||||
.and_then(|decision| decision.route_family.as_deref())
|
||||
.unwrap_or("-"),
|
||||
route_kind = control_decision
|
||||
.and_then(|decision| decision.route_kind.as_deref())
|
||||
.unwrap_or("-"),
|
||||
phase,
|
||||
timeout_ms,
|
||||
"gateway failed local execution before a candidate could be selected"
|
||||
);
|
||||
let response = build_local_http_error_response(
|
||||
trace_id,
|
||||
control_decision,
|
||||
http::StatusCode::GATEWAY_TIMEOUT,
|
||||
LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL,
|
||||
)?;
|
||||
Ok(finalize_gateway_response_with_context(
|
||||
state,
|
||||
response,
|
||||
remote_addr,
|
||||
request_context,
|
||||
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT,
|
||||
started_at,
|
||||
request_permit,
|
||||
))
|
||||
}
|
||||
|
||||
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
|
||||
match outcome {
|
||||
LocalExecutionRequestOutcome::Responded(_) => "responded",
|
||||
@@ -991,16 +1197,22 @@ pub(crate) async fn proxy_request(
|
||||
}
|
||||
let mut request_body = Some(body);
|
||||
let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) {
|
||||
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
|
||||
let body = buffer_and_normalize_request_body(
|
||||
&mut request_body,
|
||||
&mut parts.headers,
|
||||
"local proxy body buffering should own request body",
|
||||
&trace_id,
|
||||
&parts.method,
|
||||
&request_context.request_path_and_query(),
|
||||
"local_proxy",
|
||||
body_buffer_policy,
|
||||
)
|
||||
.await?;
|
||||
.await;
|
||||
match body {
|
||||
Ok(body) => Some(body),
|
||||
Err(err) => {
|
||||
return finalize_request_body_normalization_rejection(
|
||||
return finalize_request_body_buffer_rejection(
|
||||
&state,
|
||||
&request_context,
|
||||
&remote_addr,
|
||||
@@ -1171,16 +1383,22 @@ pub(crate) async fn proxy_request(
|
||||
&& request_enables_control_execute(&parts.headers);
|
||||
|
||||
let buffered_body = if should_buffer_body {
|
||||
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
|
||||
let body = buffer_and_normalize_request_body(
|
||||
&mut request_body,
|
||||
&mut parts.headers,
|
||||
"buffered auth/execution runtime path should own request body",
|
||||
&trace_id,
|
||||
&parts.method,
|
||||
&request_context.request_path_and_query(),
|
||||
"auth_execution",
|
||||
body_buffer_policy,
|
||||
)
|
||||
.await?;
|
||||
.await;
|
||||
match body {
|
||||
Ok(body) => Some(body),
|
||||
Err(err) => {
|
||||
return finalize_request_body_normalization_rejection(
|
||||
return finalize_request_body_buffer_rejection(
|
||||
&state,
|
||||
&request_context,
|
||||
&remote_addr,
|
||||
@@ -1335,14 +1553,34 @@ pub(crate) async fn proxy_request(
|
||||
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body);
|
||||
let mut local_execution_exhaustion = None;
|
||||
if stream_request {
|
||||
let stream_outcome = maybe_execute_stream_request(
|
||||
let stream_outcome = match maybe_execute_stream_request(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body,
|
||||
&trace_id,
|
||||
control_decision,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(err) => {
|
||||
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
|
||||
{
|
||||
return finalize_local_execution_planning_timeout(
|
||||
&state,
|
||||
&request_context,
|
||||
&remote_addr,
|
||||
&started_at,
|
||||
&trace_id,
|
||||
request_permit.take(),
|
||||
control_decision,
|
||||
phase,
|
||||
timeout_ms,
|
||||
);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
debug!(
|
||||
event_name = "proxy_stream_local_execute_outcome",
|
||||
log_type = "debug",
|
||||
@@ -1380,9 +1618,34 @@ pub(crate) async fn proxy_request(
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
}
|
||||
match maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
|
||||
.await?
|
||||
let sync_outcome = match maybe_execute_sync_request(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body,
|
||||
&trace_id,
|
||||
control_decision,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(err) => {
|
||||
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) {
|
||||
return finalize_local_execution_planning_timeout(
|
||||
&state,
|
||||
&request_context,
|
||||
&remote_addr,
|
||||
&started_at,
|
||||
&trace_id,
|
||||
request_permit.take(),
|
||||
control_decision,
|
||||
phase,
|
||||
timeout_ms,
|
||||
);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
match sync_outcome {
|
||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||
let execution_runtime_response = restore_redacted_sync_execution_response(
|
||||
execution_runtime_response,
|
||||
@@ -1406,15 +1669,35 @@ pub(crate) async fn proxy_request(
|
||||
LocalExecutionRequestOutcome::NoPath => {}
|
||||
}
|
||||
if parts.method != http::Method::POST {
|
||||
match maybe_execute_stream_request(
|
||||
let stream_outcome = match maybe_execute_stream_request(
|
||||
&state,
|
||||
&parts,
|
||||
buffered_body,
|
||||
&trace_id,
|
||||
control_decision,
|
||||
)
|
||||
.await?
|
||||
.await
|
||||
{
|
||||
Ok(outcome) => outcome,
|
||||
Err(err) => {
|
||||
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
|
||||
{
|
||||
return finalize_local_execution_planning_timeout(
|
||||
&state,
|
||||
&request_context,
|
||||
&remote_addr,
|
||||
&started_at,
|
||||
&trace_id,
|
||||
request_permit.take(),
|
||||
control_decision,
|
||||
phase,
|
||||
timeout_ms,
|
||||
);
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
match stream_outcome {
|
||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||
let execution_runtime_response = restore_redacted_stream_execution_response(
|
||||
execution_runtime_response,
|
||||
@@ -1977,14 +2260,17 @@ fn local_execution_runtime_miss_route_detail(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
api_key_remote_ip_allowed, diagnostic_is_auth_api_key_concurrency_limited,
|
||||
local_execution_runtime_miss_detail, restore_redacted_stream_execution_response,
|
||||
restore_redacted_sync_execution_response, GatewayControlDecision,
|
||||
LocalExecutionRuntimeMissDiagnostic,
|
||||
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
|
||||
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
|
||||
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
|
||||
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
|
||||
RequestBodyBufferPolicy,
|
||||
};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::http::{header, Response};
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::http::{header, HeaderMap, Method, Response};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
@@ -2112,6 +2398,58 @@ mod tests {
|
||||
assert!(!message.contains(&sentinel));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_body_buffer_rejects_chunked_body_when_limit_is_exceeded() {
|
||||
let mut body = Some(Body::from(Bytes::from_static(b"abcdef")));
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
let err = buffer_and_normalize_request_body(
|
||||
&mut body,
|
||||
&mut headers,
|
||||
"test owns body",
|
||||
"trace-body-large",
|
||||
&Method::POST,
|
||||
"/v1/responses",
|
||||
"test",
|
||||
RequestBodyBufferPolicy::for_tests(5, Duration::from_secs(1)),
|
||||
)
|
||||
.await
|
||||
.expect_err("body exceeding the ingress limit should fail");
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
RequestBodyBufferError::TooLarge { limit_bytes: 5 }
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_body_buffer_times_out_instead_of_waiting_forever() {
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(b"{"));
|
||||
std::future::pending::<()>().await;
|
||||
};
|
||||
let mut body = Some(Body::from_stream(stream));
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
let err = buffer_and_normalize_request_body(
|
||||
&mut body,
|
||||
&mut headers,
|
||||
"test owns body",
|
||||
"trace-body-timeout",
|
||||
&Method::POST,
|
||||
"/v1/responses",
|
||||
"test",
|
||||
RequestBodyBufferPolicy::for_tests(1024, Duration::from_millis(5)),
|
||||
)
|
||||
.await
|
||||
.expect_err("body buffering should time out");
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
RequestBodyBufferError::Timeout { timeout_ms: 5 }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
|
||||
@@ -136,12 +136,7 @@ pub(super) fn announcements_internal_error_response(detail: impl Into<String>) -
|
||||
}
|
||||
|
||||
pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
|
||||
match err {
|
||||
GatewayError::UpstreamUnavailable { message, .. }
|
||||
| GatewayError::ControlUnavailable { message, .. }
|
||||
| GatewayError::Client { message, .. }
|
||||
| GatewayError::Internal(message) => message,
|
||||
}
|
||||
err.into_message()
|
||||
}
|
||||
|
||||
pub(super) fn parse_optional_rfc3339_unix_secs(
|
||||
|
||||
Reference in New Issue
Block a user