feat(security): harden gateway boundaries and usage policies

Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
@@ -92,22 +92,19 @@ impl AccountSelfCheckWorkerConfig {
enum AccountSelfCheckOutcome {
Success {
status_code: Option<u16>,
message: Option<String>,
},
Blocked {
status_code: Option<u16>,
message: String,
},
AutoRemoved {
status_code: Option<u16>,
message: String,
},
Failed {
status_code: Option<u16>,
message: String,
category: &'static str,
},
Skipped {
message: String,
category: &'static str,
},
}
@@ -132,13 +129,12 @@ impl AccountSelfCheckOutcome {
}
}
fn message(&self) -> Option<&str> {
fn category(&self) -> &'static str {
match self {
Self::Success { message, .. } => message.as_deref(),
Self::Blocked { message, .. }
| Self::AutoRemoved { message, .. }
| Self::Failed { message, .. }
| Self::Skipped { message, .. } => Some(message.as_str()),
Self::Success { .. } => "quota_refresh_succeeded",
Self::Blocked { .. } => "account_blocked",
Self::AutoRemoved { .. } => "account_auto_removed",
Self::Failed { category, .. } | Self::Skipped { category } => category,
}
}
}
@@ -370,13 +366,13 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
let Some(payload) = payload else {
return AccountSelfCheckOutcome::Failed {
status_code: None,
message: "quota refresh returned no payload".to_string(),
category: "missing_payload",
};
};
let Some(results) = payload.get("results").and_then(Value::as_array) else {
return AccountSelfCheckOutcome::Failed {
status_code: None,
message: "quota refresh returned no result list".to_string(),
category: "invalid_payload",
};
};
let Some(item) = results.iter().find(|item| {
@@ -386,7 +382,7 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
}) else {
return AccountSelfCheckOutcome::Failed {
status_code: None,
message: "quota refresh result missing key".to_string(),
category: "missing_key_result",
};
};
@@ -404,34 +400,34 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
.get("message")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
.filter(|value| !value.is_empty());
let auto_removed = item
.get("auto_removed")
.and_then(Value::as_bool)
.unwrap_or(false);
if status == "success" {
return AccountSelfCheckOutcome::Success {
status_code,
message,
};
return AccountSelfCheckOutcome::Success { status_code };
}
if auto_removed {
return AccountSelfCheckOutcome::AutoRemoved {
status_code,
message: message.unwrap_or_else(|| "已自动删除".to_string()),
};
return AccountSelfCheckOutcome::AutoRemoved { status_code };
}
if quota_result_status_is_blocked(&status, status_code, message.as_deref()) {
return AccountSelfCheckOutcome::Blocked {
status_code,
message: message.unwrap_or_else(|| status.clone()),
};
if quota_result_status_is_blocked(&status, status_code, message) {
return AccountSelfCheckOutcome::Blocked { status_code };
}
AccountSelfCheckOutcome::Failed {
status_code,
message: message.unwrap_or_else(|| status.clone()),
category: quota_refresh_failure_category(&status, status_code),
}
}
fn quota_refresh_failure_category(status: &str, status_code: Option<u16>) -> &'static str {
match status_code {
Some(408 | 504) => "upstream_timed_out",
Some(429) => "rate_limited",
Some(500..=599) => "upstream_error",
_ if matches!(status, "unsupported" | "not_supported") => "unsupported",
_ => "quota_refresh_failed",
}
}
@@ -546,21 +542,7 @@ async fn record_score_probe_result_for_key(
succeeded,
hard_state,
probe_status,
score_reason_patch: Some(json!({
"last_probe": {
"source": "account_self_check",
"status": outcome.score_status(),
"status_code": outcome.status_code(),
"message": outcome.message()
},
"last_self_check": {
"source": "account_self_check",
"status": outcome.score_status(),
"status_code": outcome.status_code(),
"message": outcome.message(),
"attempted_at": attempted_at
}
})),
score_reason_patch: Some(score_reason_patch_for_outcome(outcome, attempted_at)),
};
if let Err(err) = state.data.record_pool_member_probe_result(result).await {
debug!(
@@ -572,6 +554,24 @@ async fn record_score_probe_result_for_key(
}
}
fn score_reason_patch_for_outcome(outcome: &AccountSelfCheckOutcome, attempted_at: u64) -> Value {
json!({
"last_probe": {
"source": "account_self_check",
"status": outcome.score_status(),
"status_code": outcome.status_code(),
"category": outcome.category()
},
"last_self_check": {
"source": "account_self_check",
"status": outcome.score_status(),
"status_code": outcome.status_code(),
"category": outcome.category(),
"attempted_at": attempted_at
}
})
}
fn endpoint_for_self_check(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
@@ -579,8 +579,21 @@ fn endpoint_for_self_check(
provider_quota_refresh_endpoint_for_provider(provider_type, endpoints, true)
}
fn gateway_error_message(err: GatewayError) -> String {
err.into_message()
fn gateway_error_category(err: &GatewayError) -> &'static str {
match err {
GatewayError::UpstreamUnavailable { .. } => "upstream_unavailable",
GatewayError::ControlUnavailable { .. } => "control_unavailable",
GatewayError::LocalExecutionPlanningTimeout { .. } => "planning_timed_out",
GatewayError::AdmissionTimeout { .. } => "admission_timed_out",
GatewayError::Client { status, .. } if status.as_u16() == 429 => "rate_limited",
GatewayError::Client { status, .. } if status.is_server_error() => "upstream_error",
GatewayError::Client { .. } => "request_rejected",
GatewayError::PlanUsageLimited(_) => "plan_usage_limited",
GatewayError::LastActiveAdminUpdateDenied | GatewayError::LastActiveAdminDeleteDenied => {
"operation_rejected"
}
GatewayError::Internal(_) => "internal_error",
}
}
fn update_summary_from_outcome(
@@ -727,7 +740,7 @@ pub(crate) async fn perform_account_self_check_once_with_config(
Ok(outcome) => outcome,
Err(err) => AccountSelfCheckOutcome::Failed {
status_code: None,
message: gateway_error_message(err),
category: gateway_error_category(&err),
},
};
record_score_probe_result_for_key(state, &provider.id, &key.id, now_ts, &outcome).await;
@@ -797,7 +810,12 @@ pub(crate) fn spawn_account_self_check_worker(
#[cfg(test)]
mod tests {
use super::select_account_self_check_key_ids;
use super::{
gateway_error_category, quota_payload_result_for_key, score_reason_patch_for_outcome,
select_account_self_check_key_ids,
};
use crate::GatewayError;
use serde_json::json;
use std::collections::BTreeMap;
#[test]
@@ -813,4 +831,38 @@ mod tests {
assert_eq!(selected, vec!["never".to_string(), "stale".to_string()]);
}
#[test]
fn account_self_check_patch_does_not_persist_upstream_message() {
let secret = "Authorization: Bearer quota-secret from /srv/aether/credentials.json";
let outcome = quota_payload_result_for_key(
"key-secret-regression",
Some(json!({
"results": [{
"key_id": "key-secret-regression",
"status": "error",
"status_code": 502,
"message": secret
}]
})),
);
let patch = score_reason_patch_for_outcome(&outcome, 1_777_000_000);
let serialized = patch.to_string();
assert_eq!(patch["last_self_check"]["category"], "upstream_error");
assert!(patch["last_self_check"].get("message").is_none());
assert!(!serialized.contains("quota-secret"));
assert!(!serialized.contains("/srv/aether/credentials.json"));
}
#[test]
fn account_self_check_gateway_error_category_drops_internal_details() {
let error = GatewayError::Internal(
"postgresql://admin:[email protected]/aether".to_string(),
);
assert_eq!(gateway_error_category(&error), "internal_error");
assert!(!gateway_error_category(&error).contains("database-secret"));
}
}
@@ -15,6 +15,15 @@ const CLEANUP_RUN_HISTORY_KEY: &str = "admin_cleanup_run_history";
const CLEANUP_RUN_HISTORY_LIMIT: usize = 50;
const REQUEST_BODY_PROGRESS_UPDATE_BATCHES: usize = 10;
const CLEANUP_ERROR_INVALID_CONFIGURATION: &str = "invalid_configuration";
const CLEANUP_ERROR_INVALID_INPUT: &str = "invalid_input";
const CLEANUP_ERROR_POSTGRES: &str = "postgres";
const CLEANUP_ERROR_REDIS: &str = "redis";
const CLEANUP_ERROR_SQL: &str = "sql";
const CLEANUP_ERROR_TIMED_OUT: &str = "timed_out";
const CLEANUP_ERROR_UNEXPECTED_VALUE: &str = "unexpected_value";
const CLEANUP_ERROR_INTERNAL: &str = "internal_error";
pub(crate) const USAGE_CLEANUP_KIND: &str = "usage_cleanup";
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct AdminCleanupRunRecord {
@@ -30,6 +39,73 @@ pub(crate) struct AdminCleanupRunRecord {
pub(crate) error: Option<String>,
}
pub(super) fn cleanup_data_layer_error_category(error: &DataLayerError) -> &'static str {
match error {
DataLayerError::InvalidConfiguration(_) => CLEANUP_ERROR_INVALID_CONFIGURATION,
DataLayerError::InvalidInput(_) => CLEANUP_ERROR_INVALID_INPUT,
DataLayerError::Postgres(_) => CLEANUP_ERROR_POSTGRES,
DataLayerError::Redis(_) => CLEANUP_ERROR_REDIS,
DataLayerError::Sql(_) => CLEANUP_ERROR_SQL,
DataLayerError::TimedOut(_) => CLEANUP_ERROR_TIMED_OUT,
DataLayerError::UnexpectedValue(_) => CLEANUP_ERROR_UNEXPECTED_VALUE,
}
}
fn cleanup_gateway_error_category(error: &GatewayError) -> &'static str {
match error {
GatewayError::UpstreamUnavailable { .. } => "upstream_unavailable",
GatewayError::ControlUnavailable { .. } => "control_unavailable",
GatewayError::LocalExecutionPlanningTimeout { .. } => "planning_timed_out",
GatewayError::AdmissionTimeout { .. } => "admission_timed_out",
GatewayError::Client { status, .. } if status.is_server_error() => "upstream_error",
GatewayError::Client { .. } => "request_rejected",
GatewayError::PlanUsageLimited(_) => "plan_usage_limited",
GatewayError::LastActiveAdminUpdateDenied | GatewayError::LastActiveAdminDeleteDenied => {
"operation_rejected"
}
GatewayError::Internal(_) => CLEANUP_ERROR_INTERNAL,
}
}
fn normalize_stored_cleanup_error(error: &str) -> &'static str {
let error = error.trim();
match error {
CLEANUP_ERROR_INVALID_CONFIGURATION => CLEANUP_ERROR_INVALID_CONFIGURATION,
CLEANUP_ERROR_INVALID_INPUT => CLEANUP_ERROR_INVALID_INPUT,
CLEANUP_ERROR_POSTGRES => CLEANUP_ERROR_POSTGRES,
CLEANUP_ERROR_REDIS => CLEANUP_ERROR_REDIS,
CLEANUP_ERROR_SQL => CLEANUP_ERROR_SQL,
CLEANUP_ERROR_TIMED_OUT => CLEANUP_ERROR_TIMED_OUT,
CLEANUP_ERROR_UNEXPECTED_VALUE => CLEANUP_ERROR_UNEXPECTED_VALUE,
CLEANUP_ERROR_INTERNAL => CLEANUP_ERROR_INTERNAL,
"upstream_unavailable" => "upstream_unavailable",
"control_unavailable" => "control_unavailable",
"planning_timed_out" => "planning_timed_out",
"admission_timed_out" => "admission_timed_out",
"upstream_error" => "upstream_error",
"request_rejected" => "request_rejected",
"plan_usage_limited" => "plan_usage_limited",
"operation_rejected" => "operation_rejected",
_ if error.starts_with("invalid configuration:") => CLEANUP_ERROR_INVALID_CONFIGURATION,
_ if error.starts_with("invalid input:") => CLEANUP_ERROR_INVALID_INPUT,
_ if error.starts_with("postgres error:") => CLEANUP_ERROR_POSTGRES,
_ if error.starts_with("redis error:") => CLEANUP_ERROR_REDIS,
_ if error.starts_with("sql error:") => CLEANUP_ERROR_SQL,
_ if error.starts_with("operation timed out:") => CLEANUP_ERROR_TIMED_OUT,
_ if error.starts_with("unexpected database value:") => CLEANUP_ERROR_UNEXPECTED_VALUE,
_ => CLEANUP_ERROR_INTERNAL,
}
}
fn normalize_cleanup_run_record(mut record: AdminCleanupRunRecord) -> AdminCleanupRunRecord {
record.error = record
.error
.as_deref()
.map(normalize_stored_cleanup_error)
.map(str::to_string);
record
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AdminCleanupTaskKind {
Config,
@@ -225,7 +301,7 @@ pub(crate) async fn record_failed_cleanup_run(
.unwrap_or(u64::MAX),
),
summary: json!({}),
error: Some(error.to_string()),
error: Some(cleanup_data_layer_error_category(error).to_string()),
};
if let Err(err) = record_cleanup_run(data, record).await {
warn!(error = %err, kind, "failed to record failed cleanup run");
@@ -277,7 +353,7 @@ async fn run_request_body_cleanup_task(
batch_size,
&total,
Some(started_at),
Some(err.to_string()),
Some(cleanup_data_layer_error_category(&err).to_string()),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record request body cleanup failure");
@@ -350,7 +426,7 @@ async fn run_admin_system_purge_task(
kind.failure_message().to_string(),
json!({}),
Some(started_at),
Some(format!("{err:?}")),
Some(cleanup_gateway_error_category(&err).to_string()),
);
if let Err(record_err) = record_cleanup_run(&data, failed).await {
warn!(error = %record_err, "failed to record admin system purge task failure");
@@ -493,6 +569,7 @@ pub(crate) async fn record_admin_cleanup_run(
data: &GatewayDataState,
record: AdminCleanupRunRecord,
) -> Result<(), DataLayerError> {
let record = normalize_cleanup_run_record(record);
let mut records = list_admin_cleanup_run_records(data).await?;
records.retain(|existing| existing.id != record.id);
records.insert(0, record);
@@ -522,5 +599,50 @@ fn parse_cleanup_run_records(value: Value) -> Vec<AdminCleanupRunRecord> {
.into_iter()
.flat_map(|items| items.iter())
.filter_map(|item| serde_json::from_value::<AdminCleanupRunRecord>(item.clone()).ok())
.map(normalize_cleanup_run_record)
.collect()
}
#[cfg(test)]
mod tests {
use super::{
cleanup_data_layer_error_category, parse_cleanup_run_records, AdminCleanupRunRecord,
};
use aether_data_contracts::DataLayerError;
use serde_json::json;
#[test]
fn cleanup_error_categories_do_not_include_data_layer_details() {
let error = DataLayerError::Postgres(
"connection failed for postgresql://admin:[email protected]/aether"
.to_string(),
);
assert_eq!(cleanup_data_layer_error_category(&error), "postgres");
assert!(!cleanup_data_layer_error_category(&error).contains("database-secret"));
}
#[test]
fn cleanup_history_read_projection_removes_legacy_error_details() {
let secret = "postgres error: password=database-secret path=/srv/aether/private.db";
let stored = AdminCleanupRunRecord {
id: "cleanup-secret-regression".to_string(),
kind: "usage_cleanup".to_string(),
trigger: "manual".to_string(),
status: "failed".to_string(),
message: "请求记录手动清理失败".to_string(),
started_at_unix_secs: 1,
completed_at_unix_secs: Some(2),
duration_ms: Some(10),
summary: json!({}),
error: Some(secret.to_string()),
};
let records = parse_cleanup_run_records(json!([stored]));
let serialized = serde_json::to_string(&records).expect("cleanup records should serialize");
assert_eq!(records[0].error.as_deref(), Some("postgres"));
assert!(!serialized.contains("database-secret"));
assert!(!serialized.contains("/srv/aether/private.db"));
}
}
@@ -289,7 +289,7 @@ mod tests {
assert_eq!(responses.max_retries, Some(9));
assert_eq!(
responses.proxy,
Some(json!({"url": "http://proxy.internal:8080"}))
Some(json!({"url": "http://proxy.internal:8080/"}))
);
assert_eq!(
responses
@@ -312,7 +312,7 @@ mod tests {
assert_eq!(live.max_retries, Some(6));
assert_eq!(
live.proxy,
Some(json!({"url": "http://voice-proxy.internal:8080"}))
Some(json!({"url": "http://voice-proxy.internal:8080/"}))
);
assert_eq!(
live.config
@@ -128,7 +128,7 @@ pub(crate) async fn perform_oauth_token_refresh_once(
Ok(None) => {
summary.skipped = summary.skipped.saturating_add(1);
}
Err(err) => {
Err(_) => {
summary.failed = summary.failed.saturating_add(1);
warn!(
event_name = "oauth_token_refresh_failed",
@@ -136,7 +136,6 @@ pub(crate) async fn perform_oauth_token_refresh_once(
worker = "oauth_token_refresh",
provider_id = %provider.id,
key_id = %key.id,
error = ?err,
"gateway oauth token auto refresh failed"
);
}
@@ -1042,6 +1042,17 @@ async fn refresh_provider_probe_keys(
.await
}
fn pool_quota_probe_worker_error_score_reason(_error: &GatewayError) -> Value {
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "worker_error",
"category": "worker_error",
"message": "Provider quota probe worker failed"
}
})
}
fn update_summary_from_payload(
summary: &mut PoolQuotaProbeRunSummary,
selected_count: usize,
@@ -1449,20 +1460,14 @@ async fn perform_pool_quota_probe_for_provider(
now_ts,
false,
Some(PoolMemberHardState::Cooldown),
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "worker_error",
"message": format!("{err:?}")
}
}),
pool_quota_probe_worker_error_score_reason(&err),
)
.await;
warn!(
provider_id = %provider_short_id,
provider_type,
key_id,
error = ?err,
error_category = "provider_quota_probe_failed",
"gateway pool quota probe failed"
);
}
@@ -1742,9 +1747,12 @@ pub(crate) fn spawn_pool_quota_probe_worker(
interval.tick().await;
loop {
interval.tick().await;
if let Err(err) = perform_pool_quota_probe_once_with_config(&state, config).await {
if perform_pool_quota_probe_once_with_config(&state, config)
.await
.is_err()
{
warn!(
error = ?err,
error_category = "pool_quota_probe_worker_failed",
"gateway pool quota probe worker tick failed"
);
}
@@ -1758,6 +1766,30 @@ mod tests {
use super::*;
use serde_json::json;
#[test]
fn worker_error_score_reason_drops_runtime_error_details() {
let secret = "postgresql://admin:[email protected]/aether; Authorization: Bearer quota-secret; https://user:[email protected]/quota?q=secret";
let patch =
pool_quota_probe_worker_error_score_reason(&GatewayError::Internal(secret.to_string()));
let serialized = patch.to_string();
assert_eq!(patch["last_probe"]["status"], "worker_error");
assert_eq!(patch["last_probe"]["category"], "worker_error");
assert_eq!(
patch["last_probe"]["message"],
"Provider quota probe worker failed"
);
for sensitive in [
"db-secret",
"quota-secret",
"user:password",
"db.internal",
"upstream.test",
] {
assert!(!serialized.contains(sensitive));
}
}
fn key(
id: &str,
provider_id: &str,
@@ -199,7 +199,14 @@ async fn run_provider_quota_alert_for_provider(
let Some(total_available) = extract_total_available(&payload) else {
warn!(
provider_id = %provider_id,
payload = %payload,
payload_status = payload
.get("status")
.and_then(serde_json::Value::as_str)
.unwrap_or("unknown"),
action_type = payload
.get("action_type")
.and_then(serde_json::Value::as_str)
.unwrap_or("unknown"),
"provider quota alert skipped because balance payload has no total_available"
);
write_checked_runtime_state_without_balance(state, &provider_id, now_unix_secs).await;
@@ -53,6 +53,7 @@ pub(super) async fn cleanup_stale_proxy_nodes_once(
);
let mutation = ProxyNodeTunnelStatusMutation {
node_id: node.id.clone(),
expected_tunnel_generation: Some(node.tunnel_generation.clone()),
connected: false,
conn_count: 0,
detail: Some(detail),
@@ -27,6 +27,8 @@ struct ProxyUpgradeRolloutPlan {
#[serde(default)]
skipped_node_ids: Vec<String>,
#[serde(default)]
skipped_node_generations: std::collections::BTreeMap<String, String>,
#[serde(default)]
tracked_nodes: Vec<ProxyUpgradeRolloutTrackedNode>,
}
@@ -39,6 +41,8 @@ pub(crate) struct ProxyUpgradeRolloutProbeConfig {
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
struct ProxyUpgradeRolloutTrackedNode {
node_id: String,
#[serde(default)]
tunnel_generation: String,
dispatched_at_unix_secs: u64,
version_confirmed_at_unix_secs: Option<u64>,
#[serde(default)]
@@ -51,6 +55,7 @@ struct ProxyUpgradeRolloutTrackedNode {
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ProxyUpgradeRolloutPendingProbe {
pub(crate) node_id: String,
pub(crate) tunnel_generation: String,
pub(crate) url: String,
pub(crate) timeout_secs: u64,
}
@@ -167,6 +172,7 @@ struct RolloutSnapshot {
completed: Vec<String>,
pending: Vec<String>,
pending_conflicts: Vec<String>,
pending_conflict_nodes: Vec<StoredProxyNode>,
skipped: Vec<String>,
available: Vec<StoredProxyNode>,
ready_to_finalize: Vec<StoredProxyNode>,
@@ -245,6 +251,11 @@ pub(crate) async fn start_proxy_upgrade_rollout(
.filter(|_| preserve_existing)
.map(|plan| plan.skipped_node_ids.clone())
.unwrap_or_default(),
skipped_node_generations: existing
.as_ref()
.filter(|_| preserve_existing)
.map(|plan| plan.skipped_node_generations.clone())
.unwrap_or_default(),
tracked_nodes: existing
.as_ref()
.filter(|_| preserve_existing)
@@ -307,6 +318,7 @@ pub(crate) async fn collect_proxy_upgrade_rollout_probes(
}
Some(ProxyUpgradeRolloutPendingProbe {
node_id: tracked.node_id,
tunnel_generation: tracked.tunnel_generation,
url: probe.url.clone(),
timeout_secs: probe.timeout_secs,
})
@@ -430,10 +442,11 @@ pub(crate) async fn clear_proxy_upgrade_rollout_conflicts(
}
let mut cleared_node_ids = Vec::new();
for node_id in snapshot.pending_conflicts {
for node in snapshot.pending_conflict_nodes {
let Some(updated) = data
.update_proxy_node_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: node_id.clone(),
node_id: node.id,
expected_tunnel_generation: Some(node.tunnel_generation),
node_name: None,
allowed_ports: None,
log_level: None,
@@ -484,6 +497,8 @@ pub(crate) async fn skip_proxy_upgrade_rollout_node(
plan.skipped_node_ids.sort();
plan.skipped_node_ids.dedup();
}
plan.skipped_node_generations
.insert(node_id.to_string(), node.tunnel_generation.clone());
plan.tracked_nodes
.retain(|tracked| tracked.node_id != node_id);
plan.updated_at_unix_secs = now;
@@ -495,6 +510,7 @@ pub(crate) async fn skip_proxy_upgrade_rollout_node(
let _ = data
.update_proxy_node_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: node_id.to_string(),
expected_tunnel_generation: Some(node.tunnel_generation),
node_name: None,
allowed_ports: None,
log_level: None,
@@ -554,6 +570,7 @@ pub(crate) async fn restore_proxy_upgrade_rollout_skipped_nodes(
}
plan.skipped_node_ids.clear();
plan.skipped_node_generations.clear();
plan.updated_at_unix_secs = now;
save_proxy_upgrade_rollout_plan(data, &plan).await?;
@@ -584,14 +601,15 @@ pub(crate) async fn retry_proxy_upgrade_rollout_node(
let Some(mut plan) = load_proxy_upgrade_rollout_plan(data).await? else {
return Ok(None);
};
let Some(_node) = data.find_proxy_node(node_id).await? else {
let Some(node) = data.find_proxy_node(node_id).await? else {
return Ok(None);
};
let now = now_unix_secs();
let _ = data
let Some(updated) = data
.update_proxy_node_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: node_id.to_string(),
expected_tunnel_generation: Some(node.tunnel_generation),
node_name: None,
allowed_ports: None,
log_level: None,
@@ -599,13 +617,18 @@ pub(crate) async fn retry_proxy_upgrade_rollout_node(
scheduling_state: None,
upgrade_to: Some(Some(plan.version.clone())),
})
.await?;
.await?
else {
return Ok(None);
};
plan.skipped_node_ids.retain(|id| id != node_id);
plan.skipped_node_generations.remove(node_id);
plan.tracked_nodes
.retain(|tracked| tracked.node_id != node_id);
plan.tracked_nodes.push(ProxyUpgradeRolloutTrackedNode {
node_id: node_id.to_string(),
tunnel_generation: updated.tunnel_generation,
dispatched_at_unix_secs: now,
version_confirmed_at_unix_secs: None,
traffic_confirmed_at_unix_secs: None,
@@ -659,6 +682,7 @@ async fn advance_proxy_upgrade_rollout(
let _ = data
.update_proxy_node_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: node.id.clone(),
expected_tunnel_generation: Some(node.tunnel_generation.clone()),
node_name: None,
allowed_ports: None,
log_level: None,
@@ -711,11 +735,12 @@ async fn advance_proxy_upgrade_rollout(
return Ok(summary);
}
let mut updated_node_ids = Vec::with_capacity(selected.len());
let mut updated_nodes = Vec::with_capacity(selected.len());
for node in selected {
let Some(updated) = data
.update_proxy_node_remote_config(&ProxyNodeRemoteConfigMutation {
node_id: node.id.clone(),
expected_tunnel_generation: Some(node.tunnel_generation),
node_name: None,
allowed_ports: None,
log_level: None,
@@ -727,28 +752,30 @@ async fn advance_proxy_upgrade_rollout(
else {
continue;
};
updated_node_ids.push(updated.id);
updated_nodes.push((updated.id, updated.tunnel_generation));
}
plan.last_dispatched_at_unix_secs = Some(now);
plan.updated_at_unix_secs = now;
plan.tracked_nodes
.extend(
updated_node_ids
.iter()
.cloned()
.map(|node_id| ProxyUpgradeRolloutTrackedNode {
node_id,
dispatched_at_unix_secs: now,
version_confirmed_at_unix_secs: None,
traffic_confirmed_at_unix_secs: None,
confirm_failed_requests: None,
confirm_dns_failures: None,
confirm_stream_errors: None,
}),
);
.extend(updated_nodes.iter().map(|(node_id, tunnel_generation)| {
ProxyUpgradeRolloutTrackedNode {
node_id: node_id.clone(),
tunnel_generation: tunnel_generation.clone(),
dispatched_at_unix_secs: now,
version_confirmed_at_unix_secs: None,
traffic_confirmed_at_unix_secs: None,
confirm_failed_requests: None,
confirm_dns_failures: None,
confirm_stream_errors: None,
}
}));
save_proxy_upgrade_rollout_plan(data, &plan).await?;
let updated_node_ids = updated_nodes
.into_iter()
.map(|(node_id, _)| node_id)
.collect::<Vec<_>>();
summary.updated = updated_node_ids.len();
summary.skipped = summary.skipped.saturating_sub(summary.updated);
summary.node_ids = updated_node_ids.clone();
@@ -785,7 +812,12 @@ fn build_rollout_snapshot(
snapshot.online_eligible_total = snapshot.online_eligible_total.saturating_add(1);
}
if skipped_node_ids.contains(node.id.as_str()) {
if skipped_node_ids.contains(node.id.as_str())
&& plan
.skipped_node_generations
.get(node.id.as_str())
.is_some_and(|generation| generation == &node.tunnel_generation)
{
snapshot.skipped.push(node.id.clone());
tracked_by_node_id.remove(node.id.as_str());
continue;
@@ -794,7 +826,13 @@ fn build_rollout_snapshot(
let reported_version = proxy_reported_version(node.proxy_metadata.as_ref());
let pending_target = remote_config_upgrade_target(node.remote_config.as_ref());
if let Some(mut tracked) = tracked_by_node_id.remove(node.id.as_str()) {
if let Some(mut tracked) = tracked_by_node_id
.remove(node.id.as_str())
.filter(|tracked| {
!tracked.tunnel_generation.is_empty()
&& tracked.tunnel_generation == node.tunnel_generation
})
{
snapshot.remaining_total = snapshot.remaining_total.saturating_add(1);
if reported_version.as_deref() == Some(target_version.as_str()) {
@@ -845,6 +883,7 @@ fn build_rollout_snapshot(
snapshot.pending.push(node.id.clone());
} else {
snapshot.pending_conflicts.push(node.id.clone());
snapshot.pending_conflict_nodes.push(node);
}
continue;
}
@@ -905,9 +944,24 @@ async fn save_proxy_upgrade_rollout_plan(
Ok(())
}
/// Records rollout traffic only when the callback is bound to the exact tunnel
/// generation that was dispatched. Callers handling a live connection should
/// use [`record_proxy_upgrade_traffic_success_for_generation`].
pub(crate) async fn record_proxy_upgrade_traffic_success(
data: &GatewayDataState,
node_id: &str,
) -> Result<bool, DataLayerError> {
let Some(node) = data.find_proxy_node(node_id).await? else {
return Ok(false);
};
record_proxy_upgrade_traffic_success_for_generation(data, node_id, &node.tunnel_generation)
.await
}
pub(crate) async fn record_proxy_upgrade_traffic_success_for_generation(
data: &GatewayDataState,
node_id: &str,
tunnel_generation: &str,
) -> Result<bool, DataLayerError> {
if !data.has_system_config_store() {
return Ok(false);
@@ -916,13 +970,17 @@ pub(crate) async fn record_proxy_upgrade_traffic_success(
let Some(mut plan) = load_proxy_upgrade_rollout_plan(data).await? else {
return Ok(false);
};
let Some(node) = data.find_proxy_node(node_id).await? else {
return Ok(false);
};
if tunnel_generation.is_empty() || node.tunnel_generation != tunnel_generation {
return Ok(false);
}
let now = now_unix_secs();
let Some(tracked) = plan
.tracked_nodes
.iter_mut()
.find(|tracked| tracked.node_id == node_id)
else {
let Some(tracked) = plan.tracked_nodes.iter_mut().find(|tracked| {
tracked.node_id == node_id && tracked.tunnel_generation == tunnel_generation
}) else {
return Ok(false);
};
let Some(version_confirmed_at_unix_secs) = tracked.version_confirmed_at_unix_secs else {
@@ -7,17 +7,19 @@ use tracing::{info, warn};
use crate::data::GatewayDataState;
use crate::{AppState, GatewayError};
use super::cleanup_runs::cleanup_data_layer_error_category;
use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_once,
cleanup_expired_gemini_file_mappings_once, cleanup_proxy_node_metrics_once,
cleanup_request_candidates_once, cleanup_stale_pending_requests_once,
cleanup_stale_proxy_nodes_once, collect_proxy_upgrade_rollout_probes, now_unix_secs,
perform_db_maintenance_once, perform_manual_usage_cleanup_once, perform_provider_checkin_once,
perform_automatic_usage_cleanup_once, perform_db_maintenance_once,
perform_manual_usage_cleanup_once, perform_provider_checkin_once,
perform_stats_aggregation_once, perform_stats_hourly_aggregation_once,
perform_usage_cleanup_once, perform_wallet_daily_usage_aggregation_once,
record_admin_cleanup_run, record_completed_cleanup_run, record_failed_cleanup_run,
record_proxy_upgrade_traffic_success, summarize_database_pool, AdminCleanupRunRecord,
ManualUsageCleanupOptions,
perform_wallet_daily_usage_aggregation_once, record_admin_cleanup_run,
record_completed_cleanup_run, record_failed_cleanup_run,
record_proxy_upgrade_traffic_success_for_generation, summarize_database_pool,
AdminCleanupRunRecord, ManualUsageCleanupOptions,
};
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
@@ -120,14 +122,19 @@ pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(
.await
{
Ok(status) if (200..300).contains(&status) => {
let _ = record_proxy_upgrade_traffic_success(&state.data, &probe.node_id).await?;
let _ = record_proxy_upgrade_traffic_success_for_generation(
&state.data,
&probe.node_id,
&probe.tunnel_generation,
)
.await?;
probe_recorded = true;
info!(
event_name = "proxy_upgrade_rollout_probe_succeeded",
log_type = "ops",
worker = "proxy_upgrade_rollout",
node_id = %probe.node_id,
url = %probe.url,
probe_origin = %crate::handlers::shared::security_log_url_origin(&probe.url),
status,
"gateway confirmed proxy upgrade health probe"
);
@@ -138,7 +145,7 @@ pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(
log_type = "ops",
worker = "proxy_upgrade_rollout",
node_id = %probe.node_id,
url = %probe.url,
probe_origin = %crate::handlers::shared::security_log_url_origin(&probe.url),
status,
"gateway proxy upgrade health probe returned non-success status"
);
@@ -149,7 +156,7 @@ pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(
log_type = "ops",
worker = "proxy_upgrade_rollout",
node_id = %probe.node_id,
url = %probe.url,
probe_origin = %crate::handlers::shared::security_log_url_origin(&probe.url),
error = %error,
"gateway proxy upgrade health probe failed"
);
@@ -237,7 +244,7 @@ pub(super) async fn run_stats_aggregation_once(
pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
let started_at_unix_secs = now_unix_secs();
let started_at = Instant::now();
let summary = match perform_usage_cleanup_once(data).await {
let summary = match perform_automatic_usage_cleanup_once(data).await {
Ok(summary) => summary,
Err(err) => {
record_failed_cleanup_run(
@@ -265,6 +272,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
"cost_reservations_deleted": summary.cost_reservations_deleted,
"request_admissions_deleted": summary.request_admissions_deleted,
}),
format!(
"请求记录自动清理完成,影响 {} 项",
@@ -275,6 +284,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted)
.saturating_add(summary.cost_reservations_deleted)
.saturating_add(summary.request_admissions_deleted)
),
)
.await;
@@ -284,6 +295,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
|| summary.header_cleaned > 0
|| summary.keys_cleaned > 0
|| summary.records_deleted > 0
|| summary.cost_reservations_deleted > 0
|| summary.request_admissions_deleted > 0
{
info!(
event_name = "usage_cleanup_completed",
@@ -295,6 +308,8 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
header_cleaned = summary.header_cleaned,
keys_cleaned = summary.keys_cleaned,
records_deleted = summary.records_deleted,
cost_reservations_deleted = summary.cost_reservations_deleted,
request_admissions_deleted = summary.request_admissions_deleted,
"gateway finished usage cleanup"
);
}
@@ -384,6 +399,8 @@ pub(crate) async fn run_manual_usage_cleanup_once(
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
"cost_reservations_deleted": summary.cost_reservations_deleted,
"request_admissions_deleted": summary.request_admissions_deleted,
"mode": options.mode.as_str(),
"requested_older_than_days": options.requested_older_than_days,
"targets": options.targets,
@@ -401,6 +418,8 @@ pub(crate) async fn run_manual_usage_cleanup_once(
requested_older_than_days = options.requested_older_than_days,
actor_user_id = actor_user_id.as_deref(),
total_affected = total,
cost_reservations_deleted = summary.cost_reservations_deleted,
request_admissions_deleted = summary.request_admissions_deleted,
"gateway finished manual usage cleanup"
);
Ok(summary)
@@ -476,7 +495,7 @@ async fn run_manual_usage_cleanup_task(
None,
actor_user_id.as_deref(),
),
error: Some(err.to_string()),
error: Some(cleanup_data_layer_error_category(&err).to_string()),
};
if let Err(record_err) = record_admin_cleanup_run(&data, record).await {
warn!(error = %record_err, "failed to record manual usage cleanup failure");
@@ -496,6 +515,8 @@ fn usage_cleanup_total(
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted)
.saturating_add(summary.cost_reservations_deleted)
.saturating_add(summary.request_admissions_deleted)
}
fn manual_usage_cleanup_start_message(options: ManualUsageCleanupOptions) -> String {
@@ -549,6 +570,8 @@ fn manual_usage_cleanup_progress_summary(
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
"cost_reservations_deleted": summary.cost_reservations_deleted,
"request_admissions_deleted": summary.request_admissions_deleted,
"total": usage_cleanup_total(summary),
"actor_user_id": actor_user_id,
})
@@ -564,7 +587,7 @@ impl std::fmt::Display for ManualUsageCleanupError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::AlreadyRunning => f.write_str("a usage cleanup run is already in progress"),
Self::DataLayer(err) => write!(f, "{err}"),
Self::DataLayer(_) => f.write_str("usage cleanup data operation failed"),
}
}
}
@@ -597,20 +620,49 @@ pub(super) fn run_pool_monitor_once(data: &GatewayDataState) {
);
}
pub(super) async fn run_pending_cleanup_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
let summary = cleanup_stale_pending_requests_once(data).await?;
if summary.failed > 0 || summary.recovered > 0 {
pub(super) async fn run_pending_cleanup_once(app: &AppState) -> Result<(), DataLayerError> {
let data = &app.data;
let cleanup_result = cleanup_stale_pending_requests_once(data).await;
let referral_result = if data.has_referral_data_backend() {
Some(app.reconcile_referral_rewards_once().await)
} else {
None
};
let summary = cleanup_result.as_ref().copied().unwrap_or_default();
let referral_summary = referral_result
.as_ref()
.and_then(|result| result.as_ref().ok().copied())
.unwrap_or_default();
if summary.failed > 0
|| summary.recovered > 0
|| cleanup_result.is_err()
|| referral_result.as_ref().is_some_and(Result::is_err)
|| referral_summary.order_attempted > 0
|| referral_summary.reward_attempted > 0
|| referral_summary.reversal_attempted > 0
{
info!(
event_name = "pending_cleanup_completed",
log_type = "ops",
worker = "pending_cleanup",
failed = summary.failed,
recovered = summary.recovered,
referral_order_attempted = referral_summary.order_attempted,
referral_order_repaired = referral_summary.order_repaired,
referral_reward_attempted = referral_summary.reward_attempted,
referral_reward_applied = referral_summary.reward_applied,
referral_reversal_attempted = referral_summary.reversal_attempted,
referral_reversal_applied = referral_summary.reversal_applied,
referral_deferred = referral_summary.deferred,
"gateway cleaned stale pending and streaming requests"
);
}
if let Err(error) = cleanup_result {
return Err(error);
}
if let Some(Err(error)) = referral_result {
return Err(DataLayerError::UnexpectedValue(error.into_message()));
}
Ok(())
}
@@ -2,10 +2,16 @@ use std::collections::{HashSet, VecDeque};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::proxy_nodes::{
bucket_start_unix_secs, InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeWriteRepository, StoredProxyNode,
};
use aether_data::repository::settlement::{
InMemorySettlementRepository, ReserveUsagePolicyCostInput, ReserveUsagePolicyRequestInput,
SettlementWriteRepository, UsagePolicyCostWindow, UsagePolicyRequestWindow,
};
use aether_runtime::bounded_queue;
use axum::extract::ws::Message;
use chrono::{DateTime, Utc};
@@ -18,7 +24,8 @@ use super::{
cleanup_proxy_node_metrics_once, cleanup_stale_proxy_nodes_once, inspect_proxy_upgrade_rollout,
next_daily_run_after, next_db_maintenance_run_after, next_stats_aggregation_run_after,
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
pending_cleanup_timeout_minutes, perform_automatic_usage_cleanup_once,
perform_manual_usage_cleanup_once, plan_pending_cleanup_batch, provider_checkin_schedule,
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
@@ -167,6 +174,29 @@ fn sample_connected_proxy_node(
)
}
const ACTIVE_PROBE_TUNNEL_TEST_PSK: &str = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
const ACTIVE_PROBE_TUNNEL_TEST_GENERATION: &str = "active-probe-test-generation-1";
fn active_probe_tunnel_metadata(version: &str) -> serde_json::Value {
json!({
"version": version,
"tunnel_security": {
"mode": TUNNEL_SECURITY_NON_TLS_REQUIRED,
"encryption_key": ACTIVE_PROBE_TUNNEL_TEST_PSK,
}
})
}
async fn recv_tunnel_test_frame(
proxy_rx: &mut aether_runtime::BoundedQueueReceiver<Message>,
description: &str,
) -> Message {
tokio::time::timeout(Duration::from_secs(5), proxy_rx.recv())
.await
.unwrap_or_else(|_| panic!("timed out waiting for {description}"))
.unwrap_or_else(|| panic!("proxy channel closed before {description}"))
}
#[tokio::test]
async fn stale_proxy_node_cleanup_marks_timed_out_tunnel_offline() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
@@ -230,6 +260,7 @@ async fn proxy_upgrade_rollout_advances_next_wave_after_version_health_confirmat
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-alpha".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
@@ -267,6 +298,7 @@ async fn proxy_upgrade_rollout_advances_next_wave_after_version_health_confirmat
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-zeta".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
@@ -378,6 +410,7 @@ async fn proxy_upgrade_rollout_blocks_next_wave_after_post_upgrade_transport_err
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-alpha".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
@@ -406,6 +439,7 @@ async fn proxy_upgrade_rollout_blocks_next_wave_after_post_upgrade_transport_err
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-alpha".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
@@ -436,9 +470,10 @@ async fn proxy_upgrade_rollout_blocks_next_wave_after_post_upgrade_transport_err
#[tokio::test]
async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_confirmation() {
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000);
let mut alpha = sample_connected_proxy_node("node-alpha", 30, 1_800_000_000)
.with_tunnel_generation(ACTIVE_PROBE_TUNNEL_TEST_GENERATION.to_string());
alpha.name = "alpha".to_string();
alpha.proxy_metadata = Some(json!({"version": "1.0.0"}));
alpha.proxy_metadata = Some(active_probe_tunnel_metadata("1.0.0"));
alpha.remote_config = None;
alpha.config_version = 0;
@@ -450,7 +485,8 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository))
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new())
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data.clone());
@@ -472,6 +508,7 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-alpha".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: None,
active_connections: Some(2),
total_requests_delta: Some(1),
@@ -479,7 +516,7 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
failed_requests_delta: Some(0),
dns_failures_delta: Some(0),
stream_errors_delta: Some(0),
proxy_metadata: Some(json!({"version": "2.0.0"})),
proxy_metadata: Some(active_probe_tunnel_metadata("2.0.0")),
proxy_version: Some("2.0.0".to_string()),
})
.await
@@ -488,9 +525,8 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
let tunnel_state = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_state
.hub
.register_proxy(Arc::new(crate::tunnel::TunnelProxyConn::new(
tunnel_state.hub.register_proxy(Arc::new(
crate::tunnel::TunnelProxyConn::new(
700,
"node-alpha".to_string(),
"Node Alpha".to_string(),
@@ -498,14 +534,18 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
proxy_close_tx,
16,
2,
)));
)
.with_tunnel_generation(ACTIVE_PROBE_TUNNEL_TEST_GENERATION.to_string())
.with_authenticated_key(ACTIVE_PROBE_TUNNEL_TEST_PSK.to_string()),
));
let responder_hub = tunnel_state.hub.clone();
let responder = tokio::spawn(async move {
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_headers =
match recv_tunnel_test_frame(&mut proxy_rx, "probe headers frame").await {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header = crate::tunnel::tunnel_protocol::FrameHeader::parse(&request_headers)
.expect("probe request headers should parse");
assert_eq!(
@@ -513,7 +553,7 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
crate::tunnel::tunnel_protocol::REQUEST_HEADERS
);
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
let request_body = match recv_tunnel_test_frame(&mut proxy_rx, "probe body frame").await {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
@@ -594,6 +634,106 @@ async fn spawn_usage_cleanup_worker_skips_when_usage_writer_unavailable() {
assert!(spawn_usage_cleanup_worker(state).is_none());
}
#[tokio::test]
async fn spawn_usage_cleanup_worker_starts_for_settlement_writer_only() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::disabled().with_settlement_writer_for_tests(Arc::new(
InMemorySettlementRepository::default(),
)),
);
let handle = spawn_usage_cleanup_worker(state)
.expect("settlement ledger cleanup should have a worker without a usage writer");
handle.abort();
}
#[tokio::test]
async fn automatic_usage_cleanup_collects_ledgers_when_usage_cleanup_is_disabled() {
let repository = Arc::new(InMemorySettlementRepository::default());
repository
.reserve_usage_policy_request(ReserveUsagePolicyRequestInput {
request_id: "expired-request".to_string(),
subject_id: "user-1".to_string(),
event_token: "expired-event".to_string(),
admitted_at_unix_secs: 1,
retain_until_unix_secs: 2,
windows: vec![UsagePolicyRequestWindow {
starts_at_unix_secs: 0,
ends_at_unix_secs: 2,
limit_requests: 10,
}],
})
.await
.expect("request admission should be seeded");
repository
.reserve_usage_policy_cost(ReserveUsagePolicyCostInput {
request_id: "expired-cost".to_string(),
subject_id: "user-1".to_string(),
reservation_token: "expired-reservation".to_string(),
admitted_at_unix_secs: 1,
reserved_cost_units: 1,
reservation_expires_at_unix_secs: 2,
retain_until_unix_secs: 2,
windows: vec![UsagePolicyCostWindow {
window_id: "expired-window".to_string(),
starts_at_unix_secs: 0,
ends_at_unix_secs: 2,
limit_cost_units: 10,
}],
})
.await
.expect("cost reservation should be seeded");
let data = GatewayDataState::disabled()
.with_settlement_writer_for_tests(repository)
.with_system_config_values_for_tests([
("enable_auto_cleanup".to_string(), json!(false)),
("cleanup_batch_size".to_string(), json!(1)),
]);
let summary = perform_automatic_usage_cleanup_once(&data)
.await
.expect("automatic ledger cleanup should succeed");
assert_eq!(summary.cost_reservations_deleted, 1);
assert_eq!(summary.request_admissions_deleted, 1);
}
#[tokio::test]
async fn manual_usage_cleanup_does_not_collect_policy_ledgers() {
let repository = Arc::new(InMemorySettlementRepository::default());
repository
.reserve_usage_policy_request(ReserveUsagePolicyRequestInput {
request_id: "manual-request".to_string(),
subject_id: "user-1".to_string(),
event_token: "manual-event".to_string(),
admitted_at_unix_secs: 1,
retain_until_unix_secs: 2,
windows: vec![UsagePolicyRequestWindow {
starts_at_unix_secs: 0,
ends_at_unix_secs: 2,
limit_requests: 10,
}],
})
.await
.expect("request admission should be seeded");
let data = GatewayDataState::disabled().with_settlement_writer_for_tests(repository.clone());
let summary =
perform_manual_usage_cleanup_once(&data, super::ManualUsageCleanupOptions::policy())
.await
.expect("manual usage cleanup should succeed");
assert_eq!(summary.request_admissions_deleted, 0);
assert_eq!(
repository
.cleanup_usage_policy_request_admissions(u64::MAX, 10)
.await
.expect("direct cleanup should find the retained admission"),
1
);
}
#[tokio::test]
async fn spawn_wallet_daily_usage_aggregation_worker_skips_when_wallet_daily_usage_backend_unavailable(
) {
@@ -851,6 +991,7 @@ async fn proxy_node_metrics_cleanup_deletes_expired_buckets_in_batches() {
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: node_id.to_string(),
expected_tunnel_generation: None,
heartbeat_interval: Some(30),
active_connections: Some(i32::try_from(idx + 1).unwrap()),
total_requests_delta: None,
@@ -915,6 +1056,7 @@ async fn proxy_node_metrics_cleanup_respects_auto_cleanup_toggle() {
repository
.apply_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-metrics-disabled".to_string(),
expected_tunnel_generation: None,
heartbeat_interval: Some(30),
active_connections: Some(1),
total_requests_delta: None,
@@ -67,6 +67,14 @@ pub(super) async fn perform_usage_cleanup_once(
perform_usage_cleanup_once_with_override(data, None, true).await
}
pub(super) async fn perform_automatic_usage_cleanup_once(
data: &GatewayDataState,
) -> Result<UsageCleanupSummary, DataLayerError> {
let mut summary = perform_usage_cleanup_once(data).await?;
cleanup_usage_policy_ledgers(data, &mut summary).await?;
Ok(summary)
}
pub(super) async fn perform_usage_cleanup_once_with_override(
data: &GatewayDataState,
override_older_than: Option<chrono::Duration>,
@@ -128,6 +136,41 @@ async fn perform_usage_cleanup_once_with_options(
.await
}
async fn cleanup_usage_policy_ledgers(
data: &GatewayDataState,
summary: &mut UsageCleanupSummary,
) -> Result<(), DataLayerError> {
if !data.has_settlement_writer() {
return Ok(());
}
let settings = usage_cleanup_settings(data).await?;
let now_unix_secs = u64::try_from(Utc::now().timestamp()).unwrap_or(0);
let cleanup_batch_size = settings.batch_size.max(1);
// Bound each maintenance run while still draining faster than one day's normal growth.
for _ in 0..32 {
let deleted = data
.cleanup_usage_policy_cost_reservations(now_unix_secs, cleanup_batch_size)
.await?;
summary.cost_reservations_deleted =
summary.cost_reservations_deleted.saturating_add(deleted);
if deleted < cleanup_batch_size {
break;
}
}
for _ in 0..32 {
let deleted = data
.cleanup_usage_policy_request_admissions(now_unix_secs, cleanup_batch_size)
.await?;
summary.request_admissions_deleted =
summary.request_admissions_deleted.saturating_add(deleted);
if deleted < cleanup_batch_size {
break;
}
}
Ok(())
}
pub(crate) async fn preview_manual_usage_cleanup(
data: &GatewayDataState,
options: ManualUsageCleanupOptions,
@@ -279,7 +279,7 @@ pub(crate) fn spawn_stats_aggregation_worker(app: AppState) -> Option<tokio::tas
}
pub(crate) fn spawn_usage_cleanup_worker(app: AppState) -> Option<tokio::task::JoinHandle<()>> {
if !app.data.has_usage_writer() {
if !app.data.has_usage_writer() && !app.data.has_settlement_writer() {
return None;
}
@@ -595,7 +595,7 @@ pub(crate) fn spawn_gemini_file_mapping_cleanup_worker(
}
pub(crate) fn spawn_pending_cleanup_worker(app: AppState) -> Option<tokio::task::JoinHandle<()>> {
if !app.data.has_usage_writer() {
if !app.data.has_usage_writer() && !app.data.has_referral_data_backend() {
return None;
}
@@ -603,8 +603,8 @@ pub(crate) fn spawn_pending_cleanup_worker(app: AppState) -> Option<tokio::task:
app,
crate::task_runtime::TASK_KEY_PENDING_CLEANUP,
|app| async move {
let data = app.data;
if let Err(err) = run_pending_cleanup_once(&data).await {
let data = app.data.clone();
if let Err(err) = run_pending_cleanup_once(&app).await {
log_maintenance_worker_failure("pending_cleanup", "startup", &err);
}
let mut interval = tokio::time::interval(PENDING_CLEANUP_INTERVAL);
@@ -617,7 +617,7 @@ pub(crate) fn spawn_pending_cleanup_worker(app: AppState) -> Option<tokio::task:
{
continue;
}
if let Err(err) = run_pending_cleanup_once(&data).await {
if let Err(err) = run_pending_cleanup_once(&app).await {
log_maintenance_worker_failure("pending_cleanup", "tick", &err);
}
}