i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #513 from mayrainnn/fix/request-body-content-encoding

Browse Source 
feat: normalize compressed request bodies
This commit is contained in:
fawney19
2026-05-20 01:05:59 +08:00
committed by GitHub
parent b1a3a26815 8f4f4d2d82
commit d1ad1815f7
12 changed files with 729 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/Cargo.toml
View File

@@ -67,6 +67,7 @@ uuid.workspace = true
webpki-roots.workspace = true
wreq.workspace = true
wreq-util.workspace = true
zstd.workspace = true
[target.'cfg(not(target_env = "msvc"))'.dependencies]
tikv-jemallocator = "0.6"

40
apps/aether-gateway/src/ai_serving/api.rs
View File

@@ -70,10 +70,13 @@ pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,
body_bytes: &axum::body::Bytes,
) -> Option<(serde_json::Value, Option<String>)> {
aether_ai_formats::api::parse_direct_request_body(
is_json_request(&parts.headers),
body_bytes.as_ref(),
)
let is_json_request = is_json_request(&parts.headers);
let body_bytes = if is_json_request {
crate::ai_serving::decoded_request_body_bytes(&parts.headers, body_bytes.as_ref()).ok()?
} else {
std::borrow::Cow::Borrowed(body_bytes.as_ref())
};
aether_ai_formats::api::parse_direct_request_body(is_json_request, body_bytes.as_ref())
}
pub(crate) fn resolve_execution_runtime_stream_plan_kind(
@@ -215,3 +218,32 @@ fn value_has_non_empty_text(value: Option<&serde_json::Value>) -> bool {
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::parse_direct_request_body;
use axum::body::Bytes;
use axum::http::{header, Request};
#[test]
fn parse_direct_request_body_reads_zstd_encoded_json_body() {
let (parts, _) = Request::builder()
.method("POST")
.uri("/v1/responses")
.header(header::CONTENT_TYPE, "application/json")
.header(header::CONTENT_ENCODING, "zstd")
.body(())
.expect("request should build")
.into_parts();
let encoded =
zstd::stream::encode_all(br#"{"model":"gpt-5.4","stream":true}"#.as_slice(), 0)
.expect("zstd body should encode");
let (body_json, body_base64) =
parse_direct_request_body(&parts, &Bytes::from(encoded)).expect("body should parse");
assert_eq!(body_json["model"].as_str(), Some("gpt-5.4"));
assert_eq!(body_json["stream"].as_bool(), Some(true));
assert!(body_base64.is_none());
}
}

7
apps/aether-gateway/src/ai_serving/mod.rs
View File

@@ -150,6 +150,13 @@ pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
crate::headers::is_json_request(headers)
}
pub(crate) fn decoded_request_body_bytes<'a>(
headers: &http::HeaderMap,
body_bytes: &'a [u8],
) -> Result<std::borrow::Cow<'a, [u8]>, crate::headers::RequestBodyNormalizationError> {
crate::headers::decoded_request_body_bytes(headers, body_bytes)
}
pub(crate) fn tls_fingerprint_from_headers(headers: &http::HeaderMap) -> Option<serde_json::Value> {
crate::headers::tls_fingerprint_from_headers(headers)
}

8
apps/aether-gateway/src/ai_serving/planner/common.rs
View File

@@ -31,7 +31,13 @@ pub(crate) fn parse_direct_request_body(
parts: &http::request::Parts,
body_bytes: &Bytes,
) -> Option<(serde_json::Value, Option<String>)> {
parse_direct_request_body_impl(is_json_request(&parts.headers), body_bytes.as_ref())
let is_json_request = is_json_request(&parts.headers);
let body_bytes = if is_json_request {
crate::ai_serving::decoded_request_body_bytes(&parts.headers, body_bytes.as_ref()).ok()?
} else {
std::borrow::Cow::Borrowed(body_bytes.as_ref())
};
parse_direct_request_body_impl(is_json_request, body_bytes.as_ref())
}
pub(crate) fn force_upstream_streaming_for_provider(

1
apps/aether-gateway/src/constants.rs
View File

@@ -32,6 +32,7 @@ pub(crate) const API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS: u64 = 150;
pub(crate) const API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS: u64 = 10;
pub(crate) const EXECUTION_PATH_LOCAL_AUTH_DENIED: &str = "local_auth_denied";
pub(crate) const EXECUTION_PATH_LOCAL_RATE_LIMITED: &str = "local_rate_limited";
pub(crate) const EXECUTION_PATH_LOCAL_INVALID_REQUEST: &str = "local_invalid_request";
pub(crate) const EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND: &str = "local_route_not_found";
pub(crate) const EXECUTION_PATH_LOCAL_OVERLOADED: &str = "local_overloaded";
pub(crate) const EXECUTION_PATH_DISTRIBUTED_OVERLOADED: &str = "distributed_overloaded";

44
apps/aether-gateway/src/control/auth/credentials.rs
View File

@@ -7,7 +7,7 @@ use url::form_urlencoded;
use crate::{
ai_serving::extract_gemini_model_from_path,
headers::{header_value_str, is_json_request},
headers::{decoded_request_body_bytes, header_value_str, is_json_request},
};
use super::super::GatewayControlDecision;
@@ -31,7 +31,8 @@ pub(crate) fn extract_requested_model(
if !is_json_request(headers) || body.is_empty() {
return None;
}
serde_json::from_slice::<serde_json::Value>(body)
let body = decoded_request_body_bytes(headers, body.as_ref()).ok()?;
serde_json::from_slice::<serde_json::Value>(body.as_ref())
.ok()
.and_then(|payload| {
payload
@@ -356,15 +357,50 @@ pub(super) fn current_unix_secs() -> u64 {
#[cfg(test)]
mod tests {
use super::{
build_auth_context_cache_key, extract_request_credentials, GatewayCredentialCarrier,
GatewayPrimaryCredential, GatewayTrustedAdminHeaders, GatewayTrustedAuthHeaders,
build_auth_context_cache_key, extract_request_credentials, extract_requested_model,
GatewayCredentialCarrier, GatewayPrimaryCredential, GatewayTrustedAdminHeaders,
GatewayTrustedAuthHeaders,
};
use crate::control::GatewayControlDecision;
use axum::body::Bytes;
use axum::http::{self, Uri};
fn uri(path: &str) -> Uri {
path.parse().expect("uri should parse")
}
#[test]
fn extract_requested_model_reads_zstd_encoded_json_body() {
let decision = GatewayControlDecision::synthetic(
"/v1/responses",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
);
let mut headers = http::HeaderMap::new();
headers.insert(
http::header::CONTENT_TYPE,
http::HeaderValue::from_static("application/json"),
);
headers.insert(
http::header::CONTENT_ENCODING,
http::HeaderValue::from_static("zstd"),
);
let encoded =
zstd::stream::encode_all(br#"{"model":"gpt-5.4","input":"hello"}"#.as_slice(), 0)
.expect("zstd body should encode");
let requested_model = extract_requested_model(
&decision,
&uri("/v1/responses"),
&headers,
&Bytes::from(encoded),
);
assert_eq!(requested_model.as_deref(), Some("gpt-5.4"));
}
#[test]
fn selects_openai_bearer_as_provider_api_key() {
let mut headers = http::HeaderMap::new();

12
apps/aether-gateway/src/control/auth/gate.rs
View File

@@ -94,6 +94,7 @@ pub(crate) async fn request_model_local_rejection(
decision,
auth_context,
requested_model.as_deref(),
headers,
body,
)
.await
@@ -104,6 +105,7 @@ async fn balance_capacity_rejection(
decision: &GatewayControlDecision,
auth_context: &GatewayControlAuthContext,
requested_model: Option<&str>,
headers: &http::HeaderMap,
body: &Bytes,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
if auth_context.api_key_is_standalone {
@@ -146,7 +148,8 @@ async fn balance_capacity_rejection(
return Ok(None);
};
let Some(estimated_cost_usd) =
estimate_request_cost_upper_bound_usd(state, decision, requested_model, body).await?
estimate_request_cost_upper_bound_usd(state, decision, requested_model, headers, body)
.await?
else {
return Ok(None);
};
@@ -173,6 +176,7 @@ async fn estimate_request_cost_upper_bound_usd(
state: &AppState,
decision: &GatewayControlDecision,
requested_model: &str,
headers: &http::HeaderMap,
body: &Bytes,
) -> Result<Option<f64>, GatewayError> {
let Some(api_format) = decision
@@ -183,7 +187,11 @@ async fn estimate_request_cost_upper_bound_usd(
else {
return Ok(None);
};
let body_json = serde_json::from_slice::<serde_json::Value>(body).ok();
let body = crate::headers::decoded_request_body_bytes(headers, body.as_ref()).ok();
let Some(body) = body else {
return Ok(None);
};
let body_json = serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok();
let Some(input_tokens) = body_json
.as_ref()
.map(estimate_json_tokens)

41
apps/aether-gateway/src/handlers/proxy/finalize.rs
View File

@@ -17,6 +17,7 @@ use tracing::{info, trace, warn};
pub(super) fn request_wants_stream(
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
body: &axum::body::Bytes,
) -> bool {
if request_context
@@ -34,7 +35,11 @@ pub(super) fn request_wants_stream(
{
return false;
}
serde_json::from_slice::<serde_json::Value>(body)
let body = match crate::headers::decoded_request_body_bytes(headers, body.as_ref()) {
Ok(body) => body,
Err(_) => return false,
};
serde_json::from_slice::<serde_json::Value>(body.as_ref())
.ok()
.and_then(|value| value.get("stream").and_then(|stream| stream.as_bool()))
.unwrap_or(false)
@@ -258,11 +263,11 @@ pub(super) fn finalize_gateway_response_with_context(
#[cfg(test)]
mod tests {
use super::finalize_gateway_response;
use crate::control::GatewayControlDecision;
use super::{finalize_gateway_response, request_wants_stream};
use crate::control::{GatewayControlDecision, GatewayPublicRequestContext};
use crate::AppState;
use axum::body::Body;
use axum::http::{Method, Response, StatusCode};
use axum::body::{Body, Bytes};
use axum::http::{HeaderMap, HeaderValue, Method, Response, StatusCode};
use std::sync::{Arc, Mutex};
use std::time::Instant;
use tracing_subscriber::filter::LevelFilter;
@@ -306,6 +311,32 @@ mod tests {
}
}
#[test]
fn request_wants_stream_reads_zstd_encoded_json_body() {
let request_context = GatewayPublicRequestContext {
trace_id: "trace-zstd-stream".to_string(),
request_method: Method::POST,
request_path: "/v1/responses".to_string(),
request_query_string: None,
request_content_type: Some("application/json".to_string()),
host_header: None,
control_decision: None,
};
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("zstd"),
);
let encoded = zstd::stream::encode_all(br#"{"stream":true}"#.as_slice(), 0)
.expect("zstd body should encode");
assert!(request_wants_stream(
&request_context,
&headers,
&Bytes::from(encoded),
));
}
#[test]
fn finalize_gateway_response_logs_sanitized_path_and_query() {
let state = AppState::new().expect("gateway state should build");

147
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -21,12 +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_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_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,
};
@@ -51,7 +52,7 @@ use crate::handlers::shared::{
};
use crate::headers::{
extract_or_generate_trace_id, request_origin_from_headers_and_remote_addr,
should_skip_request_header,
should_skip_request_header, RequestBodyNormalizationError,
};
use crate::router::RequestAdmissionError;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
@@ -92,6 +93,69 @@ 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(
trace_id: &str,
request_context: &GatewayPublicRequestContext,
error: &RequestBodyNormalizationError,
) -> Result<Response<Body>, GatewayError> {
warn!(
event_name = "frontdoor_request_body_normalization_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"
);
build_local_http_error_response(
trace_id,
request_context.control_decision.as_ref(),
error.http_status(),
error.client_message().as_str(),
)
}
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(
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,
) -> Result<Response<Body>, GatewayError> {
let response =
build_request_body_normalization_error_response(trace_id, request_context, error)?;
Ok(finalize_gateway_response_with_context(
state,
response,
remote_addr,
request_context,
EXECUTION_PATH_LOCAL_INVALID_REQUEST,
started_at,
request_permit,
))
}
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
match outcome {
LocalExecutionRequestOutcome::Responded(_) => "responded",
@@ -330,8 +394,11 @@ async fn maybe_forward_public_request_to_tunnel_owner(
) else {
return Ok(None);
};
let body_json =
buffered_body.and_then(|body| serde_json::from_slice::<serde_json::Value>(body).ok());
let body_json = buffered_body.and_then(|body| {
let body =
crate::headers::decoded_request_body_bytes(&parts.headers, body.as_ref()).ok()?;
serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok()
});
let client_session_affinity =
crate::client_session_affinity::client_session_affinity_from_parts(
parts,
@@ -461,6 +528,7 @@ async fn maybe_forward_public_request_to_tunnel_owner(
let mut response = build_sync_aware_affinity_forward_response(
request_context,
&parts.headers,
buffered_body,
decision,
upstream_response,
@@ -720,6 +788,7 @@ fn build_stream_sse_proxy_response(
async fn build_sync_aware_affinity_forward_response(
request_context: &GatewayPublicRequestContext,
request_headers: &http::HeaderMap,
buffered_body: Option<&Bytes>,
decision: &GatewayControlDecision,
upstream_response: reqwest::Response,
@@ -727,7 +796,7 @@ async fn build_sync_aware_affinity_forward_response(
let Some(buffered_body) = buffered_body else {
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
};
let stream_request = request_wants_stream(request_context, buffered_body);
let stream_request = request_wants_stream(request_context, request_headers, buffered_body);
let upstream_is_sse = upstream_response_is_sse(upstream_response.headers());
if (!stream_request && !upstream_is_sse) || (stream_request && upstream_is_sse) {
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
@@ -943,16 +1012,26 @@ 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) {
Some(
to_bytes(
request_body
.take()
.expect("local proxy body buffering should own request body"),
usize::MAX,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?,
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"local proxy body buffering should own request body",
)
.await?;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
&err,
);
}
}
} else {
None
};
@@ -1113,16 +1192,26 @@ pub(crate) async fn proxy_request(
&& request_enables_control_execute(&parts.headers);
let buffered_body = if should_buffer_body {
Some(
to_bytes(
request_body
.take()
.expect("buffered auth/execution runtime path should own request body"),
usize::MAX,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?,
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"buffered auth/execution runtime path should own request body",
)
.await?;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
&err,
);
}
}
} else {
None
};
@@ -1264,7 +1353,7 @@ pub(crate) async fn proxy_request(
let buffered_body = buffered_body
.as_ref()
.expect("execution runtime/control auth gate should have buffered request body");
let stream_request = request_wants_stream(&request_context, buffered_body);
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(

474
apps/aether-gateway/src/headers.rs
View File

@@ -1,9 +1,26 @@
use std::{collections::BTreeMap, net::SocketAddr};
use std::{borrow::Cow, collections::BTreeMap, fmt, io::Read, net::SocketAddr, sync::LazyLock};
use crate::constants::*;
use axum::body::Bytes;
use flate2::read::{DeflateDecoder, GzDecoder, ZlibDecoder};
use serde_json::{Map, Value};
use uuid::Uuid;
const DEFAULT_MAX_REQUEST_BODY_MB: u64 = 64;
const MAX_REQUEST_BODY_MB_ENV: &str = "AETHER_MAX_REQUEST_BODY_MB";
/// Upper bound applied to a request body after Content-Encoding decoding, and to
/// uncompressed bodies as-is. Guards against decompression bombs and oversized
/// request allocations. Overridable via `AETHER_MAX_REQUEST_BODY_MB`.
static MAX_REQUEST_BODY_BYTES: LazyLock<u64> = LazyLock::new(|| {
std::env::var(MAX_REQUEST_BODY_MB_ENV)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_MAX_REQUEST_BODY_MB)
.saturating_mul(1024 * 1024)
});
pub(crate) fn extract_or_generate_trace_id(headers: &http::HeaderMap) -> String {
header_value_str(headers, TRACE_ID_HEADER).unwrap_or_else(|| Uuid::new_v4().to_string())
}
@@ -157,6 +174,216 @@ pub(crate) fn is_json_request(headers: &http::HeaderMap) -> bool {
.unwrap_or(false)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum RequestBodyNormalizationError {
UnsupportedContentEncoding(String),
DecodeFailed { encoding: String, reason: String },
DecompressedBodyTooLarge { encoding: String, limit_bytes: u64 },
RequestBodyTooLarge { limit_bytes: u64 },
}
impl RequestBodyNormalizationError {
pub(crate) fn client_message(&self) -> String {
match self {
Self::UnsupportedContentEncoding(encoding) => {
format!("Unsupported request Content-Encoding: {encoding}")
}
Self::DecodeFailed { encoding, .. } => {
format!("Failed to decode request body with Content-Encoding: {encoding}")
}
Self::DecompressedBodyTooLarge {
encoding,
limit_bytes,
} => format!(
"Decoded request body with Content-Encoding {encoding} exceeds {limit_bytes} bytes"
),
Self::RequestBodyTooLarge { limit_bytes } => {
format!("Request body exceeds {limit_bytes} bytes")
}
}
}
pub(crate) fn http_status(&self) -> http::StatusCode {
match self {
Self::DecompressedBodyTooLarge { .. } | Self::RequestBodyTooLarge { .. } => {
http::StatusCode::PAYLOAD_TOO_LARGE
}
Self::UnsupportedContentEncoding(_) | Self::DecodeFailed { .. } => {
http::StatusCode::BAD_REQUEST
}
}
}
}
impl fmt::Display for RequestBodyNormalizationError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::UnsupportedContentEncoding(encoding) => {
write!(f, "unsupported request Content-Encoding: {encoding}")
}
Self::DecodeFailed { encoding, reason } => {
write!(
f,
"failed to decode request body with Content-Encoding {encoding}: {reason}"
)
}
Self::DecompressedBodyTooLarge {
encoding,
limit_bytes,
} => write!(
f,
"decoded request body with Content-Encoding {encoding} exceeds {limit_bytes} bytes"
),
Self::RequestBodyTooLarge { limit_bytes } => {
write!(f, "request body exceeds {limit_bytes} bytes")
}
}
}
}
impl std::error::Error for RequestBodyNormalizationError {}
pub(crate) fn normalize_request_body_headers_and_bytes(
headers: &mut http::HeaderMap,
body_bytes: Bytes,
) -> Result<Bytes, RequestBodyNormalizationError> {
let body_was_encoded = !request_content_encodings(headers).is_empty();
let decoded = decoded_request_body_bytes(headers, body_bytes.as_ref())?;
if !body_was_encoded {
return Ok(body_bytes);
}
headers.remove(http::header::CONTENT_ENCODING);
headers.remove(http::header::CONTENT_LENGTH);
Ok(Bytes::from(decoded.into_owned()))
}
/// Rejects a request whose declared `Content-Length` already exceeds the body
/// limit, before the body is buffered into memory. Chunked or length-less
/// requests pass this check and stay bounded by the post-decode guard instead.
pub(crate) fn check_request_content_length(
headers: &http::HeaderMap,
) -> Result<(), RequestBodyNormalizationError> {
let limit = *MAX_REQUEST_BODY_BYTES;
let declared = header_value_str(headers, http::header::CONTENT_LENGTH.as_str())
.and_then(|value| value.trim().parse::<u64>().ok());
if declared.is_some_and(|value| value > limit) {
return Err(RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit });
}
Ok(())
}
pub(crate) fn decoded_request_body_bytes<'a>(
headers: &http::HeaderMap,
body_bytes: &'a [u8],
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
let encodings = request_content_encodings(headers);
if encodings.is_empty() {
let limit = *MAX_REQUEST_BODY_BYTES;
if body_bytes.len() as u64 > limit {
return Err(RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit });
}
return Ok(Cow::Borrowed(body_bytes));
}
let mut decoded = body_bytes.to_vec();
for encoding in encodings.iter().rev() {
decoded = decode_single_request_body(encoding, decoded.as_slice())?;
}
Ok(Cow::Owned(decoded))
}
fn request_content_encodings(headers: &http::HeaderMap) -> Vec<String> {
header_value_str(headers, http::header::CONTENT_ENCODING.as_str())
.map(|value| {
value
.split(',')
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_ascii_lowercase)
.filter(|value| value != "identity")
.collect()
})
.unwrap_or_default()
}
fn decode_single_request_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
match encoding {
"gzip" | "x-gzip" => decode_gzip_body(encoding, body_bytes),
"deflate" => decode_deflate_body(encoding, body_bytes),
"zstd" => decode_zstd_body(encoding, body_bytes),
_ => Err(RequestBodyNormalizationError::UnsupportedContentEncoding(
encoding.to_string(),
)),
}
}
fn decode_gzip_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut decoder = GzDecoder::new(body_bytes);
read_request_decoder_to_end(encoding, &mut decoder)
}
fn decode_deflate_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut zlib_decoder = ZlibDecoder::new(body_bytes);
match read_request_decoder_to_end(encoding, &mut zlib_decoder) {
Ok(decoded) => Ok(decoded),
Err(err @ RequestBodyNormalizationError::DecompressedBodyTooLarge { .. }) => Err(err),
Err(zlib_error) => {
let mut raw_decoder = DeflateDecoder::new(body_bytes);
read_request_decoder_to_end(encoding, &mut raw_decoder).map_err(|raw_error| {
RequestBodyNormalizationError::DecodeFailed {
encoding: encoding.to_string(),
reason: format!("{zlib_error}; raw deflate fallback failed: {raw_error}"),
}
})
}
}
}
fn decode_zstd_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut decoder = zstd::stream::read::Decoder::new(body_bytes).map_err(|err| {
RequestBodyNormalizationError::DecodeFailed {
encoding: encoding.to_string(),
reason: err.to_string(),
}
})?;
read_request_decoder_to_end(encoding, &mut decoder)
}
fn read_request_decoder_to_end(
encoding: &str,
decoder: &mut impl Read,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let limit = *MAX_REQUEST_BODY_BYTES;
let mut limited = decoder.take(limit.saturating_add(1));
let mut out = Vec::new();
limited
.read_to_end(&mut out)
.map_err(|err| RequestBodyNormalizationError::DecodeFailed {
encoding: encoding.to_string(),
reason: err.to_string(),
})?;
if out.len() as u64 > limit {
return Err(RequestBodyNormalizationError::DecompressedBodyTooLarge {
encoding: encoding.to_string(),
limit_bytes: limit,
});
}
Ok(out)
}
pub(crate) fn header_equals(
headers: &reqwest::header::HeaderMap,
key: &'static str,
@@ -172,12 +399,20 @@ pub(crate) fn header_equals(
#[cfg(test)]
mod tests {
use super::{
decoded_request_body_bytes, normalize_request_body_headers_and_bytes,
request_origin_from_headers, request_origin_from_headers_and_remote_addr,
tls_fingerprint_from_headers, RequestOrigin,
tls_fingerprint_from_headers, RequestBodyNormalizationError, RequestOrigin,
};
use flate2::{
write::{DeflateEncoder, GzEncoder, ZlibEncoder},
Compression,
};
use http::{HeaderMap, HeaderValue};
use serde_json::json;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::{
io::Write,
net::{IpAddr, Ipv4Addr, SocketAddr},
};
#[test]
fn request_origin_prefers_first_forwarded_for_ip() {
@@ -201,6 +436,239 @@ mod tests {
);
}
#[test]
fn decoded_request_body_bytes_decodes_zstd() {
let payload = br#"{"model":"gpt-5.4"}"#;
let encoded =
zstd::stream::encode_all(payload.as_slice(), 0).expect("zstd body should encode");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("zstd"),
);
let decoded =
decoded_request_body_bytes(&headers, encoded.as_slice()).expect("body should decode");
assert_eq!(decoded.as_ref(), payload);
}
#[test]
fn decoded_request_body_bytes_decodes_x_gzip() {
let payload = br#"{"model":"gpt-5.4"}"#;
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(payload).expect("gzip body should write");
let encoded = encoder.finish().expect("gzip body should finish");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("x-gzip"),
);
let decoded =
decoded_request_body_bytes(&headers, encoded.as_slice()).expect("body should decode");
assert_eq!(decoded.as_ref(), payload);
}
#[test]
fn decoded_request_body_bytes_decodes_zlib_wrapped_deflate() {
let payload = br#"{"model":"gpt-5.4"}"#;
let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
encoder
.write_all(payload)
.expect("deflate body should write");
let encoded = encoder.finish().expect("deflate body should finish");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("deflate"),
);
let decoded =
decoded_request_body_bytes(&headers, encoded.as_slice()).expect("body should decode");
assert_eq!(decoded.as_ref(), payload);
}
#[test]
fn decoded_request_body_bytes_decodes_raw_deflate_fallback() {
let payload = br#"{"model":"gpt-5.4"}"#;
let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default());
encoder
.write_all(payload)
.expect("deflate body should write");
let encoded = encoder.finish().expect("deflate body should finish");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("deflate"),
);
let decoded =
decoded_request_body_bytes(&headers, encoded.as_slice()).expect("body should decode");
assert_eq!(decoded.as_ref(), payload);
}
#[test]
fn decoded_request_body_bytes_decodes_multiple_chained_encodings() {
let payload = br#"{"model":"gpt-5.4"}"#;
let mut gzip_encoder = GzEncoder::new(Vec::new(), Compression::default());
gzip_encoder
.write_all(payload)
.expect("gzip body should write");
let gzipped = gzip_encoder.finish().expect("gzip body should finish");
let encoded =
zstd::stream::encode_all(gzipped.as_slice(), 0).expect("zstd body should encode");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("gzip, zstd"),
);
let decoded =
decoded_request_body_bytes(&headers, encoded.as_slice()).expect("body should decode");
assert_eq!(decoded.as_ref(), payload);
}
#[test]
fn decoded_request_body_bytes_rejects_corrupt_encoded_body() {
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("zstd"),
);
let err = decoded_request_body_bytes(&headers, br#"{"model":"gpt-5.4"}"#.as_slice())
.expect_err("corrupt body should fail");
assert!(matches!(
err,
RequestBodyNormalizationError::DecodeFailed { .. }
));
}
#[test]
fn normalize_request_body_headers_and_bytes_clears_encoding_headers() {
let payload = br#"{"model":"gpt-5.4"}"#;
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(payload).expect("gzip body should write");
let encoded = encoder.finish().expect("gzip body should finish");
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("x-gzip"),
);
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_static("999"),
);
let decoded = normalize_request_body_headers_and_bytes(
&mut headers,
axum::body::Bytes::from(encoded),
)
.expect("body should normalize");
assert_eq!(decoded.as_ref(), payload);
assert!(!headers.contains_key(http::header::CONTENT_ENCODING));
assert!(!headers.contains_key(http::header::CONTENT_LENGTH));
}
#[test]
fn decoded_request_body_bytes_rejects_unsupported_encoding() {
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_ENCODING,
HeaderValue::from_static("br"),
);
let err = decoded_request_body_bytes(&headers, br#"{"model":"gpt-5.4"}"#.as_slice())
.expect_err("unsupported encoding should fail");
assert_eq!(
err,
RequestBodyNormalizationError::UnsupportedContentEncoding("br".to_string())
);
}
#[test]
fn decoded_request_body_bytes_rejects_oversized_uncompressed_body() {
let limit = *super::MAX_REQUEST_BODY_BYTES;
let oversized = vec![b'a'; limit as usize + 1];
let headers = HeaderMap::new();
let err = decoded_request_body_bytes(&headers, oversized.as_slice())
.expect_err("oversized uncompressed body should fail");
assert_eq!(
err,
RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit }
);
}
#[test]
fn check_request_content_length_rejects_oversized_declared_length() {
let limit = *super::MAX_REQUEST_BODY_BYTES;
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_str(&(limit + 1).to_string()).expect("length header should build"),
);
let err = super::check_request_content_length(&headers)
.expect_err("oversized declared length should fail");
assert_eq!(
err,
RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit }
);
}
#[test]
fn request_body_normalization_error_maps_http_status() {
assert_eq!(
RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: 1 }.http_status(),
http::StatusCode::PAYLOAD_TOO_LARGE
);
assert_eq!(
RequestBodyNormalizationError::DecompressedBodyTooLarge {
encoding: "zstd".to_string(),
limit_bytes: 1,
}
.http_status(),
http::StatusCode::PAYLOAD_TOO_LARGE
);
assert_eq!(
RequestBodyNormalizationError::UnsupportedContentEncoding("br".to_string())
.http_status(),
http::StatusCode::BAD_REQUEST
);
assert_eq!(
RequestBodyNormalizationError::DecodeFailed {
encoding: "gzip".to_string(),
reason: "bad".to_string(),
}
.http_status(),
http::StatusCode::BAD_REQUEST
);
}
#[test]
fn check_request_content_length_allows_missing_or_within_limit() {
let empty = HeaderMap::new();
assert!(super::check_request_content_length(&empty).is_ok());
let mut headers = HeaderMap::new();
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_static("1024"),
);
assert!(super::check_request_content_length(&headers).is_ok());
}
#[test]
fn request_origin_uses_real_ip_after_empty_forwarded_for_segments() {
let mut headers = HeaderMap::new();