feat(proxy): record tunnel stability metrics

This commit is contained in:
fawney19
2026-05-08 22:03:17 +08:00
parent 9a84a6ff6c
commit a703acd1fe
81 changed files with 7229 additions and 1110 deletions

View File

@@ -92,6 +92,7 @@ fn maybe_build_local_openai_image_sync_finalize_response(
payload.report_kind.as_str(),
payload.status_code,
payload.report_context.as_ref(),
payload.body_json.as_ref(),
payload.body_base64.as_deref(),
)
.map_err(GatewayError::from)?

View File

@@ -75,6 +75,18 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
}
impl<'a> LocalExecutionCandidateAttemptSource<'a> {
pub(crate) fn from_static_attempts_for_image_bridge(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> Self {
let mut items = VecDeque::new();
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: VecDeque::from(attempts),
});
}
Self { items }
}
pub(crate) async fn next_attempt(&mut self) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = self.items.front_mut()?;

View File

@@ -398,15 +398,11 @@ pub(crate) fn candidate_auth_channel_skip_reason(
) -> Option<&'static str> {
let request_auth_channel = normalize_request_auth_channel(request_auth_channel?)?;
let upstream_auth_channel = resolve_transport_request_auth_channel(transport)?;
if request_auth_channel == upstream_auth_channel
|| provider_runtime_policy(&transport.provider.provider_type)
.allow_auth_channel_mismatch_by_default
|| allow_auth_channel_mismatch_for_format(transport)
{
None
} else {
Some("auth_channel_mismatch")
if request_auth_channel == upstream_auth_channel {
return None;
}
auth_channel_mismatch_is_explicitly_disabled_for_format(transport)
.then_some("auth_channel_mismatch")
}
fn normalize_request_auth_channel(value: &str) -> Option<&'static str> {
@@ -452,19 +448,22 @@ fn resolve_transport_auth_type_for_endpoint_format(
.unwrap_or(default_auth_type)
}
fn allow_auth_channel_mismatch_for_format(transport: &GatewayProviderTransportSnapshot) -> bool {
fn auth_channel_mismatch_is_explicitly_disabled_for_format(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
transport
let Some(items) = transport
.key
.allow_auth_channel_mismatch_formats
.as_ref()
.and_then(serde_json::Value::as_array)
.is_some_and(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
})
else {
return false;
};
!items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
}
pub(crate) async fn read_candidate_transport_snapshot(
@@ -585,11 +584,11 @@ mod tests {
}
#[test]
fn auth_channel_gate_skips_mismatched_raw_secret_auth() {
fn auth_channel_gate_allows_mismatched_raw_secret_auth_by_default() {
let transport = sample_transport("bearer");
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}
@@ -603,6 +602,16 @@ mod tests {
);
}
#[test]
fn auth_channel_gate_blocks_explicitly_disabled_mismatch_format() {
let mut transport = sample_transport("bearer");
transport.key.allow_auth_channel_mismatch_formats = Some(json!(["openai:responses"]));
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
);
}
#[test]
fn auth_channel_gate_treats_cli_oauth_provider_as_bearer_like() {
let mut transport = sample_transport("oauth");
@@ -613,7 +622,7 @@ mod tests {
);
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
None
);
}

View File

@@ -159,6 +159,40 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
preselect_local_execution_candidates_for_api_formats_with_serving(
state,
client_api_format,
requested_model,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
use_api_format_alias_match,
key_mode,
candidate_api_formats,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_serving(
state: PlannerAppState<'_>,
client_api_format: &str,
requested_model: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
) -> Result<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
GatewayError,
> {
let mut model_directive_enabled_api_formats = BTreeSet::new();
for api_format in &candidate_api_formats {
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(

View File

@@ -7,6 +7,7 @@ use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
@@ -284,7 +285,13 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
return Ok(None);
};
match build_passthrough_sync_plan_from_decision(self.parts, payload) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_sync_plan_from_decision(self.parts, self.body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(self.parts, payload)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -320,7 +327,13 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
return Ok(None);
};
match build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_stream_plan_from_decision(self.parts, self.body_json, payload)
} else {
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
};
match built {
Ok(value) => Ok(value),
Err(err) => {
warn!(
@@ -508,7 +521,13 @@ async fn build_local_sync_plan_and_reports(
continue;
};
match build_passthrough_sync_plan_from_decision(parts, payload) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_sync_plan_from_decision(parts, body_json, payload)
} else {
build_passthrough_sync_plan_from_decision(parts, payload)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
@@ -578,7 +597,13 @@ async fn build_local_stream_plan_and_reports(
continue;
};
match build_standard_stream_plan_from_decision(parts, body_json, payload, false) {
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
let built = if provider_api_format == "gemini:generate_content" {
build_gemini_stream_plan_from_decision(parts, body_json, payload)
} else {
build_standard_stream_plan_from_decision(parts, body_json, payload, false)
};
match built {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {

View File

@@ -10,7 +10,7 @@ use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
};
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
use crate::{AiExecutionDecision, AppState};
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
use super::request::resolve_local_openai_image_candidate_payload_parts;
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
@@ -47,8 +47,12 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
} = attempt;
let candidate = eligible.candidate;
let transport = resolved.transport;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let provider_api_format = resolved.provider_api_format.clone();
let needs_conversion = provider_api_format != spec_metadata.api_format;
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
provider_api_format.as_str(),
);
let proxy = planner_state
.app()
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
@@ -84,41 +88,47 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
.get("stream")
.and_then(serde_json::Value::as_bool)
.unwrap_or(spec_metadata.require_streaming);
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion: false,
extra_fields,
});
let report_context = append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
attempt_identity,
model: &resolved.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
needs_conversion,
extra_fields,
}),
execution_strategy,
conversion_mode,
spec_metadata.api_format,
provider_api_format.as_str(),
);
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
@@ -137,7 +147,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_method: Some(parts.method.to_string()),
auth_header: Some(resolved.auth_header),
auth_value: Some(resolved.auth_value),
provider_api_format: spec_metadata.api_format.to_string(),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: resolved.requested_model,
mapped_model: resolved.mapped_model,

View File

@@ -9,14 +9,17 @@ use crate::ai_serving::planner::candidate_preparation::{
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::transport::{
build_openai_image_headers, build_openai_image_upstream_url,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput,
build_standard_provider_request_headers, openai_image_transport_unsupported_reason,
resolve_openai_image_auth, ProviderOpenAiImageHeadersInput,
StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
build_chatgpt_web_image_request_body, build_openai_image_provider_request_body,
default_model_for_openai_image_operation, normalize_openai_image_request,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, PlannerAppState,
build_chatgpt_web_image_request_body,
build_gemini_image_request_body_from_openai_image_request,
build_openai_image_provider_request_body, default_model_for_openai_image_operation,
normalize_openai_image_request, request_conversion_direct_auth, CandidateFailureDiagnostic,
GatewayProviderTransportSnapshot, PlannerAppState, RequestConversionKind,
};
use crate::AppState;
@@ -35,6 +38,7 @@ pub(super) struct LocalOpenAiImageCandidatePayloadParts {
pub(super) auth_value: String,
pub(super) requested_model: String,
pub(super) mapped_model: String,
pub(super) provider_api_format: String,
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) upstream_url: String,
@@ -54,6 +58,21 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if provider_api_format == "gemini:generate_content" {
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
state,
parts,
body_json,
body_base64,
trace_id,
input,
attempt,
spec,
)
.await;
}
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, spec_metadata.api_format)
@@ -70,7 +89,6 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
@@ -216,6 +234,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
auth_value,
requested_model,
mapped_model,
provider_api_format: spec_metadata.api_format.to_string(),
provider_request_headers,
provider_request_body,
upstream_url,
@@ -223,6 +242,190 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
})
}
async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
body_json: &Value,
body_base64: Option<&str>,
trace_id: &str,
input: &LocalOpenAiImageDecisionInput,
attempt: &LocalOpenAiImageCandidateAttempt,
spec: LocalOpenAiImageSpec,
) -> Option<LocalOpenAiImageCandidatePayloadParts> {
let spec_metadata = local_openai_image_spec_metadata(spec);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = "gemini:generate_content";
let prepared_candidate = match prepare_header_authenticated_candidate(
PlannerAppState::new(state),
transport,
candidate,
request_conversion_direct_auth(transport, RequestConversionKind::ToGeminiStandard),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "openai_image_to_gemini_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_openai_image_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(normalized_request) = normalize_openai_image_request(parts, body_json, body_base64)
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_request_normalize",
),
)
.await;
return None;
};
let Some(mut converted) = build_gemini_image_request_body_from_openai_image_request(
&normalized_request,
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_request_body",
),
)
.await;
return None;
};
converted.body_json =
match crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers(
converted.body_json,
transport.endpoint.body_rules.as_ref(),
body_json,
&parts.headers,
) {
Some(body) => body,
None => {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
CandidateFailureDiagnostic::provider_request_body_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_body_rules",
),
)
.await;
return None;
}
};
let upstream_is_stream = spec_metadata.require_streaming;
let Some(upstream_url) = crate::ai_serving::planner::standard::build_standard_upstream_url(
parts,
transport,
&converted.mapped_model,
provider_api_format,
upstream_is_stream,
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_url",
),
)
.await;
return None;
};
let Some(resolved_headers) =
build_standard_provider_request_headers(StandardProviderRequestHeadersInput {
transport,
provider_api_format,
same_format: false,
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
extra_headers: &BTreeMap::new(),
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
upstream_is_stream,
})
else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
spec_metadata.api_format,
provider_api_format,
"openai_image_to_gemini_headers",
),
)
.await;
return None;
};
Some(LocalOpenAiImageCandidatePayloadParts {
transport: Arc::clone(transport),
auth_header: resolved_headers.auth_header,
auth_value: resolved_headers.auth_value,
requested_model: converted.requested_model,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_headers: resolved_headers.headers,
provider_request_body: converted.body_json,
upstream_url,
input_summary: converted.summary_json,
})
}
fn chatgpt_web_image_internal_url(base_url: &str) -> String {
let base_url = base_url.trim().trim_end_matches('/');
let base_url = if base_url.is_empty() {

View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_serving::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
@@ -19,9 +20,9 @@ use crate::ai_serving::planner::materialization_policy::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::{
extract_pool_sticky_session_token, resolve_local_decision_execution_runtime_auth_context,
CandidateFailureDiagnostic, ExecutionRuntimeAuthContext, GatewayControlDecision,
PlannerAppState,
extract_pool_sticky_session_token, request_candidate_api_formats,
resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
@@ -88,52 +89,69 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
api_format: &str,
decision_kind: &str,
) -> Option<Vec<LocalOpenAiImageCandidateAttempt>> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return None;
let candidate_api_formats = image_candidate_api_formats(api_format);
let mut attempts = Vec::new();
for candidate_api_format in candidate_api_formats {
let matches_client_format = candidate_api_format == api_format;
let planner_state = PlannerAppState::new(state);
let (mut candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
};
if !matches_client_format {
candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
}
};
attempts.extend(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
);
}
Some(
materialize_local_openai_image_candidate_attempts(
planner_state,
trace_id,
input,
body_json,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
api_format,
)
.await,
)
Some(attempts)
}
pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
@@ -145,29 +163,56 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
decision_kind: &str,
) -> Result<Option<(LocalOpenAiImageCandidateAttemptSource<'a>, usize)>, GatewayError> {
let planner_state = PlannerAppState::new(state);
let (candidates, preselection_skipped) = match planner_state
.list_selectable_candidates_with_skip_reasons(
api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok(candidates) => candidates,
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
return Ok(None);
let mut candidates = Vec::new();
let mut preselection_skipped = Vec::new();
for candidate_api_format in image_candidate_api_formats(api_format) {
let matches_client_format = candidate_api_format == api_format;
match planner_state
.list_selectable_candidates_with_skip_reasons(
candidate_api_format,
&input.requested_model,
false,
input.required_capabilities.as_ref(),
matches_client_format.then_some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await
{
Ok((mut format_candidates, mut format_skipped)) => {
if !matches_client_format {
format_candidates.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
candidate,
false,
)
});
format_skipped.retain(|candidate| {
auth_snapshot_allows_cross_format_candidate(
&input.auth_snapshot,
&input.requested_model,
&candidate.candidate,
false,
)
});
}
candidates.append(&mut format_candidates);
preselection_skipped.append(&mut format_skipped);
}
Err(err) => {
warn!(
trace_id = %trace_id,
decision_kind,
api_format = candidate_api_format,
error = ?err,
"gateway local openai image decision scheduler selection failed"
);
continue;
}
}
};
}
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
@@ -198,23 +243,34 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
extra_data: None,
})
.collect(),
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -254,23 +310,34 @@ async fn materialize_local_openai_image_candidate_attempts(
persistence_policy,
candidates,
preselection_skipped,
LocalCandidateResolutionMode::Standard,
LocalCandidateResolutionMode::WithoutTransportPairGate,
|eligible| {
Some(build_local_execution_candidate_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: api_format,
provider_api_format: eligible.provider_api_format.as_str(),
client_api_format: api_format,
extra_fields: serde_json::Map::new(),
},
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
skipped_candidate.extra_data =
Some(build_local_execution_candidate_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
api_format,
provider_api_format.as_str(),
api_format,
serde_json::Map::new(),
));
@@ -282,6 +349,14 @@ async fn materialize_local_openai_image_candidate_attempts(
outcome.attempts
}
fn image_candidate_api_formats(api_format: &str) -> Vec<&'static str> {
if api_format.trim().eq_ignore_ascii_case("openai:image") {
vec!["openai:image", "gemini:generate_content"]
} else {
request_candidate_api_formats(api_format, false)
}
}
pub(super) async fn mark_skipped_local_openai_image_candidate(
state: &AppState,
input: &LocalOpenAiImageDecisionInput,

View File

@@ -11,6 +11,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_source::{
preselect_local_execution_candidates_for_api_formats_with_serving,
preselect_local_execution_candidates_with_serving, LocalCandidatePreselectionKeyMode,
};
use crate::ai_serving::planner::common::extract_requested_model_from_request;
@@ -23,7 +24,8 @@ use crate::ai_serving::planner::materialization_policy::{
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::{
ai_local_execution_contract_for_formats, extract_pool_sticky_session_token,
resolve_local_decision_execution_runtime_auth_context, GatewayControlDecision, PlannerAppState,
gemini_request_is_image_generation, resolve_local_decision_execution_runtime_auth_context,
GatewayControlDecision, PlannerAppState,
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::{AppState, GatewayError};
@@ -86,6 +88,8 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
spec: LocalStandardSpec,
) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
let spec_metadata = local_standard_spec_metadata(spec);
let is_gemini_image_bridge = spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
@@ -105,6 +109,17 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
let (candidates, skipped_candidates) = maybe_append_gemini_image_openai_image_preselection(
state,
planner_state,
trace_id,
input,
body_json,
spec,
preselection.candidates,
preselection.skipped_candidates,
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
planner_state,
trace_id,
@@ -116,9 +131,13 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
candidates,
skipped_candidates,
if is_gemini_image_bridge {
LocalCandidateResolutionMode::WithoutTransportPairGate
} else {
LocalCandidateResolutionMode::Standard
},
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
@@ -188,7 +207,18 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
Ok(
if spec_metadata.api_format == "gemini:generate_content"
&& gemini_request_is_image_generation(body_json)
{
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, input, body_json, spec)
.await?;
let source =
LocalExecutionCandidateAttemptSource::from_static_attempts_for_image_bridge(attempts);
return Ok((source, candidate_count));
}
let (source, candidate_count) =
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
@@ -253,6 +283,51 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
skipped_candidate
},
)
.await,
)
.await;
Ok((source, candidate_count))
}
#[allow(clippy::too_many_arguments)]
async fn maybe_append_gemini_image_openai_image_preselection(
state: &AppState,
planner_state: PlannerAppState<'_>,
trace_id: &str,
input: &LocalStandardDecisionInput,
body_json: &serde_json::Value,
spec: LocalStandardSpec,
mut candidates: Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
mut skipped_candidates: Vec<
crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate,
>,
) -> Result<
(
Vec<aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate>,
Vec<crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate>,
),
GatewayError,
> {
let spec_metadata = local_standard_spec_metadata(spec);
if spec_metadata.api_format != "gemini:generate_content"
|| !gemini_request_is_image_generation(body_json)
{
return Ok((candidates, skipped_candidates));
}
let image_preselection = preselect_local_execution_candidates_for_api_formats_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
input.client_session_affinity.as_ref(),
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
vec!["openai:image".to_string()],
)
.await?;
candidates.extend(image_preselection.candidates);
skipped_candidates.extend(image_preselection.skipped_candidates);
let _ = (state, trace_id);
Ok((candidates, skipped_candidates))
}

View File

@@ -18,10 +18,13 @@ use crate::ai_serving::transport::kiro::{
KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_kiro_cross_format_upstream_url, build_standard_provider_request_headers,
StandardProviderRequestHeadersInput,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
openai_image_transport_unsupported_reason, resolve_openai_image_auth,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
};
use crate::ai_serving::{
build_openai_image_request_body_from_gemini_image_request, gemini_request_is_image_generation,
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use crate::AppState;
@@ -59,6 +62,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = attempt.eligible.provider_api_format.as_str();
if spec_metadata.api_format == "gemini:generate_content"
&& provider_api_format == "openai:image"
&& gemini_request_is_image_generation(body_json)
{
return resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state, parts, trace_id, body_json, input, attempt,
)
.await;
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
@@ -331,6 +343,141 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
})
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
) -> Option<LocalStandardCandidatePayloadParts> {
let client_api_format = "gemini:generate_content";
let provider_api_format = "openai:image";
let planner_state = crate::ai_serving::PlannerAppState::new(state);
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
if let Some(skip_reason) =
openai_image_transport_unsupported_reason(transport, provider_api_format)
{
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
let prepared_candidate = match prepare_header_authenticated_candidate(
planner_state,
transport,
candidate,
resolve_openai_image_auth(transport),
OauthPreparationContext {
trace_id,
api_format: provider_api_format,
operation: "gemini_image_to_openai_image_candidate_request",
},
)
.await
{
Ok(prepared) => prepared,
Err(skip_reason) => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
skip_reason,
)
.await;
return None;
}
};
let Some(converted) = build_openai_image_request_body_from_gemini_image_request(
body_json,
parts.uri.path(),
&prepared_candidate.mapped_model,
) else {
mark_skipped_local_standard_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
body_json,
client_api_format,
provider_api_format,
),
)
.await;
return None;
};
let upstream_is_stream = true;
let upstream_url = build_openai_image_upstream_url(transport, None);
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &converted.body_json,
original_request_body: body_json,
})
else {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"gemini_image_to_openai_image_headers",
),
)
.await;
return None;
};
apply_codex_openai_responses_special_headers(
&mut provider_request_headers,
&converted.body_json,
&parts.headers,
transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
Some(LocalStandardCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: converted.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: converted.body_json,
provider_request_headers,
upstream_url,
upstream_is_stream,
envelope_name: None,
transport: Arc::clone(transport),
})
}
#[allow(clippy::too_many_arguments)]
async fn build_kiro_cross_format_payload_parts(
state: &AppState,

View File

@@ -10,14 +10,20 @@ pub(crate) use aether_ai_formats::api::{
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives,
build_generated_tool_call_id, build_kiro_final_message_sse_events,
build_kiro_initial_sse_events, build_kiro_stream_error_sse_events,
build_local_openai_chat_request_body,
build_gemini_image_request_body_from_openai_image_request,
build_gemini_image_response_from_openai_image_response,
build_gemini_image_response_from_openai_responses_image_response, build_generated_tool_call_id,
build_kiro_final_message_sse_events, build_kiro_initial_sse_events,
build_kiro_stream_error_sse_events, build_local_openai_chat_request_body,
build_local_openai_chat_request_body_with_model_directives,
build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives,
build_local_success_background_report, build_local_success_conversion_background_report,
build_openai_image_provider_request_body, build_openai_responses_response,
build_openai_image_provider_body_from_response_stream_sync_body,
build_openai_image_provider_request_body,
build_openai_image_request_body_from_gemini_image_request,
build_openai_image_response_from_gemini_response,
build_openai_image_response_from_response_stream_sync_body, build_openai_responses_response,
build_standard_request_body, build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives,
@@ -37,10 +43,10 @@ pub(crate) use aether_ai_formats::api::{
encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider,
implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
is_matching_stream_http_request, is_matching_stream_request, is_openai_image_stream_request,
is_openai_responses_family_format, is_openai_responses_format, kiro_crc32,
map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output,
gemini_request_is_image_generation, implicit_sync_finalize_report_kind,
is_core_error_finalize_kind, is_matching_stream_http_request, is_matching_stream_request,
is_openai_image_stream_request, is_openai_responses_family_format, is_openai_responses_format,
kiro_crc32, map_claude_stop_reason, map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_gemini_budget, maybe_bridge_standard_sync_json_to_stream,
maybe_build_ai_surface_stream_rewriter,
maybe_build_openai_chat_cross_format_sync_product_from_normalized_payload,
@@ -72,6 +78,7 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request, stream_body_contains_error_event,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
@@ -79,15 +86,16 @@ pub(crate) use aether_ai_formats::api::{
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState,
ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter,
GeminiProviderState, KiroToClaudeCliStreamState, LocalCoreSyncErrorKind, LocalGeminiFilesSpec,
LocalOpenAiImageSpec, LocalOpenAiResponsesSpec, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily, LocalVideoCreateSpec,
NormalizedOpenAiImageRequest, OpenAIChatClientEmitter, OpenAIChatProviderState,
OpenAIResponsesClientEmitter, OpenAIResponsesProviderState, OpenAiImageOperation,
OpenAiImageResponseFormat, OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct,
ProviderAdaptationDescriptor, ProviderAdaptationSurface, ProviderPrivateStreamNormalizer,
RequestConversionKind, StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState,
LocalCoreSyncErrorKind, LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
LocalStandardSourceMode, LocalStandardSpec, LocalSyncReportParts, LocalVideoCreateFamily,
LocalVideoCreateSpec, NormalizedOpenAiImageRequest, OpenAIChatClientEmitter,
OpenAIChatProviderState, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
OpenAiImageOperation, OpenAiImageRequestForGemini, OpenAiImageResponseFormat,
OpenAiImageStreamState, OpenAiImageSyncFinalizeProduct, ProviderAdaptationDescriptor,
ProviderAdaptationSurface, ProviderPrivateStreamNormalizer, RequestConversionKind,
StandardCrossFormatSyncProduct, StandardSyncFinalizeNormalizedProduct,
StreamingStandardFormatMatrix, SyncChatResponseConversionKind, SyncCliResponseConversionKind,
SyncToStreamBridgeOutcome, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, CLAUDE_CHAT_STREAM_PLAN_KIND,
CLAUDE_CHAT_STREAM_SUCCESS_REPORT_KIND, CLAUDE_CHAT_SYNC_ERROR_REPORT_KIND,

View File

@@ -113,6 +113,34 @@ pub(super) fn classify_admin_operations_family_route(
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& matches!(
normalized_path,
"/api/admin/proxy-nodes/metrics/fleet" | "/api/admin/proxy-nodes/metrics/fleet/"
)
{
Some(classified(
"admin_proxy",
"proxy_nodes_manage",
"list_fleet_metrics",
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/proxy-nodes/")
&& normalized_path_no_trailing.ends_with("/metrics")
&& normalized_path_no_trailing["/api/admin/proxy-nodes/".len()..]
.split('/')
.count()
== 2
{
Some(classified(
"admin_proxy",
"proxy_nodes_manage",
"list_node_metrics",
"admin:proxy_nodes",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/proxy-nodes/")
&& normalized_path_no_trailing["/api/admin/proxy-nodes/".len()..]

View File

@@ -100,6 +100,24 @@ fn classifies_admin_proxy_nodes_detail_as_admin_proxy_route() {
);
}
#[test]
fn classifies_admin_proxy_node_metrics_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(
http::Method::GET,
"/api/admin/proxy-nodes/node-1/metrics?from=1700000000&to=1700003600&step=1m",
"list_node_metrics",
);
}
#[test]
fn classifies_admin_proxy_fleet_metrics_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(
http::Method::GET,
"/api/admin/proxy-nodes/metrics/fleet?from=1700000000&to=1700003600&step=1m",
"list_fleet_metrics",
);
}
#[test]
fn classifies_admin_proxy_nodes_delete_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(

View File

@@ -6,7 +6,8 @@ use super::{
RegenerateManagementTokenSecret, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
StoredLdapModuleConfig, StoredManagementToken, StoredManagementTokenListPage,
StoredManagementTokenWithUser, StoredOAuthProviderConfig, StoredOAuthProviderModuleConfig,
StoredProxyNode, StoredProxyNodeEvent, StoredUserAuthRecord, StoredUserOAuthLinkSummary,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, StoredUserAuthRecord, StoredUserOAuthLinkSummary,
StoredUserPreferenceRecord, StoredUserSessionRecord, StoredWalletSnapshot,
UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
};
@@ -979,6 +980,56 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_proxy_node_events_filtered(
&self,
node_id: &str,
query: &super::ProxyNodeEventQuery,
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_node_events_filtered(node_id, query)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_proxy_node_metrics(
&self,
node_id: &str,
step: super::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyNodeMetricsBucket>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_proxy_fleet_metrics(
&self,
step: super::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
match &self.proxy_node_reader {
Some(repository) => {
repository
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn register_proxy_node(
&self,
mutation: &ProxyNodeRegistrationMutation,
@@ -1018,6 +1069,21 @@ impl GatewayDataState {
}
}
pub(crate) async fn cleanup_proxy_node_metrics(
&self,
retain_1m_from_unix_secs: u64,
retain_1h_from_unix_secs: u64,
) -> Result<super::ProxyNodeMetricsCleanupSummary, DataLayerError> {
match &self.proxy_node_writer {
Some(repository) => {
repository
.cleanup_proxy_node_metrics(retain_1m_from_unix_secs, retain_1h_from_unix_secs)
.await
}
None => Ok(super::ProxyNodeMetricsCleanupSummary::default()),
}
}
pub(crate) async fn apply_proxy_node_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,

View File

@@ -42,10 +42,12 @@ use aether_data::repository::oauth_providers::{
UpsertOAuthProviderConfigRecord,
};
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyNode, StoredProxyNodeEvent,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket,
};
pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
use aether_data::repository::users::{

View File

@@ -104,11 +104,7 @@ pub(crate) fn normalize_allow_auth_channel_mismatch_formats(
normalized.push(serde_json::Value::String(canonical));
}
}
if normalized.is_empty() {
Ok(None)
} else {
Ok(Some(serde_json::Value::Array(normalized)))
}
Ok(Some(serde_json::Value::Array(normalized)))
}
pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String> {
@@ -179,9 +175,9 @@ fn normalize_json_like_object(
#[cfg(test)]
mod tests {
use super::{
normalize_api_format_json_object_keys, normalize_api_format_list, normalize_auth_type,
normalize_auth_type_by_format, normalize_pool_advanced_config,
normalize_provider_type_input, validate_vertex_api_formats,
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_pool_advanced_config, normalize_provider_type_input, validate_vertex_api_formats,
};
use serde_json::json;
@@ -280,6 +276,36 @@ mod tests {
);
}
#[test]
fn normalize_allow_auth_channel_mismatch_formats_preserves_explicit_empty_array() {
assert_eq!(
normalize_allow_auth_channel_mismatch_formats(
Some(Vec::new()),
"allow_auth_channel_mismatch_formats",
&["claude:messages".to_string()],
)
.expect("empty array should normalize"),
Some(json!([]))
);
}
#[test]
fn normalize_allow_auth_channel_mismatch_formats_normalizes_and_dedupes_values() {
assert_eq!(
normalize_allow_auth_channel_mismatch_formats(
Some(vec![
"claude:messages".to_string(),
"CLAUDE:MESSAGES".to_string(),
" claude:messages ".to_string(),
]),
"allow_auth_channel_mismatch_formats",
&["claude:messages".to_string()],
)
.expect("format list should normalize"),
Some(json!(["claude:messages"]))
);
}
#[test]
fn validate_vertex_api_formats_uses_canonical_message_formats() {
assert!(validate_vertex_api_formats(

View File

@@ -3,8 +3,10 @@ use crate::handlers::shared::unix_secs_to_rfc3339;
use crate::maintenance::{inspect_proxy_upgrade_rollout, ProxyUpgradeRolloutStatus};
use crate::GatewayError;
use aether_admin::system::{
build_admin_proxy_node_event_payload, build_admin_proxy_node_events_payload_response,
build_admin_proxy_node_payload, build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_fleet_metrics_payload_response, build_admin_proxy_node_event_payload,
build_admin_proxy_node_events_payload_response,
build_admin_proxy_node_metrics_payload_response, build_admin_proxy_node_payload,
build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_nodes_invalid_status_response, build_admin_proxy_nodes_list_payload_response,
build_admin_proxy_nodes_not_found_response,
};
@@ -84,6 +86,30 @@ impl<'a> AdminAppState<'a> {
pub(crate) async fn build_admin_proxy_node_events_response(
&self,
node_id: &str,
query: &aether_data::repository::proxy_nodes::ProxyNodeEventQuery,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
return Ok(build_admin_proxy_nodes_data_unavailable_response());
}
if self.find_proxy_node(node_id).await?.is_none() {
return Ok(build_admin_proxy_nodes_not_found_response());
}
let items = self
.app
.list_proxy_node_events_filtered(node_id, query)
.await?
.into_iter()
.map(|event| build_admin_proxy_node_event_payload(&event))
.collect::<Vec<_>>();
Ok(build_admin_proxy_node_events_payload_response(items))
}
pub(crate) async fn build_admin_proxy_node_metrics_response(
&self,
node_id: &str,
step: aether_data::repository::proxy_nodes::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
@@ -93,12 +119,37 @@ impl<'a> AdminAppState<'a> {
return Ok(build_admin_proxy_nodes_not_found_response());
}
let items = self
.list_proxy_node_events(node_id, limit)
.await?
.into_iter()
.map(|event| build_admin_proxy_node_event_payload(&event))
.collect::<Vec<_>>();
Ok(build_admin_proxy_node_events_payload_response(items))
.app
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await?;
Ok(build_admin_proxy_node_metrics_payload_response(
step,
from_unix_secs,
to_unix_secs,
items,
))
}
pub(crate) async fn build_admin_proxy_fleet_metrics_response(
&self,
step: aether_data::repository::proxy_nodes::ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Response<Body>, GatewayError> {
if !self.has_proxy_node_reader() {
return Ok(build_admin_proxy_nodes_data_unavailable_response());
}
let items = self
.app
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await?;
Ok(build_admin_proxy_fleet_metrics_payload_response(
step,
from_unix_secs,
to_unix_secs,
items,
))
}
pub(crate) async fn unregister_proxy_node(

View File

@@ -10,12 +10,14 @@ use crate::maintenance::{
};
use crate::GatewayError;
use aether_admin::system::{
admin_proxy_node_event_node_id_from_path, build_admin_proxy_node_payload,
build_admin_proxy_nodes_data_unavailable_response, build_admin_proxy_nodes_not_found_response,
admin_proxy_node_event_node_id_from_path, admin_proxy_node_metrics_node_id_from_path,
build_admin_proxy_node_payload, build_admin_proxy_nodes_data_unavailable_response,
build_admin_proxy_nodes_not_found_response,
};
use aether_contracts::tunnel::{
TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use aether_data::repository::proxy_nodes::{ProxyNodeEventQuery, ProxyNodeMetricsStep};
use axum::{
body::{Body, Bytes},
http,
@@ -148,6 +150,9 @@ const DEFAULT_PROXY_CONNECTIVITY_PROBE_URL: &str = "https://www.cloudflare.com/c
const PROXY_CONNECTIVITY_TIMEOUT_SECS: u64 = 10;
const TUNNEL_RELAY_ENVELOPE_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
const MAX_PROXY_CONNECTIVITY_RESPONSE_BYTES: usize = 64 * 1024;
const PROXY_NODE_METRICS_MAX_POINTS: usize = 50_000;
const PROXY_NODE_METRICS_1M_MAX_WINDOW_SECS: u64 = 30 * 24 * 60 * 60;
const PROXY_NODE_METRICS_1H_MAX_WINDOW_SECS: u64 = 365 * 24 * 60 * 60;
#[cfg(test)]
fn manual_proxy_connectivity_probe_url_override() -> &'static std::sync::RwLock<Option<String>> {
@@ -244,13 +249,53 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
let limit = query_param_value(request_context.query_string(), "limit")
.and_then(|value| value.parse::<usize>().ok())
.filter(|value| *value > 0 && *value <= 200)
.unwrap_or(50);
let query = match parse_proxy_node_event_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_node_events_response(node_id, limit)
.build_admin_proxy_node_events_response(node_id, &query)
.await?,
));
}
if decision.route_kind.as_deref() == Some("list_node_metrics")
&& request_context.method() == http::Method::GET
{
let Some(node_id) = admin_proxy_node_metrics_node_id_from_path(request_context.path())
else {
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
let (step, from_unix_secs, to_unix_secs, limit) =
match parse_proxy_node_metrics_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_node_metrics_response(
node_id,
step,
from_unix_secs,
to_unix_secs,
limit,
)
.await?,
));
}
if decision.route_kind.as_deref() == Some("list_fleet_metrics")
&& request_context.method() == http::Method::GET
{
let (step, from_unix_secs, to_unix_secs, limit) =
match parse_proxy_node_metrics_query(request_context.query_string()) {
Ok(query) => query,
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
state
.build_admin_proxy_fleet_metrics_response(step, from_unix_secs, to_unix_secs, limit)
.await?,
));
}
@@ -1976,6 +2021,86 @@ fn validate_optional_object(value: Option<&Value>, field: &str) -> Result<(), Re
Ok(())
}
fn parse_proxy_node_event_query(
query: Option<&str>,
) -> Result<ProxyNodeEventQuery, Response<Body>> {
let limit = query_param_value(query, "limit")
.map(|value| parse_query_u64("limit", &value))
.transpose()?
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0 && *value <= 200)
.unwrap_or(50);
let from_unix_secs = query_param_value(query, "from")
.map(|value| parse_query_u64("from", &value))
.transpose()?;
let to_unix_secs = query_param_value(query, "to")
.map(|value| parse_query_u64("to", &value))
.transpose()?;
if from_unix_secs
.zip(to_unix_secs)
.is_some_and(|(from, to)| from > to)
{
return Err(bad_request_response("from 不能大于 to"));
}
let event_type = query_param_value(query, "event_type")
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
Ok(ProxyNodeEventQuery {
limit,
from_unix_secs,
to_unix_secs,
event_type,
})
}
fn parse_proxy_node_metrics_query(
query: Option<&str>,
) -> Result<(ProxyNodeMetricsStep, u64, u64, usize), Response<Body>> {
let step = match query_param_value(query, "step")
.unwrap_or_else(|| "1m".to_string())
.trim()
{
"1m" => ProxyNodeMetricsStep::OneMinute,
"1h" => ProxyNodeMetricsStep::OneHour,
_ => return Err(bad_request_response("step 仅支持 1m 或 1h")),
};
let from_unix_secs = query_param_value(query, "from")
.ok_or_else(|| bad_request_response("from 为必填 Unix 秒时间戳"))?;
let from_unix_secs = parse_query_u64("from", &from_unix_secs)?;
let to_unix_secs = query_param_value(query, "to")
.ok_or_else(|| bad_request_response("to 为必填 Unix 秒时间戳"))?;
let to_unix_secs = parse_query_u64("to", &to_unix_secs)?;
if from_unix_secs > to_unix_secs {
return Err(bad_request_response("from 不能大于 to"));
}
let window_secs = to_unix_secs.saturating_sub(from_unix_secs);
let max_window_secs = match step {
ProxyNodeMetricsStep::OneMinute => PROXY_NODE_METRICS_1M_MAX_WINDOW_SECS,
ProxyNodeMetricsStep::OneHour => PROXY_NODE_METRICS_1H_MAX_WINDOW_SECS,
};
if window_secs > max_window_secs {
return Err(bad_request_response(match step {
ProxyNodeMetricsStep::OneMinute => "1m 最大查询窗口为 30 天",
ProxyNodeMetricsStep::OneHour => "1h 最大查询窗口为 365 天",
}));
}
let points = window_secs / step.bucket_size_secs() + 1;
let limit = usize::try_from(points)
.ok()
.filter(|value| *value > 0 && *value <= PROXY_NODE_METRICS_MAX_POINTS)
.ok_or_else(|| bad_request_response("查询点数过多"))?;
Ok((step, from_unix_secs, to_unix_secs, limit))
}
fn parse_query_u64(field: &str, value: &str) -> Result<u64, Response<Body>> {
value
.parse::<u64>()
.map_err(|_| bad_request_response(format!("{field} 必须是非负 Unix 秒时间戳")))
}
fn bad_request_response(detail: impl Into<String>) -> Response<Body> {
(
http::StatusCode::BAD_REQUEST,

View File

@@ -11,13 +11,14 @@ pub(crate) use runtime::{
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_request_candidate_cleanup_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
};

View File

@@ -20,6 +20,8 @@ mod pending_cleanup;
mod pool_quota_probe;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
mod proxy_node_metrics_cleanup;
#[path = "runtime/proxy_node_staleness.rs"]
mod proxy_node_staleness;
#[path = "runtime/proxy_upgrade_rollout.rs"]
@@ -59,6 +61,7 @@ pub(crate) use pool_quota_probe::{
PoolQuotaProbeWorkerConfig,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;
use proxy_upgrade_rollout::*;
pub(crate) use proxy_upgrade_rollout::{
@@ -90,6 +93,8 @@ const AUDIT_LOG_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const GEMINI_FILE_MAPPING_CLEANUP_INTERVAL: Duration = Duration::from_secs(60 * 60);
const PENDING_CLEANUP_INTERVAL: Duration = Duration::from_secs(5 * 60);
const PROXY_NODE_STALE_SWEEP_INTERVAL: Duration = Duration::from_secs(5);
const PROXY_NODE_METRICS_CLEANUP_HOUR: u32 = 2;
const PROXY_NODE_METRICS_CLEANUP_MINUTE: u32 = 10;
const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;

View File

@@ -0,0 +1,20 @@
use aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary;
use aether_data_contracts::DataLayerError;
use crate::data::GatewayDataState;
use super::now_unix_secs;
const PROXY_NODE_METRICS_1M_RETENTION_SECS: u64 = 30 * 24 * 60 * 60;
const PROXY_NODE_METRICS_1H_RETENTION_SECS: u64 = 180 * 24 * 60 * 60;
pub(super) async fn cleanup_proxy_node_metrics_once(
data: &GatewayDataState,
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
let now = now_unix_secs();
data.cleanup_proxy_node_metrics(
now.saturating_sub(PROXY_NODE_METRICS_1M_RETENTION_SECS),
now.saturating_sub(PROXY_NODE_METRICS_1H_RETENTION_SECS),
)
.await
}

View File

@@ -6,10 +6,10 @@ use crate::{AppState, GatewayError};
use super::{
advance_proxy_upgrade_rollout_once, cleanup_audit_logs_once,
cleanup_expired_gemini_file_mappings_once, cleanup_request_candidates_once,
cleanup_stale_pending_requests_once, cleanup_stale_proxy_nodes_once,
collect_proxy_upgrade_rollout_probes, perform_db_maintenance_once,
perform_provider_checkin_once, perform_stats_aggregation_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,
perform_db_maintenance_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_proxy_upgrade_traffic_success,
summarize_database_pool,
@@ -61,6 +61,23 @@ pub(super) async fn run_proxy_node_stale_cleanup_once(
Ok(())
}
pub(super) async fn run_proxy_node_metrics_cleanup_once(
data: &GatewayDataState,
) -> Result<(), DataLayerError> {
let summary = cleanup_proxy_node_metrics_once(data).await?;
if summary.deleted_1m_rows > 0 || summary.deleted_1h_rows > 0 {
info!(
event_name = "proxy_node_metrics_cleanup_completed",
log_type = "ops",
worker = "proxy_node_metrics_cleanup",
deleted_1m_rows = summary.deleted_1m_rows,
deleted_1h_rows = summary.deleted_1h_rows,
"gateway deleted expired proxy node metrics buckets"
);
}
Ok(())
}
pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(), DataLayerError> {
let mut summary = advance_proxy_upgrade_rollout_once(&state.data).await?;
let probes = collect_proxy_upgrade_rollout_probes(&state.data).await?;

View File

@@ -12,12 +12,14 @@ use super::{
maintenance_timezone, parse_hhmm_time, perform_oauth_token_refresh_once,
provider_checkin_schedule, run_audit_cleanup_once, run_db_maintenance_once,
run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once, run_pool_monitor_once,
run_provider_checkin_once, run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_request_candidate_cleanup_once, run_stats_aggregation_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_STALE_SWEEP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE, PROXY_NODE_STALE_SWEEP_INTERVAL,
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
@@ -292,6 +294,30 @@ pub(crate) fn spawn_proxy_node_stale_cleanup_worker(
}))
}
pub(crate) fn spawn_proxy_node_metrics_cleanup_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if !data.has_proxy_node_writer() {
return None;
}
let timezone = maintenance_timezone();
Some(tokio::spawn(async move {
loop {
tokio::time::sleep(duration_until_next_daily_run(
Utc::now(),
timezone,
PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE,
))
.await;
if let Err(err) = run_proxy_node_metrics_cleanup_once(&data).await {
log_maintenance_worker_failure("proxy_node_metrics_cleanup", "tick", &err);
}
}
}))
}
pub(crate) fn spawn_proxy_upgrade_rollout_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {

View File

@@ -4,8 +4,10 @@ use std::sync::Mutex as StdMutex;
use std::time::Duration;
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsStep, ProxyNodeTrafficMutation,
ProxyNodeTunnelStatusMutation, StoredProxyFleetMetricsBucket, StoredProxyNode,
StoredProxyNodeEvent, StoredProxyNodeMetricsBucket,
};
use aether_http::{build_http_client, HttpClientConfig};
use aether_runtime::{
@@ -46,6 +48,7 @@ use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
use crate::maintenance::spawn_proxy_upgrade_rollout_worker;
use crate::maintenance::spawn_request_candidate_cleanup_worker;
@@ -593,6 +596,44 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_events_filtered(
&self,
node_id: &str,
query: &ProxyNodeEventQuery,
) -> Result<Vec<StoredProxyNodeEvent>, GatewayError> {
self.data
.list_proxy_node_events_filtered(node_id, query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_node_metrics(
&self,
node_id: &str,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyNodeMetricsBucket>, GatewayError> {
self.data
.list_proxy_node_metrics(node_id, step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_proxy_fleet_metrics(
&self,
step: ProxyNodeMetricsStep,
from_unix_secs: u64,
to_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredProxyFleetMetricsBucket>, GatewayError> {
self.data
.list_proxy_fleet_metrics(step, from_unix_secs, to_unix_secs, limit)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn register_proxy_node(
&self,
mutation: &aether_data::repository::proxy_nodes::ProxyNodeRegistrationMutation,
@@ -630,6 +671,18 @@ impl AppState {
.map_err(|err| std::io::Error::other(err.to_string()))
}
pub(crate) async fn cleanup_proxy_node_metrics(
&self,
retain_1m_from_unix_secs: u64,
retain_1h_from_unix_secs: u64,
) -> Result<aether_data::repository::proxy_nodes::ProxyNodeMetricsCleanupSummary, GatewayError>
{
self.data
.cleanup_proxy_node_metrics(retain_1m_from_unix_secs, retain_1h_from_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn apply_proxy_node_heartbeat(
&self,
mutation: &ProxyNodeHeartbeatMutation,
@@ -979,6 +1032,9 @@ impl AppState {
if let Some(handle) = spawn_proxy_node_stale_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_node_metrics_cleanup_worker(self.data.clone()) {
tasks.push(handle);
}
if let Some(handle) = spawn_proxy_upgrade_rollout_worker(self.clone()) {
tasks.push(handle);
}

View File

@@ -19,6 +19,721 @@ use aether_data_contracts::repository::provider_catalog::{
use base64::Engine as _;
use sha2::{Digest, Sha256};
#[tokio::test]
async fn gateway_converts_openai_image_sync_to_gemini_image_provider() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
auth_header_value: String,
has_model_field: bool,
prompt: String,
response_modalities: Vec<String>,
image_size: String,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai", "google"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(serde_json::json!(["openai", "google"])),
Some(serde_json::json!(["openai:image"])),
Some(serde_json::json!(["gpt-image-2"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-gemini-image-bridge-1".to_string(),
provider_name: "google".to_string(),
provider_type: "google".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-gemini-image-bridge-1".to_string(),
endpoint_api_format: "gemini:generate_content".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-gemini-image-bridge-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["gemini:generate_content".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({
"gemini:generate_content": 1
})),
model_id: "model-gemini-image-bridge-1".to_string(),
global_model_id: "global-model-gemini-image-bridge-1".to_string(),
global_model_name: "gpt-image-2".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gemini-2.5-flash-image-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-flash-image-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-gemini-image-bridge-1".to_string(),
"google".to_string(),
Some("https://generativelanguage.googleapis.com".to_string()),
"google".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-gemini-image-bridge-1".to_string(),
"provider-gemini-image-bridge-1".to_string(),
"gemini:generate_content".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
Some(serde_json::json!([
{"action":"drop","path":"model"}
])),
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-gemini-image-bridge-1".to_string(),
"provider-gemini-image-bridge-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:generate_content"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-gemini-image")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:generate_content": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let body_json = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or_else(|| json!({}));
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
auth_header_value: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-key"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
has_model_field: body_json.get("model").is_some(),
prompt: body_json
.get("contents")
.and_then(|value| value.get(0))
.and_then(|value| value.get("parts"))
.and_then(|value| value.get(0))
.and_then(|value| value.get("text"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
response_modalities: body_json
.get("generationConfig")
.and_then(|value| value.get("responseModalities"))
.and_then(|value| value.as_array())
.map(|items| {
items
.iter()
.filter_map(|item| item.as_str().map(ToOwned::to_owned))
.collect()
})
.unwrap_or_default(),
image_size: body_json
.get("generationConfig")
.and_then(|value| value.get("imageSize"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "trace-openai-image-to-gemini-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"modelVersion": "gemini-2.5-flash-image-upstream",
"usageMetadata": {
"promptTokenCount": 11,
"candidatesTokenCount": 22,
"totalTokenCount": 33
},
"candidates": [{
"content": {
"role": "model",
"parts": [
{"text": "revised kite prompt"},
{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}
]
},
"finishReason": "STOP"
}]
}
},
"telemetry": {
"elapsed_ms": 37
}
}))
}
}),
);
let client_api_key = "sk-client-openai-image-to-gemini";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
sample_auth_snapshot(
"key-openai-image-client-bridge-1",
"user-openai-image-bridge-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::new(InMemoryRequestCandidateRepository::default()),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/images/generations"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
format!("Bearer {client_api_key}"),
)
.header(TRACE_ID_HEADER, "trace-openai-image-to-gemini-123")
.body(
"{\"model\":\"gpt-image-2\",\"prompt\":\"Draw a red kite\",\"size\":\"1024x1024\",\"response_format\":\"b64_json\"}",
)
.send()
.await
.expect("request should succeed");
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK, "{response_body}");
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(response_json["data"][0]["b64_json"], "aGVsbG8=");
assert_eq!(
response_json["data"][0]["revised_prompt"],
"revised kite prompt"
);
assert_eq!(response_json["model"], "gemini-2.5-flash-image-upstream");
assert_eq!(response_json["usage"]["input_tokens"], 11);
assert_eq!(response_json["usage"]["output_tokens"], 22);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime sync should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-openai-image-to-gemini-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-upstream:generateContent"
);
assert_eq!(
seen_execution_runtime_request.auth_header_value,
"sk-upstream-gemini-image"
);
assert!(!seen_execution_runtime_request.has_model_field);
assert_eq!(seen_execution_runtime_request.prompt, "Draw a red kite");
assert_eq!(
seen_execution_runtime_request.response_modalities,
vec!["TEXT".to_string(), "IMAGE".to_string()]
);
assert_eq!(seen_execution_runtime_request.image_size, "1024x1024");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_converts_gemini_image_sync_to_openai_image_provider() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeSyncRequest {
trace_id: String,
url: String,
authorization: String,
model: String,
action: String,
prompt: String,
image_url: String,
request_stream: bool,
}
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["gemini", "openai"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-image"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800_i64),
Some(serde_json::json!(["gemini", "openai"])),
Some(serde_json::json!(["gemini:generate_content"])),
Some(serde_json::json!(["gemini-image"])),
)
.expect("auth snapshot should build")
}
fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-image-bridge-1".to_string(),
provider_name: "openai".to_string(),
provider_type: "openai".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: "endpoint-openai-image-bridge-1".to_string(),
endpoint_api_format: "openai:image".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("image".to_string()),
endpoint_is_active: true,
key_id: "key-openai-image-bridge-1".to_string(),
key_name: "prod".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:image".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 5,
key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})),
model_id: "model-openai-image-bridge-1".to_string(),
global_model_id: "global-model-openai-image-bridge-1".to_string(),
global_model_name: "gemini-image".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-image-2-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-image-2-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:image".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-image-bridge-1".to_string(),
"openai".to_string(),
Some("https://api.openai.com".to_string()),
"openai".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-image-bridge-1".to_string(),
"provider-openai-image-bridge-1".to_string(),
"openai:image".to_string(),
Some("openai".to_string()),
Some("image".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-image-bridge-1".to_string(),
"provider-openai-image-bridge-1".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:image"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-openai-image")
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"openai:image": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let body_json = payload
.get("body")
.and_then(|value| value.get("json_body"))
.cloned()
.unwrap_or_else(|| json!({}));
let content = body_json
.get("input")
.and_then(|value| value.get(0))
.and_then(|value| value.get("content"))
.cloned()
.unwrap_or_else(|| json!([]));
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
trace_id: parts
.headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
model: body_json
.get("model")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
action: body_json
.get("tools")
.and_then(|value| value.get(0))
.and_then(|value| value.get("action"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
prompt: content
.as_array()
.into_iter()
.flatten()
.find(|item| {
item.get("type").and_then(|value| value.as_str()) == Some("input_text")
})
.and_then(|item| item.get("text"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
image_url: content
.as_array()
.into_iter()
.flatten()
.find(|item| {
item.get("type").and_then(|value| value.as_str()) == Some("input_image")
})
.and_then(|item| item.get("image_url"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
request_stream: body_json
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(true),
});
Json(json!({
"request_id": "trace-gemini-image-to-openai-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "resp_img_bridge_123",
"object": "response",
"model": "gpt-image-2-upstream",
"status": "completed",
"usage": {
"input_tokens": 3,
"output_tokens": 4,
"total_tokens": 7
},
"output": [{
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"revised_prompt": "converted gemini prompt",
"result": "aGVsbG8="
}]
}
},
"telemetry": {
"elapsed_ms": 43
}
}))
}
}),
);
let client_api_key = "client-gemini-image-to-openai";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(client_api_key)),
sample_auth_snapshot(
"key-gemini-image-client-bridge-1",
"user-gemini-image-bridge-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::new(InMemoryRequestCandidateRepository::default()),
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/models/gemini-image:generateContent?key={client_api_key}"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-gemini-image-to-openai-123")
.body(
"{\"generationConfig\":{\"responseModalities\":[\"TEXT\",\"IMAGE\"]},\"contents\":[{\"role\":\"user\",\"parts\":[{\"text\":\"Change the background\"},{\"inlineData\":{\"mimeType\":\"image/png\",\"data\":\"aGVsbG8=\"}}]}]}",
)
.send()
.await
.expect("request should succeed");
let response_status = response.status();
let response_body = response.text().await.expect("body should read");
assert_eq!(response_status, StatusCode::OK, "{response_body}");
let response_json: serde_json::Value =
serde_json::from_str(&response_body).expect("body should parse");
assert_eq!(response_json["modelVersion"], "gpt-image-2-upstream");
assert_eq!(
response_json["candidates"][0]["content"]["parts"][0]["text"],
"converted gemini prompt"
);
assert_eq!(
response_json["candidates"][0]["content"]["parts"][1]["inlineData"]["mimeType"],
"image/png"
);
assert_eq!(
response_json["candidates"][0]["content"]["parts"][1]["inlineData"]["data"],
"aGVsbG8="
);
assert_eq!(response_json["usageMetadata"]["promptTokenCount"], 3);
assert_eq!(response_json["usageMetadata"]["candidatesTokenCount"], 4);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime sync should be captured");
assert_eq!(
seen_execution_runtime_request.trace_id,
"trace-gemini-image-to-openai-123"
);
assert_eq!(
seen_execution_runtime_request.url,
"https://api.openai.com/v1/responses"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-upstream-openai-image"
);
assert_eq!(seen_execution_runtime_request.model, "gpt-image-2-upstream");
assert_eq!(seen_execution_runtime_request.action, "edit");
assert_eq!(
seen_execution_runtime_request.prompt,
"Change the background"
);
assert_eq!(
seen_execution_runtime_request.image_url,
"data:image/png;base64,aGVsbG8="
);
assert!(!seen_execution_runtime_request.request_stream);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_executes_codex_image_sync_via_local_decision_gate_after_oauth_refresh() {
#[derive(Debug, Clone)]

View File

@@ -1,4 +1,5 @@
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::repository::management_tokens::InMemoryManagementTokenRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -1442,6 +1443,7 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
node_id: "node-1".to_string(),
event_type: "connected".to_string(),
detail: Some("older".to_string()),
event_metadata: None,
created_at_unix_ms: Some(1_710_000_000),
},
StoredProxyNodeEvent {
@@ -1449,6 +1451,7 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
node_id: "node-1".to_string(),
event_type: "disconnected".to_string(),
detail: Some("newer".to_string()),
event_metadata: None,
created_at_unix_ms: Some(1_710_000_100),
},
],
@@ -1490,6 +1493,150 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reports_proxy_node_metrics_and_filters_events_locally() {
let proxy_node_repository =
Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-1",
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
proxy_node_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system time should be after epoch")
.as_secs();
let heartbeat_response = client
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-1",
"heartbeat_id": 91,
"heartbeat_interval": 30,
"active_connections": 7,
"proxy_metadata": {
"tunnel_metrics": {
"connect_errors": 3,
"disconnects": 1,
"error_events_total": 1,
"ws_in_bytes": 1000,
"ws_out_bytes": 2000,
"ws_in_frames": 10,
"ws_out_frames": 20,
"heartbeat_rtt_last_ms": 42
},
"recent_tunnel_errors": [{
"timestamp_unix_secs": now_unix_secs,
"category": "tcp_connect_timeout",
"message": "tunnel TCP connect timeout"
}]
},
"proxy_version": "2.0.0"
}))
.send()
.await
.expect("heartbeat request should succeed");
assert_eq!(heartbeat_response.status(), StatusCode::OK);
let from = now_unix_secs.saturating_sub(120);
let to = now_unix_secs.saturating_add(120);
let metrics_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=1m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("metrics request should succeed");
assert_eq!(metrics_response.status(), StatusCode::OK);
let metrics_payload: serde_json::Value = metrics_response
.json()
.await
.expect("metrics json should parse");
assert_eq!(metrics_payload["step"], "1m");
assert_eq!(metrics_payload["summary"]["samples"], 1);
assert_eq!(metrics_payload["summary"]["uptime_samples"], 1);
assert_eq!(metrics_payload["summary"]["active_connections_max"], 7);
assert_eq!(metrics_payload["summary"]["heartbeat_rtt_ms_avg"], 42.0);
assert_eq!(metrics_payload["summary"]["connect_errors_delta"], 3);
assert_eq!(metrics_payload["summary"]["ws_out_frames_delta"], 20);
let metric_items = metrics_payload["items"]
.as_array()
.expect("metrics items should be array");
assert_eq!(metric_items.len(), 1);
assert_eq!(metric_items[0]["node_id"], "node-1");
assert!(metric_items[0]["bucket_start"].is_string());
let fleet_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/metrics/fleet?from={from}&to={to}&step=1m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("fleet metrics request should succeed");
assert_eq!(fleet_response.status(), StatusCode::OK);
let fleet_payload: serde_json::Value = fleet_response
.json()
.await
.expect("fleet json should parse");
assert_eq!(fleet_payload["summary"]["samples"], 1);
assert_eq!(fleet_payload["summary"]["error_events_delta"], 1);
let events_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/events?from={from}&to={to}&event_type=tunnel_err"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("events request should succeed");
assert_eq!(events_response.status(), StatusCode::OK);
let events_payload: serde_json::Value = events_response
.json()
.await
.expect("events json should parse");
let event_items = events_payload["items"]
.as_array()
.expect("event items should be array");
assert_eq!(event_items.len(), 1);
assert_eq!(event_items[0]["event_type"], "tunnel_err");
assert_eq!(
event_items[0]["event_metadata"]["category"],
"tcp_connect_timeout"
);
let invalid_metrics_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/metrics?from={from}&to={to}&step=5m"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("invalid metrics request should succeed");
assert_eq!(invalid_metrics_response.status(), StatusCode::BAD_REQUEST);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_updates_proxy_node_config_and_dispatches_upgrade_targets_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));

View File

@@ -110,21 +110,24 @@ sudo aether-proxy uninstall
| `--upstream-pool-idle-timeout-secs` | `AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT_SECS` | `300` | 连接池空闲超时(秒) |
| `--upstream-tcp-keepalive-secs` | `AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE_SECS` | `60` | TCP keepalive0 关闭) |
| `--upstream-tcp-nodelay` | `AETHER_PROXY_UPSTREAM_TCP_NODELAY` | `true` | 启用 TCP_NODELAY |
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理,支持 `http://``socks5://``socks5h://` |
| `--upstream-proxy-url` | `AETHER_PROXY_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理 |
| `--redirect-replay-budget-bytes` | `AETHER_PROXY_REDIRECT_REPLAY_BUDGET_BYTES` | `5M` | 307/308 请求体重放的预读预算,支持 `K/M/G``0` 表示禁用 body replay buffering |
`upstream_proxy_url` 只影响 `aether-proxy` 访问 OpenAI、Claude、Gemini 等 provider 的上游请求,不影响节点回连 Aether 服务器的 WebSocket tunnel。配合 WARP sidecar 时可填写:
出口代理支持 `http://``socks5://``socks5h://`。配合 WARP sidecar 时可填写:
```toml
upstream_proxy_url = "socks5h://microwarp:1080"
```
如果需要让 Aether 管理 API 和 WebSocket tunnel 也走代理,使用 `aether_proxy_url`
#### Aether API 客户端
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--aether-request-timeout-secs` | `AETHER_PROXY_AETHER_REQUEST_TIMEOUT_SECS` | `10` | 请求总超时(秒) |
| `--aether-connect-timeout-secs` | `AETHER_PROXY_AETHER_CONNECT_TIMEOUT_SECS` | `10` | 建连超时(秒) |
| `--aether-proxy-url` | `AETHER_PROXY_AETHER_PROXY_URL` | 空 | Aether 注册、心跳和 WebSocket tunnel 回连使用的出口代理(默认不走代理) |
| `--aether-retry-max-attempts` | `AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS` | `3` | 最大重试次数 |
#### DNS 与安全
@@ -154,6 +157,15 @@ upstream_proxy_url = "socks5h://microwarp:1080"
-`systemd``OpenRC` 安装时默认会额外打开文件日志到 `/var/log/aether-proxy`
- OpenRC 安装时,`aether-proxy logs` 实际读取 `/var/log/aether-proxy/current.log``/var/log/aether-proxy/error.log`;这些文件通常需要用 `sudo aether-proxy logs` 查看
### 隧道健康上报Heartbeat
proxy 会在心跳 `proxy_metadata` 中主动上报隧道稳定性指标,便于后端直接入库/告警:
- `proxy_metadata.tunnel_metrics`:建连尝试/成功/失败、断开次数、累计在线时长、心跳 RTT、WebSocket 收发帧与字节等。
- `proxy_metadata.recent_tunnel_errors`:最近隧道异常事件(时间戳、类别、错误摘要,环形缓冲)。
说明:仅主连接(`conn=0`)发送 heartbeat避免多条 tunnel 重复上报同一份全局指标。
### 多服务器配置
`aether-proxy.toml` 中使用 `[[servers]]` 配置 Aether 服务器。即使只有一个服务器,也必须写成一个 `[[servers]]` 条目;旧的顶层单服务器写法已不再支持。

View File

@@ -17,7 +17,7 @@ use crate::config::{Config, ServerEntry, TunnelPoolSizing};
use crate::net;
use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig};
use crate::state::{AppState, ProxyMetrics, ServerContext};
use crate::state::{AppState, ProxyMetrics, ServerContext, TunnelMetrics};
use crate::upstream_client;
use crate::{hardware, target_filter, tunnel};
@@ -80,6 +80,14 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
server_count = servers.len(),
"aether-proxy starting (tunnel mode)"
);
if let Some(proxy_url) = config.effective_aether_proxy_url() {
if let Ok(proxy) = crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url) {
info!(
aether_proxy_url = %proxy.redacted_url(),
"Aether control and tunnel egress proxy configured"
);
}
}
if let Some(proxy_url) = config
.upstream_proxy_url
.as_deref()
@@ -457,6 +465,7 @@ fn build_server_context(
dynamic: Arc::new(ArcSwap::from_pointee(dynamic)),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
tunnel_metrics: Arc::new(TunnelMetrics::new()),
})
}
@@ -960,6 +969,7 @@ mod tests {
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_proxy_url: None,
aether_retry_max_attempts: 1,
aether_retry_base_delay_ms: 50,
aether_retry_max_delay_ms: 100,

View File

@@ -339,6 +339,11 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_AETHER_HTTP2", default_value_t = true)]
pub aether_http2: bool,
/// Optional egress proxy used for Aether API registration and WebSocket tunnel reconnects.
/// Supported schemes: http, socks5, socks5h.
#[arg(long, env = "AETHER_PROXY_AETHER_PROXY_URL")]
pub aether_proxy_url: Option<String>,
/// Aether API retry attempts (including initial)
#[arg(
long,
@@ -680,12 +685,11 @@ impl Config {
if self.upstream_connect_timeout_secs == 0 {
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
}
if let Some(proxy_url) = self
.upstream_proxy_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Some(proxy_url) = normalized_proxy_url(&self.aether_proxy_url) {
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
.map_err(|err| anyhow::anyhow!("aether_proxy_url invalid: {err}"))?;
}
if let Some(proxy_url) = normalized_proxy_url(&self.upstream_proxy_url) {
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
.map_err(|err| anyhow::anyhow!("upstream_proxy_url invalid: {err}"))?;
}
@@ -740,6 +744,10 @@ impl Config {
Ok(Duration::from_millis(self.tunnel_stale_timeout_ms))
}
pub fn effective_aether_proxy_url(&self) -> Option<&str> {
normalized_proxy_url(&self.aether_proxy_url)
}
pub fn resolve_tunnel_pool_sizing(
&self,
hw_info: &HardwareInfo,
@@ -850,6 +858,8 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_http2: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_proxy_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_retry_max_attempts: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_retry_base_delay_ms: Option<u64>,
@@ -929,8 +939,7 @@ impl ConfigFile {
/// Load from a TOML file.
pub fn load(path: &Path) -> anyhow::Result<Self> {
let content = std::fs::read_to_string(path)?;
reject_removed_config_keys(&content)?;
Ok(toml::from_str(&content)?)
parse_config_file_content(&content)
}
/// Save to a TOML file.
@@ -1006,6 +1015,7 @@ impl ConfigFile {
);
set!("AETHER_PROXY_AETHER_TCP_NODELAY", self.aether_tcp_nodelay);
set!("AETHER_PROXY_AETHER_HTTP2", self.aether_http2);
set!("AETHER_PROXY_AETHER_PROXY_URL", self.aether_proxy_url);
set!(
"AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS",
self.aether_retry_max_attempts
@@ -1124,6 +1134,59 @@ impl ConfigFile {
}
}
fn parse_config_file_content(content: &str) -> anyhow::Result<ConfigFile> {
reject_removed_config_keys(content)?;
let mut value: toml::Value = toml::from_str(content)?;
promote_server_scoped_upstream_proxy_url(&mut value)?;
Ok(value.try_into()?)
}
fn normalized_proxy_url(value: &Option<String>) -> Option<&str> {
value
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn promote_server_scoped_upstream_proxy_url(value: &mut toml::Value) -> anyhow::Result<()> {
const KEY: &str = "upstream_proxy_url";
let Some(root) = value.as_table_mut() else {
return Ok(());
};
let mut promoted = root.get(KEY).cloned();
let Some(servers) = root.get_mut("servers").and_then(toml::Value::as_array_mut) else {
return Ok(());
};
for (index, server) in servers.iter_mut().enumerate() {
let Some(table) = server.as_table_mut() else {
continue;
};
let Some(server_value) = table.remove(KEY) else {
continue;
};
match promoted.as_ref() {
Some(existing) if existing != &server_value => {
anyhow::bail!(
"conflicting upstream_proxy_url values: top-level value and [[servers]] entry {} differ; configure it once at the top level",
index + 1
);
}
Some(_) => {}
None => promoted = Some(server_value),
}
}
if let Some(promoted) = promoted {
root.insert(KEY.to_string(), promoted);
}
Ok(())
}
fn reject_removed_config_keys(content: &str) -> anyhow::Result<()> {
let value: toml::Value = toml::from_str(content)?;
let Some(table) = value.as_table() else {
@@ -1233,6 +1296,92 @@ mod tests {
);
}
#[test]
fn config_file_deserializes_aether_proxy_url() {
let cfg: ConfigFile = toml::from_str("aether_proxy_url = \"socks5h://127.0.0.1:1080\"")
.expect("proxy URL toml");
assert_eq!(
cfg.aether_proxy_url.as_deref(),
Some("socks5h://127.0.0.1:1080")
);
}
#[test]
fn aether_proxy_url_requires_explicit_opt_in() {
let default_direct = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--upstream-proxy-url",
"socks5h://127.0.0.1:1080",
]);
assert_eq!(default_direct.effective_aether_proxy_url(), None);
let explicit = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--upstream-proxy-url",
"socks5h://127.0.0.1:1080",
"--aether-proxy-url",
"http://127.0.0.1:8080",
]);
assert_eq!(
explicit.effective_aether_proxy_url(),
Some("http://127.0.0.1:8080")
);
}
#[test]
fn config_file_load_accepts_server_scoped_upstream_proxy_url() {
let cfg = parse_config_file_content(
r#"
[[servers]]
aether_url = "https://aether.example.com"
upstream_proxy_url = "socks5://127.0.0.1:1080"
management_token = "ae_test"
node_name = "proxy-test"
"#,
)
.expect("server-scoped proxy URL should be promoted");
assert_eq!(
cfg.upstream_proxy_url.as_deref(),
Some("socks5://127.0.0.1:1080")
);
assert_eq!(cfg.servers.len(), 1);
assert_eq!(cfg.servers[0].aether_url, "https://aether.example.com");
}
#[test]
fn config_file_load_rejects_conflicting_server_scoped_upstream_proxy_url() {
let error = parse_config_file_content(
r#"
upstream_proxy_url = "socks5://127.0.0.1:1080"
[[servers]]
aether_url = "https://aether.example.com"
upstream_proxy_url = "socks5://127.0.0.1:1081"
management_token = "ae_test"
node_name = "proxy-test"
"#,
)
.expect_err("conflicting proxy URLs should be rejected");
assert!(
error.to_string().contains("conflicting upstream_proxy_url"),
"error should mention the conflicting key"
);
}
#[test]
fn config_file_rejects_removed_tunnel_seconds_keys() {
let error = reject_removed_config_keys("tunnel_ping_interval_secs = 5")

View File

@@ -1,4 +1,11 @@
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::time::Duration;
use base64::Engine;
use socket2::{SockRef, TcpKeepalive};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use url::Url;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -114,6 +121,303 @@ impl UpstreamProxyConfig {
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct ProxyConnectOptions {
pub connect_timeout: Duration,
pub tcp_nodelay: bool,
pub tcp_keepalive: Option<Duration>,
}
pub(crate) async fn connect_target_via_proxy(
proxy: &UpstreamProxyConfig,
target_host: &str,
target_port: u16,
options: ProxyConnectOptions,
) -> io::Result<TcpStream> {
let mut tcp = connect_proxy_tcp(
proxy,
options.connect_timeout,
options.tcp_nodelay,
options.tcp_keepalive,
)
.await?;
match proxy.scheme() {
UpstreamProxyScheme::Http => {
http_connect(&mut tcp, &target_authority(target_host, target_port), proxy).await?;
}
UpstreamProxyScheme::Socks5 | UpstreamProxyScheme::Socks5h => {
socks5_connect(&mut tcp, proxy, target_host, target_port).await?;
}
}
Ok(tcp)
}
pub(crate) async fn connect_proxy_tcp(
proxy: &UpstreamProxyConfig,
connect_timeout: Duration,
tcp_nodelay: bool,
tcp_keepalive: Option<Duration>,
) -> io::Result<TcpStream> {
let resolved = tokio::time::timeout(
connect_timeout,
tokio::net::lookup_host((proxy.host(), proxy.port())),
)
.await
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "proxy DNS timeout"))?
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
let mut last_error = None;
for addr in resolved {
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
Ok(Ok(stream)) => {
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
return Ok(stream);
}
Ok(Err(error)) => last_error = Some(error),
Err(_) => {
last_error = Some(io::Error::new(
io::ErrorKind::TimedOut,
format!("proxy connect timeout: {addr}"),
));
}
}
}
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses")))
}
fn configure_tcp_stream(
stream: &TcpStream,
tcp_nodelay: bool,
tcp_keepalive: Option<Duration>,
) -> io::Result<()> {
stream.set_nodelay(tcp_nodelay)?;
if let Some(keepalive) = tcp_keepalive {
let keepalive = TcpKeepalive::new().with_time(keepalive);
SockRef::from(stream).set_tcp_keepalive(&keepalive)?;
}
Ok(())
}
pub(crate) async fn http_connect(
stream: &mut TcpStream,
target_authority: &str,
proxy: &UpstreamProxyConfig,
) -> io::Result<()> {
let mut request = format!(
"CONNECT {target_authority} HTTP/1.1\r\nHost: {target_authority}\r\nProxy-Connection: Keep-Alive\r\n"
);
if let Some(auth) = proxy.basic_auth_header() {
request.push_str("Proxy-Authorization: ");
request.push_str(&auth);
request.push_str("\r\n");
}
request.push_str("\r\n");
stream.write_all(request.as_bytes()).await?;
stream.flush().await?;
let mut response = Vec::with_capacity(1024);
let mut chunk = [0u8; 1024];
loop {
if response.len() >= 16 * 1024 {
return Err(io::Error::other("proxy CONNECT response too large"));
}
let n = stream.read(&mut chunk).await?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"proxy closed during CONNECT",
));
}
response.extend_from_slice(&chunk[..n]);
if response.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let status_line_end = response
.windows(2)
.position(|window| window == b"\r\n")
.ok_or_else(|| io::Error::other("proxy CONNECT response missing status line"))?;
let status_line = std::str::from_utf8(&response[..status_line_end])
.map_err(|_| io::Error::other("proxy CONNECT status line is not UTF-8"))?;
let status = status_line.split_whitespace().nth(1).unwrap_or_default();
if status == "200" {
Ok(())
} else {
Err(io::Error::other(format!(
"proxy CONNECT failed: {status_line}"
)))
}
}
pub(crate) async fn socks5_connect(
stream: &mut TcpStream,
proxy: &UpstreamProxyConfig,
target_host: &str,
target_port: u16,
) -> io::Result<()> {
let requires_auth = proxy.username().is_some();
if requires_auth {
stream.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
} else {
stream.write_all(&[0x05, 0x01, 0x00]).await?;
}
let mut method_response = [0u8; 2];
stream.read_exact(&mut method_response).await?;
if method_response[0] != 0x05 {
return Err(io::Error::other("invalid SOCKS5 method response"));
}
match method_response[1] {
0x00 => {}
0x02 => socks5_authenticate(stream, proxy).await?,
0xff => return Err(io::Error::other("SOCKS5 proxy rejected all auth methods")),
method => {
return Err(io::Error::other(format!(
"SOCKS5 proxy selected unsupported auth method 0x{method:02x}"
)))
}
}
let address = socks5_target_address(target_host, target_port, proxy.uses_remote_dns()).await?;
stream.write_all(&address).await?;
let mut response = [0u8; 4];
stream.read_exact(&mut response).await?;
if response[0] != 0x05 {
return Err(io::Error::other("invalid SOCKS5 connect response"));
}
if response[1] != 0x00 {
return Err(io::Error::other(format!(
"SOCKS5 connect failed: {}",
socks5_reply_message(response[1])
)));
}
match response[3] {
0x01 => {
let mut ignored = [0u8; 4 + 2];
stream.read_exact(&mut ignored).await?;
}
0x03 => {
let mut len = [0u8; 1];
stream.read_exact(&mut len).await?;
let mut ignored = vec![0u8; len[0] as usize + 2];
stream.read_exact(&mut ignored).await?;
}
0x04 => {
let mut ignored = [0u8; 16 + 2];
stream.read_exact(&mut ignored).await?;
}
atyp => {
return Err(io::Error::other(format!(
"SOCKS5 proxy returned unsupported address type 0x{atyp:02x}"
)))
}
}
Ok(())
}
async fn socks5_authenticate(
stream: &mut TcpStream,
proxy: &UpstreamProxyConfig,
) -> io::Result<()> {
let username = proxy.username().unwrap_or_default().as_bytes();
let password = proxy.password().unwrap_or_default().as_bytes();
if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize {
return Err(io::Error::other(
"SOCKS5 username/password must be at most 255 bytes",
));
}
let mut request = Vec::with_capacity(username.len() + password.len() + 3);
request.push(0x01);
request.push(username.len() as u8);
request.extend_from_slice(username);
request.push(password.len() as u8);
request.extend_from_slice(password);
stream.write_all(&request).await?;
let mut response = [0u8; 2];
stream.read_exact(&mut response).await?;
if response[0] != 0x01 || response[1] != 0x00 {
return Err(io::Error::other("SOCKS5 username/password auth failed"));
}
Ok(())
}
pub(crate) async fn socks5_target_address(
target_host: &str,
target_port: u16,
remote_dns: bool,
) -> io::Result<Vec<u8>> {
let mut request = vec![0x05, 0x01, 0x00];
if let Ok(ip) = target_host.parse::<IpAddr>() {
push_socks5_ip_address(&mut request, ip);
} else if remote_dns {
let host = target_host.as_bytes();
if host.len() > u8::MAX as usize {
return Err(io::Error::other("SOCKS5 target hostname is too long"));
}
request.push(0x03);
request.push(host.len() as u8);
request.extend_from_slice(host);
} else {
let mut resolved = tokio::net::lookup_host((target_host, target_port))
.await
.map_err(|err| io::Error::other(format!("SOCKS5 target DNS failed: {err}")))?;
let addr = resolved
.next()
.ok_or_else(|| io::Error::other("SOCKS5 target DNS returned no addresses"))?;
push_socks5_socket_address(&mut request, addr);
}
request.extend_from_slice(&target_port.to_be_bytes());
Ok(request)
}
fn push_socks5_socket_address(request: &mut Vec<u8>, addr: SocketAddr) {
push_socks5_ip_address(request, addr.ip());
}
fn push_socks5_ip_address(request: &mut Vec<u8>, ip: IpAddr) {
match ip {
IpAddr::V4(ip) => {
request.push(0x01);
request.extend_from_slice(&ip.octets());
}
IpAddr::V6(ip) => {
request.push(0x04);
request.extend_from_slice(&ip.octets());
}
}
}
fn socks5_reply_message(reply: u8) -> &'static str {
match reply {
0x01 => "general failure",
0x02 => "connection not allowed",
0x03 => "network unreachable",
0x04 => "host unreachable",
0x05 => "connection refused",
0x06 => "TTL expired",
0x07 => "command not supported",
0x08 => "address type not supported",
_ => "unknown error",
}
}
pub(crate) fn target_authority(host: &str, port: u16) -> String {
if host.contains(':') && !host.starts_with('[') {
format!("[{host}]:{port}")
} else {
format!("{host}:{port}")
}
}
fn non_empty_url_part(value: &str) -> Option<String> {
let value = value.trim();
if value.is_empty() {

View File

@@ -67,6 +67,7 @@ impl AetherClient {
tcp_nodelay: config.aether_tcp_nodelay,
http2_adaptive_window: config.aether_http2,
user_agent: Some(format!("aether-proxy/{}", env!("CARGO_PKG_VERSION"))),
proxy_url: config.effective_aether_proxy_url().map(str::to_string),
..HttpClientConfig::default()
})
.expect("failed to create HTTP client");

View File

@@ -93,15 +93,6 @@ impl ServerTab {
required: true,
help: "Node name for identification in Aether dashboard",
},
Field {
label: "Upstream Proxy",
key: "upstream_proxy_url",
value: String::new(),
kind: FieldKind::Text,
required: false,
help:
"Optional provider egress proxy, e.g. http://127.0.0.1:8080 or socks5h://127.0.0.1:1080",
},
],
}
}
@@ -115,12 +106,6 @@ impl ServerTab {
}
tab
}
fn set_field_value(&mut self, key: &str, value: &str) {
if let Some(field) = self.fields.iter_mut().find(|field| field.key == key) {
field.value = value.to_string();
}
}
}
// -- App state ----------------------------------------------------------------
@@ -153,6 +138,15 @@ impl App {
server_tabs: vec![ServerTab::new()],
active_tab: 0,
global_fields: vec![
Field {
label: "Egress Proxy",
key: "upstream_proxy_url",
value: String::new(),
kind: FieldKind::Text,
required: false,
help:
"Optional egress proxy for Aether tunnel/API and provider requests, e.g. http://127.0.0.1:8080 or socks5h://127.0.0.1:1080",
},
Field {
label: "Install Service",
key: "install_service",
@@ -282,6 +276,7 @@ impl App {
"allow_private_targets" => cfg.allow_private_targets.map(|v| v.to_string()),
"heartbeat_interval" => cfg.heartbeat_interval.map(|v| v.to_string()),
"redirect_replay_budget_bytes" => cfg.redirect_replay_budget_bytes.clone(),
"upstream_proxy_url" => cfg.upstream_proxy_url.clone(),
_ => None,
};
if let Some(v) = val {
@@ -296,11 +291,6 @@ impl App {
} else {
self.server_tabs = servers.iter().map(ServerTab::from_entry).collect();
}
if let Some(proxy_url) = cfg.upstream_proxy_url.as_deref() {
for tab in &mut self.server_tabs {
tab.set_field_value("upstream_proxy_url", proxy_url);
}
}
self.active_tab = 0;
self.selected = 0;
self.scroll_offset = 0;
@@ -322,12 +312,6 @@ impl App {
.filter(|v| !v.is_empty())
}
fn get_shared_server_field(&self, key: &str) -> Option<String> {
self.server_tabs
.iter()
.find_map(|tab| Self::get_tab(tab, key))
}
fn toggle_enabled(&self, key: &str) -> bool {
self.get_global(key).as_deref() == Some("true")
}
@@ -373,7 +357,7 @@ impl App {
}
fn parse_optional_upstream_proxy_url(&self) -> anyhow::Result<Option<String>> {
let Some(raw) = self.get_shared_server_field("upstream_proxy_url") else {
let Some(raw) = self.get_global("upstream_proxy_url") else {
return Ok(None);
};
let trimmed = raw.trim();
@@ -381,7 +365,7 @@ impl App {
return Ok(None);
}
UpstreamProxyConfig::parse(trimmed)
.map_err(|err| anyhow::anyhow!("upstream proxy URL invalid: {err}"))?;
.map_err(|err| anyhow::anyhow!("egress proxy URL invalid: {err}"))?;
Ok(Some(trimmed.to_string()))
}
@@ -613,11 +597,7 @@ impl App {
}
// -- Add / remove server --
KeyCode::Char('+') | KeyCode::Char('a') => {
let upstream_proxy_url = self.get_shared_server_field("upstream_proxy_url");
let mut tab = ServerTab::new();
if let Some(proxy_url) = upstream_proxy_url.as_deref() {
tab.set_field_value("upstream_proxy_url", proxy_url);
}
let tab = ServerTab::new();
self.server_tabs.push(tab);
self.active_tab = self.server_tabs.len() - 1;
self.selected = 0;
@@ -694,14 +674,7 @@ impl App {
fn commit_edit_buffer(&mut self) -> bool {
if self.validate_edit() {
let key = self.selected_field().key;
if key == "upstream_proxy_url" {
for tab in &mut self.server_tabs {
tab.set_field_value(key, &self.edit_buffer);
}
} else {
self.selected_field_mut().value = self.edit_buffer.clone();
}
self.selected_field_mut().value = self.edit_buffer.clone();
self.modified = true;
self.mode = Mode::Normal;
true
@@ -1124,23 +1097,20 @@ mod tests {
}
#[test]
fn new_app_places_upstream_proxy_under_node_name() {
fn new_app_places_egress_proxy_in_global_fields() {
let app = sample_app();
let keys: Vec<&str> = app.server_tabs[0]
let server_keys: Vec<&str> = app.server_tabs[0]
.fields
.iter()
.map(|field| field.key)
.collect();
let global_keys: Vec<&str> = app.global_fields.iter().map(|field| field.key).collect();
assert_eq!(
keys,
vec![
"aether_url",
"management_token",
"node_name",
"upstream_proxy_url"
]
server_keys,
vec!["aether_url", "management_token", "node_name"]
);
assert_eq!(global_keys.first().copied(), Some("upstream_proxy_url"));
}
#[test]
@@ -1172,7 +1142,7 @@ mod tests {
set_global_field(&mut app, "allow_private_targets", "true");
set_global_field(&mut app, "heartbeat_interval", "45");
set_global_field(&mut app, "redirect_replay_budget_bytes", "6m");
set_server_field(&mut app, "upstream_proxy_url", "socks5h://127.0.0.1:1080");
set_global_field(&mut app, "upstream_proxy_url", "socks5h://127.0.0.1:1080");
let cfg = app.to_config().expect("config should serialize");
assert_eq!(cfg.allow_private_targets, Some(true));
@@ -1194,14 +1164,14 @@ mod tests {
}
#[test]
fn to_config_rejects_invalid_upstream_proxy_url() {
fn to_config_rejects_invalid_egress_proxy_url() {
let mut app = sample_app();
set_server_field(&mut app, "upstream_proxy_url", "ftp://proxy.example");
set_global_field(&mut app, "upstream_proxy_url", "ftp://proxy.example");
let error = app
.to_config()
.expect_err("invalid upstream proxy should be rejected");
assert!(error.to_string().contains("upstream proxy URL"));
.expect_err("invalid egress proxy should be rejected");
assert!(error.to_string().contains("egress proxy URL"));
}
#[test]

View File

@@ -1,8 +1,10 @@
//! Shared application state passed to all subsystems.
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_runtime::{AdmissionPermit, ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot};
use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
@@ -50,6 +52,8 @@ pub struct ServerContext {
pub active_connections: Arc<AtomicU64>,
/// Per-server request/latency metrics.
pub metrics: Arc<ProxyMetrics>,
/// Per-server tunnel stability/traffic metrics.
pub tunnel_metrics: Arc<TunnelMetrics>,
}
/// Aggregate metrics for reporting to Aether.
@@ -83,6 +87,221 @@ impl ProxyMetrics {
}
}
const RECENT_TUNNEL_ERROR_CAPACITY: usize = 64;
const TUNNEL_ERROR_CATEGORY_MAX_CHARS: usize = 48;
const TUNNEL_ERROR_MESSAGE_MAX_CHARS: usize = 320;
#[derive(Debug, Clone, serde::Serialize)]
pub struct TunnelErrorEvent {
pub timestamp_unix_secs: u64,
pub category: String,
pub message: String,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct TunnelMetricsSnapshot {
pub connect_attempts: u64,
pub connect_successes: u64,
pub connect_errors: u64,
pub disconnects: u64,
pub last_connected_at_unix_secs: u64,
pub last_disconnected_at_unix_secs: u64,
pub last_connected_duration_ms: u64,
pub connected_duration_total_ms: u64,
pub heartbeat_sent: u64,
pub heartbeat_ack: u64,
pub heartbeat_rtt_last_ms: u64,
pub heartbeat_rtt_total_ms: u64,
pub ws_in_frames: u64,
pub ws_in_bytes: u64,
pub ws_out_frames: u64,
pub ws_out_bytes: u64,
pub error_events_total: u64,
}
impl TunnelMetricsSnapshot {
pub fn heartbeat_rtt_avg_ms(self) -> Option<f64> {
if self.heartbeat_ack == 0 {
None
} else {
Some(self.heartbeat_rtt_total_ms as f64 / self.heartbeat_ack as f64)
}
}
}
pub struct TunnelMetrics {
connect_attempts: AtomicU64,
connect_successes: AtomicU64,
connect_errors: AtomicU64,
disconnects: AtomicU64,
last_connected_at_unix_secs: AtomicU64,
last_disconnected_at_unix_secs: AtomicU64,
last_connected_duration_ms: AtomicU64,
connected_duration_total_ms: AtomicU64,
heartbeat_sent: AtomicU64,
heartbeat_ack: AtomicU64,
heartbeat_rtt_last_ms: AtomicU64,
heartbeat_rtt_total_ms: AtomicU64,
ws_in_frames: AtomicU64,
ws_in_bytes: AtomicU64,
ws_out_frames: AtomicU64,
ws_out_bytes: AtomicU64,
error_events_total: AtomicU64,
recent_errors: Mutex<VecDeque<TunnelErrorEvent>>,
}
impl TunnelMetrics {
pub fn new() -> Self {
Self {
connect_attempts: AtomicU64::new(0),
connect_successes: AtomicU64::new(0),
connect_errors: AtomicU64::new(0),
disconnects: AtomicU64::new(0),
last_connected_at_unix_secs: AtomicU64::new(0),
last_disconnected_at_unix_secs: AtomicU64::new(0),
last_connected_duration_ms: AtomicU64::new(0),
connected_duration_total_ms: AtomicU64::new(0),
heartbeat_sent: AtomicU64::new(0),
heartbeat_ack: AtomicU64::new(0),
heartbeat_rtt_last_ms: AtomicU64::new(0),
heartbeat_rtt_total_ms: AtomicU64::new(0),
ws_in_frames: AtomicU64::new(0),
ws_in_bytes: AtomicU64::new(0),
ws_out_frames: AtomicU64::new(0),
ws_out_bytes: AtomicU64::new(0),
error_events_total: AtomicU64::new(0),
recent_errors: Mutex::new(VecDeque::with_capacity(RECENT_TUNNEL_ERROR_CAPACITY)),
}
}
pub fn record_connect_attempt(&self) {
self.connect_attempts.fetch_add(1, Ordering::Release);
}
pub fn record_connect_success(&self) {
self.connect_successes.fetch_add(1, Ordering::Release);
self.last_connected_at_unix_secs
.store(now_unix_secs(), Ordering::Release);
}
pub fn record_connect_error(&self) {
self.connect_errors.fetch_add(1, Ordering::Release);
}
pub fn record_disconnect(&self, connected_for: Duration) {
let duration_ms = duration_to_millis_u64(connected_for);
self.disconnects.fetch_add(1, Ordering::Release);
self.last_disconnected_at_unix_secs
.store(now_unix_secs(), Ordering::Release);
self.last_connected_duration_ms
.store(duration_ms, Ordering::Release);
self.connected_duration_total_ms
.fetch_add(duration_ms, Ordering::Release);
}
pub fn record_heartbeat_sent(&self) {
self.heartbeat_sent.fetch_add(1, Ordering::Release);
}
pub fn record_heartbeat_ack(&self, rtt: Duration) {
let rtt_ms = duration_to_millis_u64(rtt);
self.heartbeat_ack.fetch_add(1, Ordering::Release);
self.heartbeat_rtt_last_ms.store(rtt_ms, Ordering::Release);
self.heartbeat_rtt_total_ms
.fetch_add(rtt_ms, Ordering::Release);
}
pub fn record_ws_incoming_frame(&self, payload_len: usize) {
self.ws_in_frames.fetch_add(1, Ordering::Release);
self.ws_in_bytes.fetch_add(
u64::try_from(payload_len).unwrap_or(u64::MAX),
Ordering::Release,
);
}
pub fn record_ws_outgoing_frame(&self, payload_len: usize) {
self.ws_out_frames.fetch_add(1, Ordering::Release);
self.ws_out_bytes.fetch_add(
u64::try_from(payload_len).unwrap_or(u64::MAX),
Ordering::Release,
);
}
pub fn record_error(&self, category: &str, message: &str) {
self.error_events_total.fetch_add(1, Ordering::Release);
let event = TunnelErrorEvent {
timestamp_unix_secs: now_unix_secs(),
category: normalize_error_field(category, TUNNEL_ERROR_CATEGORY_MAX_CHARS, "unknown"),
message: normalize_error_field(message, TUNNEL_ERROR_MESSAGE_MAX_CHARS, "n/a"),
};
let mut recent_errors = match self.recent_errors.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
if recent_errors.len() >= RECENT_TUNNEL_ERROR_CAPACITY {
recent_errors.pop_front();
}
recent_errors.push_back(event);
}
pub fn recent_errors(&self, limit: usize) -> Vec<TunnelErrorEvent> {
if limit == 0 {
return Vec::new();
}
let recent_errors = match self.recent_errors.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
let start = recent_errors.len().saturating_sub(limit);
recent_errors.iter().skip(start).cloned().collect()
}
pub fn snapshot(&self) -> TunnelMetricsSnapshot {
TunnelMetricsSnapshot {
connect_attempts: self.connect_attempts.load(Ordering::Acquire),
connect_successes: self.connect_successes.load(Ordering::Acquire),
connect_errors: self.connect_errors.load(Ordering::Acquire),
disconnects: self.disconnects.load(Ordering::Acquire),
last_connected_at_unix_secs: self.last_connected_at_unix_secs.load(Ordering::Acquire),
last_disconnected_at_unix_secs: self
.last_disconnected_at_unix_secs
.load(Ordering::Acquire),
last_connected_duration_ms: self.last_connected_duration_ms.load(Ordering::Acquire),
connected_duration_total_ms: self.connected_duration_total_ms.load(Ordering::Acquire),
heartbeat_sent: self.heartbeat_sent.load(Ordering::Acquire),
heartbeat_ack: self.heartbeat_ack.load(Ordering::Acquire),
heartbeat_rtt_last_ms: self.heartbeat_rtt_last_ms.load(Ordering::Acquire),
heartbeat_rtt_total_ms: self.heartbeat_rtt_total_ms.load(Ordering::Acquire),
ws_in_frames: self.ws_in_frames.load(Ordering::Acquire),
ws_in_bytes: self.ws_in_bytes.load(Ordering::Acquire),
ws_out_frames: self.ws_out_frames.load(Ordering::Acquire),
ws_out_bytes: self.ws_out_bytes.load(Ordering::Acquire),
error_events_total: self.error_events_total.load(Ordering::Acquire),
}
}
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
fn duration_to_millis_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn normalize_error_field(value: &str, max_chars: usize, fallback: &str) -> String {
let normalized = value.split_whitespace().collect::<Vec<_>>().join(" ");
if normalized.is_empty() {
return fallback.to_string();
}
normalized.chars().take(max_chars).collect()
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ProxyAdmissionError {
#[error("proxy stream admission saturated at {limit} for gate {gate}")]

View File

@@ -1,7 +1,7 @@
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::sync::watch;
@@ -10,6 +10,7 @@ use tokio_tungstenite::tungstenite::http;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tracing::{debug, warn};
use crate::egress_proxy::{connect_target_via_proxy, ProxyConnectOptions, UpstreamProxyConfig};
use crate::state::{AppState, ServerContext};
use super::{dispatcher, heartbeat, writer};
@@ -71,14 +72,7 @@ pub async fn connect_and_run(
.config
.tunnel_connect_timeout()
.expect("validated config should resolve tunnel connect timeout");
let tcp_stream = tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})??;
let tcp_stream = connect_tunnel_tcp(state, host, port, connect_timeout).await?;
// Configure TCP parameters via socket2
configure_tcp_socket(&tcp_stream, state);
@@ -133,6 +127,8 @@ pub async fn connect_and_run(
ping_interval_ms = ping_interval.as_millis(),
"tunnel connected"
);
server.tunnel_metrics.record_connect_success();
let connected_at = Instant::now();
// NOTE: reconnect_attempts reset is handled by the caller (mod.rs)
// based on how long the connection stayed alive.
@@ -141,7 +137,11 @@ pub async fn connect_and_run(
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
// Spawn writer task (with WebSocket ping keepalive)
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
let (frame_tx, mut writer_handle) = writer::spawn_writer_with_metrics(
ws_sink,
ping_interval,
Some(Arc::clone(&server.tunnel_metrics)),
);
let drain_signal = spawn_drain_signal(conn_idx, frame_tx.clone(), drain.clone());
// Spawn heartbeat task (only for primary connection to avoid
@@ -174,8 +174,13 @@ pub async fn connect_and_run(
drain.clone(),
) => {
match result {
Ok(()) => TunnelOutcome::Disconnected,
Err(e) => return Err(e),
Ok(()) => Ok(TunnelOutcome::Disconnected),
Err(e) => {
server
.tunnel_metrics
.record_error("dispatcher_error", &e.to_string());
Err(e)
}
}
}
writer_result = &mut writer_handle => {
@@ -184,16 +189,22 @@ pub async fn connect_and_run(
Err(e) => {
if e.is_panic() {
tracing::error!(error = %e, "writer task panicked, triggering reconnect");
server
.tunnel_metrics
.record_error("writer_task_panic", &e.to_string());
} else {
warn!(error = %e, "writer task cancelled, triggering reconnect");
server
.tunnel_metrics
.record_error("writer_task_cancelled", &e.to_string());
}
}
}
TunnelOutcome::Disconnected
Ok(TunnelOutcome::Disconnected)
}
_ = shutdown.changed() => {
debug!("shutdown during tunnel dispatch");
TunnelOutcome::Shutdown
Ok(TunnelOutcome::Shutdown)
}
};
@@ -213,8 +224,12 @@ pub async fn connect_and_run(
let _ = tokio::time::timeout(Duration::from_secs(35), writer_handle).await;
}
server
.tunnel_metrics
.record_disconnect(connected_at.elapsed());
debug!("tunnel disconnected");
Ok(outcome)
outcome
}
fn spawn_drain_signal(
@@ -246,6 +261,56 @@ fn spawn_drain_signal(
})
}
async fn connect_tunnel_tcp(
state: &Arc<AppState>,
host: &str,
port: u16,
connect_timeout: Duration,
) -> Result<TcpStream, anyhow::Error> {
if let Some(proxy_url) = state.config.effective_aether_proxy_url() {
let proxy = UpstreamProxyConfig::parse(proxy_url)
.map_err(|err| anyhow::anyhow!("aether proxy URL invalid: {err}"))?;
debug!(
proxy_url = %proxy.redacted_url(),
host = %host,
port = port,
"connecting tunnel via Aether egress proxy"
);
return tokio::time::timeout(
connect_timeout,
connect_target_via_proxy(
&proxy,
host,
port,
ProxyConnectOptions {
connect_timeout,
tcp_nodelay: state.config.tunnel_tcp_nodelay,
tcp_keepalive: (state.config.tunnel_tcp_keepalive_secs > 0)
.then(|| Duration::from_secs(state.config.tunnel_tcp_keepalive_secs)),
},
),
)
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel proxy TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})?
.map_err(anyhow::Error::from);
}
tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})?
.map_err(anyhow::Error::from)
}
/// Configure TCP keepalive and NODELAY on an established socket.
fn configure_tcp_socket(stream: &TcpStream, state: &Arc<AppState>) {
let sock_ref = socket2::SockRef::from(stream);

View File

@@ -84,6 +84,10 @@ where
stale_ms = stale_timeout.as_millis(),
"tunnel connection stale, no data received"
);
server.tunnel_metrics.record_error(
"stale_timeout",
&format!("no tunnel frame received for {}ms", stale_timeout.as_millis()),
);
break None;
}
};
@@ -92,6 +96,9 @@ where
Ok(m) => m,
Err(e) => {
error!(error = %e, "WebSocket read error");
server
.tunnel_metrics
.record_error("ws_read_error", &e.to_string());
break Some(e);
}
};
@@ -100,7 +107,10 @@ where
last_data_at = tokio::time::Instant::now();
let data = match msg {
Message::Binary(data) => Bytes::from(data),
Message::Binary(data) => {
server.tunnel_metrics.record_ws_incoming_frame(data.len());
Bytes::from(data)
}
Message::Ping(_) => continue,
Message::Pong(_) => continue,
Message::Close(_) => {
@@ -114,6 +124,9 @@ where
Ok(f) => f,
Err(e) => {
warn!(error = %e, "failed to decode frame");
server
.tunnel_metrics
.record_error("frame_decode_error", &e.to_string());
continue;
}
};
@@ -223,6 +236,10 @@ where
if is_end || dispatch != StreamDispatchStatus::Delivered {
streams.remove(&sid);
if dispatch == StreamDispatchStatus::TimedOut {
server.tunnel_metrics.record_error(
"stream_dispatch_timeout",
&format!("request body dispatch timed out for stream {}", sid),
);
try_send_stream_error(
&frame_tx,
sid,

View File

@@ -8,6 +8,7 @@ use std::time::UNIX_EPOCH;
use bytes::Bytes;
use tokio::sync::watch;
use tokio::time::Instant;
use tracing::{debug, info, warn};
use crate::registration::client::RemoteConfig;
@@ -60,6 +61,13 @@ struct HeartbeatSnapshot {
stream_errors: u64,
}
#[derive(Debug, Clone, Copy)]
struct PendingHeartbeat {
heartbeat_id: u64,
snapshot: HeartbeatSnapshot,
sent_at: Option<Instant>,
}
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
pub fn spawn(
state: Arc<AppState>,
@@ -76,7 +84,7 @@ pub fn spawn(
// At most one in-flight heartbeat snapshot is tracked at a time.
// Snapshot is only cleared after receiving an ACK, which avoids losing
// interval counters when ACK/frame delivery is temporarily unstable.
let mut pending: Option<(u64, HeartbeatSnapshot)> = None;
let mut pending: Option<PendingHeartbeat> = None;
let mut next_heartbeat_id: u64 = 1;
let heartbeat_session_id = format!(
"{}-{}",
@@ -93,8 +101,8 @@ pub fn spawn(
loop {
tokio::select! {
_ = tokio::time::sleep(current_interval) => {
let (heartbeat_id, snapshot) = if let Some((id, snap)) = pending {
(id, snap)
let pending_entry = if let Some(entry) = pending {
entry
} else {
let snap = collect_snapshot(&server);
let id = next_heartbeat_id;
@@ -102,24 +110,34 @@ pub fn spawn(
if next_heartbeat_id == 0 {
next_heartbeat_id = 1;
}
pending = Some((id, snap));
(id, snap)
let entry = PendingHeartbeat {
heartbeat_id: id,
snapshot: snap,
sent_at: None,
};
pending = Some(entry);
entry
};
let payload = build_heartbeat_payload(
&state,
&server,
&heartbeat_session_id,
heartbeat_id,
snapshot
pending_entry.heartbeat_id,
pending_entry.snapshot
).await;
let frame = Frame::control(MsgType::HeartbeatData, payload);
if frame_tx.send(frame).await.is_err() {
if let Some((_, snap)) = pending.take() {
restore_snapshot(&server, snap);
if let Some(entry) = pending.take() {
restore_snapshot(&server, entry.snapshot);
}
break; // Writer closed
}
server.tunnel_metrics.record_heartbeat_sent();
if let Some(mut entry) = pending {
entry.sent_at = Some(Instant::now());
pending = Some(entry);
}
debug!("sent heartbeat data");
// Re-read interval from dynamic config (remote config may have
@@ -142,8 +160,11 @@ pub fn spawn(
heartbeat_id: ack_id,
upgrade_to,
} => {
if let Some((pending_id, _)) = pending {
if ack_id == pending_id {
if let Some(entry) = pending {
if ack_id == entry.heartbeat_id {
if let Some(sent_at) = entry.sent_at {
server.tunnel_metrics.record_heartbeat_ack(sent_at.elapsed());
}
pending = None;
}
}
@@ -154,8 +175,8 @@ pub fn spawn(
}
_ = shutdown.changed() => {
debug!("heartbeat task shutting down");
if let Some((_, snap)) = pending.take() {
restore_snapshot(&server, snap);
if let Some(entry) = pending.take() {
restore_snapshot(&server, entry.snapshot);
}
break;
}
@@ -217,6 +238,8 @@ async fn build_heartbeat_payload(
snapshot: HeartbeatSnapshot,
) -> Bytes {
let node_id = server.node_id.read().unwrap().clone();
let tunnel_snapshot = server.tunnel_metrics.snapshot();
let recent_errors = server.tunnel_metrics.recent_errors(8);
let avg_latency_ms = if snapshot.requests > 0 {
Some(snapshot.latency_ns as f64 / snapshot.requests as f64 / 1_000_000.0)
@@ -268,6 +291,26 @@ async fn build_heartbeat_payload(
"proxy_metadata": {
"version": CURRENT_VERSION,
"admission": admission,
"tunnel_metrics": {
"connect_attempts": tunnel_snapshot.connect_attempts,
"connect_successes": tunnel_snapshot.connect_successes,
"connect_errors": tunnel_snapshot.connect_errors,
"disconnects": tunnel_snapshot.disconnects,
"last_connected_at_unix_secs": tunnel_snapshot.last_connected_at_unix_secs,
"last_disconnected_at_unix_secs": tunnel_snapshot.last_disconnected_at_unix_secs,
"last_connected_duration_ms": tunnel_snapshot.last_connected_duration_ms,
"connected_duration_total_ms": tunnel_snapshot.connected_duration_total_ms,
"heartbeat_sent": tunnel_snapshot.heartbeat_sent,
"heartbeat_ack": tunnel_snapshot.heartbeat_ack,
"heartbeat_rtt_last_ms": tunnel_snapshot.heartbeat_rtt_last_ms,
"heartbeat_rtt_avg_ms": tunnel_snapshot.heartbeat_rtt_avg_ms(),
"ws_in_frames": tunnel_snapshot.ws_in_frames,
"ws_in_bytes": tunnel_snapshot.ws_in_bytes,
"ws_out_frames": tunnel_snapshot.ws_out_frames,
"ws_out_bytes": tunnel_snapshot.ws_out_bytes,
"error_events_total": tunnel_snapshot.error_events_total,
},
"recent_tunnel_errors": recent_errors,
},
});
@@ -277,6 +320,9 @@ async fn build_heartbeat_payload(
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
if payload.is_empty() {
warn!("received empty heartbeat ACK");
server
.tunnel_metrics
.record_error("heartbeat_ack_empty", "received empty heartbeat ACK");
return AckDecision::Ignore;
}
@@ -303,6 +349,9 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
}
Err(e) => {
warn!(error = %e, "failed to parse heartbeat ACK");
server
.tunnel_metrics
.record_error("heartbeat_ack_parse", &e.to_string());
AckDecision::Ignore
}
}
@@ -373,7 +422,7 @@ mod tests {
use super::{handle_ack, AckDecision};
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::{ProxyMetrics, ServerContext};
use crate::state::{ProxyMetrics, ServerContext, TunnelMetrics};
fn sample_server() -> Arc<ServerContext> {
let config = Arc::new(crate::config::Config::parse_from([
@@ -399,6 +448,7 @@ mod tests {
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
tunnel_metrics: Arc::new(TunnelMetrics::new()),
})
}

View File

@@ -76,6 +76,7 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "tunnel drained, exiting slot");
return;
}
server.tunnel_metrics.record_connect_attempt();
let started_at = Instant::now();
match client::connect_and_run(state, server, conn_idx, &mut shutdown, drain.clone()).await {
Ok(client::TunnelOutcome::Shutdown) => {
@@ -86,6 +87,10 @@ pub async fn run(
debug!(server = %server.server_label, conn = conn_idx, "tunnel disconnected, reconnecting");
}
Err(e) => {
server.tunnel_metrics.record_connect_error();
server
.tunnel_metrics
.record_error("tunnel_connect_error", &e.to_string());
error!(server = %server.server_label, conn = conn_idx, error = %e, "tunnel connection error, reconnecting");
}
}
@@ -237,7 +242,7 @@ mod tests {
use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::{AppState as ProxyAppState, ProxyMetrics, ServerContext};
use crate::state::{AppState as ProxyAppState, ProxyMetrics, ServerContext, TunnelMetrics};
use crate::target_filter::DnsCache;
use crate::tunnel::protocol;
use crate::upstream_client;
@@ -478,6 +483,7 @@ mod tests {
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
tunnel_metrics: Arc::new(TunnelMetrics::new()),
})
}
@@ -498,6 +504,7 @@ mod tests {
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_proxy_url: None,
aether_retry_max_attempts: 1,
aether_retry_base_delay_ms: 50,
aether_retry_max_delay_ms: 100,

View File

@@ -1532,7 +1532,7 @@ mod tests {
use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::ProxyMetrics;
use crate::state::{ProxyMetrics, TunnelMetrics};
use crate::target_filter::DnsCache;
use crate::tunnel::client::build_tls_config;
@@ -2335,6 +2335,7 @@ mod tests {
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
tunnel_metrics: Arc::new(TunnelMetrics::new()),
})
}
@@ -2355,6 +2356,7 @@ mod tests {
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_proxy_url: None,
aether_retry_max_attempts: 3,
aether_retry_base_delay_ms: 200,
aether_retry_max_delay_ms: 2_000,

View File

@@ -5,9 +5,10 @@
//! periodic WebSocket Ping frames to keep the connection alive through
//! intermediary proxies (Nginx, Cloudflare, etc.).
use std::sync::Arc;
use std::time::Duration;
use aether_contracts::tunnel::MsgType;
use aether_contracts::tunnel::{MsgType, HEADER_SIZE};
#[cfg(test)]
use aether_runtime::QueueSnapshot;
use aether_runtime::{bounded_queue, BoundedQueueSender, QueueSendError};
@@ -16,6 +17,8 @@ use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, trace};
use crate::state::TunnelMetrics;
use super::protocol::Frame;
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
@@ -77,7 +80,20 @@ impl FrameSender {
///
/// `ping_interval` controls WebSocket-level Ping frequency (typically 15s).
/// This keeps the connection alive through intermediary proxies/load-balancers.
pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
#[cfg(test)]
pub fn spawn_writer<S>(sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
spawn_writer_with_metrics(sink, ping_interval, None)
}
/// Spawn the writer task with optional tunnel metrics instrumentation.
pub fn spawn_writer_with_metrics<S>(
mut sink: S,
ping_interval: Duration,
tunnel_metrics: Option<Arc<TunnelMetrics>>,
) -> (FrameSender, JoinHandle<()>)
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
@@ -93,7 +109,7 @@ where
loop {
if let Ok(frame) = high_rx.try_recv() {
if !write_frame(&mut sink, frame).await {
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
break;
}
continue;
@@ -107,7 +123,7 @@ where
frame = high_rx.recv(), if high_open => {
match frame {
Some(frame) => {
if !write_frame(&mut sink, frame).await {
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
break;
}
}
@@ -117,6 +133,9 @@ where
_ = ping_ticker.tick(), if high_open || normal_open => {
if let Err(e) = sink.send(Message::Ping(vec![])).await {
error!(error = %e, "failed to send WebSocket ping");
if let Some(metrics) = tunnel_metrics.as_deref() {
metrics.record_error("ws_ping_error", &e.to_string());
}
break;
}
trace!("sent WebSocket ping");
@@ -124,7 +143,7 @@ where
frame = normal_rx.recv(), if normal_open => {
match frame {
Some(frame) => {
if !write_frame(&mut sink, frame).await {
if !write_frame(&mut sink, frame, tunnel_metrics.as_deref()).await {
break;
}
}
@@ -156,15 +175,22 @@ fn classify_frame_priority(frame: &Frame) -> FramePriority {
}
}
async fn write_frame<S>(sink: &mut S, frame: Frame) -> bool
async fn write_frame<S>(sink: &mut S, frame: Frame, tunnel_metrics: Option<&TunnelMetrics>) -> bool
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
let data = frame.encode();
let wire_len = data.len().max(HEADER_SIZE);
if let Err(e) = sink.send(Message::Binary(data.into())).await {
error!(error = %e, "failed to write frame to WebSocket");
if let Some(metrics) = tunnel_metrics {
metrics.record_error("ws_write_error", &e.to_string());
}
return false;
}
if let Some(metrics) = tunnel_metrics {
metrics.record_ws_outgoing_frame(wire_len);
}
true
}

View File

@@ -2,7 +2,7 @@ use std::collections::HashMap;
use std::convert::Infallible;
use std::future::Future;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
@@ -28,14 +28,15 @@ use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use rustls::pki_types::ServerName;
use rustls::ClientConfig;
use socket2::{SockRef, TcpKeepalive};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;
use tower_service::Service;
use crate::config::Config;
use crate::egress_proxy::{UpstreamProxyConfig, UpstreamProxyScheme};
use crate::egress_proxy::{
connect_proxy_tcp, http_connect, socks5_connect, ProxyConnectOptions, UpstreamProxyConfig,
UpstreamProxyScheme,
};
use crate::target_filter::{self, DnsCache};
type BoxError = Box<dyn std::error::Error + Send + Sync>;
@@ -263,12 +264,6 @@ pub struct InstrumentedConnector {
tcp_keepalive: Option<Duration>,
}
struct ProxyConnectOptions {
connect_timeout: Duration,
tcp_nodelay: bool,
tcp_keepalive: Option<Duration>,
}
impl Service<Uri> for InstrumentedConnector {
type Response = TimedConn;
type Error = BoxError;
@@ -402,262 +397,6 @@ async fn connect_via_proxy(
}
}
async fn connect_proxy_tcp(
proxy: &UpstreamProxyConfig,
connect_timeout: Duration,
tcp_nodelay: bool,
tcp_keepalive: Option<Duration>,
) -> io::Result<TcpStream> {
let resolved = tokio::time::timeout(
connect_timeout,
tokio::net::lookup_host((proxy.host(), proxy.port())),
)
.await
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "proxy DNS timeout"))?
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
let mut last_error = None;
for addr in resolved {
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
Ok(Ok(stream)) => {
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
return Ok(stream);
}
Ok(Err(error)) => last_error = Some(error),
Err(_) => {
last_error = Some(io::Error::new(
io::ErrorKind::TimedOut,
format!("proxy connect timeout: {addr}"),
));
}
}
}
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses")))
}
fn configure_tcp_stream(
stream: &TcpStream,
tcp_nodelay: bool,
tcp_keepalive: Option<Duration>,
) -> io::Result<()> {
stream.set_nodelay(tcp_nodelay)?;
if let Some(keepalive) = tcp_keepalive {
let keepalive = TcpKeepalive::new().with_time(keepalive);
SockRef::from(stream).set_tcp_keepalive(&keepalive)?;
}
Ok(())
}
async fn http_connect(
stream: &mut TcpStream,
target_authority: &str,
proxy: &UpstreamProxyConfig,
) -> io::Result<()> {
let mut request = format!(
"CONNECT {target_authority} HTTP/1.1\r\nHost: {target_authority}\r\nProxy-Connection: Keep-Alive\r\n"
);
if let Some(auth) = proxy.basic_auth_header() {
request.push_str("Proxy-Authorization: ");
request.push_str(&auth);
request.push_str("\r\n");
}
request.push_str("\r\n");
stream.write_all(request.as_bytes()).await?;
stream.flush().await?;
let mut response = Vec::with_capacity(1024);
let mut chunk = [0u8; 1024];
loop {
if response.len() >= 16 * 1024 {
return Err(io::Error::other("proxy CONNECT response too large"));
}
let n = stream.read(&mut chunk).await?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"proxy closed during CONNECT",
));
}
response.extend_from_slice(&chunk[..n]);
if response.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let status_line_end = response
.windows(2)
.position(|window| window == b"\r\n")
.ok_or_else(|| io::Error::other("proxy CONNECT response missing status line"))?;
let status_line = std::str::from_utf8(&response[..status_line_end])
.map_err(|_| io::Error::other("proxy CONNECT status line is not UTF-8"))?;
let status = status_line.split_whitespace().nth(1).unwrap_or_default();
if status == "200" {
Ok(())
} else {
Err(io::Error::other(format!(
"proxy CONNECT failed: {status_line}"
)))
}
}
async fn socks5_connect(
stream: &mut TcpStream,
proxy: &UpstreamProxyConfig,
target_host: &str,
target_port: u16,
) -> io::Result<()> {
let requires_auth = proxy.username().is_some();
if requires_auth {
stream.write_all(&[0x05, 0x02, 0x00, 0x02]).await?;
} else {
stream.write_all(&[0x05, 0x01, 0x00]).await?;
}
let mut method_response = [0u8; 2];
stream.read_exact(&mut method_response).await?;
if method_response[0] != 0x05 {
return Err(io::Error::other("invalid SOCKS5 method response"));
}
match method_response[1] {
0x00 => {}
0x02 => socks5_authenticate(stream, proxy).await?,
0xff => return Err(io::Error::other("SOCKS5 proxy rejected all auth methods")),
method => {
return Err(io::Error::other(format!(
"SOCKS5 proxy selected unsupported auth method 0x{method:02x}"
)))
}
}
let address = socks5_target_address(target_host, target_port, proxy.uses_remote_dns()).await?;
stream.write_all(&address).await?;
let mut response = [0u8; 4];
stream.read_exact(&mut response).await?;
if response[0] != 0x05 {
return Err(io::Error::other("invalid SOCKS5 connect response"));
}
if response[1] != 0x00 {
return Err(io::Error::other(format!(
"SOCKS5 connect failed: {}",
socks5_reply_message(response[1])
)));
}
match response[3] {
0x01 => {
let mut ignored = [0u8; 4 + 2];
stream.read_exact(&mut ignored).await?;
}
0x03 => {
let mut len = [0u8; 1];
stream.read_exact(&mut len).await?;
let mut ignored = vec![0u8; len[0] as usize + 2];
stream.read_exact(&mut ignored).await?;
}
0x04 => {
let mut ignored = [0u8; 16 + 2];
stream.read_exact(&mut ignored).await?;
}
atyp => {
return Err(io::Error::other(format!(
"SOCKS5 proxy returned unsupported address type 0x{atyp:02x}"
)))
}
}
Ok(())
}
async fn socks5_authenticate(
stream: &mut TcpStream,
proxy: &UpstreamProxyConfig,
) -> io::Result<()> {
let username = proxy.username().unwrap_or_default().as_bytes();
let password = proxy.password().unwrap_or_default().as_bytes();
if username.len() > u8::MAX as usize || password.len() > u8::MAX as usize {
return Err(io::Error::other(
"SOCKS5 username/password must be at most 255 bytes",
));
}
let mut request = Vec::with_capacity(username.len() + password.len() + 3);
request.push(0x01);
request.push(username.len() as u8);
request.extend_from_slice(username);
request.push(password.len() as u8);
request.extend_from_slice(password);
stream.write_all(&request).await?;
let mut response = [0u8; 2];
stream.read_exact(&mut response).await?;
if response[0] != 0x01 || response[1] != 0x00 {
return Err(io::Error::other("SOCKS5 username/password auth failed"));
}
Ok(())
}
async fn socks5_target_address(
target_host: &str,
target_port: u16,
remote_dns: bool,
) -> io::Result<Vec<u8>> {
let mut request = vec![0x05, 0x01, 0x00];
if let Ok(ip) = target_host.parse::<IpAddr>() {
push_socks5_ip_address(&mut request, ip);
} else if remote_dns {
let host = target_host.as_bytes();
if host.len() > u8::MAX as usize {
return Err(io::Error::other("SOCKS5 target hostname is too long"));
}
request.push(0x03);
request.push(host.len() as u8);
request.extend_from_slice(host);
} else {
let mut resolved = tokio::net::lookup_host((target_host, target_port))
.await
.map_err(|err| io::Error::other(format!("SOCKS5 target DNS failed: {err}")))?;
let addr = resolved
.next()
.ok_or_else(|| io::Error::other("SOCKS5 target DNS returned no addresses"))?;
push_socks5_socket_address(&mut request, addr);
}
request.extend_from_slice(&target_port.to_be_bytes());
Ok(request)
}
fn push_socks5_socket_address(request: &mut Vec<u8>, addr: SocketAddr) {
push_socks5_ip_address(request, addr.ip());
}
fn push_socks5_ip_address(request: &mut Vec<u8>, ip: IpAddr) {
match ip {
IpAddr::V4(ip) => {
request.push(0x01);
request.extend_from_slice(&ip.octets());
}
IpAddr::V6(ip) => {
request.push(0x04);
request.extend_from_slice(&ip.octets());
}
}
}
fn socks5_reply_message(reply: u8) -> &'static str {
match reply {
0x01 => "general failure",
0x02 => "connection not allowed",
0x03 => "network unreachable",
0x04 => "host unreachable",
0x05 => "connection refused",
0x06 => "TTL expired",
0x07 => "command not supported",
0x08 => "address type not supported",
_ => "unknown error",
}
}
fn uri_host(uri: &Uri) -> Result<String, io::Error> {
uri.host()
.map(|host| {
@@ -934,8 +673,11 @@ mod tests {
use clap::Parser;
use http_body_util::BodyExt;
use hyper::Response;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use crate::egress_proxy::socks5_target_address;
#[test]
fn fresh_connection_uses_connector_breakdown() {
let mut response = Response::new(());