mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 00:17:45 +08:00
feat: optimize usage body viewing and provider card layout
This commit is contained in:
@@ -1593,6 +1593,19 @@ impl GatewayDataState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_body_payload(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::usage::StoredUsageBodyPayload>,
|
||||
DataLayerError,
|
||||
> {
|
||||
match &self.usage_reader {
|
||||
Some(repository) => repository.read_body_payload(body_ref).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
|
||||
@@ -17,7 +17,10 @@ use aether_admin::observability::usage::{
|
||||
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
|
||||
admin_usage_provider_key_name, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageBodyField};
|
||||
use aether_data_contracts::repository::usage::{
|
||||
canonical_usage_body_ref_for, StoredRequestUsageAudit, StoredUsageBodyPayload,
|
||||
UsageBodyCaptureState, UsageBodyField, MAX_DECOMPRESSED_USAGE_JSON_BYTES,
|
||||
};
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -28,9 +31,59 @@ use serde_json::{json, Value};
|
||||
use std::collections::BTreeMap;
|
||||
use tokio::try_join;
|
||||
|
||||
#[derive(Default)]
|
||||
struct AdminUsageDetailBodyValue {
|
||||
value: Option<Value>,
|
||||
load_failed: bool,
|
||||
error_code: Option<&'static str>,
|
||||
}
|
||||
|
||||
impl AdminUsageDetailBodyValue {
|
||||
fn resolved(
|
||||
item: &StoredRequestUsageAudit,
|
||||
field: UsageBodyField,
|
||||
value: Option<Value>,
|
||||
) -> Self {
|
||||
let missing = value.is_none()
|
||||
&& item
|
||||
.body_capture_result(field, item.body_value(field))
|
||||
.available;
|
||||
Self {
|
||||
value,
|
||||
error_code: missing.then_some("missing"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_usage_body_load_error_code(error: &GatewayError) -> &'static str {
|
||||
if let GatewayError::Internal(message) = error {
|
||||
if message.contains("decompressed usage json exceeds ")
|
||||
|| message.contains("encoded usage json exceeds ")
|
||||
{
|
||||
return "too_large";
|
||||
}
|
||||
if message.contains("failed to decompress usage json:")
|
||||
|| message.contains("failed to parse decompressed usage json:")
|
||||
{
|
||||
return "decode_failed";
|
||||
}
|
||||
}
|
||||
"storage_unavailable"
|
||||
}
|
||||
|
||||
async fn resolve_admin_usage_detail_field(
|
||||
state: &AdminAppState<'_>,
|
||||
item: &StoredRequestUsageAudit,
|
||||
field: UsageBodyField,
|
||||
selected_field: Option<UsageBodyField>,
|
||||
) -> AdminUsageDetailBodyValue {
|
||||
if selected_field.is_some_and(|selected| selected != field) {
|
||||
return AdminUsageDetailBodyValue::default();
|
||||
}
|
||||
if field == UsageBodyField::RequestBody {
|
||||
resolve_admin_usage_detail_request_body(state, item).await
|
||||
} else {
|
||||
resolve_admin_usage_detail_body_value(state, item, field).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn resolve_admin_usage_detail_request_body(
|
||||
@@ -38,10 +91,7 @@ async fn resolve_admin_usage_detail_request_body(
|
||||
item: &StoredRequestUsageAudit,
|
||||
) -> AdminUsageDetailBodyValue {
|
||||
match admin_usage_resolve_request_capture_body_for_item(state, item, None).await {
|
||||
Ok(body) => AdminUsageDetailBodyValue {
|
||||
value: body,
|
||||
load_failed: false,
|
||||
},
|
||||
Ok(body) => AdminUsageDetailBodyValue::resolved(item, UsageBodyField::RequestBody, body),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
error = ?err,
|
||||
@@ -52,7 +102,9 @@ async fn resolve_admin_usage_detail_request_body(
|
||||
);
|
||||
let value = admin_usage_resolve_request_capture_body(item, None);
|
||||
AdminUsageDetailBodyValue {
|
||||
load_failed: value.is_none(),
|
||||
error_code: value
|
||||
.is_none()
|
||||
.then(|| admin_usage_body_load_error_code(&err)),
|
||||
value,
|
||||
}
|
||||
}
|
||||
@@ -66,10 +118,7 @@ async fn resolve_admin_usage_detail_body_value(
|
||||
) -> AdminUsageDetailBodyValue {
|
||||
let inline_body = item.body_value(field);
|
||||
match admin_usage_resolve_body_value(state, item, inline_body, field).await {
|
||||
Ok(body) => AdminUsageDetailBodyValue {
|
||||
value: body,
|
||||
load_failed: false,
|
||||
},
|
||||
Ok(body) => AdminUsageDetailBodyValue::resolved(item, field, body),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
error = ?err,
|
||||
@@ -80,13 +129,139 @@ async fn resolve_admin_usage_detail_body_value(
|
||||
);
|
||||
let value = inline_body.cloned();
|
||||
AdminUsageDetailBodyValue {
|
||||
load_failed: value.is_none(),
|
||||
error_code: value
|
||||
.is_none()
|
||||
.then(|| admin_usage_body_load_error_code(&err)),
|
||||
value,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_admin_usage_raw_body(
|
||||
state: &AdminAppState<'_>,
|
||||
item: &StoredRequestUsageAudit,
|
||||
field: UsageBodyField,
|
||||
) -> Result<Option<StoredUsageBodyPayload>, GatewayError> {
|
||||
if matches!(
|
||||
item.body_state(field),
|
||||
Some(
|
||||
UsageBodyCaptureState::Disabled
|
||||
| UsageBodyCaptureState::Unavailable
|
||||
| UsageBodyCaptureState::None
|
||||
)
|
||||
) {
|
||||
return Ok(None);
|
||||
}
|
||||
let inline_body = item.body_value(field);
|
||||
let prefer_inline = matches!(
|
||||
item.body_state(field),
|
||||
Some(UsageBodyCaptureState::Inline | UsageBodyCaptureState::Truncated)
|
||||
) && inline_body.is_some();
|
||||
if !prefer_inline {
|
||||
if let Some(body_ref) = item
|
||||
.body_ref(field)
|
||||
.and_then(|reference| canonical_usage_body_ref_for(reference, &item.request_id, field))
|
||||
{
|
||||
if let Some(payload) = state.read_request_usage_body_payload(&body_ref).await? {
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
}
|
||||
}
|
||||
let fallback = inline_body.cloned().or_else(|| {
|
||||
(field == UsageBodyField::RequestBody)
|
||||
.then(|| admin_usage_resolve_request_capture_body(item, None))
|
||||
.flatten()
|
||||
});
|
||||
fallback
|
||||
.map(|value| {
|
||||
serde_json::to_vec(&value)
|
||||
.map(StoredUsageBodyPayload::Json)
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn build_admin_usage_raw_body_response(
|
||||
state: &AdminAppState<'_>,
|
||||
item: &StoredRequestUsageAudit,
|
||||
field: UsageBodyField,
|
||||
) -> Response<Body> {
|
||||
let result = read_admin_usage_raw_body(state, item, field).await;
|
||||
let mut response = match result {
|
||||
Ok(Some(payload)) => admin_usage_raw_payload_response(payload),
|
||||
Ok(None) => admin_usage_raw_body_error(http::StatusCode::NOT_FOUND, "missing"),
|
||||
Err(error) => {
|
||||
tracing::warn!(error = ?error, usage_id = %item.id, field = field.as_storage_field(), "failed to read admin usage raw body");
|
||||
let code = admin_usage_body_load_error_code(&error);
|
||||
admin_usage_raw_body_error(
|
||||
if code == "too_large" {
|
||||
http::StatusCode::PAYLOAD_TOO_LARGE
|
||||
} else {
|
||||
http::StatusCode::SERVICE_UNAVAILABLE
|
||||
},
|
||||
code,
|
||||
)
|
||||
}
|
||||
};
|
||||
let headers = response.headers_mut();
|
||||
headers.insert(
|
||||
http::header::CACHE_CONTROL,
|
||||
http::HeaderValue::from_static("no-store, no-transform"),
|
||||
);
|
||||
headers.insert(
|
||||
"x-content-type-options",
|
||||
http::HeaderValue::from_static("nosniff"),
|
||||
);
|
||||
headers.insert(
|
||||
"x-aether-body-field",
|
||||
http::HeaderValue::from_static(field.as_storage_field()),
|
||||
);
|
||||
if let Ok(value) = http::HeaderValue::from_str(&item.id) {
|
||||
headers.insert("x-aether-usage-id", value);
|
||||
}
|
||||
attach_admin_audit_response(
|
||||
response,
|
||||
"admin_usage_detail_viewed",
|
||||
"view_usage_detail",
|
||||
"usage_record",
|
||||
&item.id,
|
||||
)
|
||||
}
|
||||
|
||||
fn admin_usage_raw_payload_response(payload: StoredUsageBodyPayload) -> Response<Body> {
|
||||
let (encoding, bytes, limit) = match payload {
|
||||
StoredUsageBodyPayload::Gzip(bytes) => (
|
||||
"gzip",
|
||||
bytes,
|
||||
MAX_DECOMPRESSED_USAGE_JSON_BYTES + 1024 * 1024,
|
||||
),
|
||||
StoredUsageBodyPayload::Json(bytes) => ("json", bytes, MAX_DECOMPRESSED_USAGE_JSON_BYTES),
|
||||
};
|
||||
if bytes.len() > limit {
|
||||
admin_usage_raw_body_error(http::StatusCode::PAYLOAD_TOO_LARGE, "too_large")
|
||||
} else {
|
||||
(
|
||||
[
|
||||
("content-type", "application/octet-stream"),
|
||||
("content-encoding", "identity"),
|
||||
("x-aether-body-encoding", encoding),
|
||||
],
|
||||
bytes,
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_usage_raw_body_error(status: http::StatusCode, code: &'static str) -> Response<Body> {
|
||||
(
|
||||
status,
|
||||
[("x-aether-body-error", code)],
|
||||
Json(json!({ "body_load_error_code": code })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
@@ -212,6 +387,25 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
"include_bodies",
|
||||
true,
|
||||
);
|
||||
let body_field =
|
||||
match request_context
|
||||
.request_query_string
|
||||
.as_deref()
|
||||
.and_then(|query| {
|
||||
url::form_urlencoded::parse(query.as_bytes())
|
||||
.find(|(key, _)| key == "body_field")
|
||||
.map(|(_, value)| value.into_owned())
|
||||
}) {
|
||||
Some(value) => {
|
||||
match UsageBodyField::from_storage_field(value.trim()) {
|
||||
Some(field) if include_bodies => Some(field),
|
||||
_ => return Ok(Some(admin_usage_bad_request_response(
|
||||
"body_field 必须是有效的正文字段,且 include_bodies 必须为 true",
|
||||
))),
|
||||
}
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
let Some(item) = state.find_request_usage_by_id(&usage_id).await? else {
|
||||
return Ok(Some(
|
||||
@@ -223,6 +417,27 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
));
|
||||
};
|
||||
|
||||
let body_format = request_context
|
||||
.request_query_string
|
||||
.as_deref()
|
||||
.and_then(|query| {
|
||||
url::form_urlencoded::parse(query.as_bytes())
|
||||
.find(|(key, _)| key == "body_format")
|
||||
.map(|(_, value)| value.into_owned())
|
||||
});
|
||||
if let Some(format) = body_format {
|
||||
if format != "raw" || body_field.is_none() {
|
||||
return Ok(Some(admin_usage_bad_request_response(
|
||||
"body_format=raw 必须指定 body_field",
|
||||
)));
|
||||
}
|
||||
if let Some(field) = body_field {
|
||||
return Ok(Some(
|
||||
build_admin_usage_raw_body_response(state, &item, field).await,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let user_ids = item.user_id.clone().into_iter().collect::<Vec<_>>();
|
||||
let (users_by_id, provider_key_names, api_key_names): (
|
||||
BTreeMap<String, aether_data::repository::users::StoredUserSummary>,
|
||||
@@ -250,23 +465,32 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
}
|
||||
}
|
||||
let mut body_load_errors = serde_json::Map::new();
|
||||
let request_body = if include_bodies {
|
||||
let mut body_load_error_codes = serde_json::Map::new();
|
||||
let mut request_body = if include_bodies {
|
||||
let (request_body, provider_request_body, response_body, client_response_body) = tokio::join!(
|
||||
resolve_admin_usage_detail_request_body(state, &item),
|
||||
resolve_admin_usage_detail_body_value(
|
||||
resolve_admin_usage_detail_field(
|
||||
state,
|
||||
&item,
|
||||
UsageBodyField::RequestBody,
|
||||
body_field
|
||||
),
|
||||
resolve_admin_usage_detail_field(
|
||||
state,
|
||||
&item,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
body_field,
|
||||
),
|
||||
resolve_admin_usage_detail_body_value(
|
||||
resolve_admin_usage_detail_field(
|
||||
state,
|
||||
&item,
|
||||
UsageBodyField::ResponseBody,
|
||||
body_field,
|
||||
),
|
||||
resolve_admin_usage_detail_body_value(
|
||||
resolve_admin_usage_detail_field(
|
||||
state,
|
||||
&item,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
body_field,
|
||||
),
|
||||
);
|
||||
for (field, resolved) in [
|
||||
@@ -275,20 +499,25 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
(UsageBodyField::ResponseBody, &response_body),
|
||||
(UsageBodyField::ClientResponseBody, &client_response_body),
|
||||
] {
|
||||
if resolved.load_failed {
|
||||
if let Some(error_code) = resolved.error_code {
|
||||
body_load_errors.insert(field.as_storage_field().to_string(), json!(true));
|
||||
body_load_error_codes
|
||||
.insert(field.as_storage_field().to_string(), json!(error_code));
|
||||
}
|
||||
}
|
||||
detail_item.provider_request_body = provider_request_body.value;
|
||||
detail_item.response_body = response_body.value;
|
||||
detail_item.client_response_body = client_response_body.value;
|
||||
if body_field.is_none_or(|field| field == UsageBodyField::ProviderRequestBody) {
|
||||
detail_item.provider_request_body = provider_request_body.value;
|
||||
}
|
||||
if body_field.is_none_or(|field| field == UsageBodyField::ResponseBody) {
|
||||
detail_item.response_body = response_body.value;
|
||||
}
|
||||
if body_field.is_none_or(|field| field == UsageBodyField::ClientResponseBody) {
|
||||
detail_item.client_response_body = client_response_body.value;
|
||||
}
|
||||
request_body.value
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if include_bodies {
|
||||
// request_body 已通过 request capture 解析;其余 detached body 在上方并行加载。
|
||||
}
|
||||
let default_headers = admin_usage_curl_headers();
|
||||
let mut payload = build_admin_usage_detail_payload(
|
||||
&detail_item,
|
||||
@@ -297,15 +526,33 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
state.has_auth_user_data_reader(),
|
||||
state.has_auth_api_key_data_reader(),
|
||||
provider_key_name.as_deref(),
|
||||
include_bodies,
|
||||
request_body,
|
||||
include_bodies && body_field.is_none(),
|
||||
if body_field.is_none() {
|
||||
request_body.take()
|
||||
} else {
|
||||
None
|
||||
},
|
||||
&default_headers,
|
||||
);
|
||||
if let Some(field) = body_field {
|
||||
payload[field.as_storage_field()] = match field {
|
||||
UsageBodyField::RequestBody => request_body,
|
||||
UsageBodyField::ProviderRequestBody => detail_item.provider_request_body.take(),
|
||||
UsageBodyField::ResponseBody => detail_item.response_body.take(),
|
||||
UsageBodyField::ClientResponseBody => detail_item.client_response_body.take(),
|
||||
}
|
||||
.unwrap_or(Value::Null);
|
||||
}
|
||||
payload["body_load_errors"] = if include_bodies && !body_load_errors.is_empty() {
|
||||
Value::Object(body_load_errors)
|
||||
} else {
|
||||
Value::Null
|
||||
};
|
||||
payload["body_load_error_codes"] = if body_load_error_codes.is_empty() {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::Object(body_load_error_codes)
|
||||
};
|
||||
|
||||
return Ok(Some(attach_admin_audit_response(
|
||||
Json(payload).into_response(),
|
||||
@@ -320,3 +567,61 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::admin_usage_body_load_error_code;
|
||||
use crate::GatewayError;
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_raw_body_does_not_decode_or_reencode_stored_bytes() {
|
||||
use super::{admin_usage_raw_payload_response, StoredUsageBodyPayload};
|
||||
for (payload, encoding, expected) in [
|
||||
(
|
||||
StoredUsageBodyPayload::Gzip(vec![31, 139, 8, 0, 1]),
|
||||
"gzip",
|
||||
vec![31, 139, 8, 0, 1],
|
||||
),
|
||||
(
|
||||
StoredUsageBodyPayload::Json(b"{ \"untouched\" : true }".to_vec()),
|
||||
"json",
|
||||
b"{ \"untouched\" : true }".to_vec(),
|
||||
),
|
||||
] {
|
||||
let response = admin_usage_raw_payload_response(payload);
|
||||
assert_eq!(response.headers()["content-encoding"], "identity");
|
||||
assert_eq!(response.headers()["x-aether-body-encoding"], encoding);
|
||||
let bytes = axum::body::to_bytes(response.into_body(), 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(bytes.as_ref(), expected.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_load_errors_expose_safe_codes_instead_of_internal_messages() {
|
||||
for (message, expected) in [
|
||||
(
|
||||
"unexpected database value: decompressed usage json exceeds 67108864 bytes",
|
||||
"too_large",
|
||||
),
|
||||
(
|
||||
"failed to decompress usage json: invalid gzip header",
|
||||
"decode_failed",
|
||||
),
|
||||
(
|
||||
"failed to parse decompressed usage json: invalid JSON",
|
||||
"decode_failed",
|
||||
),
|
||||
(
|
||||
"postgres error: private connection details",
|
||||
"storage_unavailable",
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
admin_usage_body_load_error_code(&GatewayError::Internal(message.to_string())),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,6 +234,20 @@ impl<'a> AdminAppState<'a> {
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn read_request_usage_body_payload(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::usage::StoredUsageBodyPayload>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.app
|
||||
.data
|
||||
.read_request_usage_body_payload(body_ref)
|
||||
.await
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn build_api_format_health_monitor_payload(
|
||||
&self,
|
||||
lookback_hours: u64,
|
||||
|
||||
@@ -211,9 +211,19 @@ fn apply_sensitive_route_cache_policy(
|
||||
return;
|
||||
}
|
||||
|
||||
let preserve_no_transform = headers
|
||||
.get_all(http::header::CACHE_CONTROL)
|
||||
.iter()
|
||||
.filter_map(|value| value.to_str().ok())
|
||||
.flat_map(|value| value.split(','))
|
||||
.any(|directive| directive.trim().eq_ignore_ascii_case("no-transform"));
|
||||
headers.insert(
|
||||
http::header::CACHE_CONTROL,
|
||||
HeaderValue::from_static("no-store"),
|
||||
HeaderValue::from_static(if preserve_no_transform {
|
||||
"no-store, no-transform"
|
||||
} else {
|
||||
"no-store"
|
||||
}),
|
||||
);
|
||||
headers.insert(http::header::PRAGMA, HeaderValue::from_static("no-cache"));
|
||||
}
|
||||
@@ -354,6 +364,29 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raw_body_no_transform_survives_sensitive_cache_policy() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.append(
|
||||
http::header::CACHE_CONTROL,
|
||||
HeaderValue::from_static("public, max-age=3600"),
|
||||
);
|
||||
headers.append(
|
||||
http::header::CACHE_CONTROL,
|
||||
HeaderValue::from_static(" No-Transform "),
|
||||
);
|
||||
apply_sensitive_route_cache_policy(
|
||||
&mut headers,
|
||||
"/api/admin/usage/usage-1?body_format=raw",
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
headers[http::header::CACHE_CONTROL],
|
||||
"no-store, no-transform"
|
||||
);
|
||||
assert_eq!(headers[http::header::PRAGMA], "no-cache");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn authenticated_user_data_responses_are_never_cacheable() {
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
@@ -31,7 +31,11 @@ fn apply_frontdoor_cors_headers(
|
||||
);
|
||||
headers.insert(
|
||||
http::header::ACCESS_CONTROL_EXPOSE_HEADERS,
|
||||
HeaderValue::from_static("*"),
|
||||
HeaderValue::from_static(if headers.contains_key("x-aether-body-field") {
|
||||
"*, X-Aether-Body-Encoding, X-Aether-Body-Field, X-Aether-Body-Error, X-Aether-Usage-Id"
|
||||
} else {
|
||||
"*"
|
||||
}),
|
||||
);
|
||||
if let Some(value) = requested_headers {
|
||||
if let Ok(value) = HeaderValue::from_str(value) {
|
||||
@@ -109,3 +113,36 @@ pub(crate) async fn frontdoor_cors_middleware(
|
||||
);
|
||||
response
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn raw_body_headers_are_explicitly_exposed_for_credentialed_requests() {
|
||||
let cors =
|
||||
FrontdoorCorsConfig::new(vec!["https://console.example".to_string()], true).unwrap();
|
||||
let mut headers = http::HeaderMap::new();
|
||||
headers.insert(
|
||||
"x-aether-body-field",
|
||||
HeaderValue::from_static("request_body"),
|
||||
);
|
||||
apply_frontdoor_cors_headers(&mut headers, &cors, "https://console.example", None);
|
||||
let exposed = headers[http::header::ACCESS_CONTROL_EXPOSE_HEADERS]
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_ascii_lowercase();
|
||||
for name in [
|
||||
"x-aether-body-encoding",
|
||||
"x-aether-body-field",
|
||||
"x-aether-body-error",
|
||||
"x-aether-usage-id",
|
||||
] {
|
||||
assert!(exposed.split(',').any(|header| header.trim() == name));
|
||||
}
|
||||
assert_eq!(
|
||||
headers[http::header::ACCESS_CONTROL_ALLOW_CREDENTIALS],
|
||||
"true"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2366,6 +2366,265 @@ async fn gateway_handles_admin_usage_detail_with_ref_backed_bodies() {
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
fn sample_selective_body_usage() -> StoredRequestUsageAudit {
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-selected-body",
|
||||
"req-selected-body",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.request_body = Some(json!({ "marker": "request_body" }));
|
||||
usage.provider_request_body = Some(json!({ "marker": "provider_request_body" }));
|
||||
usage.response_body = Some(json!({ "marker": "response_body" }));
|
||||
usage.client_response_body = Some(json!({ "marker": "client_response_body" }));
|
||||
usage
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_returns_only_the_selected_body() {
|
||||
let fields = [
|
||||
"request_body",
|
||||
"provider_request_body",
|
||||
"response_body",
|
||||
"client_response_body",
|
||||
];
|
||||
for detached in [false, true] {
|
||||
let usage = sample_selective_body_usage();
|
||||
let repository = if detached {
|
||||
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage])
|
||||
} else {
|
||||
InMemoryUsageReadRepository::seed(vec![usage])
|
||||
};
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
repository,
|
||||
)));
|
||||
for selected in fields {
|
||||
let response = local_admin_usage_response(
|
||||
&state, http::Method::GET,
|
||||
&format!("/api/admin/usage/usage-selected-body?include_bodies=true&body_field={selected}"),
|
||||
None,
|
||||
).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should collect");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&bytes).expect("json should parse");
|
||||
for field in fields {
|
||||
assert_eq!(payload[format!("has_{field}")], true);
|
||||
if field == selected {
|
||||
assert_eq!(payload[field]["marker"], selected);
|
||||
} else {
|
||||
assert!(
|
||||
payload[field].is_null(),
|
||||
"unselected {field} must not be returned"
|
||||
);
|
||||
}
|
||||
}
|
||||
assert!(payload["body_load_errors"].is_null());
|
||||
assert!(payload["body_load_error_codes"].is_null());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_validates_body_field_selection() {
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
InMemoryUsageReadRepository::seed(vec![sample_selective_body_usage()]),
|
||||
)));
|
||||
for query in [
|
||||
"body_field=unknown",
|
||||
"body_field=",
|
||||
"include_bodies=false&body_field=request_body",
|
||||
"body_format=raw",
|
||||
"body_format=unknown&body_field=request_body",
|
||||
"body_format=&body_field=request_body",
|
||||
] {
|
||||
let response = local_admin_usage_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
&format!("/api/admin/usage/usage-selected-body?{query}"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
response.status(),
|
||||
StatusCode::BAD_REQUEST,
|
||||
"invalid query: {query}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_raw_reads_only_selected_body_and_is_not_cacheable() {
|
||||
for detached in [false, true] {
|
||||
let repository = if detached {
|
||||
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![
|
||||
sample_selective_body_usage(),
|
||||
])
|
||||
} else {
|
||||
InMemoryUsageReadRepository::seed(vec![sample_selective_body_usage()])
|
||||
};
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(Arc::new(repository)),
|
||||
);
|
||||
for field in [
|
||||
"request_body",
|
||||
"provider_request_body",
|
||||
"response_body",
|
||||
"client_response_body",
|
||||
] {
|
||||
let response = local_admin_usage_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
&format!("/api/admin/usage/usage-selected-body?body_field={field}&body_format=raw"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.headers()["cache-control"],
|
||||
"no-store, no-transform"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers()["x-aether-usage-id"],
|
||||
"usage-selected-body"
|
||||
);
|
||||
assert_eq!(response.headers()["x-aether-body-field"], field);
|
||||
assert_eq!(response.headers()["x-aether-body-encoding"], "json");
|
||||
assert!(response.extensions().get::<AdminAuditEvent>().is_some());
|
||||
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
|
||||
json!({ "marker": field })
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_raw_rejects_foreign_refs_and_disabled_capture() {
|
||||
for body_state in [
|
||||
UsageBodyCaptureState::Reference,
|
||||
UsageBodyCaptureState::Disabled,
|
||||
] {
|
||||
let mut usage = sample_selective_body_usage();
|
||||
usage.response_body = None;
|
||||
usage.response_body_ref = Some("usage://request/foreign-request/response_body".to_string());
|
||||
usage.response_body_state = Some(body_state);
|
||||
let mut foreign = sample_selective_body_usage();
|
||||
foreign.id = "foreign-usage".to_string();
|
||||
foreign.request_id = "foreign-request".to_string();
|
||||
foreign.response_body = Some(json!({ "secret": "must not leak" }));
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage, foreign]),
|
||||
)),
|
||||
);
|
||||
let response = local_admin_usage_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/usage/usage-selected-body?body_field=response_body&body_format=raw",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
assert_eq!(response.headers()["x-aether-body-error"], "missing");
|
||||
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
assert!(!String::from_utf8_lossy(&bytes).contains("secret"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_raw_preserves_authorization_and_binary_headers() {
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![
|
||||
sample_selective_body_usage(),
|
||||
]),
|
||||
)),
|
||||
);
|
||||
let gateway =
|
||||
build_router_with_state(state).layer(tower_http::compression::CompressionLayer::new());
|
||||
let (url, server) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
let endpoint = format!(
|
||||
"{url}/api/admin/usage/usage-selected-body?body_field=response_body&body_format=raw"
|
||||
);
|
||||
let unauthorized = client.get(&endpoint).send().await.unwrap();
|
||||
assert!(matches!(
|
||||
unauthorized.status(),
|
||||
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
|
||||
));
|
||||
let response = admin_request(client.get(&endpoint))
|
||||
.header("accept-encoding", "gzip")
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()["content-encoding"], "identity");
|
||||
assert_eq!(response.headers()["x-aether-body-encoding"], "json");
|
||||
assert_eq!(
|
||||
response.headers()["cache-control"],
|
||||
"no-store, no-transform"
|
||||
);
|
||||
let bytes = response.bytes().await.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&bytes).unwrap(),
|
||||
json!({ "marker": "response_body" })
|
||||
);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_usage_detail_isolates_missing_body_errors() {
|
||||
let mut usage = sample_selective_body_usage();
|
||||
usage.request_body = None;
|
||||
usage.request_body_ref = Some("usage://request/req-selected-body/request_body".to_string());
|
||||
usage.request_body_state = Some(UsageBodyCaptureState::Reference);
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(Arc::new(
|
||||
InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage]),
|
||||
)));
|
||||
for selected in ["response_body", "request_body"] {
|
||||
let response = local_admin_usage_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
&format!("/api/admin/usage/usage-selected-body?body_field={selected}"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should collect");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&bytes).expect("json should parse");
|
||||
if selected == "request_body" {
|
||||
assert_eq!(payload["body_load_errors"]["request_body"], true);
|
||||
assert_eq!(payload["body_load_error_codes"]["request_body"], "missing");
|
||||
assert!(payload["request_body"].is_null());
|
||||
} else {
|
||||
assert_eq!(payload["response_body"]["marker"], "response_body");
|
||||
assert!(payload["body_load_errors"].is_null());
|
||||
assert!(payload["body_load_error_codes"].is_null());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_resolves_admin_usage_detail_when_inline_state_has_body_ref() {
|
||||
let (_upstream_url, upstream_hits, upstream_handle) =
|
||||
|
||||
Reference in New Issue
Block a user