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

fix: restore security hardening compatibility and validation

Browse Source 
Restore authorized rule reveal, explicit full HTTP capture and retention, video task business fields, and valid payment URLs. Add opt-in credential preservation for trusted recovery, fix frontend type contracts and async races, and eliminate PostgreSQL test fixture resource leaks. Document audit coverage and successful fmt and CI-scoped Clippy checks.
This commit is contained in:
elky
2026-09-07 21:14:27 +08:00
parent a5c3699ae9
commit a90d564931
191 changed files with 6785 additions and 1643 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/usage/cleanup.rs
+744 -118
View File
File diff suppressed because it is too large Load Diff
crates/aether-data/adapters/postgres/src/usage/mod.rs
+59 -39
View File
@@ -1374,8 +1374,8 @@ WHERE u.request_id = ANY($1)
"#;
const UPSERT_USAGE_ROUTING_SNAPSHOT_SQL: &str =
include_str!("queries/upsert_usage_routing_snapshot_sql.sql");
#[cfg(test)]
const UPSERT_USAGE_HTTP_AUDIT_SQL: &str = include_str!("queries/upsert_usage_http_audit_sql.sql");
const UPSERT_USAGE_BODY_BLOB_SQL: &str = include_str!("queries/upsert_usage_body_blob_sql.sql");
const UPSERT_USAGE_SETTLEMENT_PRICING_SNAPSHOT_SQL: &str =
include_str!("queries/upsert_usage_settlement_pricing_snapshot_sql.sql");
@@ -13362,19 +13362,36 @@ async fn sync_usage_body_blob_storage<'e, E>(
executor: E,
request_id: &str,
field: UsageBodyField,
_value: Option<&Value>,
_storage: &UsageBodyStorage,
_clear_existing: bool,
value: Option<&Value>,
storage: &UsageBodyStorage,
clear_existing: bool,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
let body_ref = usage_body_ref(request_id, field);
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
if clear_existing {
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
} else if let Some(payload_gzip) = storage.detached_blob_bytes.as_ref() {
sqlx::query(UPSERT_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.bind(request_id)
.bind(field.as_storage_field())
.bind(payload_gzip)
.execute(executor)
.await
.map_postgres_err()?;
} else if value.is_some() {
sqlx::query(DELETE_USAGE_BODY_BLOB_SQL)
.bind(&body_ref)
.execute(executor)
.await
.map_postgres_err()?;
}
Ok(())
}
@@ -13383,43 +13400,46 @@ async fn sync_usage_http_audit_storage<'e, E>(
request_id: &str,
headers: &UsageHttpAuditHeaders<'_>,
refs: &UsageHttpAuditRefs,
_states: &UsageHttpAuditStates,
states: &UsageHttpAuditStates,
body_capture_mode: &str,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if headers.any_present() || refs.any_present() || body_capture_mode != "none" {
return Err(DataLayerError::InvalidInput(
"usage HTTP capture persistence is disabled".to_string(),
));
if !headers.any_present()
&& !refs.any_present()
&& !states.any_present()
&& body_capture_mode == "none"
{
return Ok(());
}
sqlx::query(
r#"
WITH deleted_audit AS (
DELETE FROM usage_http_audits WHERE request_id = $1
)
UPDATE usage
SET request_headers = NULL,
request_body = NULL,
provider_request_headers = NULL,
provider_request_body = NULL,
response_headers = NULL,
response_body = NULL,
client_response_headers = NULL,
client_response_body = NULL,
request_body_compressed = NULL,
provider_request_body_compressed = NULL,
response_body_compressed = NULL,
client_response_body_compressed = NULL
WHERE request_id = $1
"#,
)
.bind(request_id)
.execute(executor)
.await
.map_postgres_err()?;
sqlx::query(UPSERT_USAGE_HTTP_AUDIT_SQL)
.bind(request_id)
.bind(headers.request_headers_json)
.bind(headers.provider_request_headers_json)
.bind(headers.response_headers_json)
.bind(headers.client_response_headers_json)
.bind(refs.request_body_ref.as_deref())
.bind(refs.provider_request_body_ref.as_deref())
.bind(refs.response_body_ref.as_deref())
.bind(refs.client_response_body_ref.as_deref())
.bind(usage_body_capture_state_bind_text(
states.request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.provider_request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.response_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.client_response_body_state,
))
.bind(body_capture_mode)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
crates/aether-data/adapters/postgres/src/usage/tests.rs
+204 -18
View File
@@ -187,6 +187,152 @@ async fn pending_batch_is_opt_in_and_rejects_non_pending_before_connecting() {
.contains("pending usage batch requires pending status"));
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_full_http_capture_round_trips_for_direct_and_batch_writes() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: std::env::var("AETHER_TEST_DATABASE_URL").unwrap(),
min_connections: 1,
max_connections: 2,
acquire_timeout_ms: 10_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.unwrap();
let repository = SqlxUsageReadRepository::new(factory.connect_lazy().unwrap());
crate::run_migrations(repository.pool()).await.unwrap();
for batch in [false, true] {
let request_id = format!("req-full-capture-{}", uuid::Uuid::new_v4().simple());
let now_unix_secs = Utc::now().timestamp() as u64;
let mut pending = fast_clear_usage_record(
&request_id,
"full-capture-test",
now_unix_secs,
false,
UsageBodyCaptureState::Inline,
None,
);
pending.request_headers =
Some(json!({"content-type": "application/json", "authorization": "Bearer private"}));
pending.request_body =
Some(json!({"messages": [{"role": "user", "content": "original request"}]}));
pending.request_body_state = Some(UsageBodyCaptureState::Inline);
pending.provider_request_body = Some(json!({"input": "provider request"}));
pending.response_body = Some(json!("pending response"));
pending.response_body_state = Some(UsageBodyCaptureState::Inline);
pending.client_response_body = Some(json!("pending client response"));
pending.client_response_body_state = Some(UsageBodyCaptureState::Inline);
if batch {
repository
.upsert_pending_many(vec![pending.clone()])
.await
.unwrap();
} else {
repository.upsert(pending.clone()).await.unwrap();
}
for (field, expected) in [
(UsageBodyField::RequestBody, pending.request_body.as_ref()),
(
UsageBodyField::ProviderRequestBody,
pending.provider_request_body.as_ref(),
),
(UsageBodyField::ResponseBody, pending.response_body.as_ref()),
(
UsageBodyField::ClientResponseBody,
pending.client_response_body.as_ref(),
),
] {
assert_eq!(
repository
.resolve_body_ref(&usage_body_ref(&request_id, field))
.await
.unwrap()
.as_ref(),
expected,
"batch={batch}, field={field:?}"
);
}
let mut terminal = fast_clear_usage_record(
&request_id,
"full-capture-test",
now_unix_secs,
true,
UsageBodyCaptureState::None,
None,
);
terminal.provider_request_body_state = None;
terminal.response_headers =
Some(json!({"content-type": "text/event-stream", "set-cookie": "private"}));
terminal.response_body = Some(json!(format!(
"data: {}\n\ndata: [DONE]\n\n",
"streamed text".repeat(8192)
)));
terminal.response_body_state = Some(UsageBodyCaptureState::Inline);
terminal.client_response_body = Some(json!({"output": "final response"}));
terminal.client_response_body_state = Some(UsageBodyCaptureState::Inline);
repository.upsert(terminal.clone()).await.unwrap();
let stored = repository
.find_by_request_id_shallow(&request_id)
.await
.unwrap()
.unwrap();
assert_eq!(
stored.request_headers,
Some(json!({"content-type": "application/json", "authorization": "[redacted]"}))
);
assert_eq!(
stored.response_headers,
Some(json!({"content-type": "text/event-stream", "set-cookie": "[redacted]"}))
);
for (field, expected) in [
(UsageBodyField::RequestBody, pending.request_body.as_ref()),
(
UsageBodyField::ProviderRequestBody,
pending.provider_request_body.as_ref(),
),
(
UsageBodyField::ResponseBody,
terminal.response_body.as_ref(),
),
(
UsageBodyField::ClientResponseBody,
terminal.client_response_body.as_ref(),
),
] {
assert_eq!(
stored.body_state(field),
Some(UsageBodyCaptureState::Reference)
);
assert_eq!(
stored.body_ref(field),
Some(usage_body_ref(&request_id, field).as_str())
);
assert_eq!(
repository
.resolve_body_ref(stored.body_ref(field).unwrap())
.await
.unwrap()
.as_ref(),
expected,
"batch={batch}, field={field:?}"
);
}
let legacy_content_present: bool = sqlx::query_scalar("SELECT request_body IS NOT NULL OR request_headers IS NOT NULL OR response_body IS NOT NULL FROM usage WHERE request_id = $1")
.bind(&request_id).fetch_one(repository.pool()).await.unwrap();
assert!(!legacy_content_present);
sqlx::query("DELETE FROM usage WHERE request_id = $1")
.bind(&request_id)
.execute(repository.pool())
.await
.unwrap();
}
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_stale_terminal_event_is_a_full_transaction_noop() {
@@ -253,7 +399,7 @@ async fn live_stale_terminal_event_is_a_full_transaction_noop() {
.unwrap(),
);
let settlement_before = sqlx::query(
"SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
"SELECT billing_status, billing_total_cost_usd::DOUBLE PRECISION AS billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
@@ -318,7 +464,7 @@ async fn live_stale_terminal_event_is_a_full_transaction_noop() {
.unwrap(),
);
let settlement_after = sqlx::query(
"SELECT billing_status, billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
"SELECT billing_status, billing_total_cost_usd::DOUBLE PRECISION AS billing_total_cost_usd FROM usage_settlement_snapshots WHERE request_id = $1",
)
.bind(&request_id)
.fetch_one(repository.pool())
@@ -542,7 +688,7 @@ async fn live_pending_batch_persists_auxiliary_state_and_preserves_terminal_conf
.fetch_one(repository.pool())
.await
.expect("HTTP audit count should be readable");
assert_eq!(http_count, 0);
assert_eq!(http_count, 1);
let blob_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*)::BIGINT FROM usage_body_blobs WHERE request_id = $1",
)
@@ -550,7 +696,23 @@ async fn live_pending_batch_persists_auxiliary_state_and_preserves_terminal_conf
.fetch_one(repository.pool())
.await
.expect("body blob count should be readable");
assert_eq!(blob_count, 0);
assert_eq!(blob_count, 4);
let captured = repository
.find_by_request_id_shallow(&rich_request_id)
.await
.unwrap()
.unwrap();
assert_eq!(
captured.request_headers,
Some(json!({"x-request": "[redacted]"}))
);
assert_eq!(
repository
.resolve_body_ref(captured.body_ref(UsageBodyField::RequestBody).unwrap())
.await
.unwrap(),
Some(json!({"messages": [{"role": "user", "content": "hello"}]}))
);
let routing = sqlx::query(
"SELECT candidate_id, candidate_index, selected_provider_api_key_id FROM usage_routing_snapshots WHERE request_id = $1",
@@ -667,7 +829,7 @@ async fn live_pending_batch_and_terminal_upserts_count_each_provider_request_onc
let suffix = uuid::Uuid::new_v4().simple().to_string();
let provider_name = format!("pending-terminal-race-provider-{suffix}");
let provider_key_id = format!("pending-terminal-race-key-{suffix}");
let provider_key_id = uuid::Uuid::new_v4().to_string();
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let request_ids = (0..REQUESTS)
.map(|index| format!("req-pending-terminal-race-{index}-{suffix}"))
@@ -793,7 +955,7 @@ async fn live_first_byte_fast_path_is_atomic_and_preserves_terminal_state() {
let existing_request_id = format!("req-first-byte-existing-{suffix}");
let metadata_fill_request_id = format!("req-first-byte-metadata-fill-{suffix}");
let provider_name = format!("first-byte-fast-{suffix}");
let missing_provider_key_id = format!("key-first-byte-missing-{suffix}");
let missing_provider_key_id = uuid::Uuid::new_v4().to_string();
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let mut missing_first_byte = first_byte_usage_record(
@@ -1064,7 +1226,7 @@ async fn live_first_byte_reads_provider_contribution_after_waiting_for_canonical
let suffix = uuid::Uuid::new_v4().simple().to_string();
let request_id = format!("req-first-byte-lock-snapshot-{suffix}");
let provider_name = format!("first-byte-lock-snapshot-{suffix}");
let provider_key_id = format!("key-first-byte-lock-snapshot-{suffix}");
let provider_key_id = uuid::Uuid::new_v4().to_string();
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let mut pending = first_byte_usage_record(
&request_id,
@@ -1209,14 +1371,14 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
let request_b = format!("req-first-byte-batch-b-{suffix}");
let request_missing = format!("req-first-byte-batch-missing-{suffix}");
let request_terminal = format!("req-first-byte-batch-terminal-{suffix}");
let missing_provider_key_id = format!("key-first-byte-batch-missing-{suffix}");
let missing_provider_key_id = uuid::Uuid::new_v4().to_string();
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
let mut pending_a = first_byte_usage_record(
&request_a,
&provider_name,
now_unix_secs,
Some(json!({"seed": "a"})),
Some(json!({"trace_id": "seed-a"})),
);
pending_a.status = "pending".to_string();
pending_a.first_byte_time_ms = None;
@@ -1241,7 +1403,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
);
terminal.is_stream = Some(true);
terminal.first_byte_time_ms = Some(44);
terminal.request_metadata = Some(json!({"terminal": true}));
terminal.request_metadata = Some(json!({"trace_id": "terminal"}));
repository
.upsert(pending_a)
@@ -1260,7 +1422,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
&request_a,
&provider_name,
now_unix_secs + 1,
Some(json!({"incoming": "a"})),
Some(json!({"trace_id": "incoming-a"})),
);
first_a.first_byte_time_ms = Some(30);
first_a.has_format_conversion = None;
@@ -1272,7 +1434,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
&request_b,
&provider_name,
now_unix_secs + 1,
Some(json!({"incoming": "b"})),
Some(json!({"trace_id": "incoming-b"})),
);
first_b.has_format_conversion = Some(true);
@@ -1280,7 +1442,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
&request_terminal,
&provider_name,
now_unix_secs + 2,
Some(json!({"late": true})),
Some(json!({"trace_id": "late"})),
);
late_terminal.first_byte_time_ms = Some(3);
late_terminal.has_format_conversion = None;
@@ -1288,7 +1450,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
&request_missing,
&provider_name,
now_unix_secs + 1,
Some(json!({"incoming": "missing"})),
Some(json!({"trace_id": "incoming-missing"})),
);
first_missing.provider_api_key_id = Some(missing_provider_key_id.clone());
@@ -1345,7 +1507,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
row_a
.try_get::<Option<serde_json::Value>, _>("request_metadata")
.unwrap(),
Some(json!({"seed": "a"})),
Some(json!({"trace_id": "seed-a"})),
"existing metadata remains authoritative"
);
@@ -1364,7 +1526,7 @@ async fn live_first_byte_batch_preserves_duplicate_order_and_terminal_guards() {
row_b
.try_get::<Option<serde_json::Value>, _>("request_metadata")
.unwrap(),
Some(json!({"incoming": "b"}))
Some(json!({"trace_id": "incoming-b"}))
);
let row_terminal = rows
@@ -2044,7 +2206,7 @@ async fn live_provider_performance_grouping_sets_matches_separate_queries() {
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and a populated PostgreSQL database"]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_dashboard_daily_breakdown_uses_canonical_covering_read_path() {
let database_url = std::env::var("AETHER_TEST_DATABASE_URL")
.expect("AETHER_TEST_DATABASE_URL must point at the test database");
@@ -2062,6 +2224,20 @@ async fn live_dashboard_daily_breakdown_uses_canonical_covering_read_path() {
let repository =
SqlxUsageReadRepository::new(factory.connect_lazy().expect("lazy pool should build"));
let until = Utc::now().timestamp().max(0) as u64;
crate::run_migrations(repository.pool()).await.unwrap();
let request_id = format!("daily-breakdown-{}", uuid::Uuid::new_v4().simple());
let provider_name = format!("daily-provider-{}", uuid::Uuid::new_v4().simple());
repository
.upsert(fast_clear_usage_record(
&request_id,
&provider_name,
until.saturating_sub(60),
true,
UsageBodyCaptureState::None,
None,
))
.await
.unwrap();
let started = std::time::Instant::now();
let rows = repository
.list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery {
@@ -2077,7 +2253,17 @@ async fn live_dashboard_daily_breakdown_uses_canonical_covering_read_path() {
started.elapsed(),
rows.len()
);
assert!(!rows.is_empty());
let seeded = rows
.iter()
.find(|row| row.provider == provider_name)
.unwrap();
assert_eq!(seeded.requests, 1);
assert_eq!(seeded.total_tokens, 2);
sqlx::query("DELETE FROM \"usage\" WHERE request_id = $1")
.bind(&request_id)
.execute(repository.pool())
.await
.unwrap();
}
#[tokio::test]