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

feat(proxy): 重构 Proxy 节点管理与隧道系统

Browse Source 
- 重构 proxy_nodes 管理端，支持节点注册、心跳、隧道生命周期管理
- 增强 tunnel 嵌入式 hub 和隧道协议
- 重构 aether-proxy 配置、隧道客户端、心跳和调度机制
- 调整 admin OAuth/配额/导入等处理器的参数传递
- 扩展数据迁移模块
- 补充 proxy nodes、OAuth、配额、系统导入等测试
- 更新前端 proxy nodes 视图和 API
This commit is contained in:
fawney19
2026-04-14 22:51:02 +08:00
parent fb31928e44
commit a4e7ac1df6
55 changed files with 3477 additions and 892 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
apps/aether-gateway/src/control/route/admin/operations_families.rs
View File

@@ -5,7 +5,7 @@ use super::{classified, ClassifiedRoute};
pub(super) fn classify_admin_operations_family_route(
method: &http::Method,
normalized_path: &str,
_normalized_path_no_trailing: &str,
normalized_path_no_trailing: &str,
) -> Option<ClassifiedRoute> {
if method == http::Method::GET
&& matches!(
@@ -113,6 +113,24 @@ pub(super) fn classify_admin_operations_family_route(
"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()..]
.split('/')
.count()
== 1
&& !matches!(
&normalized_path_no_trailing["/api/admin/proxy-nodes/".len()..],
"register" | "heartbeat" | "unregister" | "manual" | "upgrade" | "test-url"
)
{
Some(classified(
"admin_proxy",
"proxy_nodes_manage",
"get_node",
"admin:proxy_nodes",
false,
))
} else if method == http::Method::POST
&& matches!(
normalized_path,

9
apps/aether-gateway/src/control/tests/admin_proxy_nodes.rs
View File

@@ -91,6 +91,15 @@ fn classifies_admin_proxy_nodes_manual_update_as_admin_proxy_route() {
);
}
#[test]
fn classifies_admin_proxy_nodes_detail_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(
http::Method::GET,
"/api/admin/proxy-nodes/node-1",
"get_node",
);
}
#[test]
fn classifies_admin_proxy_nodes_delete_as_admin_proxy_route() {
assert_proxy_nodes_admin_route(

28
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -793,7 +793,7 @@ mod tests {
use axum::{Json, Router};
use serde_json::json;
use super::DirectSyncExecutionRuntime;
use super::{build_client, DirectSyncExecutionRuntime, ExecutionTransportControls};
use crate::frontdoor_loop_guard::{
frontdoor_self_loop_public_ai_path, gateway_frontdoor_self_loop_guard_error_with_port,
gateway_frontdoor_self_loop_guard_matches_with_port,
@@ -842,6 +842,32 @@ mod tests {
);
}
#[test]
fn direct_sync_execution_runtime_builds_clients_for_socks_proxy_urls() {
let timeouts = ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
};
for proxy_url in ["socks5://127.0.0.1:1080", "socks5h://127.0.0.1:1080"] {
build_client(
Some(&timeouts),
Some(&aether_contracts::ProxySnapshot {
enabled: Some(true),
mode: Some("socks".into()),
node_id: None,
label: Some("manual-proxy".into()),
url: Some(proxy_url.to_string()),
extra: None,
}),
None,
ExecutionTransportControls::default(),
)
.unwrap_or_else(|err| panic!("client should build for {proxy_url}: {err}"));
}
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),

2
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -31,3 +31,5 @@ pub(crate) use self::request::{
AdminAppState, AdminRequestContext, AdminRouteRequest, AdminRouteResponse, AdminRouteResult,
};
pub(crate) use self::routes::maybe_build_local_admin_response;
#[cfg(test)]
pub(crate) use self::system::override_proxy_connectivity_probe_url_for_tests;

25
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
View File

@@ -8,9 +8,11 @@ use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider
use crate::handlers::admin::provider::oauth::provisioning::build_provider_oauth_auth_config_from_token_payload;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
};
@@ -116,7 +118,18 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
.list_provider_catalog_endpoints_by_provider_ids(&[provider_id.to_string()])
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id,
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let mut results = Vec::with_capacity(entries.len());
let mut success = 0usize;
let mut failed = 0usize;
@@ -126,7 +139,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
state,
template,
entry.refresh_token.as_str(),
proxy_node_id,
request_proxy.clone(),
)
.await
{
@@ -189,7 +202,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -232,7 +245,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?

32
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs
View File

@@ -2,9 +2,11 @@ use super::parse::{AdminProviderOAuthBatchImportEntry, AdminProviderOAuthBatchIm
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::decode_jwt_claims;
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
@@ -14,6 +16,7 @@ use aether_admin::provider::oauth::{
build_kiro_batch_import_key_name, coerce_admin_provider_oauth_import_str,
parse_admin_provider_oauth_kiro_batch_import_entries,
};
use aether_contracts::ProxySnapshot;
use serde_json::{json, Map, Value};
use std::collections::BTreeSet;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -102,7 +105,7 @@ fn admin_provider_oauth_kiro_refresh_error_detail(
async fn refresh_admin_provider_oauth_kiro_auth_config(
state: &AdminAppState<'_>,
auth_config: &AdminKiroAuthConfig,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
social_refresh_base_url: Option<&str>,
idc_refresh_base_url: Option<&str>,
) -> Result<AdminKiroAuthConfig, String> {
@@ -164,7 +167,7 @@ async fn refresh_admin_provider_oauth_kiro_auth_config(
"grantType": "refresh_token",
})),
None,
proxy_node_id,
proxy.clone(),
)
.await
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
@@ -268,7 +271,7 @@ async fn refresh_admin_provider_oauth_kiro_auth_config(
.unwrap_or_default(),
})),
None,
proxy_node_id,
proxy,
)
.await
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
@@ -352,7 +355,18 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(&[provider_id.to_string()])
.await?;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id,
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let social_refresh_base_url =
admin_provider_oauth_kiro_refresh_base_url_override(state, "kiro_social_refresh");
let idc_refresh_base_url =
@@ -392,7 +406,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
refreshed_auth_config = match refresh_admin_provider_oauth_kiro_auth_config(
state,
&refreshed_auth_config,
proxy_node_id,
request_proxy.clone(),
social_refresh_base_url.as_deref(),
idc_refresh_base_url.as_deref(),
)
@@ -477,7 +491,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
refreshed_auth_config.expires_at,
)
.await?
@@ -511,7 +525,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&access_token,
&auth_config,
&provider_oauth_active_api_formats(&endpoints),
key_proxy.clone(),
None,
refreshed_auth_config.expires_at,
)
.await?

19
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
View File

@@ -1,5 +1,6 @@
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::quota::codex::refresh_codex_provider_quota_locally;
use super::super::super::runtime::provider_oauth_runtime_endpoint_for_provider;
use super::super::super::state::{
admin_provider_oauth_template, enrich_admin_provider_oauth_auth_config,
is_fixed_provider_type_for_provider_oauth, json_non_empty_string, json_u64_value,
@@ -123,6 +124,22 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
"该 Provider 不支持 OAuth 授权",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
key.proxy.as_ref(),
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_code(
@@ -130,7 +147,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{

32
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/provider.rs
View File

@@ -2,10 +2,11 @@ use super::super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use super::super::super::runtime::refresh_provider_oauth_account_state_after_update;
use super::super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
is_fixed_provider_type_for_provider_oauth,
@@ -114,6 +115,21 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
"该 Provider 不支持 OAuth 授权",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_code(
@@ -121,7 +137,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
@@ -138,11 +154,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(payload.proxy_node_id.as_deref());
let duplicate = match state
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
.await
@@ -163,7 +175,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -204,7 +216,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?

24
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/authorize.rs
View File

@@ -1,5 +1,6 @@
use super::session::AdminProviderOAuthDeviceAuthorizePayload;
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::runtime::provider_oauth_runtime_endpoint_for_provider;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, current_unix_secs,
default_kiro_device_start_url, generate_provider_oauth_nonce, json_non_empty_string,
@@ -69,6 +70,21 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
"设备授权仅支持 Kiro provider",
));
}
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let region = normalize_kiro_device_region(Some(payload.region.as_str())).ok_or_else(|| {
build_internal_control_error_response(http::StatusCode::BAD_REQUEST, "region 格式无效")
@@ -85,11 +101,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
};
let client_registration = match state
.register_admin_kiro_device_oidc_client(
&region,
&start_url,
payload.proxy_node_id.as_deref(),
)
.register_admin_kiro_device_oidc_client(&region, &start_url, request_proxy.clone())
.await
{
Ok(payload) => payload,
@@ -114,7 +126,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
&client_id,
&client_secret,
&start_url,
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{

34
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
View File

@@ -5,9 +5,11 @@ use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, build_kiro_device_key_name,
current_unix_secs, decode_jwt_claims, json_non_empty_string, json_u64_value,
@@ -125,6 +127,21 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
"Provider 不存在",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
session.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_result = match state
.poll_admin_kiro_device_token(
@@ -132,7 +149,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&session.client_id,
&session.client_secret,
&session.device_code,
session.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
@@ -252,12 +269,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}
};
let key_proxy = provider_oauth_key_proxy_value(session.proxy_node_id.as_deref());
let api_formats = provider_oauth_active_api_formats(
&state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?,
);
let api_formats = provider_oauth_active_api_formats(&endpoints);
let mut replaced = false;
let persisted_key = if let Some(existing_key) = duplicate {
replaced = true;
@@ -266,7 +278,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
Some(expires_at),
)
.await?
@@ -288,7 +300,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
Some(expires_at),
)
.await?

36
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
View File

@@ -2,10 +2,11 @@ use super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::errors::build_internal_control_error_response;
use super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use super::super::runtime::refresh_provider_oauth_account_state_after_update;
use super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
exchange_admin_provider_oauth_refresh_token, is_fixed_provider_type_for_provider_oauth,
@@ -96,13 +97,24 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_admin_provider_oauth_backend_unavailable_response());
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_refresh_token(
template,
refresh_token_input,
proxy_node_id.as_deref(),
)
.exchange_admin_provider_oauth_refresh_token(template, refresh_token_input, request_proxy)
.await
{
Ok(payload) => payload,
@@ -124,11 +136,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
}
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id.as_deref());
let duplicate = match state
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
.await
@@ -149,7 +157,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -189,7 +197,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?

2
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/request.rs
View File

@@ -79,7 +79,7 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, endpoints)
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints)
else {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,

32
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
View File

@@ -1,13 +1,14 @@
use super::shared::{
coerce_json_f64, coerce_json_string, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
coerce_json_f64, coerce_json_string, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_antigravity_usage_response;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -31,6 +32,14 @@ async fn execute_antigravity_quota_plan(
.or_insert_with(|| "antigravity".to_string());
let body = json!({ "project": project_id });
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("antigravity-quota:{}", transport.key.id),
candidate_id: None,
@@ -56,20 +65,9 @@ async fn execute_antigravity_quota_plan(
client_api_format: "gemini:chat".to_string(),
provider_api_format: "antigravity:fetch_available_models".to_string(),
model_name: Some("fetchAvailableModels".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "antigravity").await

30
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
View File

@@ -1,9 +1,12 @@
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
use super::super::shared::{
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
ProviderQuotaExecutionOutcome,
};
use super::parse::normalize_codex_plan_type;
use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use std::collections::BTreeMap;
pub(super) fn build_codex_refresh_headers(
@@ -58,6 +61,14 @@ pub(super) async fn execute_codex_quota_plan(
transport: &AdminGatewayProviderTransportSnapshot,
headers: BTreeMap<String, String>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("codex-quota:{}", transport.key.id),
candidate_id: None,
@@ -79,20 +90,9 @@ pub(super) async fn execute_codex_quota_plan(
client_api_format: "openai:cli".to_string(),
provider_api_format: "openai:cli".to_string(),
model_name: Some("codex-wham-usage".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "codex").await
}

26
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs
View File

@@ -1,3 +1,4 @@
use super::super::shared::default_provider_quota_execution_timeouts;
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
use crate::handlers::admin::provider::shared::payloads::{
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
@@ -6,7 +7,7 @@ use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use std::collections::BTreeMap;
use url::form_urlencoded;
use uuid::Uuid;
@@ -66,6 +67,14 @@ pub(super) async fn execute_kiro_quota_plan(
transport: &AdminGatewayProviderTransportSnapshot,
auth: &AdminKiroRequestAuth,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("kiro-quota:{}", transport.key.id),
candidate_id: None,
@@ -87,20 +96,9 @@ pub(super) async fn execute_kiro_quota_plan(
client_api_format: "claude:cli".to_string(),
provider_api_format: "kiro:usage".to_string(),
model_name: Some("kiro-usage-limits".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "kiro").await

27
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -4,16 +4,37 @@ use crate::handlers::admin::provider::shared::payloads::{
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionResult};
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
const PROVIDER_QUOTA_DEFAULT_TIMEOUT_MS: u64 = 30_000;
const PROVIDER_QUOTA_PROXY_TIMEOUT_MS: u64 = 60_000;
pub(super) enum ProviderQuotaExecutionOutcome {
Response(ExecutionResult),
Failure(String),
}
pub(super) fn default_provider_quota_execution_timeouts(
proxy: Option<&ProxySnapshot>,
) -> ExecutionTimeouts {
let timeout_ms = if proxy.is_some() {
PROVIDER_QUOTA_PROXY_TIMEOUT_MS
} else {
PROVIDER_QUOTA_DEFAULT_TIMEOUT_MS
};
ExecutionTimeouts {
connect_ms: Some(timeout_ms),
read_ms: Some(timeout_ms),
write_ms: Some(timeout_ms),
pool_ms: Some(timeout_ms),
total_ms: Some(timeout_ms),
..ExecutionTimeouts::default()
}
}
pub(super) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
@@ -127,6 +148,9 @@ pub(super) async fn execute_provider_quota_plan(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let proxy_source = state
.resolve_transport_proxy_source_with_tunnel_affinity(transport)
.await;
let proxy_url_present = plan
.proxy
.as_ref()
@@ -138,6 +162,7 @@ pub(super) async fn execute_provider_quota_plan(
endpoint_id = %transport.endpoint.id,
url = %plan.url,
tls_profile = ?plan.tls_profile.as_deref(),
proxy_source = ?proxy_source,
proxy_node_id = ?proxy_node_id,
proxy_url_present,
error = %error,

54
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
View File

@@ -9,28 +9,34 @@ use aether_data_contracts::repository::provider_catalog::{
pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
provider_type: &str,
endpoints: Vec<StoredProviderCatalogEndpoint>,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
let provider_type = provider_type.trim().to_ascii_lowercase();
match provider_type.as_str() {
"codex" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:cli")
}),
"antigravity" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& (endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:chat")
|| endpoint
"codex" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:cli"))
}),
.eq_ignore_ascii_case("openai:cli")
})
.cloned(),
"antigravity" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& (endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:chat")
|| endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:cli"))
})
.cloned(),
"kiro" => endpoints
.iter()
.find(|endpoint| {
@@ -41,8 +47,16 @@ pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
.eq_ignore_ascii_case("claude:cli")
})
.cloned()
.or_else(|| endpoints.into_iter().find(|endpoint| endpoint.is_active)),
_ => endpoints.into_iter().find(|endpoint| endpoint.is_active),
.or_else(|| {
endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned()
}),
_ => endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned(),
}
}
@@ -59,7 +73,7 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, endpoints)
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints)
else {
return Ok((false, None));
};

35
apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs
View File

@@ -3,16 +3,25 @@ use super::super::errors::{
};
use super::json_non_empty_string;
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
use aether_contracts::ProxySnapshot;
use axum::{body::Body, http, response::Response};
use url::form_urlencoded;
fn provider_oauth_transport_error_detail(prefix: &str, error: &str) -> String {
let error = error.trim();
if error.is_empty() {
return prefix.to_string();
}
format!("{prefix}: {error}")
}
pub(crate) async fn exchange_admin_provider_oauth_code(
state: &AdminAppState<'_>,
template: AdminProviderOAuthTemplate,
code: &str,
state_nonce: &str,
pkce_verifier: Option<&str>,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
let response = if template.provider_type == "claude_code" {
@@ -63,7 +72,7 @@ pub(crate) async fn exchange_admin_provider_oauth_code(
Some("application/json"),
Some(serde_json::Value::Object(body)),
None,
proxy_node_id,
proxy.clone(),
)
.await
} else {
@@ -100,12 +109,15 @@ pub(crate) async fn exchange_admin_provider_oauth_code(
Some("application/x-www-form-urlencoded"),
None,
Some(form_body),
proxy_node_id,
proxy.clone(),
)
.await
}
.map_err(|_| {
build_internal_control_error_response(http::StatusCode::BAD_REQUEST, "token exchange 失败")
.map_err(|error| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
provider_oauth_transport_error_detail("token exchange 失败", &error),
)
})?;
if !response.status.is_success() {
@@ -134,7 +146,7 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
state: &AdminAppState<'_>,
template: AdminProviderOAuthTemplate,
refresh_token: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
let scope = template.scopes.join(" ");
@@ -175,7 +187,7 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
Some("application/json"),
Some(serde_json::Value::Object(body)),
None,
proxy_node_id,
proxy.clone(),
)
.await
} else {
@@ -211,14 +223,17 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
Some("application/x-www-form-urlencoded"),
None,
Some(form_body),
proxy_node_id,
proxy.clone(),
)
.await
}
.map_err(|_| {
.map_err(|error| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Refresh Token 验证失败: token exchange 失败",
provider_oauth_transport_error_detail(
"Refresh Token 验证失败: token exchange 失败",
&error,
),
)
})?;

13
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
View File

@@ -28,9 +28,6 @@ pub(crate) fn normalize_pool_advanced_config(
serde_json::Value::Null => Ok(None),
// `pool_advanced: {}` still means "enable pool mode with defaults".
serde_json::Value::Object(map) => Ok(Some(serde_json::Value::Object(map))),
// Backward compatibility for older boolean payloads.
serde_json::Value::Bool(true) => Ok(Some(serde_json::json!({}))),
serde_json::Value::Bool(false) => Ok(None),
_ => Err("pool_advanced 必须是 JSON 对象".to_string()),
}
}
@@ -78,14 +75,14 @@ mod tests {
}
#[test]
fn normalize_pool_advanced_accepts_legacy_booleans() {
fn normalize_pool_advanced_rejects_legacy_booleans() {
assert_eq!(
normalize_pool_advanced_config(Some(json!(true))).expect("true should normalize"),
Some(json!({}))
normalize_pool_advanced_config(Some(json!(true))).unwrap_err(),
"pool_advanced 必须是 JSON 对象"
);
assert_eq!(
normalize_pool_advanced_config(Some(json!(false))).expect("false should normalize"),
None
normalize_pool_advanced_config(Some(json!(false))).unwrap_err(),
"pool_advanced 必须是 JSON 对象"
);
}
}

81
apps/aether-gateway/src/handlers/admin/request/provider/oauth.rs
View File

@@ -1,7 +1,7 @@
use super::*;
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTimeouts, RequestBody,
ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot, RequestBody,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
};
use aether_data::repository::provider_oauth::{
@@ -21,6 +21,7 @@ use url::Url;
const KIRO_IDC_AMZ_USER_AGENT: &str =
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
const ADMIN_PROVIDER_OAUTH_TIMEOUT_MS: u64 = 30_000;
const ADMIN_PROVIDER_OAUTH_PROXY_TIMEOUT_MS: u64 = 60_000;
pub(crate) struct AdminProviderOAuthHttpResponse {
pub(crate) status: http::StatusCode,
@@ -132,7 +133,7 @@ impl<'a> AdminAppState<'a> {
code: &str,
state_nonce: &str,
pkce_verifier: Option<&str>,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
crate::handlers::admin::provider::oauth::state::exchange_admin_provider_oauth_code(
self,
@@ -140,7 +141,7 @@ impl<'a> AdminAppState<'a> {
code,
state_nonce,
pkce_verifier,
proxy_node_id,
proxy,
)
.await
}
@@ -149,13 +150,13 @@ impl<'a> AdminAppState<'a> {
&self,
template: AdminProviderOAuthTemplate,
refresh_token: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
crate::handlers::admin::provider::oauth::state::exchange_admin_provider_oauth_refresh_token(
self,
template,
refresh_token,
proxy_node_id,
proxy,
)
.await
}
@@ -315,7 +316,7 @@ impl<'a> AdminAppState<'a> {
&self,
region: &str,
start_url: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let payload = post_kiro_device_oidc_json(
self,
@@ -337,7 +338,7 @@ impl<'a> AdminAppState<'a> {
],
"issuerUrl": start_url,
}),
proxy_node_id,
proxy,
)
.await?;
if payload
@@ -364,7 +365,7 @@ impl<'a> AdminAppState<'a> {
client_id: &str,
client_secret: &str,
start_url: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let payload = post_kiro_device_oidc_json(
self,
@@ -375,7 +376,7 @@ impl<'a> AdminAppState<'a> {
"clientSecret": client_secret,
"startUrl": start_url,
}),
proxy_node_id,
proxy,
)
.await?;
if payload
@@ -402,7 +403,7 @@ impl<'a> AdminAppState<'a> {
client_id: &str,
client_secret: &str,
device_code: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
post_kiro_device_oidc_json(
self,
@@ -414,11 +415,35 @@ impl<'a> AdminAppState<'a> {
"grantType": "urn:ietf:params:oauth:grant-type:device_code",
"deviceCode": device_code,
}),
proxy_node_id,
proxy,
)
.await
}
pub(crate) async fn resolve_admin_provider_oauth_operation_proxy_snapshot(
&self,
temporary_proxy_node_id: Option<&str>,
configured_proxies: &[Option<&serde_json::Value>],
) -> Option<ProxySnapshot> {
if let Some(snapshot) = self
.resolve_admin_proxy_node_snapshot(temporary_proxy_node_id)
.await
{
return Some(snapshot);
}
for proxy in configured_proxies {
if let Some(snapshot) = self
.app
.resolve_configured_proxy_snapshot_with_tunnel_affinity(*proxy)
.await
{
return Some(snapshot);
}
}
self.app.resolve_system_proxy_snapshot().await
}
pub(crate) async fn find_duplicate_provider_oauth_key(
&self,
provider_id: &str,
@@ -504,7 +529,7 @@ async fn post_kiro_device_oidc_json(
endpoint_key: &str,
default_url: String,
body: serde_json::Value,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let url = state.provider_oauth_token_url(endpoint_key, &default_url);
let host = Url::parse(&url)
@@ -539,7 +564,7 @@ async fn post_kiro_device_oidc_json(
Some("application/json"),
Some(body),
None,
proxy_node_id,
proxy,
)
.await
.map_err(|_| {
@@ -570,7 +595,7 @@ impl<'a> AdminAppState<'a> {
content_type: Option<&str>,
json_body: Option<serde_json::Value>,
body_bytes: Option<Vec<u8>>,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<AdminProviderOAuthHttpResponse, String> {
let body = if let Some(json_body) = json_body {
RequestBody::from_json(json_body)
@@ -581,6 +606,7 @@ impl<'a> AdminAppState<'a> {
body_ref: None,
}
};
let timeout_ms = admin_provider_oauth_timeout_ms(proxy.as_ref());
let plan = ExecutionPlan {
request_id: request_id.to_string(),
candidate_id: None,
@@ -601,21 +627,14 @@ impl<'a> AdminAppState<'a> {
client_api_format: "provider_oauth:exchange".to_string(),
provider_api_format: "provider_oauth:exchange".to_string(),
model_name: Some("oauth-exchange".to_string()),
proxy: if proxy_node_id
.map(str::trim)
.is_some_and(|value| !value.is_empty())
{
self.resolve_admin_proxy_node_snapshot(proxy_node_id).await
} else {
self.app.resolve_system_proxy_snapshot().await
},
proxy,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(ADMIN_PROVIDER_OAUTH_TIMEOUT_MS),
read_ms: Some(ADMIN_PROVIDER_OAUTH_TIMEOUT_MS),
write_ms: Some(ADMIN_PROVIDER_OAUTH_TIMEOUT_MS),
pool_ms: Some(ADMIN_PROVIDER_OAUTH_TIMEOUT_MS),
total_ms: Some(ADMIN_PROVIDER_OAUTH_TIMEOUT_MS),
connect_ms: Some(timeout_ms),
read_ms: Some(timeout_ms),
write_ms: Some(timeout_ms),
pool_ms: Some(timeout_ms),
total_ms: Some(timeout_ms),
..ExecutionTimeouts::default()
}),
};
@@ -632,6 +651,14 @@ impl<'a> AdminAppState<'a> {
}
}
fn admin_provider_oauth_timeout_ms(proxy: Option<&ProxySnapshot>) -> u64 {
if proxy.is_some() {
ADMIN_PROVIDER_OAUTH_PROXY_TIMEOUT_MS
} else {
ADMIN_PROVIDER_OAUTH_TIMEOUT_MS
}
}
fn maybe_insert_host_header(
mut headers: reqwest::header::HeaderMap,
host: &str,

166
apps/aether-gateway/src/handlers/admin/request/provider/transport.rs
View File

@@ -82,6 +82,15 @@ impl<'a> AdminAppState<'a> {
.await
}
pub(crate) async fn resolve_transport_proxy_source_with_tunnel_affinity(
&self,
transport: &AdminGatewayProviderTransportSnapshot,
) -> Option<&'static str> {
self.app
.resolve_transport_proxy_source_with_tunnel_affinity(transport)
.await
}
pub(crate) fn fixed_provider_template(
&self,
provider_type: &str,
@@ -115,15 +124,14 @@ impl<'a> AdminAppState<'a> {
return Some(snapshot);
}
if explicit_node_id.is_none() {
if let Some(snapshot) = self.app.resolve_system_proxy_snapshot().await {
return Some(snapshot);
}
let proxy = connector_config
.and_then(|config| config.get("proxy"))
.and_then(admin_provider_transport_proxy_snapshot);
if proxy.is_some() {
return proxy;
}
connector_config
.and_then(|config| config.get("proxy"))
.and_then(admin_provider_transport_legacy_proxy_snapshot)
self.app.resolve_system_proxy_snapshot().await
}
pub(crate) async fn resolve_admin_proxy_node_snapshot(
@@ -256,65 +264,46 @@ impl<'a> AdminAppState<'a> {
}
}
fn admin_provider_transport_legacy_proxy_snapshot(value: &Value) -> Option<ProxySnapshot> {
match value {
Value::String(proxy_url) => {
let proxy_url = proxy_url.trim();
if proxy_url.is_empty() {
return None;
}
Some(ProxySnapshot {
enabled: Some(true),
mode: admin_provider_transport_proxy_mode(Some(proxy_url)),
node_id: None,
label: None,
url: Some(proxy_url.to_string()),
extra: None,
})
}
Value::Object(object) => {
if object.get("enabled").and_then(Value::as_bool) == Some(false) {
return None;
}
let proxy_url = object
.get("url")
.or_else(|| object.get("proxy_url"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let username = object
.get("username")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let password = object
.get("password")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
Some(ProxySnapshot {
enabled: Some(true),
mode: object
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| admin_provider_transport_proxy_mode(Some(proxy_url))),
node_id: None,
label: object
.get("label")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
url: admin_provider_transport_inject_proxy_auth(proxy_url, username, password)
.or_else(|| Some(proxy_url.to_string())),
extra: None,
})
}
_ => None,
fn admin_provider_transport_proxy_snapshot(value: &Value) -> Option<ProxySnapshot> {
let object = value.as_object()?;
if object.get("enabled").and_then(Value::as_bool) == Some(false) {
return None;
}
let proxy_url = object
.get("url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let username = object
.get("username")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let password = object
.get("password")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
Some(ProxySnapshot {
enabled: Some(true),
mode: object
.get("mode")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| admin_provider_transport_proxy_mode(Some(proxy_url))),
node_id: None,
label: object
.get("label")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
url: admin_provider_transport_inject_proxy_auth(proxy_url, username, password)
.or_else(|| Some(proxy_url.to_string())),
extra: None,
})
}
fn admin_provider_transport_inject_proxy_auth(
@@ -354,3 +343,50 @@ fn admin_provider_transport_string_field(config: &Map<String, Value>, key: &str)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
#[cfg(test)]
mod tests {
use aether_contracts::ProxySnapshot;
use serde_json::json;
use super::admin_provider_transport_proxy_snapshot;
#[test]
fn connector_proxy_snapshot_requires_object_value() {
assert_eq!(
admin_provider_transport_proxy_snapshot(&json!("http://proxy.example:8080")),
None
);
}
#[test]
fn connector_proxy_snapshot_requires_url_field() {
assert_eq!(
admin_provider_transport_proxy_snapshot(&json!({
"proxy_url": "http://proxy.example:8080"
})),
None
);
}
#[test]
fn connector_proxy_snapshot_keeps_current_object_shape() {
assert_eq!(
admin_provider_transport_proxy_snapshot(&json!({
"url": "http://proxy.example:8080",
"username": "alice",
"password": "secret",
"mode": "http",
"label": "manual"
})),
Some(ProxySnapshot {
enabled: Some(true),
mode: Some("http".to_string()),
node_id: None,
label: Some("manual".to_string()),
url: Some("http://alice:secret@proxy.example:8080/".to_string()),
extra: None,
})
);
}
}

125
apps/aether-gateway/src/handlers/admin/request/system/import.rs
View File

@@ -46,6 +46,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
const ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES: usize = 10 * 1024 * 1024;
const MIN_ADMIN_SYSTEM_IMPORT_VERSION: (u32, u32) = (2, 2);
fn invalid_request(detail: impl Into<String>) -> (http::StatusCode, Value) {
(
@@ -342,19 +343,33 @@ struct ImportedWalletTarget {
updated_at_unix_secs: Option<u64>,
}
fn imported_users_export_is_legacy(version: Option<&Value>) -> bool {
fn imported_system_export_version(version: Option<&Value>) -> Result<(u32, u32), String> {
let Some(Value::String(version)) = version else {
return true;
return Err("version 必须是 x.y 字符串".to_string());
};
let version = version.trim();
if version.is_empty() {
return Err("version 必须是 x.y 字符串".to_string());
}
let mut parts = version.split('.');
let Some(major) = parts.next().and_then(|value| value.parse::<u32>().ok()) else {
return true;
return Err("version 必须是 x.y 字符串".to_string());
};
let Some(minor) = parts.next().and_then(|value| value.parse::<u32>().ok()) else {
return true;
return Err("version 必须是 x.y 字符串".to_string());
};
(major, minor) < (1, 3)
Ok((major, minor))
}
fn validate_imported_system_export_version(version: Option<&Value>) -> Result<(), String> {
let parsed = imported_system_export_version(version)?;
if parsed < MIN_ADMIN_SYSTEM_IMPORT_VERSION {
return Err(format!(
"version {}.{} 已不再支持；仅支持 2.2+ 导出格式",
parsed.0, parsed.1
));
}
Ok(())
}
fn imported_object_field<'a>(
@@ -385,11 +400,6 @@ fn imported_optional_bool(value: Option<&Value>) -> Result<Option<bool>, String>
match value {
None | Some(Value::Null) => Ok(None),
Some(Value::Bool(value)) => Ok(Some(*value)),
Some(Value::String(raw)) => match raw.trim().to_ascii_lowercase().as_str() {
"true" => Ok(Some(true)),
"false" => Ok(Some(false)),
_ => Err("字段必须是布尔值".to_string()),
},
_ => Err("字段必须是布尔值".to_string()),
}
}
@@ -402,11 +412,6 @@ fn imported_optional_i32(value: Option<&Value>, field_name: &str) -> Result<Opti
.ok_or_else(|| format!("{field_name} 必须是整数"))
.and_then(|value| i32::try_from(value).map_err(|_| format!("{field_name} 超出范围")))
.map(Some),
Some(Value::String(raw)) => raw
.trim()
.parse::<i32>()
.map(Some)
.map_err(|_| format!("{field_name} 必须是整数")),
_ => Err(format!("{field_name} 必须是整数")),
}
}
@@ -418,11 +423,6 @@ fn imported_optional_u64(value: Option<&Value>, field_name: &str) -> Result<Opti
.as_u64()
.ok_or_else(|| format!("{field_name} 必须是非负整数"))
.map(Some),
Some(Value::String(raw)) => raw
.trim()
.parse::<u64>()
.map(Some)
.map_err(|_| format!("{field_name} 必须是非负整数")),
_ => Err(format!("{field_name} 必须是非负整数")),
}
}
@@ -435,13 +435,6 @@ fn imported_optional_f64(value: Option<&Value>, field_name: &str) -> Result<Opti
.filter(|value| value.is_finite())
.ok_or_else(|| format!("{field_name} 必须是有限数值"))
.map(Some),
Some(Value::String(raw)) => raw
.trim()
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.ok_or_else(|| format!("{field_name} 必须是有限数值"))
.map(Some),
_ => Err(format!("{field_name} 必须是有限数值")),
}
}
@@ -491,16 +484,6 @@ fn imported_string_list_from_value(
.map(ToOwned::to_owned)
.collect(),
)),
Value::String(raw) => {
let trimmed = raw.trim();
if trimmed.is_empty() || trimmed.eq_ignore_ascii_case("null") {
return Ok(None);
}
if let Ok(decoded) = serde_json::from_str::<Value>(trimmed) {
return imported_string_list_from_value(Some(&decoded), field_name);
}
Ok(Some(vec![trimmed.to_string()]))
}
_ => Err(format!("{field_name} 必须是字符串列表")),
}
}
@@ -1584,7 +1567,6 @@ impl<'a> AdminAppState<'a> {
)))
}
};
let legacy_export = imported_users_export_is_legacy(root.get("version"));
let empty = Vec::new();
let users = match root.get("users") {
Some(Value::Array(items)) => items,
@@ -1614,6 +1596,8 @@ impl<'a> AdminAppState<'a> {
};
}
invalid_value!(validate_imported_system_export_version(root.get("version")));
let mut stats = AdminSystemUsersImportStats::default();
for (index, raw_user) in users.iter().enumerate() {
@@ -1883,7 +1867,7 @@ impl<'a> AdminAppState<'a> {
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(if legacy_export { 0 } else { 0 });
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
@@ -2063,7 +2047,7 @@ impl<'a> AdminAppState<'a> {
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(if legacy_export { 0 } else { 0 });
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
@@ -2321,3 +2305,64 @@ enum WalletOwner<'a> {
User(&'a str),
ApiKey(&'a str),
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
imported_optional_bool, imported_optional_f64, imported_optional_i32,
imported_optional_u64, imported_string_list_from_value,
validate_imported_system_export_version,
};
#[test]
fn import_requires_supported_export_version() {
assert!(validate_imported_system_export_version(Some(&json!("2.2"))).is_ok());
assert_eq!(
validate_imported_system_export_version(Some(&json!("2.1"))).unwrap_err(),
"version 2.1 已不再支持；仅支持 2.2+ 导出格式"
);
assert_eq!(
validate_imported_system_export_version(Some(&json!(null))).unwrap_err(),
"version 必须是 x.y 字符串"
);
}
#[test]
fn import_rejects_legacy_string_scalars() {
assert_eq!(
imported_optional_bool(Some(&json!("true"))).unwrap_err(),
"字段必须是布尔值"
);
assert_eq!(
imported_optional_i32(Some(&json!("5")), "rate_limit").unwrap_err(),
"rate_limit 必须是整数"
);
assert_eq!(
imported_optional_u64(Some(&json!("5")), "total_requests").unwrap_err(),
"total_requests 必须是非负整数"
);
assert_eq!(
imported_optional_f64(Some(&json!("1.25")), "total_cost_usd").unwrap_err(),
"total_cost_usd 必须是有限数值"
);
}
#[test]
fn import_rejects_legacy_string_lists() {
assert_eq!(
imported_string_list_from_value(Some(&json!("openai")), "allowed_providers")
.unwrap_err(),
"allowed_providers 必须是字符串列表"
);
assert_eq!(
imported_string_list_from_value(
Some(&json!("[\"openai:chat\"]")),
"allowed_api_formats"
)
.unwrap_err(),
"allowed_api_formats 必须是字符串列表"
);
}
}

2
apps/aether-gateway/src/handlers/admin/system/mod.rs
View File

@@ -6,4 +6,6 @@ mod proxy_nodes;
mod routes;
pub(super) mod shared;
#[cfg(test)]
pub(crate) use self::proxy_nodes::override_proxy_connectivity_probe_url_for_tests;
pub(super) use self::routes::maybe_build_local_admin_system_response;

457
apps/aether-gateway/src/handlers/admin/system/proxy_nodes.rs
View File

@@ -1,5 +1,6 @@
use std::time::{Duration, Instant};
use crate::execution_runtime::transport::format_upstream_request_error;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::query_param_value;
use crate::maintenance::{
@@ -12,6 +13,9 @@ 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,
};
use aether_contracts::tunnel::{
TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use axum::{
body::{Body, Bytes},
http,
@@ -21,7 +25,6 @@ use axum::{
use serde::de::DeserializeOwned;
use serde::Deserialize;
use serde_json::{json, Value};
use tokio::{net::TcpStream, time::timeout};
#[derive(Debug, Deserialize)]
struct ProxyNodeRegisterRequest {
@@ -131,6 +134,61 @@ const JSON_OBJECT_REQUIRED_DETAIL: &str = "请求体必须是合法的 JSON 对
const DEFAULT_PROXY_UPGRADE_BATCH_SIZE: usize = 1;
const DEFAULT_PROXY_UPGRADE_COOLDOWN_SECS: u64 = 60;
const DEFAULT_PROXY_UPGRADE_PROBE_TIMEOUT_SECS: u64 = 10;
const DEFAULT_PROXY_CONNECTIVITY_PROBE_URL: &str = "https://www.cloudflare.com/cdn-cgi/trace";
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;
#[cfg(test)]
fn manual_proxy_connectivity_probe_url_override() -> &'static std::sync::RwLock<Option<String>> {
static OVERRIDE: std::sync::OnceLock<std::sync::RwLock<Option<String>>> =
std::sync::OnceLock::new();
OVERRIDE.get_or_init(|| std::sync::RwLock::new(None))
}
#[cfg(test)]
fn manual_proxy_connectivity_probe_url_override_lock() -> &'static std::sync::Mutex<()> {
static LOCK: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
LOCK.get_or_init(|| std::sync::Mutex::new(()))
}
fn proxy_connectivity_probe_url() -> String {
#[cfg(test)]
if let Some(url) = manual_proxy_connectivity_probe_url_override()
.read()
.expect("probe url override lock should read")
.clone()
{
return url;
}
DEFAULT_PROXY_CONNECTIVITY_PROBE_URL.to_string()
}
#[cfg(test)]
pub(crate) struct ProxyConnectivityProbeUrlOverrideGuard(std::sync::MutexGuard<'static, ()>);
#[cfg(test)]
pub(crate) fn override_proxy_connectivity_probe_url_for_tests(
url: impl Into<String>,
) -> ProxyConnectivityProbeUrlOverrideGuard {
let guard = manual_proxy_connectivity_probe_url_override_lock()
.lock()
.expect("probe url override lock should acquire");
*manual_proxy_connectivity_probe_url_override()
.write()
.expect("probe url override lock should write") = Some(url.into());
ProxyConnectivityProbeUrlOverrideGuard(guard)
}
#[cfg(test)]
impl Drop for ProxyConnectivityProbeUrlOverrideGuard {
fn drop(&mut self) {
*manual_proxy_connectivity_probe_url_override()
.write()
.expect("probe url override lock should write") = None;
}
}
pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
state: &AdminAppState<'_>,
@@ -187,6 +245,26 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
));
}
if decision.route_kind.as_deref() == Some("get_node")
&& request_context.method() == http::Method::GET
{
if !state.has_proxy_node_reader() {
return Ok(Some(build_admin_proxy_nodes_data_unavailable_response()));
}
let Some(node_id) = admin_proxy_node_node_id_from_path(request_context.path()) else {
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
let Some(node) = state.find_proxy_node(&node_id).await? else {
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
return Ok(Some(
Json(json!({
"node": build_admin_proxy_node_detail_payload(&node),
}))
.into_response(),
));
}
if decision.route_kind.as_deref() == Some("register_node")
&& request_context.method() == http::Method::POST
{
@@ -367,7 +445,7 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
return Ok(Some(build_admin_proxy_nodes_not_found_response()));
};
return Ok(Some(
Json(test_proxy_node_connectivity(&node).await).into_response(),
Json(test_proxy_node_connectivity(state, &node).await).into_response(),
));
}
@@ -383,15 +461,7 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
Err(response) => return Ok(Some(response)),
};
return Ok(Some(
Json(
test_manual_proxy_connectivity(
&normalized.proxy_url,
normalized.host.as_str(),
normalized.port,
)
.await,
)
.into_response(),
Json(test_manual_proxy_connectivity(&normalized).await).into_response(),
));
}
@@ -716,6 +786,18 @@ struct DeletedProxyNodeCleanup {
cleared_keys: usize,
}
fn build_admin_proxy_node_detail_payload(
node: &aether_data::repository::proxy_nodes::StoredProxyNode,
) -> Value {
let mut payload = build_admin_proxy_node_payload(node);
if node.is_manual {
if let Value::Object(object) = &mut payload {
object.insert("proxy_password".to_string(), json!(node.proxy_password));
}
}
payload
}
#[derive(Debug, Clone)]
struct NormalizedManualProxyEndpoint {
proxy_url: String,
@@ -832,6 +914,7 @@ fn proxy_reference_matches_node_id(value: Option<&Value>, node_id: &str) -> bool
}
async fn test_proxy_node_connectivity(
state: &AdminAppState<'_>,
node: &aether_data::repository::proxy_nodes::StoredProxyNode,
) -> Value {
if node.is_manual {
@@ -854,12 +937,13 @@ async fn test_proxy_node_connectivity(
});
}
};
return test_manual_proxy_connectivity(
let proxy_url = proxy_url_with_auth(
&endpoint.proxy_url,
endpoint.host.as_str(),
endpoint.port,
node.proxy_username.as_deref(),
node.proxy_password.as_deref(),
)
.await;
.unwrap_or(endpoint.proxy_url);
return test_manual_proxy_connectivity(&proxy_url).await;
}
if !node.tunnel_mode {
@@ -880,39 +964,269 @@ async fn test_proxy_node_connectivity(
});
}
let probe_url = proxy_connectivity_probe_url();
match probe_tunnel_proxy_connectivity(state.app(), &node.id, &probe_url).await {
Ok(result) => {
if let Ok(status) = reqwest::StatusCode::from_u16(result.status) {
if status.is_success() {
return json!({
"success": true,
"latency_ms": result.latency_ms,
"exit_ip": parse_proxy_probe_exit_ip(&result.body),
"error": null,
});
}
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_proxy_probe_status_error(status, &result.body)),
});
}
json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": format!("代理探测返回非法状态码: {}", result.status),
})
}
Err(error) => json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&error),
}),
}
}
async fn test_manual_proxy_connectivity(proxy_url: &str) -> Value {
let probe_url = proxy_connectivity_probe_url();
test_manual_proxy_connectivity_with_probe_url(proxy_url, &probe_url).await
}
async fn test_manual_proxy_connectivity_with_probe_url(proxy_url: &str, probe_url: &str) -> Value {
let started_at = Instant::now();
let proxy = match reqwest::Proxy::all(proxy_url) {
Ok(proxy) => proxy,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
}
};
let mut builder = reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(PROXY_CONNECTIVITY_TIMEOUT_SECS))
.proxy(proxy)
.user_agent("aether-gateway/proxy-connectivity");
if proxy_url
.trim()
.to_ascii_lowercase()
.starts_with("https://")
{
builder = builder.danger_accept_invalid_certs(true);
}
let client = match builder.build() {
Ok(client) => client,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
}
};
let response = match client.get(probe_url).send().await {
Ok(response) => response,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
}
};
let status = response.status();
let body = match response.text().await {
Ok(body) => body,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
}
};
if !status.is_success() {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_proxy_probe_status_error(status, &body)),
});
}
json!({
"success": true,
"latency_ms": node.avg_latency_ms.map(|value| value.max(0.0).round() as u64),
"exit_ip": null,
"latency_ms": started_at.elapsed().as_millis() as u64,
"exit_ip": parse_proxy_probe_exit_ip(&body),
"error": null,
})
}
async fn test_manual_proxy_connectivity(_proxy_url: &str, host: &str, port: u16) -> Value {
let started_at = Instant::now();
match timeout(Duration::from_secs(5), TcpStream::connect((host, port))).await {
Ok(Ok(stream)) => {
drop(stream);
json!({
"success": true,
"latency_ms": started_at.elapsed().as_millis() as u64,
"exit_ip": null,
"error": null,
})
}
Ok(Err(error)) => json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&error.to_string()),
}),
Err(_) => json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": "连接超时",
}),
struct TunnelConnectivityProbeResult {
status: u16,
body: String,
latency_ms: u64,
}
async fn probe_tunnel_proxy_connectivity(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
) -> Result<TunnelConnectivityProbeResult, String> {
let trimmed_node_id = node_id.trim();
if trimmed_node_id.is_empty() {
return Err("proxy node id is empty".to_string());
}
if state.tunnel.has_local_proxy(trimmed_node_id) {
return probe_tunnel_proxy_connectivity_locally(state, trimmed_node_id, probe_url).await;
}
if let Some(owner) = state
.tunnel
.lookup_attachment_owner(state.data.as_ref(), trimmed_node_id)
.await
.map_err(|err| format!("lookup tunnel attachment owner failed: {err}"))?
{
if owner.gateway_instance_id != state.tunnel.local_instance_id() {
return probe_tunnel_proxy_connectivity_via_owner(
state,
trimmed_node_id,
probe_url,
&owner.relay_base_url,
&owner.gateway_instance_id,
)
.await;
}
state
.tunnel
.clear_local_attachment_if_stale(state.data.as_ref(), trimmed_node_id)
.await
.map_err(|err| format!("clear stale local tunnel attachment failed: {err}"))?;
}
probe_tunnel_proxy_connectivity_locally(state, trimmed_node_id, probe_url).await
}
async fn probe_tunnel_proxy_connectivity_locally(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
) -> Result<TunnelConnectivityProbeResult, String> {
let started_at = Instant::now();
let result = state
.tunnel
.probe_node_url_with_response(node_id, probe_url, PROXY_CONNECTIVITY_TIMEOUT_SECS)
.await?;
Ok(TunnelConnectivityProbeResult {
status: result.status,
body: result.body,
latency_ms: started_at.elapsed().as_millis() as u64,
})
}
async fn probe_tunnel_proxy_connectivity_via_owner(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
relay_base_url: &str,
owner_instance_id: &str,
) -> Result<TunnelConnectivityProbeResult, String> {
let owner_url = build_tunnel_owner_relay_url(relay_base_url, node_id)?;
let started_at = Instant::now();
let response = state
.client
.post(owner_url)
.header(
http::header::CONTENT_TYPE,
TUNNEL_RELAY_ENVELOPE_CONTENT_TYPE,
)
.header(
TUNNEL_RELAY_FORWARDED_BY_HEADER,
state.tunnel.local_instance_id(),
)
.header(TUNNEL_RELAY_OWNER_INSTANCE_HEADER, owner_instance_id)
.timeout(Duration::from_secs(PROXY_CONNECTIVITY_TIMEOUT_SECS))
.body(build_tunnel_probe_relay_envelope(probe_url)?)
.send()
.await
.map_err(|error| format!("owner tunnel relay probe failed: {error}"))?;
let status = response.status();
let body = response
.bytes()
.await
.map_err(|error| format!("failed to read owner tunnel relay probe body: {error}"))?;
if body.len() > MAX_PROXY_CONNECTIVITY_RESPONSE_BYTES {
return Err(format!(
"owner tunnel relay probe body exceeds {} bytes",
MAX_PROXY_CONNECTIVITY_RESPONSE_BYTES
));
}
Ok(TunnelConnectivityProbeResult {
status: status.as_u16(),
body: String::from_utf8_lossy(&body).to_string(),
latency_ms: started_at.elapsed().as_millis() as u64,
})
}
fn build_tunnel_probe_relay_envelope(probe_url: &str) -> Result<Vec<u8>, String> {
let meta = crate::tunnel::tunnel_protocol::RequestMeta {
method: "GET".to_string(),
url: probe_url.trim().to_string(),
headers: std::collections::HashMap::new(),
timeout: PROXY_CONNECTIVITY_TIMEOUT_SECS,
follow_redirects: Some(false),
http1_only: false,
};
let meta_bytes = serde_json::to_vec(&meta)
.map_err(|error| format!("encode tunnel probe metadata failed: {error}"))?;
let mut envelope = Vec::with_capacity(4 + meta_bytes.len());
envelope.extend_from_slice(&(meta_bytes.len() as u32).to_be_bytes());
envelope.extend_from_slice(&meta_bytes);
Ok(envelope)
}
fn build_tunnel_owner_relay_url(relay_base_url: &str, node_id: &str) -> Result<String, String> {
let mut url = url::Url::parse(relay_base_url)
.map_err(|error| format!("invalid owner relay base url: {error}"))?;
{
let mut segments = url
.path_segments_mut()
.map_err(|_| "owner relay base url cannot be a base-less URL".to_string())?;
segments.pop_if_empty();
segments.push("api");
segments.push("internal");
segments.push("tunnel");
segments.push("relay");
segments.push(node_id.trim());
}
Ok(url.to_string())
}
fn validate_register_request(
@@ -1039,10 +1353,16 @@ fn validate_manual_update_request(
fn validate_proxy_test_url_request(
input: ProxyNodeTestUrlRequest,
) -> Result<NormalizedManualProxyEndpoint, Response<Body>> {
let _ = normalize_optional_string(input.username.as_deref(), "username", 255)?;
let _ = normalize_optional_string(input.password.as_deref(), "password", 500)?;
normalize_manual_proxy_endpoint(&input.proxy_url)
) -> Result<String, Response<Body>> {
let username = normalize_optional_string(input.username.as_deref(), "username", 255)?;
let password = normalize_optional_string(input.password.as_deref(), "password", 500)?;
let endpoint = normalize_manual_proxy_endpoint(&input.proxy_url)?;
Ok(proxy_url_with_auth(
&endpoint.proxy_url,
username.as_deref(),
password.as_deref(),
)
.unwrap_or(endpoint.proxy_url))
}
fn admin_proxy_node_upgrade_action_node_id_from_path(path: &str, suffix: &str) -> Option<String> {
@@ -1425,6 +1745,53 @@ fn sanitize_proxy_error(detail: &str) -> String {
}
}
fn proxy_url_with_auth(
proxy_url: &str,
username: Option<&str>,
password: Option<&str>,
) -> Option<String> {
let username = username.map(str::trim).filter(|value| !value.is_empty())?;
let mut parsed = url::Url::parse(proxy_url).ok()?;
if parsed.set_username(username).is_err() {
return None;
}
let password = password.map(str::trim).filter(|value| !value.is_empty());
if parsed.set_password(password).is_err() {
return None;
}
Some(parsed.to_string())
}
fn parse_proxy_probe_exit_ip(body: &str) -> Option<String> {
body.lines().find_map(|line| {
let (key, value) = line.split_once('=')?;
if key.trim() != "ip" {
return None;
}
let value = value.trim();
if value.is_empty() {
return None;
}
Some(value.to_string())
})
}
fn format_proxy_probe_status_error(status: reqwest::StatusCode, body: &str) -> String {
let body = body.trim();
if body.is_empty() {
return format!("代理探测返回 HTTP {}", status.as_u16());
}
let truncated = if body.chars().count() > 200 {
let shortened: String = body.chars().take(200).collect();
format!("{shortened}...")
} else {
body.to_string()
};
format!("代理探测返回 HTTP {}: {truncated}", status.as_u16())
}
fn validate_optional_counter(value: Option<i64>, field: &str) -> Result<(), Response<Body>> {
if value.is_some_and(|value| value < 0) {
return Err(bad_request_response(format!("{field} 必须是非负整数")));

8
apps/aether-gateway/src/handlers/internal/gateway_helpers.rs
View File

@@ -351,12 +351,10 @@ pub(crate) fn gateway_error_message(error: GatewayError) -> String {
pub(crate) fn build_internal_tunnel_heartbeat_ack(
node: &StoredProxyNode,
heartbeat_id: Option<u64>,
heartbeat_id: u64,
) -> serde_json::Value {
let mut payload = serde_json::Map::new();
if let Some(heartbeat_id) = heartbeat_id {
payload.insert("heartbeat_id".to_string(), json!(heartbeat_id));
}
payload.insert("heartbeat_id".to_string(), json!(heartbeat_id));
if let Some(remote_config) = node.remote_config.as_ref() {
payload.insert("remote_config".to_string(), remote_config.clone());
payload.insert("config_version".to_string(), json!(node.config_version));
@@ -385,7 +383,7 @@ pub(crate) fn parse_internal_tunnel_heartbeat_request(
})?;
let node_id = payload.node_id.trim();
if node_id.is_empty() || node_id.len() > 36 {
if node_id.is_empty() || node_id.len() > 36 || payload.heartbeat_id == 0 {
return Err(build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"invalid heartbeat payload",

3
apps/aether-gateway/src/handlers/shared/payloads.rs
View File

@@ -4,8 +4,7 @@ use std::collections::BTreeMap;
#[derive(Debug, Deserialize)]
pub(crate) struct InternalTunnelHeartbeatRequest {
pub(crate) node_id: String,
#[serde(default)]
pub(crate) heartbeat_id: Option<u64>,
pub(crate) heartbeat_id: u64,
#[serde(default)]
pub(crate) heartbeat_interval: Option<i32>,
#[serde(default)]

4
apps/aether-gateway/src/maintenance/runtime.rs
View File

@@ -87,9 +87,9 @@ pub(super) fn postgres_error(
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(30);
const PROXY_NODE_STALE_SWEEP_INTERVAL: Duration = Duration::from_secs(5);
const PROXY_UPGRADE_ROLLOUT_INTERVAL: Duration = Duration::from_secs(15);
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 90;
const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
const POOL_MONITOR_INTERVAL: Duration = Duration::from_secs(5 * 60);
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;

80
apps/aether-gateway/src/state/proxy.rs
View File

@@ -92,26 +92,21 @@ impl AppState {
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
if let Some(snapshot) = self
.resolve_proxy_snapshot_from_config(transport.key.proxy.as_ref())
.await
{
return Some(snapshot);
}
if let Some(snapshot) = self
.resolve_proxy_snapshot_from_config(transport.provider.proxy.as_ref())
.await
{
return Some(snapshot);
}
if let Some(snapshot) = self.resolve_system_proxy_snapshot().await {
return Some(snapshot);
}
self.resolve_proxy_snapshot_from_config(transport.endpoint.proxy.as_ref())
self.resolve_transport_proxy_with_source_with_tunnel_affinity(transport)
.await
.map(|(snapshot, _)| snapshot)
}
async fn resolve_proxy_snapshot_from_config(
pub(crate) async fn resolve_transport_proxy_source_with_tunnel_affinity(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<&'static str> {
self.resolve_transport_proxy_with_source_with_tunnel_affinity(transport)
.await
.map(|(_, source)| source)
}
pub(crate) async fn resolve_configured_proxy_snapshot_with_tunnel_affinity(
&self,
raw: Option<&Value>,
) -> Option<ProxySnapshot> {
@@ -127,6 +122,37 @@ impl AppState {
proxy_snapshot_from_object(object)
}
async fn resolve_transport_proxy_with_source_with_tunnel_affinity(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<(ProxySnapshot, &'static str)> {
if let Some(snapshot) = self
.resolve_configured_proxy_snapshot_with_tunnel_affinity(transport.key.proxy.as_ref())
.await
{
return Some((snapshot, "key"));
}
if let Some(snapshot) = self
.resolve_configured_proxy_snapshot_with_tunnel_affinity(
transport.endpoint.proxy.as_ref(),
)
.await
{
return Some((snapshot, "endpoint"));
}
if let Some(snapshot) = self
.resolve_configured_proxy_snapshot_with_tunnel_affinity(
transport.provider.proxy.as_ref(),
)
.await
{
return Some((snapshot, "provider"));
}
self.resolve_system_proxy_snapshot()
.await
.map(|snapshot| (snapshot, "system"))
}
}
fn proxy_enabled(object: &Map<String, Value>) -> bool {
@@ -206,12 +232,22 @@ fn proxy_url_with_node_auth(
if parsed.set_username(username).is_err() {
return None;
}
let password = password
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_default();
if parsed.set_password(Some(password)).is_err() {
let password = password.map(str::trim).filter(|value| !value.is_empty());
if parsed.set_password(password).is_err() {
return None;
}
Some(parsed.to_string())
}
#[cfg(test)]
mod tests {
use super::proxy_url_with_node_auth;
#[test]
fn proxy_url_with_node_auth_omits_empty_password_separator() {
assert_eq!(
proxy_url_with_node_auth("socks5://proxy.example:1080", Some("alice"), None).as_deref(),
Some("socks5://alice@proxy.example:1080")
);
}
}

146
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
View File

@@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -14,7 +15,7 @@ use serde_json::json;
use super::super::super::{
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
sample_key, start_server,
sample_key, sample_proxy_node, start_server,
};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
@@ -29,6 +30,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
url: String,
authorization: String,
provider_api_format: String,
total_ms: Option<u64>,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -67,6 +69,10 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
total_ms: plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
@@ -178,6 +184,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
seen_execution_runtime_request.provider_api_format,
"openai:cli"
);
assert_eq!(seen_execution_runtime_request.total_ms, Some(30_000));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-a".to_string()])
@@ -215,6 +222,143 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_for_codex_proxy_with_extended_timeout() {
let upstream =
Router::new().route(
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
any(|_request: Request| async move {
(StatusCode::OK, Body::from("unexpected upstream hit"))
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<aether_contracts::ExecutionPlan>));
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 plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(plan.clone());
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 12.5,
"reset_after_seconds": 18000,
"reset_at": 1_900_000_000u64,
"window_minutes": 300
}
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut provider = StoredProviderCatalogProvider::new(
"provider-codex".to_string(),
"codex".to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build");
provider.proxy = Some(json!({
"node_id": "proxy-node-codex-quota",
"enabled": true
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:cli",
"https://chatgpt.com/backend-api",
)],
vec![sample_key(
"key-codex-a",
"provider-codex",
"openai:cli",
"sk-codex-123",
)],
));
let mut manual_node = sample_proxy_node("proxy-node-codex-quota");
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/refresh-quota"
))
.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("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let plan = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
plan.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-codex-quota")
);
let timeouts = plan.timeouts.expect("timeouts should exist");
assert_eq!(timeouts.connect_ms, Some(60_000));
assert_eq!(timeouts.read_ms, Some(60_000));
assert_eq!(timeouts.write_ms, Some(60_000));
assert_eq!(timeouts.pool_ms, Some(60_000));
assert_eq!(timeouts.total_ms, Some(60_000));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_admin_principal() {
#[derive(Debug, Clone)]

390
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -398,10 +398,7 @@ async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_a
.next()
.expect("persisted key should exist");
assert_eq!(persisted.auth_type, "oauth");
assert_eq!(
persisted.proxy,
Some(json!({"node_id":"proxy-node-kiro","enabled":true}))
);
assert_eq!(persisted.proxy, None);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -1133,10 +1130,7 @@ async fn gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_p
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id":"proxy-node-batch-import","enabled":true}))
);
assert_eq!(persisted.proxy, None);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -1604,10 +1598,7 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id":"proxy-node-codex-oauth","enabled":true}))
);
assert_eq!(persisted.proxy, None);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -1775,10 +1766,7 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id":"proxy-node-codex-import","enabled":true}))
);
assert_eq!(persisted.proxy, None);
let decrypted_api_key =
decrypt_python_fernet_ciphertext(DEVELOPMENT_ENCRYPTION_KEY, &persisted.encrypted_api_key)
.expect("api key should decrypt");
@@ -1819,6 +1807,12 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtim
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-codex-import"));
let timeouts = plan.timeouts.as_ref().expect("timeouts should exist");
assert_eq!(timeouts.connect_ms, Some(60_000));
assert_eq!(timeouts.read_ms, Some(60_000));
assert_eq!(timeouts.write_ms, Some(60_000));
assert_eq!(timeouts.pool_ms, Some(60_000));
assert_eq!(timeouts.total_ms, Some(60_000));
assert_eq!(plan.request_id, "provider-oauth:refresh-token");
assert_eq!(plan.method, "POST");
assert_eq!(plan.url, "https://oauth.example/oauth/token");
@@ -1920,10 +1914,7 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtim
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
assert_eq!(
keys[0].proxy,
Some(json!({"node_id":"proxy-node-codex-import","enabled":true}))
);
assert_eq!(keys[0].proxy, None);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
@@ -1932,6 +1923,136 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtim
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_provider_proxy_before_system_proxy(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-codex-provider"));
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "imported-codex-access-token",
"refresh_token": "imported-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "alice@example.com",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-codex-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut provider_node = sample_proxy_node("proxy-node-codex-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let mut system_node = sample_proxy_node("proxy-node-codex-system");
system_node.status = "online".to_string();
system_node.is_manual = true;
system_node.tunnel_mode = false;
system_node.tunnel_connected = false;
system_node.proxy_url = Some("http://proxy-system.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
provider_node,
system_node,
]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-codex-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.header(crate::constants::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")
.json(&json!({
"refresh_token": "provider-import-refresh-token",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "codex");
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].proxy, None);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
assert_eq!(
plans[0]
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-codex-provider")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_system_proxy() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
@@ -2044,6 +2165,71 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtim
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_import_refresh_token_surfaces_execution_runtime_error_detail() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|| async { StatusCode::INTERNAL_SERVER_ERROR }),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.header(crate::constants::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")
.json(&json!({
"refresh_token": "provider-import-refresh-token",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::BAD_REQUEST, "payload={payload}");
assert!(
payload["detail"]
.as_str()
.expect("detail should be string")
.contains("execution runtime returned HTTP 500"),
"payload={payload}"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -2162,10 +2348,7 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_ad
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id":"proxy-node-kiro-batch","enabled":true}))
);
assert_eq!(stored_key.proxy, None);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
@@ -2324,10 +2507,7 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id":"proxy-node-kiro-batch-runtime","enabled":true}))
);
assert_eq!(stored_key.proxy, None);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
@@ -2653,6 +2833,160 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:local-refresh-token" {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "alice@example.com",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:cli",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-provider",
"provider-codex",
"openai:cli",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let mut provider_node = sample_proxy_node("proxy-node-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let mut system_node = sample_proxy_node("proxy-node-system");
system_node.status = "online".to_string();
system_node.is_manual = true;
system_node.tunnel_mode = false;
system_node.tunnel_connected = false;
system_node.proxy_url = Some("http://proxy-system.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
provider_node,
system_node,
]));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-provider/refresh"
))
.header(crate::constants::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("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let plans = execution_plans.lock().expect("mutex should lock");
let refresh_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:local-refresh-token")
.expect("local refresh plan should exist");
assert_eq!(
refresh_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-provider")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy(
) {

237
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
View File

@@ -6,11 +6,13 @@ use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, StoredProxyNodeEvent,
};
use axum::body::Body;
use axum::extract::ws::Message;
use axum::routing::any;
use axum::{extract::Request, Router};
use base64::Engine as _;
use http::StatusCode;
use serde_json::json;
use tokio::net::TcpListener;
use tokio::sync::watch;
use super::super::{
build_router_with_state, hash_management_token, sample_endpoint, sample_key,
@@ -25,6 +27,7 @@ use crate::maintenance::{
record_proxy_upgrade_traffic_success, skip_proxy_upgrade_rollout_node,
start_proxy_upgrade_rollout,
};
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
#[tokio::test]
async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal() {
@@ -107,6 +110,49 @@ async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal(
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_full_manual_proxy_node_detail_locally_with_trusted_admin_principal() {
let mut manual_node = sample_proxy_node("proxy-node-manual");
manual_node.name = "alpha-manual".to_string();
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
manual_node.proxy_username = Some("alice".to_string());
manual_node.proxy_password = Some("supersecret".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
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 response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/proxy-node-manual"
))
.header(crate::constants::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("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["node"]["id"], "proxy-node-manual");
assert_eq!(payload["node"]["proxy_username"], "alice");
assert_eq!(payload["node"]["proxy_password"], "supersecret");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_reports_active_proxy_upgrade_rollout_in_proxy_node_list() {
let mut alpha = sample_proxy_node("node-alpha");
@@ -719,15 +765,28 @@ async fn gateway_registers_and_unregisters_proxy_nodes_locally_with_management_t
#[tokio::test]
async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let proxy_port = listener
.local_addr()
.expect("listener addr should resolve")
.port();
let accept_handle =
tokio::spawn(async move { while let Ok((_stream, _addr)) = listener.accept().await {} });
let proxy_auths = Arc::new(Mutex::new(Vec::<Option<String>>::new()));
let proxy_auths_clone = Arc::clone(&proxy_auths);
let proxy = Router::new().fallback(any(move |request: Request| {
let proxy_auths_inner = Arc::clone(&proxy_auths_clone);
async move {
proxy_auths_inner.lock().expect("mutex should lock").push(
request
.headers()
.get("proxy-authorization")
.and_then(|value| value.to_str().ok())
.map(str::to_string),
);
(
StatusCode::OK,
Body::from("fl=1234\nip=203.0.113.10\nwarp=off\n"),
)
}
}));
let (proxy_url, proxy_handle) = start_server(proxy).await;
let _probe_url_guard = crate::handlers::admin::override_proxy_connectivity_probe_url_for_tests(
"http://probe.example/cdn-cgi/trace",
);
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
let gateway = build_router_with_state(
@@ -739,7 +798,6 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let proxy_url = format!("http://127.0.0.1:{proxy_port}");
let create_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/manual"))
@@ -749,7 +807,9 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"name": "manual-node",
"proxy_url": proxy_url,
"proxy_url": proxy_url.clone(),
"username": "alice",
"password": "supersecret",
"region": "US-West"
}))
.send()
@@ -770,6 +830,8 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
assert_eq!(create_payload["node"]["is_manual"], true);
assert_eq!(create_payload["node"]["status"], "online");
assert_eq!(create_payload["node"]["proxy_url"], proxy_url);
assert_eq!(create_payload["node"]["proxy_username"], "alice");
assert_eq!(create_payload["node"]["proxy_password"], "su****et");
let test_url_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/test-url"))
@@ -778,7 +840,9 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"proxy_url": proxy_url
"proxy_url": proxy_url.clone(),
"username": "alice",
"password": "supersecret"
}))
.send()
.await
@@ -790,6 +854,7 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
.expect("json body should parse");
assert_eq!(test_url_payload["success"], true);
assert!(test_url_payload["latency_ms"].is_u64());
assert_eq!(test_url_payload["exit_ip"], "203.0.113.10");
let test_node_response = client
.post(format!(
@@ -808,6 +873,19 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
.await
.expect("json body should parse");
assert_eq!(test_node_payload["success"], true);
assert_eq!(test_node_payload["exit_ip"], "203.0.113.10");
let expected_proxy_auth = format!(
"Basic {}",
base64::engine::general_purpose::STANDARD.encode("alice:supersecret")
);
assert_eq!(
proxy_auths.lock().expect("mutex should lock").as_slice(),
[
Some(expected_proxy_auth.clone()),
Some(expected_proxy_auth.clone()),
]
);
let update_response = client
.patch(format!("{gateway_url}/api/admin/proxy-nodes/{node_id}"))
@@ -832,7 +910,7 @@ async fn gateway_creates_updates_and_tests_manual_proxy_nodes_locally() {
assert_eq!(update_payload["node"]["proxy_url"], proxy_url);
gateway_handle.abort();
accept_handle.abort();
proxy_handle.abort();
}
#[tokio::test]
@@ -869,6 +947,135 @@ async fn gateway_tests_disconnected_tunnel_proxy_nodes_locally() {
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_tests_connected_tunnel_proxy_nodes_with_active_probe() {
let _probe_url_guard = crate::handlers::admin::override_proxy_connectivity_probe_url_for_tests(
"https://probe.example/cdn-cgi/trace",
);
let mut node = sample_proxy_node("node-online");
node.status = "online".to_string();
node.tunnel_connected = true;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![node]));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
proxy_node_repository,
));
let tunnel_state = state.tunnel.app_state();
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false);
tunnel_state
.hub
.register_proxy(Arc::new(TunnelProxyConn::new(
500,
"node-online".to_string(),
"Node Online".to_string(),
proxy_tx,
proxy_close_tx,
16,
)));
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_task = tokio::spawn({
let gateway_url = gateway_url.clone();
async move {
reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/node-online/test"
))
.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
}
});
let request_headers = match proxy_rx.recv().await.expect("headers frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_header =
tunnel_protocol::FrameHeader::parse(&request_headers).expect("request header should parse");
assert_eq!(request_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
let meta_payload = tunnel_protocol::decode_payload(&request_headers, &request_header)
.expect("request header payload should decode");
let meta: tunnel_protocol::RequestMeta =
serde_json::from_slice(&meta_payload).expect("request meta should parse");
assert_eq!(meta.method, "GET");
assert_eq!(meta.url, "https://probe.example/cdn-cgi/trace");
assert_eq!(meta.follow_redirects, Some(false));
let request_body = match proxy_rx.recv().await.expect("body frame should arrive") {
Message::Binary(data) => data,
other => panic!("unexpected message: {other:?}"),
};
let request_body_header =
tunnel_protocol::FrameHeader::parse(&request_body).expect("request body should parse");
assert_eq!(request_body_header.msg_type, tunnel_protocol::REQUEST_BODY);
assert_ne!(
request_body_header.flags & tunnel_protocol::FLAG_END_STREAM,
0,
"probe body frame should close the stream"
);
let response_meta = tunnel_protocol::ResponseMeta {
status: 200,
headers: vec![("content-type".to_string(), "text/plain".to_string())],
};
let response_meta_bytes =
serde_json::to_vec(&response_meta).expect("response meta should serialize");
let mut response_headers_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_HEADERS,
0,
&response_meta_bytes,
);
tunnel_state
.hub
.handle_proxy_frame(500, &mut response_headers_frame)
.await;
let mut response_body_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::RESPONSE_BODY,
0,
b"fl=1234\nip=203.0.113.10\nwarp=off\n",
);
tunnel_state
.hub
.handle_proxy_frame(500, &mut response_body_frame)
.await;
let mut response_end_frame = tunnel_protocol::encode_frame(
request_header.stream_id,
tunnel_protocol::STREAM_END,
0,
&[],
);
tunnel_state
.hub
.handle_proxy_frame(500, &mut response_end_frame)
.await;
let response = request_task
.await
.expect("request task should complete")
.expect("test-node request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], true);
assert!(payload["latency_ms"].is_u64());
assert_eq!(payload["exit_ip"], "203.0.113.10");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_deletes_proxy_nodes_and_clears_proxy_refs_locally() {
let mut manual_node = sample_proxy_node("manual-node-1");
@@ -1164,6 +1371,7 @@ async fn gateway_updates_proxy_node_config_and_batches_upgrade_locally() {
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-online",
"heartbeat_id": 77,
"heartbeat_interval": 45,
"active_connections": 3,
"total_requests": 5,
@@ -1179,6 +1387,7 @@ async fn gateway_updates_proxy_node_config_and_batches_upgrade_locally() {
.json()
.await
.expect("json body should parse");
assert_eq!(heartbeat_payload["heartbeat_id"], 77);
assert_eq!(heartbeat_payload["config_version"], 10);
assert!(heartbeat_payload.get("upgrade_to").is_none());
assert_eq!(heartbeat_payload["remote_config"]["allowed_ports"][1], 8443);

109
apps/aether-gateway/src/tests/control/admin/system_import.rs
View File

@@ -493,7 +493,7 @@ async fn gateway_returns_503_for_admin_system_config_import_when_local_data_is_u
}
#[tokio::test]
async fn gateway_accepts_legacy_admin_system_config_import_versions() {
async fn gateway_rejects_legacy_admin_system_config_import_versions() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
@@ -519,10 +519,13 @@ async fn gateway_accepts_legacy_admin_system_config_import_versions() {
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "配置导入成功");
assert_eq!(payload["stats"]["errors"], json!([]));
let detail = payload["detail"]
.as_str()
.expect("detail should be a string");
assert!(detail.contains(&format!("不支持的配置版本: {version}")));
assert!(detail.contains("支持的版本: 2.2"));
}
gateway_handle.abort();
@@ -562,7 +565,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"version": "1.3",
"version": "2.2",
"merge_mode": "overwrite",
"users": [{
"email": "alice@example.com",
@@ -761,8 +764,40 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
}
#[tokio::test]
async fn gateway_imports_admin_system_config_fixtures_from_legacy_exports() {
for fixture in ["v20", "v21", "v22"] {
async fn gateway_imports_admin_system_config_fixture_v22() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(build_empty_admin_system_data_state()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/config/import"))
.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")
.json(&fixture_system_import_payload("v22"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "配置导入成功");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_system_config_fixtures_from_removed_legacy_exports() {
for fixture in ["v20", "v21"] {
let version = match fixture {
"v20" => "2.0",
"v21" => "2.1",
_ => unreachable!("unexpected fixture"),
};
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
@@ -783,16 +818,59 @@ async fn gateway_imports_admin_system_config_fixtures_from_legacy_exports() {
assert_eq!(
response.status(),
StatusCode::OK,
"fixture {fixture} should import"
StatusCode::BAD_REQUEST,
"fixture {fixture} should be rejected"
);
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "配置导入成功");
let detail = payload["detail"]
.as_str()
.expect("detail should be a string");
assert!(detail.contains(&format!("不支持的配置版本: {version}")));
assert!(detail.contains("支持的版本: 2.2"));
gateway_handle.abort();
}
}
#[tokio::test]
async fn gateway_rejects_legacy_user_import_string_bool_field() {
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::default());
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123")])
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/users/import"))
.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")
.json(&json!({
"version": "2.2",
"merge_mode": "overwrite",
"users": [{
"email": "legacy@example.com",
"email_verified": "true"
}]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "字段必须是布尔值");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_reports_field_path_for_invalid_admin_system_config_import_shape() {
let gateway = build_router_with_state(
@@ -835,7 +913,7 @@ async fn gateway_reports_field_path_for_invalid_admin_system_config_import_shape
}
#[tokio::test]
async fn gateway_imports_admin_system_config_with_numeric_string_prices() {
async fn gateway_rejects_admin_system_config_with_numeric_string_prices() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
@@ -860,12 +938,11 @@ async fn gateway_imports_admin_system_config_with_numeric_string_prices() {
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body: Value = response.json().await.expect("json body should parse");
assert_eq!(body["message"], "配置导入成功");
assert_eq!(body["stats"]["global_models"]["created"], json!(1));
assert_eq!(body["stats"]["providers"]["created"], json!(1));
assert_eq!(body["stats"]["models"]["created"], json!(1));
let detail = body["detail"].as_str().expect("detail should be a string");
assert!(detail.contains("配置文件格式无效"));
assert!(detail.contains("default_price_per_request"));
gateway_handle.abort();
}

31
apps/aether-gateway/src/tests/control/internal.rs
View File

@@ -76,6 +76,37 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_internal_tunnel_heartbeat_without_heartbeat_id() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-123",
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
Arc::clone(&repository),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-123",
"heartbeat_interval": 45,
"active_connections": 5
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_internal_tunnel_node_status_locally_with_loopback() {
let upstream_hits = Arc::new(Mutex::new(0usize));

83
apps/aether-gateway/src/tunnel/embedded/hub.rs
View File

@@ -85,6 +85,7 @@ pub struct ProxyConn {
next_stream_id: AtomicU32,
pub stream_count: AtomicUsize,
pub max_streams: usize,
draining: AtomicBool,
}
impl ProxyConn {
@@ -104,6 +105,7 @@ impl ProxyConn {
next_stream_id: AtomicU32::new(2),
stream_count: AtomicUsize::new(0),
max_streams,
draining: AtomicBool::new(false),
}
}
@@ -159,13 +161,21 @@ impl ProxyConn {
}
pub fn is_available(&self) -> bool {
!self.outbound.is_closing()
!self.outbound.is_closing() && !self.is_draining()
}
pub fn request_close(&self) {
self.outbound.mark_closing();
}
pub fn mark_draining(&self) -> bool {
!self.draining.swap(true, Ordering::AcqRel)
}
pub fn is_draining(&self) -> bool {
self.draining.load(Ordering::Acquire)
}
pub fn send(&self, msg: Message) -> SendStatus {
let was_closing = self.outbound.is_closing();
let status = self.outbound.send(msg);
@@ -659,10 +669,16 @@ impl HubRouter {
}
protocol::PONG => {}
protocol::GOAWAY => {
warn!(
proxy_conn_id = proxy_conn_id,
"received GOAWAY from proxy connection"
);
if let Some(pc) = self.proxy_conns_by_id.get(&proxy_conn_id) {
let first = pc.mark_draining();
if first {
warn!(
proxy_conn_id = proxy_conn_id,
node_id = %pc.node_id,
"received GOAWAY from proxy connection; marking connection draining"
);
}
}
}
_ => {
debug!(
@@ -975,6 +991,63 @@ mod tests {
assert_ne!(second_header.flags & protocol::FLAG_END_STREAM, 0);
}
#[tokio::test]
async fn goaway_marks_connection_draining_and_reroutes_new_streams() {
let hub = HubRouter::new(ControlPlaneClient::disabled());
let (proxy_one_tx, mut proxy_one_rx) = bounded_queue(8);
let (proxy_one_close_tx, _) = watch::channel(false);
let proxy_one = Arc::new(ProxyConn::new(
201,
"node-drain".to_string(),
"Node Drain".to_string(),
proxy_one_tx,
proxy_one_close_tx,
16,
));
hub.register_proxy(Arc::clone(&proxy_one));
let (proxy_two_tx, mut proxy_two_rx) = bounded_queue(8);
let (proxy_two_close_tx, _) = watch::channel(false);
let proxy_two = Arc::new(ProxyConn::new(
202,
"node-drain".to_string(),
"Node Drain".to_string(),
proxy_two_tx,
proxy_two_close_tx,
16,
));
hub.register_proxy(Arc::clone(&proxy_two));
let mut goaway = protocol::encode_goaway();
hub.handle_proxy_frame(201, &mut goaway).await;
assert!(
proxy_one.is_draining(),
"first connection should be draining"
);
assert!(
!proxy_two.is_draining(),
"second connection should remain schedulable"
);
let _stream = hub
.open_local_stream("node-drain", &build_meta())
.expect("open local stream");
assert!(
proxy_one_rx.try_recv().is_err(),
"draining connection should not receive new streams"
);
let routed = proxy_two_rx
.try_recv()
.expect("headers should route to second connection");
let routed_data = match routed {
Message::Binary(data) => data.to_vec(),
other => panic!("unexpected message: {other:?}"),
};
let header = protocol::FrameHeader::parse(&routed_data).expect("frame header");
assert_eq!(header.msg_type, protocol::REQUEST_HEADERS);
}
#[tokio::test]
async fn heartbeat_callback_failure_does_not_send_fake_ack() {
let hub = HubRouter::new(ControlPlaneClient::local(

2
apps/aether-gateway/src/tunnel/embedded/mod.rs
View File

@@ -22,7 +22,7 @@ use tracing::warn;
use crate::data::GatewayDataState;
pub use control_plane::ControlPlaneClient;
pub use hub::{ConnConfig, HubRouter, ProxyConn};
pub use hub::{ConnConfig, HubRouter, LocalBodyEvent, ProxyConn};
pub use local_relay::relay_request;
#[derive(Clone)]

96
apps/aether-gateway/src/tunnel/mod.rs
View File

@@ -46,12 +46,13 @@ pub(crate) const TUNNEL_NODE_STATUS_PATH: &str = "/api/internal/tunnel/node-stat
pub(crate) const TUNNEL_RELAY_PATH_PATTERN: &str = "/api/internal/tunnel/relay/{node_id}";
pub(crate) const TUNNEL_ROUTE_FAMILY: &str = "tunnel_manage";
const DEFAULT_PROXY_IDLE_TIMEOUT_SECS: u64 = 0;
const DEFAULT_PING_INTERVAL_SECS: u64 = 15;
const DEFAULT_PROXY_IDLE_TIMEOUT_MS: u64 = 900;
const DEFAULT_PING_INTERVAL_MS: u64 = 250;
const DEFAULT_MAX_STREAMS: usize = 2048;
const DEFAULT_OUTBOUND_QUEUE_CAPACITY: usize = 128;
const DEFAULT_ATTACHMENT_TTL_SECS: u64 = 90;
const DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES: usize = 5_242_880;
const DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES: usize = 64 * 1024;
const TUNNEL_ATTACHMENT_KEY_PREFIX: &str = "tunnel.attachments.";
const TUNNEL_ATTACHMENT_REDIS_KEY_PREFIX: &str = "tunnel:attachments:";
const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
@@ -61,8 +62,7 @@ const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
#[derive(Debug, Deserialize)]
struct InternalTunnelHeartbeatRequest {
node_id: String,
#[serde(default)]
heartbeat_id: Option<u64>,
heartbeat_id: u64,
#[serde(default)]
heartbeat_interval: Option<i32>,
#[serde(default)]
@@ -380,6 +380,12 @@ pub(crate) struct TunnelStatsSnapshot {
pub(crate) active_streams: usize,
}
#[derive(Debug, Clone)]
pub(crate) struct TunnelProbeResponse {
pub(crate) status: u16,
pub(crate) body: String,
}
impl EmbeddedTunnelState {
pub(crate) fn new() -> Self {
Self::with_data(Arc::new(GatewayDataState::disabled()))
@@ -409,8 +415,8 @@ impl EmbeddedTunnelState {
inner: TunnelAppState::new(
build_embedded_control_plane(Arc::clone(&data), attachment_directory.clone()),
ConnConfig {
ping_interval: Duration::from_secs(DEFAULT_PING_INTERVAL_SECS),
idle_timeout: Duration::from_secs(DEFAULT_PROXY_IDLE_TIMEOUT_SECS),
ping_interval: Duration::from_millis(DEFAULT_PING_INTERVAL_MS),
idle_timeout: Duration::from_millis(DEFAULT_PROXY_IDLE_TIMEOUT_MS),
outbound_queue_capacity: DEFAULT_OUTBOUND_QUEUE_CAPACITY,
},
DEFAULT_MAX_STREAMS,
@@ -451,6 +457,18 @@ impl EmbeddedTunnelState {
url: &str,
timeout_secs: u64,
) -> Result<u16, String> {
Ok(self
.probe_node_url_with_response(node_id, url, timeout_secs)
.await?
.status)
}
pub(crate) async fn probe_node_url_with_response(
&self,
node_id: &str,
url: &str,
timeout_secs: u64,
) -> Result<TunnelProbeResponse, String> {
let timeout_secs = timeout_secs.clamp(5, 60);
let meta = tunnel_protocol::RequestMeta {
method: "GET".to_string(),
@@ -469,7 +487,35 @@ impl EmbeddedTunnelState {
let response = stream
.wait_headers(Duration::from_secs(timeout_secs))
.await?;
Ok(response.status)
let Some(mut body_rx) = stream.take_body_receiver() else {
return Err("missing tunnel probe response body receiver".to_string());
};
let body = tokio::time::timeout(Duration::from_secs(timeout_secs), async {
let mut body_bytes = Vec::new();
while let Some(event) = body_rx.recv().await {
match event {
embedded::LocalBodyEvent::Chunk(chunk) => {
let next_len = body_bytes.len().saturating_add(chunk.len());
if next_len > DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES {
return Err(format!(
"tunnel probe body exceeds {} bytes",
DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES
));
}
body_bytes.extend_from_slice(&chunk);
}
embedded::LocalBodyEvent::End => break,
embedded::LocalBodyEvent::Error(error) => return Err(error),
}
}
Ok::<String, String>(String::from_utf8_lossy(&body_bytes).to_string())
})
.await
.map_err(|_| "timed out waiting for tunnel probe response body".to_string())??;
Ok(TunnelProbeResponse {
status: response.status,
body,
})
}
.await;
self.inner
@@ -510,8 +556,8 @@ impl Default for EmbeddedTunnelState {
impl fmt::Debug for EmbeddedTunnelState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EmbeddedTunnelState")
.field("proxy_idle_timeout_secs", &DEFAULT_PROXY_IDLE_TIMEOUT_SECS)
.field("ping_interval_secs", &DEFAULT_PING_INTERVAL_SECS)
.field("proxy_idle_timeout_ms", &DEFAULT_PROXY_IDLE_TIMEOUT_MS)
.field("ping_interval_ms", &DEFAULT_PING_INTERVAL_MS)
.field("max_streams", &DEFAULT_MAX_STREAMS)
.field("outbound_queue_capacity", &DEFAULT_OUTBOUND_QUEUE_CAPACITY)
.field(
@@ -879,14 +925,9 @@ async fn apply_embedded_tunnel_node_status(
.map_err(|err| format!("node status sync failed: {err}"))
}
fn build_embedded_tunnel_heartbeat_ack(
node: &StoredProxyNode,
heartbeat_id: Option<u64>,
) -> Vec<u8> {
fn build_embedded_tunnel_heartbeat_ack(node: &StoredProxyNode, heartbeat_id: u64) -> Vec<u8> {
let mut payload = serde_json::Map::new();
if let Some(heartbeat_id) = heartbeat_id {
payload.insert("heartbeat_id".to_string(), json!(heartbeat_id));
}
payload.insert("heartbeat_id".to_string(), json!(heartbeat_id));
if let Some(remote_config) = node.remote_config.as_ref() {
payload.insert("remote_config".to_string(), remote_config.clone());
payload.insert("config_version".to_string(), json!(node.config_version));
@@ -911,7 +952,7 @@ fn parse_embedded_tunnel_heartbeat_request(
.map_err(|_| "invalid heartbeat payload".to_string())?;
let node_id = payload.node_id.trim();
if node_id.is_empty() || node_id.len() > 36 {
if node_id.is_empty() || node_id.len() > 36 || payload.heartbeat_id == 0 {
return Err("invalid heartbeat payload".to_string());
}
if payload
@@ -1042,6 +1083,27 @@ mod tests {
);
}
#[tokio::test]
async fn embedded_tunnel_heartbeat_rejects_missing_heartbeat_id() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-123",
)]));
let data = GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&repository));
let error = apply_embedded_tunnel_heartbeat(
&data,
br#"{
"node_id": "node-123",
"heartbeat_interval": 45,
"active_connections": 5
}"#,
)
.await
.expect_err("heartbeat without heartbeat_id should fail");
assert_eq!(error, "invalid heartbeat payload");
}
#[tokio::test]
async fn embedded_tunnel_node_status_updates_proxy_node_repository() {
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(

10
apps/aether-proxy/Dockerfile
View File

@@ -1,10 +0,0 @@
FROM debian:bookworm-slim
ARG TARGETARCH
RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates \
&& rm -rf /var/lib/apt/lists/*
COPY build/linux-${TARGETARCH}/aether-proxy /usr/local/bin/aether-proxy
ENTRYPOINT ["aether-proxy"]

32
apps/aether-proxy/README.md
View File

@@ -10,14 +10,6 @@ Tunnel 模式下代理节点**无需对外监听端口**，仅需出站连接到
- 常规 Linux 发行版：`systemd`
- Alpine Linux：`OpenRC`
### Docker Compose 部署
```bash
cp .env.example .env
# 编辑 .env 填入 AETHER_PROXY_AETHER_URL 和 AETHER_PROXY_MANAGEMENT_TOKEN
docker compose up -d
```
### 下载预编译二进制
<!-- DOWNLOAD_TABLE_START -->
@@ -82,23 +74,30 @@ sudo aether-proxy uninstall
| `--node-name` | `AETHER_PROXY_NODE_NAME` | **必填** | 节点名称标识 |
| `--public-ip` | `AETHER_PROXY_PUBLIC_IP` | 自动检测 | 公网 IP |
| `--node-region` | `AETHER_PROXY_NODE_REGION` | 自动检测 | 地区标识 |
| `--heartbeat-interval` | `AETHER_PROXY_HEARTBEAT_INTERVAL` | `30` | 心跳间隔（秒） |
| `--heartbeat-interval` | `AETHER_PROXY_HEARTBEAT_INTERVAL` | `5` | 心跳间隔（秒） |
| `--allowed-ports` | `AETHER_PROXY_ALLOWED_PORTS` | `80,443,8080,8443` | 允许代理的目标端口 |
#### Tunnel 连接
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--tunnel-connections` | `AETHER_PROXY_TUNNEL_CONNECTIONS` | `3` | 到 Aether 的连接池大小 |
| `--tunnel-connections` | `AETHER_PROXY_TUNNEL_CONNECTIONS` | 自动（硬件估算） | 最小连接池大小；显式设置后默认固定为该值 |
| `--tunnel-connections-max` | `AETHER_PROXY_TUNNEL_CONNECTIONS_MAX` | 自动（硬件估算） | 连接池自动扩容上限；大于 `tunnel_connections` 时启用 autoscale |
| `--tunnel-max-streams` | `AETHER_PROXY_TUNNEL_MAX_STREAMS` | 自动（硬件估算） | 单连接最大并发 stream 数 |
| `--tunnel-connect-timeout-secs` | `AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT_SECS` | `15` | TCP + TLS 握手超时（秒） |
| `--tunnel-ping-interval-ms` | `AETHER_PROXY_TUNNEL_PING_INTERVAL_MS` | `250` | fast-fail 探测周期（毫秒） |
| `--tunnel-connect-timeout-ms` | `AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT_MS` | `800` | fast-reconnect 建连超时（毫秒） |
| `--tunnel-stale-timeout-ms` | `AETHER_PROXY_TUNNEL_STALE_TIMEOUT_MS` | `900` | 无入站数据断连阈值（毫秒） |
| `--tunnel-scale-check-interval-ms` | `AETHER_PROXY_TUNNEL_SCALE_CHECK_INTERVAL_MS` | `1000` | autoscale 采样周期（毫秒） |
| `--tunnel-scale-up-threshold-percent` | `AETHER_PROXY_TUNNEL_SCALE_UP_THRESHOLD_PERCENT` | `70` | 单 tunnel 占用率超过该值时扩容 |
| `--tunnel-scale-down-threshold-percent` | `AETHER_PROXY_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT` | `35` | 单 tunnel 占用率持续低于该值时允许缩容 |
| `--tunnel-scale-down-grace-secs` | `AETHER_PROXY_TUNNEL_SCALE_DOWN_GRACE_SECS` | `15` | 低负载持续时间达到该值后才回收次级 tunnel |
| `--tunnel-tcp-keepalive-secs` | `AETHER_PROXY_TUNNEL_TCP_KEEPALIVE_SECS` | `30` | TCP keepalive 初始延迟（秒） |
| `--tunnel-tcp-nodelay` | `AETHER_PROXY_TUNNEL_TCP_NODELAY` | `true` | 禁用 Nagle 算法 |
| `--tunnel-ping-interval-secs` | `AETHER_PROXY_TUNNEL_PING_INTERVAL_SECS` | `15` | WebSocket Ping 频率（秒） |
| `--tunnel-stale-timeout-secs` | `AETHER_PROXY_TUNNEL_STALE_TIMEOUT_SECS` | `45` | 无数据断连阈值（秒） |
| `--tunnel-reconnect-base-ms` | `AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS` | `500` | 指数退避基础延迟（毫秒） |
| `--tunnel-reconnect-base-ms` | `AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS` | `50` | 指数退避基础延迟（毫秒） |
| `--tunnel-reconnect-max-ms` | `AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS` | `30000` | 指数退避上限（毫秒） |
省略 `tunnel_connections` 时，proxy 会按设备能力自动计算一个基线值和扩容上限；如果显式设置了 `tunnel_connections` 但没有设置 `tunnel_connections_max`，则保持固定连接池，不自动扩缩。
#### 上游 HTTP 请求
| 参数 | 环境变量 | 默认值 | 说明 |
@@ -141,12 +140,12 @@ sudo aether-proxy uninstall
- 默认 `AETHER_PROXY_LOG_DESTINATION=stdout`，日志交给容器日志驱动或宿主机服务管理器
- 需要落盘时改成 `file` 或 `both`，并设置 `AETHER_PROXY_LOG_DIR`；setup TUI 里用 `Save Logs to File` 开关即可
- 文件日志固定写普通文本，并支持 `hourly/daily` 轮转；默认按天轮换、保留 7 天，最多保留 30 个文件
- `docker compose` 默认保持 `stdout`，避免和容器自带日志重复；以 `systemd` 或 `OpenRC` 安装时默认会额外打开文件日志到 `/var/log/aether-proxy`
- 以 `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` 查看
### 多服务器配置
在 `aether-proxy.toml` 中使用 `[[servers]]` 配置多个 Aether 服务器：
在 `aether-proxy.toml` 中使用 `[[servers]]` 配置 Aether 服务器。即使只有一个服务器，也必须写成一个 `[[servers]]` 条目；旧的顶层单服务器写法已不再支持。
```toml
[[servers]]
@@ -164,7 +163,6 @@ node_name = "jp-proxy-02"
推送 `proxy-v*` 格式的 tag，GitHub Actions 会自动：
- 编译所有平台二进制并发布到 Releases
- 构建 Docker 镜像并推送到 GHCR 和 Docker Hub
- 更新 README 中的下载链接表格
```bash

21
apps/aether-proxy/docker-compose.yml
View File

@@ -1,21 +0,0 @@
services:
aether-proxy:
image: ghcr.io/fawney19/aether-proxy:latest
container_name: aether-proxy
restart: unless-stopped
env_file:
- .env
environment:
AETHER_PROXY_LOG_JSON: ${AETHER_PROXY_LOG_JSON:-true}
AETHER_PROXY_LOG_DESTINATION: ${AETHER_PROXY_LOG_DESTINATION:-stdout}
AETHER_PROXY_LOG_DIR: ${AETHER_PROXY_LOG_DIR:-/var/log/aether-proxy}
AETHER_PROXY_LOG_ROTATION: ${AETHER_PROXY_LOG_ROTATION:-daily}
AETHER_PROXY_LOG_RETENTION_DAYS: ${AETHER_PROXY_LOG_RETENTION_DAYS:-7}
AETHER_PROXY_LOG_MAX_FILES: ${AETHER_PROXY_LOG_MAX_FILES:-30}
volumes:
- ./logs:/var/log/aether-proxy
logging:
driver: json-file
options:
max-size: "50m"
max-file: "3"

381
apps/aether-proxy/src/app.rs
View File

@@ -1,8 +1,9 @@
//! Application lifecycle: initialization, task orchestration, and shutdown.
use std::sync::atomic::AtomicU64;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Duration;
use std::time::{Duration, Instant};
use aether_http::{jittered_delay_for_retry, HttpRetryConfig};
use aether_runtime::{
@@ -14,7 +15,7 @@ use tokio::sync::{watch, Mutex};
use tokio::task::JoinHandle;
use tracing::{error, info, warn};
use crate::config::{Config, ServerEntry};
use crate::config::{Config, ServerEntry, TunnelPoolSizing};
use crate::net;
use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig};
@@ -24,6 +25,52 @@ use crate::{hardware, target_filter, tunnel};
type TaskHandles = Arc<Mutex<Vec<JoinHandle<()>>>>;
#[derive(Debug, Clone, Copy)]
struct TunnelPoolPolicy {
min_connections: usize,
max_connections: usize,
max_streams_per_tunnel: usize,
scale_check_interval: Duration,
scale_up_threshold_percent: u32,
scale_down_threshold_percent: u32,
scale_down_grace: Duration,
}
impl TunnelPoolPolicy {
fn from_config(config: &Config, sizing: TunnelPoolSizing) -> Self {
Self {
min_connections: sizing.initial_connections.max(1) as usize,
max_connections: sizing
.max_connections
.max(sizing.initial_connections)
.max(1) as usize,
max_streams_per_tunnel: config.tunnel_max_streams.unwrap_or(128).max(1) as usize,
scale_check_interval: Duration::from_millis(config.tunnel_scale_check_interval_ms),
scale_up_threshold_percent: config.tunnel_scale_up_threshold_percent,
scale_down_threshold_percent: config.tunnel_scale_down_threshold_percent,
scale_down_grace: Duration::from_secs(config.tunnel_scale_down_grace_secs),
}
}
fn scale_up_high_water_mark(&self) -> u64 {
occupancy_threshold(self.max_streams_per_tunnel, self.scale_up_threshold_percent)
}
fn scale_down_low_water_mark(&self) -> u64 {
occupancy_threshold(
self.max_streams_per_tunnel,
self.scale_down_threshold_percent,
)
}
}
struct ManagedTunnel {
slot_id: usize,
drain_tx: watch::Sender<bool>,
handle: JoinHandle<()>,
draining: bool,
}
/// Run the full application lifecycle after config has been parsed.
pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Result<()> {
config.validate()?;
@@ -63,6 +110,19 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
"auto-detected tunnel_max_streams from hardware"
);
}
let tunnel_pool_sizing = config.resolve_tunnel_pool_sizing(&hw_info)?;
let tunnel_pool_policy = TunnelPoolPolicy::from_config(&config, tunnel_pool_sizing);
info!(
tunnel_connections_initial = tunnel_pool_policy.min_connections,
tunnel_connections_max = tunnel_pool_policy.max_connections,
tunnel_max_streams = tunnel_pool_policy.max_streams_per_tunnel,
scale_check_interval_ms = tunnel_pool_policy.scale_check_interval.as_millis(),
scale_up_threshold_percent = tunnel_pool_policy.scale_up_threshold_percent,
scale_down_threshold_percent = tunnel_pool_policy.scale_down_threshold_percent,
scale_down_grace_secs = tunnel_pool_policy.scale_down_grace.as_secs(),
auto_sizing = config.tunnel_connections.is_none(),
"resolved tunnel pool policy"
);
info!(
max_concurrency = hw_info.estimated_max_concurrency,
@@ -177,15 +237,14 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
"running in tunnel mode"
);
// Spawn tunnel connections per server (pool_size connections each)
let pool_size = state.config.tunnel_connections.max(1) as usize;
// Spawn tunnel pool manager per server.
let tunnel_handles: TaskHandles = Arc::new(Mutex::new(Vec::new()));
let retry_handles: TaskHandles = Arc::new(Mutex::new(Vec::new()));
for server in server_contexts.lock().await.iter() {
spawn_tunnel_pool(
spawn_tunnel_pool_manager(
Arc::clone(&state),
Arc::clone(server),
pool_size,
tunnel_pool_policy,
shutdown_rx.clone(),
Arc::clone(&tunnel_handles),
)
@@ -200,7 +259,7 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
failed_entries,
public_ip.clone(),
hw_info.clone(),
pool_size,
tunnel_pool_policy,
shutdown_rx.clone(),
Arc::clone(&tunnel_handles),
Arc::clone(&retry_handles),
@@ -240,7 +299,7 @@ async fn spawn_registration_recovery_tasks(
failed: Vec<(String, ServerEntry)>,
public_ip: String,
hw_info: crate::hardware::HardwareInfo,
pool_size: usize,
tunnel_pool_policy: TunnelPoolPolicy,
shutdown: watch::Receiver<bool>,
tunnel_handles: TaskHandles,
retry_handles: TaskHandles,
@@ -261,7 +320,7 @@ async fn spawn_registration_recovery_tasks(
entry,
retry_public_ip,
retry_hw_info,
pool_size,
tunnel_pool_policy,
retry_shutdown,
retry_tunnels,
)
@@ -281,7 +340,7 @@ async fn retry_failed_registration(
entry: ServerEntry,
public_ip: String,
hw_info: crate::hardware::HardwareInfo,
pool_size: usize,
tunnel_pool_policy: TunnelPoolPolicy,
mut shutdown: watch::Receiver<bool>,
tunnel_handles: TaskHandles,
) {
@@ -329,10 +388,10 @@ async fn retry_failed_registration(
node_id,
);
server_contexts.lock().await.push(Arc::clone(&server));
spawn_tunnel_pool(
spawn_tunnel_pool_manager(
Arc::clone(&state),
server,
pool_size,
tunnel_pool_policy,
shutdown,
tunnel_handles,
)
@@ -386,23 +445,239 @@ fn build_server_context(
})
}
async fn spawn_tunnel_pool(
async fn spawn_tunnel_pool_manager(
state: Arc<AppState>,
server: Arc<ServerContext>,
pool_size: usize,
policy: TunnelPoolPolicy,
shutdown: watch::Receiver<bool>,
tunnel_handles: TaskHandles,
) {
let mut handles = Vec::with_capacity(pool_size);
for conn_idx in 0..pool_size {
let s = Arc::clone(&state);
let srv = Arc::clone(&server);
let rx = shutdown.clone();
handles.push(tokio::spawn(async move {
tunnel::run(&s, &srv, conn_idx, rx).await;
}));
let handle = tokio::spawn(async move {
run_tunnel_pool_manager(state, server, policy, shutdown).await;
});
tunnel_handles.lock().await.push(handle);
}
async fn run_tunnel_pool_manager(
state: Arc<AppState>,
server: Arc<ServerContext>,
policy: TunnelPoolPolicy,
mut shutdown: watch::Receiver<bool>,
) {
let mut tunnels = BTreeMap::<usize, ManagedTunnel>::new();
ensure_tunnel_capacity(
&mut tunnels,
policy.min_connections,
&policy,
&state,
&server,
&shutdown,
);
let mut ticker = tokio::time::interval(policy.scale_check_interval);
ticker.tick().await;
let mut low_load_since: Option<Instant> = None;
loop {
tokio::select! {
_ = shutdown.changed() => {
info!(server = %server.server_label, "tunnel pool manager shutting down");
break;
}
_ = ticker.tick() => {
reap_finished_tunnels(&mut tunnels).await;
let available = tunnels.values().filter(|tunnel| !tunnel.draining).count();
if available < policy.min_connections {
ensure_tunnel_capacity(
&mut tunnels,
policy.min_connections,
&policy,
&state,
&server,
&shutdown,
);
low_load_since = None;
continue;
}
let active_connections = server.active_connections.load(Ordering::Acquire);
let desired_connections = desired_tunnel_connections(active_connections, &policy);
if desired_connections > available {
ensure_tunnel_capacity(
&mut tunnels,
desired_connections,
&policy,
&state,
&server,
&shutdown,
);
info!(
server = %server.server_label,
active_connections,
available_connections = available,
target_connections = desired_connections,
"scaled tunnel pool up"
);
low_load_since = None;
continue;
}
if should_scale_down(active_connections, available, &policy) {
match low_load_since {
Some(since) if since.elapsed() >= policy.scale_down_grace => {
if request_tunnel_drain(&mut tunnels, policy.min_connections) {
info!(
server = %server.server_label,
active_connections,
available_connections = available,
"requested tunnel drain for scale-down"
);
}
low_load_since = None;
}
None => {
low_load_since = Some(Instant::now());
}
Some(_) => {}
}
} else {
low_load_since = None;
}
}
}
}
tunnel_handles.lock().await.extend(handles);
for tunnel in tunnels.values_mut() {
let _ = tunnel.drain_tx.send(true);
tunnel.draining = true;
}
while !tunnels.is_empty() {
reap_finished_tunnels(&mut tunnels).await;
if !tunnels.is_empty() {
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
}
fn ensure_tunnel_capacity(
tunnels: &mut BTreeMap<usize, ManagedTunnel>,
target_connections: usize,
policy: &TunnelPoolPolicy,
state: &Arc<AppState>,
server: &Arc<ServerContext>,
shutdown: &watch::Receiver<bool>,
) {
let target_connections = target_connections.min(policy.max_connections);
while tunnels.values().filter(|tunnel| !tunnel.draining).count() < target_connections {
let Some(slot_id) = next_available_tunnel_slot(tunnels, policy.max_connections) else {
break;
};
tunnels.insert(
slot_id,
spawn_managed_tunnel(
Arc::clone(state),
Arc::clone(server),
slot_id,
shutdown.clone(),
),
);
}
}
fn spawn_managed_tunnel(
state: Arc<AppState>,
server: Arc<ServerContext>,
slot_id: usize,
shutdown: watch::Receiver<bool>,
) -> ManagedTunnel {
let (drain_tx, drain_rx) = watch::channel(false);
let handle = tokio::spawn(async move {
tunnel::run(&state, &server, slot_id, shutdown, drain_rx).await;
});
ManagedTunnel {
slot_id,
drain_tx,
handle,
draining: false,
}
}
async fn reap_finished_tunnels(tunnels: &mut BTreeMap<usize, ManagedTunnel>) {
let finished_slots = tunnels
.iter()
.filter_map(|(slot_id, tunnel)| tunnel.handle.is_finished().then_some(*slot_id))
.collect::<Vec<_>>();
for slot_id in finished_slots {
if let Some(tunnel) = tunnels.remove(&slot_id) {
let _ = tunnel.handle.await;
}
}
}
fn next_available_tunnel_slot(
tunnels: &BTreeMap<usize, ManagedTunnel>,
max_connections: usize,
) -> Option<usize> {
(0..max_connections).find(|slot_id| !tunnels.contains_key(slot_id))
}
fn request_tunnel_drain(
tunnels: &mut BTreeMap<usize, ManagedTunnel>,
min_connections: usize,
) -> bool {
let available = tunnels.values().filter(|tunnel| !tunnel.draining).count();
if available <= min_connections {
return false;
}
let Some((_, tunnel)) = tunnels
.iter_mut()
.rev()
.find(|(_, tunnel)| tunnel.slot_id != 0 && !tunnel.draining)
else {
return false;
};
if tunnel.drain_tx.send(true).is_ok() {
tunnel.draining = true;
return true;
}
false
}
fn desired_tunnel_connections(active_connections: u64, policy: &TunnelPoolPolicy) -> usize {
let required = div_ceil_u64(active_connections.max(1), policy.scale_up_high_water_mark());
required.clamp(policy.min_connections as u64, policy.max_connections as u64) as usize
}
fn should_scale_down(
active_connections: u64,
available_connections: usize,
policy: &TunnelPoolPolicy,
) -> bool {
if available_connections <= policy.min_connections {
return false;
}
active_connections
<= (available_connections as u64)
.saturating_sub(1)
.saturating_mul(policy.scale_down_low_water_mark())
}
fn occupancy_threshold(max_streams_per_tunnel: usize, percent: u32) -> u64 {
div_ceil_u64(
(max_streams_per_tunnel as u64).saturating_mul(percent as u64),
100,
)
.max(1)
}
fn div_ceil_u64(value: u64, divisor: u64) -> u64 {
if divisor == 0 {
return value;
}
value.saturating_add(divisor.saturating_sub(1)) / divisor
}
async fn await_all_handles(handles: &TaskHandles) {
@@ -469,6 +744,8 @@ mod tests {
},
)];
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let tunnel_pool_policy =
TunnelPoolPolicy::from_config(&state.config, sample_tunnel_pool_sizing());
spawn_registration_recovery_tasks(
Arc::clone(&state),
@@ -476,7 +753,7 @@ mod tests {
failed,
"127.0.0.1".to_string(),
sample_hardware_info(),
1,
tunnel_pool_policy,
shutdown_rx.clone(),
Arc::clone(&tunnel_handles),
Arc::clone(&retry_handles),
@@ -508,6 +785,40 @@ mod tests {
gateway_handle.abort();
}
#[test]
fn desired_tunnel_connections_expands_when_load_crosses_high_water() {
let policy = TunnelPoolPolicy {
min_connections: 1,
max_connections: 6,
max_streams_per_tunnel: 1024,
scale_check_interval: Duration::from_secs(1),
scale_up_threshold_percent: 70,
scale_down_threshold_percent: 35,
scale_down_grace: Duration::from_secs(15),
};
assert_eq!(desired_tunnel_connections(1, &policy), 1);
assert_eq!(desired_tunnel_connections(2_000, &policy), 3);
assert_eq!(desired_tunnel_connections(5_000, &policy), 6);
}
#[test]
fn should_scale_down_requires_load_to_fit_remaining_tunnels() {
let policy = TunnelPoolPolicy {
min_connections: 1,
max_connections: 6,
max_streams_per_tunnel: 1024,
scale_check_interval: Duration::from_secs(1),
scale_up_threshold_percent: 70,
scale_down_threshold_percent: 35,
scale_down_grace: Duration::from_secs(15),
};
assert!(!should_scale_down(800, 3, &policy));
assert!(should_scale_down(600, 3, &policy));
assert!(!should_scale_down(200, 1, &policy));
}
async fn wait_for_registered_server(
server_contexts: &Arc<Mutex<Vec<Arc<ServerContext>>>>,
) -> Arc<ServerContext> {
@@ -595,6 +906,13 @@ mod tests {
}
}
fn sample_tunnel_pool_sizing() -> TunnelPoolSizing {
TunnelPoolSizing {
initial_connections: 1,
max_connections: 1,
}
}
fn sample_state(config: Config) -> Arc<ProxyAppState> {
let config = Arc::new(config);
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
@@ -656,13 +974,18 @@ mod tests {
log_max_files: 30,
tunnel_reconnect_base_ms: 50,
tunnel_reconnect_max_ms: 250,
tunnel_ping_interval_secs: 1,
tunnel_ping_interval_ms: 1_000,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 2,
tunnel_connect_timeout_ms: 2_000,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 5,
tunnel_connections: 1,
tunnel_stale_timeout_ms: 5_000,
tunnel_connections: Some(1),
tunnel_connections_max: Some(1),
tunnel_scale_check_interval_ms: 1_000,
tunnel_scale_up_threshold_percent: 70,
tunnel_scale_down_threshold_percent: 35,
tunnel_scale_down_grace_secs: 15,
}
}

584
apps/aether-proxy/src/config.rs
View File

@@ -1,10 +1,13 @@
use std::fmt;
use std::path::Path;
use std::time::Duration;
use aether_runtime::{FileLoggingConfig, LogDestination, LogRotation, ServiceRuntimeConfig};
use clap::Parser;
use serde::{Deserialize, Serialize};
use crate::hardware::HardwareInfo;
/// Fields that existed in 0.1.x but were removed in 0.2.0.
const LEGACY_ONLY_KEYS: &[&str] = &[
"hmac_key",
@@ -16,6 +19,12 @@ const LEGACY_ONLY_KEYS: &[&str] = &[
"tls_cert",
"tls_key",
];
const REMOVED_TUNNEL_SECONDS_KEYS: &[&str] = &[
"tunnel_ping_interval_secs",
"tunnel_connect_timeout_secs",
"tunnel_stale_timeout_secs",
];
const REMOVED_SINGLE_SERVER_KEYS: &[&str] = &["aether_url", "management_token"];
/// Fields renamed from 0.1.x `delegate_*` to 0.2.0 `upstream_*`.
const DELEGATE_TO_UPSTREAM: &[(&str, &str)] = &[
@@ -38,13 +47,31 @@ const DELEGATE_TO_UPSTREAM: &[(&str, &str)] = &[
/// Default bytes buffered before a tunnel request becomes non-replayable for
/// 307/308 redirects. Kept aligned with the current admin-side request size
/// default, but exposed as an independent proxy transport budget.
pub const DEFAULT_HEARTBEAT_INTERVAL_SECS: u64 = 30;
pub const DEFAULT_HEARTBEAT_INTERVAL_SECS: u64 = 5;
#[allow(dead_code)]
pub const DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES: usize = 5_242_880;
pub const DEFAULT_REDIRECT_REPLAY_BUDGET_HUMAN: &str = "5M";
pub const DEFAULT_LOG_RETENTION_DAYS: u64 = 7;
pub const DEFAULT_LOG_MAX_FILES: usize = 30;
pub const DEFAULT_TUNNEL_PING_INTERVAL_MS: u64 = 250;
pub const DEFAULT_TUNNEL_CONNECT_TIMEOUT_MS: u64 = 800;
pub const DEFAULT_TUNNEL_STALE_TIMEOUT_MS: u64 = 900;
pub const DEFAULT_TUNNEL_SCALE_CHECK_INTERVAL_MS: u64 = 1_000;
pub const DEFAULT_TUNNEL_SCALE_UP_THRESHOLD_PERCENT: u32 = 70;
pub const DEFAULT_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT: u32 = 35;
pub const DEFAULT_TUNNEL_SCALE_DOWN_GRACE_SECS: u64 = 15;
const AUTO_TUNNEL_CONNECTIONS_BASE_CAP: u64 = 4;
const AUTO_TUNNEL_CONNECTIONS_MAX_CAP: u64 = 8;
const TUNNEL_PING_INTERVAL_MS_ENV: &str = "AETHER_PROXY_TUNNEL_PING_INTERVAL_MS";
const TUNNEL_CONNECT_TIMEOUT_MS_ENV: &str = "AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT_MS";
const TUNNEL_STALE_TIMEOUT_MS_ENV: &str = "AETHER_PROXY_TUNNEL_STALE_TIMEOUT_MS";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TunnelPoolSizing {
pub initial_connections: u32,
pub max_connections: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum ByteSizeValue {
Text(String),
@@ -473,7 +500,7 @@ pub struct Config {
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
default_value_t = 500
default_value_t = 50
)]
pub tunnel_reconnect_base_ms: u64,
@@ -485,21 +512,25 @@ pub struct Config {
)]
pub tunnel_reconnect_max_ms: u64,
/// WebSocket tunnel ping interval in seconds
#[arg(long, env = "AETHER_PROXY_TUNNEL_PING_INTERVAL", default_value_t = 15)]
pub tunnel_ping_interval_secs: u64,
/// WebSocket tunnel ping interval in milliseconds
#[arg(
long,
env = TUNNEL_PING_INTERVAL_MS_ENV,
default_value_t = DEFAULT_TUNNEL_PING_INTERVAL_MS
)]
pub tunnel_ping_interval_ms: u64,
/// Maximum concurrent streams over tunnel (auto-detected from hardware if omitted)
#[arg(long, env = "AETHER_PROXY_TUNNEL_MAX_STREAMS")]
pub tunnel_max_streams: Option<u32>,
/// WebSocket tunnel TCP connect timeout in seconds
/// WebSocket tunnel TCP connect timeout in milliseconds
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT",
default_value_t = 15
env = TUNNEL_CONNECT_TIMEOUT_MS_ENV,
default_value_t = DEFAULT_TUNNEL_CONNECT_TIMEOUT_MS
)]
pub tunnel_connect_timeout_secs: u64,
pub tunnel_connect_timeout_ms: u64,
/// WebSocket tunnel TCP keepalive in seconds (0 disables)
#[arg(long, env = "AETHER_PROXY_TUNNEL_TCP_KEEPALIVE", default_value_t = 30)]
@@ -509,13 +540,55 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_TUNNEL_TCP_NODELAY", default_value_t = true)]
pub tunnel_tcp_nodelay: bool,
/// Tunnel connection staleness timeout in seconds (triggers reconnect if no data received)
#[arg(long, env = "AETHER_PROXY_TUNNEL_STALE_TIMEOUT", default_value_t = 45)]
pub tunnel_stale_timeout_secs: u64,
/// Tunnel connection staleness timeout in milliseconds
#[arg(
long,
env = TUNNEL_STALE_TIMEOUT_MS_ENV,
default_value_t = DEFAULT_TUNNEL_STALE_TIMEOUT_MS
)]
pub tunnel_stale_timeout_ms: u64,
/// Number of parallel WebSocket tunnel connections per server (connection pool)
#[arg(long, env = "AETHER_PROXY_TUNNEL_CONNECTIONS", default_value_t = 3)]
pub tunnel_connections: u32,
/// Minimum number of parallel WebSocket tunnel connections per server.
/// If omitted, a device-aware value is auto-detected at startup.
#[arg(long, env = "AETHER_PROXY_TUNNEL_CONNECTIONS")]
pub tunnel_connections: Option<u32>,
/// Maximum number of WebSocket tunnel connections per server.
/// When larger than `tunnel_connections`, the proxy may autoscale up to this limit.
#[arg(long, env = "AETHER_PROXY_TUNNEL_CONNECTIONS_MAX")]
pub tunnel_connections_max: Option<u32>,
/// Autoscale evaluation interval for the tunnel pool.
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_SCALE_CHECK_INTERVAL_MS",
default_value_t = DEFAULT_TUNNEL_SCALE_CHECK_INTERVAL_MS
)]
pub tunnel_scale_check_interval_ms: u64,
/// Per-tunnel occupancy percentage that triggers scale-up.
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_SCALE_UP_THRESHOLD_PERCENT",
default_value_t = DEFAULT_TUNNEL_SCALE_UP_THRESHOLD_PERCENT
)]
pub tunnel_scale_up_threshold_percent: u32,
/// Per-tunnel occupancy percentage that allows scale-down after the grace window.
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT",
default_value_t = DEFAULT_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT
)]
pub tunnel_scale_down_threshold_percent: u32,
/// Low-load grace window before a secondary tunnel is drained.
#[arg(
long,
env = "AETHER_PROXY_TUNNEL_SCALE_DOWN_GRACE_SECS",
default_value_t = DEFAULT_TUNNEL_SCALE_DOWN_GRACE_SECS
)]
pub tunnel_scale_down_grace_secs: u64,
}
impl Config {
@@ -539,22 +612,52 @@ impl Config {
anyhow::bail!("allowed_ports: port 0 is not valid");
}
}
if self.tunnel_connect_timeout_secs == 0 {
anyhow::bail!("tunnel_connect_timeout_secs must be > 0");
let tunnel_connect_timeout = self.tunnel_connect_timeout()?;
if tunnel_connect_timeout.is_zero() {
anyhow::bail!("effective tunnel connect timeout must be > 0");
}
if self.tunnel_ping_interval_secs == 0 {
anyhow::bail!("tunnel_ping_interval_secs must be > 0");
let tunnel_ping_interval = self.tunnel_ping_interval()?;
if tunnel_ping_interval.is_zero() {
anyhow::bail!("effective tunnel ping interval must be > 0");
}
if self.tunnel_stale_timeout_secs <= self.tunnel_ping_interval_secs {
let tunnel_stale_timeout = self.tunnel_stale_timeout()?;
if tunnel_stale_timeout <= tunnel_ping_interval {
anyhow::bail!(
"tunnel_stale_timeout_secs ({}) must be > tunnel_ping_interval_secs ({})",
self.tunnel_stale_timeout_secs,
self.tunnel_ping_interval_secs
"effective tunnel stale timeout ({:?}) must be > effective tunnel ping interval ({:?})",
tunnel_stale_timeout,
tunnel_ping_interval
);
}
if self.tunnel_connections == 0 {
if matches!(self.tunnel_connections, Some(0)) {
anyhow::bail!("tunnel_connections must be > 0");
}
if matches!(self.tunnel_connections_max, Some(0)) {
anyhow::bail!("tunnel_connections_max must be > 0");
}
if let (Some(min_connections), Some(max_connections)) =
(self.tunnel_connections, self.tunnel_connections_max)
{
if max_connections < min_connections {
anyhow::bail!("tunnel_connections_max must be >= tunnel_connections");
}
}
if self.tunnel_scale_check_interval_ms == 0 {
anyhow::bail!("tunnel_scale_check_interval_ms must be > 0");
}
if self.tunnel_scale_down_grace_secs == 0 {
anyhow::bail!("tunnel_scale_down_grace_secs must be > 0");
}
if !(1..=100).contains(&self.tunnel_scale_up_threshold_percent) {
anyhow::bail!("tunnel_scale_up_threshold_percent must be within 1..=100");
}
if !(1..100).contains(&self.tunnel_scale_down_threshold_percent) {
anyhow::bail!("tunnel_scale_down_threshold_percent must be within 1..100");
}
if self.tunnel_scale_down_threshold_percent >= self.tunnel_scale_up_threshold_percent {
anyhow::bail!(
"tunnel_scale_down_threshold_percent must be < tunnel_scale_up_threshold_percent"
);
}
if self.aether_retry_max_attempts == 0 {
anyhow::bail!("aether_retry_max_attempts must be >= 1");
}
@@ -600,6 +703,54 @@ impl Config {
Ok(())
}
pub fn tunnel_ping_interval(&self) -> anyhow::Result<Duration> {
Ok(Duration::from_millis(self.tunnel_ping_interval_ms))
}
pub fn tunnel_connect_timeout(&self) -> anyhow::Result<Duration> {
Ok(Duration::from_millis(self.tunnel_connect_timeout_ms))
}
pub fn tunnel_stale_timeout(&self) -> anyhow::Result<Duration> {
Ok(Duration::from_millis(self.tunnel_stale_timeout_ms))
}
pub fn resolve_tunnel_pool_sizing(
&self,
hw_info: &HardwareInfo,
) -> anyhow::Result<TunnelPoolSizing> {
let per_tunnel_capacity = u64::from(self.tunnel_max_streams.unwrap_or(128).max(1));
let estimated = hw_info.estimated_max_concurrency.max(per_tunnel_capacity);
let cpu_cap = u64::from(hw_info.cpu_cores).clamp(1, AUTO_TUNNEL_CONNECTIONS_MAX_CAP);
let auto_initial = div_ceil_u64(estimated, per_tunnel_capacity.saturating_mul(8))
.clamp(1, AUTO_TUNNEL_CONNECTIONS_BASE_CAP)
.min(cpu_cap);
let auto_max = div_ceil_u64(estimated, per_tunnel_capacity.saturating_mul(4))
.clamp(auto_initial, AUTO_TUNNEL_CONNECTIONS_MAX_CAP)
.min(cpu_cap.max(auto_initial));
let initial_connections = u64::from(self.tunnel_connections.unwrap_or(auto_initial as u32));
let max_connections = match self.tunnel_connections_max {
Some(explicit) => u64::from(explicit),
None if self.tunnel_connections.is_some() => initial_connections,
None => auto_max,
};
if max_connections < initial_connections {
anyhow::bail!(
"effective tunnel_connections_max ({max_connections}) must be >= tunnel_connections ({initial_connections})"
);
}
Ok(TunnelPoolSizing {
initial_connections: u32::try_from(initial_connections)
.expect("effective tunnel initial connections should fit in u32"),
max_connections: u32::try_from(max_connections)
.expect("effective tunnel max connections should fit in u32"),
})
}
pub fn service_runtime_config(&self) -> anyhow::Result<ServiceRuntimeConfig> {
let mut config = ServiceRuntimeConfig::new("aether-proxy", "aether_proxy=info")
.with_log_format(aether_runtime::LogFormat::Pretty)
@@ -629,6 +780,7 @@ impl Config {
/// Per-server connection config (used in multi-server TOML `[[servers]]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ServerEntry {
pub aether_url: String,
pub management_token: String,
@@ -643,11 +795,8 @@ pub struct ServerEntry {
/// Serializable config for TOML file persistence.
/// All fields are optional -- only populated values are written.
#[derive(Debug, Default, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub management_token: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub public_ip: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -717,23 +866,31 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_reconnect_max_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_ping_interval_secs: Option<u64>,
pub tunnel_ping_interval_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_max_streams: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_connect_timeout_secs: Option<u64>,
pub tunnel_connect_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_tcp_keepalive_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_tcp_nodelay: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_stale_timeout_secs: Option<u64>,
pub tunnel_stale_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_connections: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_connections_max: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_scale_check_interval_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_scale_up_threshold_percent: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_scale_down_threshold_percent: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_scale_down_grace_secs: Option<u64>,
/// Multi-server config: each entry connects to a separate Aether instance.
/// When present, top-level aether_url/management_token are ignored for
/// tunnel connections (but still injected as env for clap compatibility).
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub servers: Vec<ServerEntry>,
}
@@ -742,6 +899,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)?)
}
@@ -752,102 +910,6 @@ impl ConfigFile {
Ok(())
}
/// Detect and migrate a 0.1.x config file to 0.2.0 format in-place.
///
/// Returns `true` if migration was performed, `false` if already current.
/// The original file is backed up as `<name>.v1.bak` before rewriting.
pub fn migrate_legacy(path: &Path) -> anyhow::Result<bool> {
let content = match std::fs::read_to_string(path) {
Ok(c) => c,
Err(_) => return Ok(false),
};
let mut table: toml::map::Map<String, toml::Value> = toml::from_str(&content)?;
// Detect legacy format: presence of any 0.1.x-only key.
let is_legacy = LEGACY_ONLY_KEYS.iter().any(|k| table.contains_key(*k))
|| DELEGATE_TO_UPSTREAM
.iter()
.any(|(old, _)| table.contains_key(*old));
if !is_legacy {
return Ok(false);
}
// 1. Rename delegate_* -> upstream_* (carry over user-customized values)
for &(old, new) in DELEGATE_TO_UPSTREAM {
if let Some(val) = table.remove(old) {
table.entry(new.to_string()).or_insert(val);
}
}
// 2. Build [[servers]] from top-level aether_url + management_token + node_name
if !table.contains_key("servers") {
let aether_url = table.get("aether_url").and_then(|v| v.as_str());
let management_token = table.get("management_token").and_then(|v| v.as_str());
if let (Some(url), Some(token)) = (aether_url, management_token) {
let mut entry = toml::map::Map::new();
entry.insert("aether_url".into(), toml::Value::String(url.to_string()));
entry.insert(
"management_token".into(),
toml::Value::String(token.to_string()),
);
if let Some(name) = table.get("node_name").and_then(|v| v.as_str()) {
entry.insert("node_name".into(), toml::Value::String(name.to_string()));
}
table.insert(
"servers".into(),
toml::Value::Array(vec![toml::Value::Table(entry)]),
);
}
}
// 3. Remove top-level fields that are now in [[servers]] or obsolete
table.remove("aether_url");
table.remove("management_token");
table.remove("node_name");
for &key in LEGACY_ONLY_KEYS {
table.remove(key);
}
// 4. Backup original file (abort migration if backup fails)
let backup_path = path.with_extension("v1.bak");
std::fs::copy(path, &backup_path).map_err(|e| {
anyhow::anyhow!(
"failed to backup config before migration: {} -> {}: {}",
path.display(),
backup_path.display(),
e
)
})?;
// 5. Write migrated config
let new_content = toml::to_string_pretty(&table)?;
std::fs::write(path, &new_content)?;
eprintln!(" Config migrated from 0.1.x to 0.2.0 format.");
eprintln!(" Backup saved: {}", backup_path.display());
Ok(true)
}
/// Resolve the effective server list.
///
/// If `[[servers]]` is present, use it. Otherwise fall back to the
/// top-level `aether_url` + `management_token` as a single server.
pub fn effective_servers(&self) -> Vec<ServerEntry> {
if !self.servers.is_empty() {
return self.servers.clone();
}
match (&self.aether_url, &self.management_token) {
(Some(url), Some(token)) => vec![ServerEntry {
aether_url: url.clone(),
management_token: token.clone(),
node_name: None,
}],
_ => vec![],
}
}
/// Inject values as environment variables so clap picks them up.
///
/// Only sets variables that are **not** already present in the
@@ -874,18 +936,9 @@ impl ConfigFile {
};
}
// When top-level fields are absent, fall back to the first [[servers]]
// entry so that clap's required `aether_url` / `management_token` are
// satisfied even with the new config format.
let first_server = self.servers.first();
let aether_url = self
.aether_url
.as_deref()
.or(first_server.map(|s| s.aether_url.as_str()));
let management_token = self
.management_token
.as_deref()
.or(first_server.map(|s| s.management_token.as_str()));
let aether_url = first_server.map(|s| s.aether_url.as_str());
let management_token = first_server.map(|s| s.management_token.as_str());
let node_name = self
.node_name
.as_deref()
@@ -988,25 +1041,39 @@ impl ConfigFile {
"AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS",
self.tunnel_reconnect_max_ms
);
set!(
"AETHER_PROXY_TUNNEL_PING_INTERVAL",
self.tunnel_ping_interval_secs
);
set!(TUNNEL_PING_INTERVAL_MS_ENV, self.tunnel_ping_interval_ms);
set!("AETHER_PROXY_TUNNEL_MAX_STREAMS", self.tunnel_max_streams);
set!(
"AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT",
self.tunnel_connect_timeout_secs
TUNNEL_CONNECT_TIMEOUT_MS_ENV,
self.tunnel_connect_timeout_ms
);
set!(
"AETHER_PROXY_TUNNEL_TCP_KEEPALIVE",
self.tunnel_tcp_keepalive_secs
);
set!("AETHER_PROXY_TUNNEL_TCP_NODELAY", self.tunnel_tcp_nodelay);
set!(
"AETHER_PROXY_TUNNEL_STALE_TIMEOUT",
self.tunnel_stale_timeout_secs
);
set!(TUNNEL_STALE_TIMEOUT_MS_ENV, self.tunnel_stale_timeout_ms);
set!("AETHER_PROXY_TUNNEL_CONNECTIONS", self.tunnel_connections);
set!(
"AETHER_PROXY_TUNNEL_CONNECTIONS_MAX",
self.tunnel_connections_max
);
set!(
"AETHER_PROXY_TUNNEL_SCALE_CHECK_INTERVAL_MS",
self.tunnel_scale_check_interval_ms
);
set!(
"AETHER_PROXY_TUNNEL_SCALE_UP_THRESHOLD_PERCENT",
self.tunnel_scale_up_threshold_percent
);
set!(
"AETHER_PROXY_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT",
self.tunnel_scale_down_threshold_percent
);
set!(
"AETHER_PROXY_TUNNEL_SCALE_DOWN_GRACE_SECS",
self.tunnel_scale_down_grace_secs
);
// allowed_ports needs special handling (comma-separated)
if let Some(ref ports) = self.allowed_ports {
@@ -1022,11 +1089,70 @@ impl ConfigFile {
}
}
fn reject_removed_config_keys(content: &str) -> anyhow::Result<()> {
let value: toml::Value = toml::from_str(content)?;
let Some(table) = value.as_table() else {
return Ok(());
};
let removed_seconds = REMOVED_TUNNEL_SECONDS_KEYS
.iter()
.copied()
.filter(|key| table.contains_key(*key))
.collect::<Vec<_>>();
if !removed_seconds.is_empty() {
anyhow::bail!(
"removed tunnel config keys detected: {}. Use *_ms variants instead",
removed_seconds.join(", ")
);
}
let removed_single_server = REMOVED_SINGLE_SERVER_KEYS
.iter()
.copied()
.filter(|key| table.contains_key(*key))
.collect::<Vec<_>>();
if !removed_single_server.is_empty() {
anyhow::bail!(
"single-server top-level config keys are no longer supported: {}. Use [[servers]] entries instead",
removed_single_server.join(", ")
);
}
let removed_legacy = LEGACY_ONLY_KEYS
.iter()
.copied()
.filter(|key| table.contains_key(*key))
.chain(
DELEGATE_TO_UPSTREAM
.iter()
.map(|(old, _)| *old)
.filter(|key| table.contains_key(*key)),
)
.collect::<Vec<_>>();
if !removed_legacy.is_empty() {
anyhow::bail!(
"legacy config keys are no longer supported: {}",
removed_legacy.join(", ")
);
}
Ok(())
}
fn div_ceil_u64(value: u64, divisor: u64) -> u64 {
if divisor == 0 {
return value;
}
value.saturating_add(divisor.saturating_sub(1)) / divisor
}
#[cfg(test)]
mod tests {
use clap::CommandFactory;
use clap::{CommandFactory, Parser};
use super::*;
use crate::hardware::HardwareInfo;
#[test]
fn parse_byte_size_supports_human_units() {
@@ -1056,6 +1182,36 @@ mod tests {
assert_eq!(stringy.redirect_replay_budget_bytes.as_deref(), Some("6M"));
}
#[test]
fn config_file_rejects_removed_tunnel_seconds_keys() {
let error = reject_removed_config_keys("tunnel_ping_interval_secs = 5")
.expect_err("removed tunnel seconds keys should be rejected");
assert!(
error.to_string().contains("tunnel_ping_interval_secs"),
"error should mention removed key"
);
}
#[test]
fn config_file_rejects_top_level_single_server_keys() {
let error = reject_removed_config_keys("aether_url = \"https://example.com\"")
.expect_err("top-level single-server key should be rejected");
assert!(
error.to_string().contains("aether_url"),
"error should mention removed single-server key"
);
}
#[test]
fn config_file_rejects_legacy_keys() {
let error = reject_removed_config_keys("delegate_connect_timeout_secs = 10")
.expect_err("legacy delegate key should be rejected");
assert!(
error.to_string().contains("delegate_connect_timeout_secs"),
"error should mention removed legacy key"
);
}
#[test]
fn config_requires_node_name() {
let command = Config::command();
@@ -1067,4 +1223,130 @@ mod tests {
assert!(node_name.is_required_set());
assert!(node_name.get_default_values().is_empty());
}
#[test]
fn tunnel_fast_recovery_defaults_use_millisecond_values() {
let config = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
]);
assert_eq!(
config
.tunnel_ping_interval()
.expect("ping interval should resolve"),
Duration::from_millis(DEFAULT_TUNNEL_PING_INTERVAL_MS)
);
assert_eq!(
config
.tunnel_connect_timeout()
.expect("connect timeout should resolve"),
Duration::from_millis(DEFAULT_TUNNEL_CONNECT_TIMEOUT_MS)
);
assert_eq!(
config
.tunnel_stale_timeout()
.expect("stale timeout should resolve"),
Duration::from_millis(DEFAULT_TUNNEL_STALE_TIMEOUT_MS)
);
}
#[test]
fn tunnel_millisecond_flags_take_effect_when_explicitly_set() {
let config = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--tunnel-ping-interval-ms",
"100",
"--tunnel-connect-timeout-ms",
"200",
"--tunnel-stale-timeout-ms",
"300",
]);
assert_eq!(
config
.tunnel_ping_interval()
.expect("ping interval should resolve"),
Duration::from_millis(100)
);
assert_eq!(
config
.tunnel_connect_timeout()
.expect("connect timeout should resolve"),
Duration::from_millis(200)
);
assert_eq!(
config
.tunnel_stale_timeout()
.expect("stale timeout should resolve"),
Duration::from_millis(300)
);
}
#[test]
fn auto_tunnel_pool_sizing_uses_hardware_capacity() {
let config = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--tunnel-max-streams",
"1024",
]);
let hw = HardwareInfo {
cpu_cores: 12,
total_memory_mb: 20_480,
os_info: "test".to_string(),
fd_limit: 1_048_576,
estimated_max_concurrency: 24_000,
};
let sizing = config
.resolve_tunnel_pool_sizing(&hw)
.expect("sizing should resolve");
assert_eq!(sizing.initial_connections, 3);
assert_eq!(sizing.max_connections, 6);
}
#[test]
fn explicit_tunnel_connections_keep_fixed_pool_without_max_override() {
let config = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--tunnel-max-streams",
"512",
"--tunnel-connections",
"2",
]);
let hw = HardwareInfo {
cpu_cores: 12,
total_memory_mb: 20_480,
os_info: "test".to_string(),
fd_limit: 1_048_576,
estimated_max_concurrency: 24_000,
};
let sizing = config
.resolve_tunnel_pool_sizing(&hw)
.expect("sizing should resolve");
assert_eq!(sizing.initial_connections, 2);
assert_eq!(sizing.max_connections, 2);
}
}

37
apps/aether-proxy/src/main.rs
View File

@@ -61,12 +61,15 @@ async fn main() -> anyhow::Result<()> {
std::env::var("AETHER_PROXY_CONFIG").unwrap_or_else(|_| DEFAULT_CONFIG.to_string());
let config_path = std::path::Path::new(&config_file_path);
if config_path.exists() {
// Migrate legacy 0.1.x config to 0.2.0 format if needed
if let Err(e) = config::ConfigFile::migrate_legacy(config_path) {
eprintln!(" WARNING: config migration failed: {}", e);
}
if let Ok(file_cfg) = config::ConfigFile::load(config_path) {
file_cfg.inject_env();
match config::ConfigFile::load(config_path) {
Ok(file_cfg) => file_cfg.inject_env(),
Err(error) => {
eprintln!(
" WARNING: failed to load config {}: {}",
config_path.display(),
error
);
}
}
}
@@ -148,21 +151,19 @@ async fn run_proxy(config: Config) -> anyhow::Result<()> {
std::process::exit(1);
}
// Resolve server list: prefer [[servers]] from TOML, fall back to CLI/env single server.
// Resolve server list: if a config file exists, it must use [[servers]].
// Otherwise fall back to CLI/env single-server mode.
let config_path =
std::env::var("AETHER_PROXY_CONFIG").unwrap_or_else(|_| DEFAULT_CONFIG.to_string());
let servers = if std::path::Path::new(&config_path).exists() {
config::ConfigFile::load(std::path::Path::new(&config_path))
.ok()
.map(|f| f.effective_servers())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| {
vec![config::ServerEntry {
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: None,
}]
})
let file_cfg = config::ConfigFile::load(std::path::Path::new(&config_path))?;
if file_cfg.servers.is_empty() {
anyhow::bail!(
"config file {} must contain at least one [[servers]] entry",
config_path
);
}
file_cfg.servers.clone()
} else {
vec![config::ServerEntry {
aether_url: config.aether_url.clone(),

19
apps/aether-proxy/src/setup/tui.rs
View File

@@ -172,7 +172,7 @@ impl App {
value: DEFAULT_HEARTBEAT_INTERVAL_SECS.to_string(),
kind: FieldKind::Text,
required: false,
help: "Heartbeat interval in seconds; default is 30",
help: "Heartbeat interval in seconds; default is 5",
},
Field {
label: "Redirect Replay Budget",
@@ -264,22 +264,11 @@ impl App {
}
// Server tabs
let servers = cfg.effective_servers();
let servers = cfg.servers.clone();
if servers.is_empty() {
let mut tab = ServerTab::new();
// Single-server fallback: use top-level node_name
if let Some(ref name) = cfg.node_name {
tab.fields[2].value = name.clone();
}
self.server_tabs = vec![tab];
self.server_tabs = vec![ServerTab::new()];
} else {
self.server_tabs = servers.iter().map(ServerTab::from_entry).collect();
// For single-server mode, node_name might be in top-level only
if self.server_tabs.len() == 1 && self.server_tabs[0].fields[2].value.is_empty() {
if let Some(ref name) = cfg.node_name {
self.server_tabs[0].fields[2].value = name.clone();
}
}
}
self.active_tab = 0;
self.selected = 0;
@@ -387,7 +376,7 @@ impl App {
..ConfigFile::default()
};
// Always write [[servers]] format; old top-level fields are read-only compat
// Always write [[servers]] format.
cfg.servers = self
.server_tabs
.iter()

76
apps/aether-proxy/src/tunnel/client.rs
View File

@@ -31,6 +31,7 @@ pub async fn connect_and_run(
server: &Arc<ServerContext>,
conn_idx: usize,
shutdown: &mut watch::Receiver<bool>,
drain: watch::Receiver<bool>,
) -> Result<TunnelOutcome, anyhow::Error> {
let ws_url = build_tunnel_url(server);
info!(url = %ws_url, conn = conn_idx, "connecting tunnel");
@@ -53,8 +54,7 @@ pub async fn connect_and_run(
http::HeaderValue::from_str(&dynamic_node_name)?,
);
// Advertise per-connection max concurrent streams so the backend can
// respect the proxy's capacity limit (backward-compatible: old backends
// ignore this header).
// respect the proxy's capacity limit.
let max_streams = state.config.tunnel_max_streams.unwrap_or(128);
headers.insert("X-Tunnel-Max-Streams", http::HeaderValue::from(max_streams));
@@ -67,13 +67,16 @@ pub async fn connect_and_run(
let port = uri.port_u16().unwrap_or(if is_tls { 443 } else { 80 });
// TCP connect with timeout
let connect_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let connect_timeout = state
.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 ({}s)",
connect_timeout.as_secs()
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})??;
@@ -96,7 +99,7 @@ pub async fn connect_and_run(
max_message_size: Some(64 << 20),
..Default::default()
};
let handshake_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let handshake_timeout = connect_timeout;
let (ws_stream, _response) = tokio::time::timeout(
handshake_timeout,
tokio_tungstenite::client_async_tls_with_config(
@@ -109,16 +112,25 @@ pub async fn connect_and_run(
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel WebSocket handshake timeout ({}s)",
handshake_timeout.as_secs()
"tunnel WebSocket handshake timeout ({}ms)",
handshake_timeout.as_millis()
)
})??;
let stale_timeout = state
.config
.tunnel_stale_timeout()
.expect("validated config should resolve tunnel stale timeout");
let ping_interval = state
.config
.tunnel_ping_interval()
.expect("validated config should resolve tunnel ping interval");
info!(
conn = conn_idx,
tcp_keepalive_secs = state.config.tunnel_tcp_keepalive_secs,
tcp_nodelay = state.config.tunnel_tcp_nodelay,
connect_timeout_secs = state.config.tunnel_connect_timeout_secs,
stale_timeout_secs = state.config.tunnel_stale_timeout_secs,
connect_timeout_ms = connect_timeout.as_millis(),
stale_timeout_ms = stale_timeout.as_millis(),
ping_interval_ms = ping_interval.as_millis(),
"tunnel connected"
);
@@ -129,8 +141,8 @@ 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 ping_interval = Duration::from_secs(state.config.tunnel_ping_interval_secs);
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
let drain_signal = spawn_drain_signal(conn_idx, frame_tx.clone(), drain.clone());
// Spawn heartbeat task (only for primary connection to avoid
// resetting shared atomic metrics via swap(0))
@@ -153,7 +165,14 @@ pub async fn connect_and_run(
let state_clone = Arc::clone(state);
let server_clone = Arc::clone(server);
let outcome = tokio::select! {
result = dispatcher::run(state_clone, server_clone, ws_read, frame_tx.clone(), hb_handle) => {
result = dispatcher::run(
state_clone,
server_clone,
ws_read,
frame_tx.clone(),
hb_handle,
drain.clone(),
) => {
match result {
Ok(()) => TunnelOutcome::Disconnected,
Err(e) => return Err(e),
@@ -181,6 +200,10 @@ pub async fn connect_and_run(
// Drop our sender; the writer will exit once all stream handler clones
// are also dropped (i.e. after they finish their in-flight work).
drop(frame_tx);
if !drain_signal.is_finished() {
drain_signal.abort();
let _ = drain_signal.await;
}
// Wait for the writer task to finish with a generous timeout — the
// dispatcher already waits up to 30s for stream handlers, so 35s here
@@ -194,6 +217,35 @@ pub async fn connect_and_run(
Ok(outcome)
}
fn spawn_drain_signal(
conn_idx: usize,
frame_tx: writer::FrameSender,
mut drain: watch::Receiver<bool>,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
if !*drain.borrow() {
loop {
if drain.changed().await.is_err() {
return;
}
if *drain.borrow() {
break;
}
}
}
debug!(conn = conn_idx, "sending GOAWAY for tunnel drain");
let _ = tokio::time::timeout(
Duration::from_millis(250),
frame_tx.send(super::protocol::Frame::control(
super::protocol::MsgType::GoAway,
bytes::Bytes::new(),
)),
)
.await;
})
}
/// 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);

51
apps/aether-proxy/src/tunnel/dispatcher.rs
View File

@@ -7,6 +7,7 @@ use std::time::Duration;
use bytes::Bytes;
use futures_util::StreamExt;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, info, warn};
@@ -25,6 +26,7 @@ pub async fn run<S>(
mut ws_stream: S,
frame_tx: FrameSender,
heartbeat: HeartbeatHandle,
mut drain: watch::Receiver<bool>,
) -> Result<(), anyhow::Error>
where
S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
@@ -38,12 +40,21 @@ where
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
let mut frames_since_cleanup: u32 = 0;
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
let stale_timeout = state
.config
.tunnel_stale_timeout()
.expect("validated config should resolve tunnel stale timeout");
// Track last time we received any data to detect stale connections
let mut last_data_at = tokio::time::Instant::now();
let mut draining = *drain.borrow();
let read_err = loop {
if draining && streams.is_empty() {
info!("tunnel drained after in-flight streams completed");
break None;
}
let msg_result = tokio::select! {
msg = ws_stream.next() => {
match msg {
@@ -51,9 +62,19 @@ where
None => break None,
}
}
changed = drain.changed() => {
if changed.is_err() {
continue;
}
if *drain.borrow() {
info!("tunnel drain requested, waiting for in-flight streams");
draining = true;
}
continue;
}
_ = tokio::time::sleep_until(last_data_at + stale_timeout) => {
warn!(
stale_secs = stale_timeout.as_secs(),
stale_ms = stale_timeout.as_millis(),
"tunnel connection stale, no data received"
);
break None;
@@ -92,6 +113,24 @@ where
match frame.msg_type {
MsgType::RequestHeaders => {
if draining {
if frame_tx
.try_send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from("tunnel draining"),
))
.is_err()
{
warn!(
stream_id = frame.stream_id,
"writer channel full, StreamError dropped during drain"
);
}
continue;
}
// Decompress if the frame is gzip-compressed, then parse metadata
let payload = match decompress_if_gzip(&frame) {
Ok(p) => p,
@@ -176,6 +215,10 @@ where
let _ = tx.send(frame).await;
if is_end {
streams.remove(&sid);
if draining && streams.is_empty() {
info!("tunnel drained after request body completion");
break None;
}
}
}
}
@@ -184,6 +227,10 @@ where
// Client-side cancellation or end
if let Some(tx) = streams.remove(&frame.stream_id) {
let _ = tx.send(frame).await;
if draining && streams.is_empty() {
info!("tunnel drained after stream termination");
break None;
}
}
}

95
apps/aether-proxy/src/tunnel/heartbeat.rs
View File

@@ -24,7 +24,7 @@ static NON_ROOT_UPGRADE_WARNED: AtomicBool = AtomicBool::new(false);
enum AckDecision {
Accept {
heartbeat_id: Option<u64>,
heartbeat_id: u64,
upgrade_to: Option<String>,
},
Ignore,
@@ -143,16 +143,8 @@ pub fn spawn(
upgrade_to,
} => {
if let Some((pending_id, _)) = pending {
match ack_id {
Some(id) if id == pending_id => {
pending = None;
}
None => {
// Backward-compatible with servers that don't echo
// heartbeat_id in ACK payload yet.
pending = None;
}
_ => {}
if ack_id == pending_id {
pending = None;
}
}
maybe_trigger_upgrade(upgrade_to);
@@ -266,6 +258,7 @@ async fn build_heartbeat_payload(
"node_id": node_id,
"heartbeat_session_id": heartbeat_session_id,
"heartbeat_id": heartbeat_id,
"heartbeat_interval": server.dynamic.load().heartbeat_interval,
"active_connections": server.active_connections.load(Ordering::Acquire),
"total_requests": snapshot.requests,
"avg_latency_ms": avg_latency_ms,
@@ -283,10 +276,8 @@ async fn build_heartbeat_payload(
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
if payload.is_empty() {
return AckDecision::Accept {
heartbeat_id: None,
upgrade_to: None,
};
warn!("received empty heartbeat ACK");
return AckDecision::Ignore;
}
#[derive(serde::Deserialize)]
@@ -295,8 +286,7 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: u64,
#[serde(default)]
heartbeat_id: Option<u64>,
heartbeat_id: u64,
#[serde(default)]
upgrade_to: Option<String>,
}
@@ -371,3 +361,74 @@ fn maybe_trigger_upgrade(version: Option<String>) {
}
});
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock};
use arc_swap::ArcSwap;
use clap::Parser;
use super::{handle_ack, AckDecision};
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::{ProxyMetrics, ServerContext};
fn sample_server() -> Arc<ServerContext> {
let config = Arc::new(crate::config::Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
]));
Arc::new(ServerContext {
server_label: "heartbeat-test".to_string(),
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: config.node_name.clone(),
node_id: Arc::new(RwLock::new("node-123".to_string())),
aether_client: Arc::new(AetherClient::new(
&config,
&config.aether_url,
&config.management_token,
)),
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
})
}
#[test]
fn heartbeat_ack_requires_heartbeat_id() {
let server = sample_server();
let decision = handle_ack(
&server,
br#"{"config_version":1,"remote_config":{"heartbeat_interval":9}}"#,
);
assert!(matches!(decision, AckDecision::Ignore));
assert_eq!(server.dynamic.load().heartbeat_interval, 5);
}
#[test]
fn heartbeat_ack_applies_remote_config_with_heartbeat_id() {
let server = sample_server();
let decision = handle_ack(
&server,
br#"{"heartbeat_id":7,"config_version":1,"remote_config":{"heartbeat_interval":9}}"#,
);
assert!(matches!(
decision,
AckDecision::Accept {
heartbeat_id: 7,
upgrade_to: None
}
));
assert_eq!(server.dynamic.load().heartbeat_interval, 9);
}
}

44
apps/aether-proxy/src/tunnel/mod.rs
View File

@@ -36,10 +36,16 @@ pub async fn run(
server: &Arc<ServerContext>,
conn_idx: usize,
mut shutdown: watch::Receiver<bool>,
mut drain: watch::Receiver<bool>,
) {
info!(server = %server.server_label, conn = conn_idx, "starting tunnel");
let reconnect_salt = compute_connection_salt(server, conn_idx);
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested before startup");
return;
}
let startup_delay = compute_startup_stagger(conn_idx, reconnect_salt);
if !startup_delay.is_zero() {
info!(
@@ -54,14 +60,24 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during startup stagger");
return;
}
_ = drain.changed() => {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested during startup stagger");
return;
}
}
}
}
let mut consecutive_failures: u32 = 0;
loop {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drained, exiting slot");
return;
}
let started_at = Instant::now();
match client::connect_and_run(state, server, conn_idx, &mut shutdown).await {
match client::connect_and_run(state, server, conn_idx, &mut shutdown, drain.clone()).await {
Ok(client::TunnelOutcome::Shutdown) => {
info!(server = %server.server_label, conn = conn_idx, "tunnel shut down gracefully");
return;
@@ -78,6 +94,10 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested, not reconnecting");
return;
}
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drained after disconnect");
return;
}
// Reset backoff after a stable session to keep recovery snappy when
// failures are only occasional.
@@ -108,6 +128,12 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during reconnect wait");
return;
}
_ = drain.changed() => {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested during reconnect wait");
return;
}
}
}
}
}
@@ -266,8 +292,9 @@ mod tests {
let proxy_task = tokio::spawn({
let state = Arc::clone(&state);
let server = Arc::clone(&server);
let (_drain_tx, drain_rx) = watch::channel(false);
async move {
run(&state, &server, 0, shutdown_rx).await;
run(&state, &server, 0, shutdown_rx, drain_rx).await;
}
});
@@ -486,13 +513,18 @@ mod tests {
log_max_files: 30,
tunnel_reconnect_base_ms: 50,
tunnel_reconnect_max_ms: 250,
tunnel_ping_interval_secs: 1,
tunnel_ping_interval_ms: 1_000,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 2,
tunnel_connect_timeout_ms: 2_000,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 5,
tunnel_connections: 1,
tunnel_stale_timeout_ms: 5_000,
tunnel_connections: Some(1),
tunnel_connections_max: Some(1),
tunnel_scale_check_interval_ms: 1_000,
tunnel_scale_up_threshold_percent: 70,
tunnel_scale_down_threshold_percent: 35,
tunnel_scale_down_grace_secs: 15,
}
}

13
apps/aether-proxy/src/tunnel/stream_handler.rs
View File

@@ -1594,13 +1594,18 @@ mod tests {
log_max_files: 30,
tunnel_reconnect_base_ms: 500,
tunnel_reconnect_max_ms: 30_000,
tunnel_ping_interval_secs: 15,
tunnel_ping_interval_ms: 15_000,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 15,
tunnel_connect_timeout_ms: 15_000,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 45,
tunnel_connections: 1,
tunnel_stale_timeout_ms: 45_000,
tunnel_connections: Some(1),
tunnel_connections_max: Some(1),
tunnel_scale_check_interval_ms: 1_000,
tunnel_scale_up_threshold_percent: 70,
tunnel_scale_down_threshold_percent: 35,
tunnel_scale_down_grace_secs: 15,
}
}