mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'origin/main'
This commit is contained in:
@@ -177,6 +177,17 @@ pub(super) fn classify_admin_endpoints_family_route(
|
||||
"admin:endpoints_manage",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/endpoints/providers/")
|
||||
&& normalized_path.ends_with("/key-balance")
|
||||
{
|
||||
Some(classified(
|
||||
"admin_proxy",
|
||||
"endpoints_manage",
|
||||
"query_key_balance",
|
||||
"admin:endpoints_manage",
|
||||
false,
|
||||
))
|
||||
} else if method == http::Method::POST
|
||||
&& normalized_path.starts_with("/api/admin/endpoints/providers/")
|
||||
&& normalized_path.ends_with("/keys")
|
||||
|
||||
@@ -380,6 +380,20 @@ fn classifies_admin_refresh_provider_quota_as_admin_proxy_route() {
|
||||
assert!(!decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_query_provider_key_balance_as_admin_proxy_route() {
|
||||
let headers = http::HeaderMap::new();
|
||||
let uri: Uri = "/api/admin/endpoints/providers/provider-newapi/key-balance"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
|
||||
.expect("decision should resolve");
|
||||
assert_eq!(decision.route_class.as_deref(), Some("admin_proxy"));
|
||||
assert_eq!(decision.route_family.as_deref(), Some("endpoints_manage"));
|
||||
assert_eq!(decision.route_kind.as_deref(), Some("query_key_balance"));
|
||||
assert!(!decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_refresh_provider_quota_buffers_request_body_for_key_selection() {
|
||||
let headers = headers(&[]);
|
||||
@@ -399,6 +413,25 @@ fn admin_refresh_provider_quota_buffers_request_body_for_key_selection() {
|
||||
assert!(local_proxy_route_requires_buffered_body(&context));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_query_provider_key_balance_buffers_request_body_for_key_secret() {
|
||||
let headers = headers(&[]);
|
||||
let uri: Uri = "/api/admin/endpoints/providers/provider-newapi/key-balance"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision = classify_control_route(&http::Method::POST, &uri, &headers)
|
||||
.expect("decision should resolve");
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-key-balance",
|
||||
&http::Method::POST,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
|
||||
assert!(local_proxy_route_requires_buffered_body(&context));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_default_body_rules_as_admin_proxy_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod balance;
|
||||
mod mutations;
|
||||
mod quota;
|
||||
mod reads;
|
||||
@@ -18,6 +19,10 @@ pub(crate) async fn maybe_build_local_admin_endpoints_keys_response(
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = balance::maybe_handle(state, request_context, request_body).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = mutations::maybe_handle(state, request_context, request_body).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -32,6 +32,68 @@ pub(super) fn admin_provider_ops_is_valid_action_type(action_type: &str) -> bool
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_ops_saved_connector_credentials(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
) -> serde_json::Map<String, serde_json::Value> {
|
||||
admin_provider_ops_decrypted_credentials(
|
||||
state,
|
||||
admin_provider_ops_config_object(provider)
|
||||
.and_then(admin_provider_ops_connector_object)
|
||||
.and_then(|connector| connector.get("credentials")),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn admin_provider_ops_query_balance_response_for_credentials(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
architecture_id: &str,
|
||||
base_url: &str,
|
||||
provider_ops_config: &serde_json::Map<String, serde_json::Value>,
|
||||
connector_config: &serde_json::Map<String, serde_json::Value>,
|
||||
credentials: &serde_json::Map<String, serde_json::Value>,
|
||||
request_config: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
) -> serde_json::Value {
|
||||
let architecture_id = normalize_architecture_id(architecture_id);
|
||||
let Some(architecture) = get_architecture(architecture_id) else {
|
||||
return responses::admin_provider_ops_action_not_supported(
|
||||
"query_balance",
|
||||
ADMIN_PROVIDER_OPS_ACTION_RUST_ONLY_MESSAGE,
|
||||
);
|
||||
};
|
||||
let headers = match build_headers(architecture.architecture_id, connector_config, credentials) {
|
||||
Ok(headers) => headers,
|
||||
Err(message) => {
|
||||
return responses::admin_provider_ops_action_not_configured("query_balance", message);
|
||||
}
|
||||
};
|
||||
let Some(action_config) = resolve_action_config(
|
||||
architecture_id,
|
||||
provider_ops_config,
|
||||
"query_balance",
|
||||
request_config,
|
||||
) else {
|
||||
return responses::admin_provider_ops_action_not_supported(
|
||||
"query_balance",
|
||||
ADMIN_PROVIDER_OPS_ACTION_RUST_ONLY_MESSAGE,
|
||||
);
|
||||
};
|
||||
|
||||
query_balance::admin_provider_ops_run_query_balance_action(
|
||||
state,
|
||||
provider_id,
|
||||
provider,
|
||||
&architecture,
|
||||
base_url,
|
||||
&action_config,
|
||||
&headers,
|
||||
credentials,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn admin_provider_ops_local_action_response(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
|
||||
@@ -111,11 +111,14 @@ pub(super) async fn admin_provider_ops_run_query_balance_action(
|
||||
|
||||
if status != http::StatusCode::OK {
|
||||
let cookie_auth = architecture.query_balance_cookie_auth_errors;
|
||||
let new_api_token_auth = architecture.architecture_id == "new_api";
|
||||
return match status {
|
||||
http::StatusCode::UNAUTHORIZED => admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
if cookie_auth {
|
||||
if new_api_token_auth {
|
||||
"访问令牌无效,请使用 New API 个人安全设置里的访问令牌"
|
||||
} else if cookie_auth {
|
||||
"Cookie 已失效,请重新配置"
|
||||
} else {
|
||||
"认证失败"
|
||||
@@ -125,7 +128,9 @@ pub(super) async fn admin_provider_ops_run_query_balance_action(
|
||||
http::StatusCode::FORBIDDEN => admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
if cookie_auth {
|
||||
if new_api_token_auth {
|
||||
"访问令牌无效或无权限,请使用 New API 个人安全设置里的访问令牌"
|
||||
} else if cookie_auth {
|
||||
"Cookie 已失效或无权限"
|
||||
} else {
|
||||
"无权限访问"
|
||||
|
||||
@@ -8,7 +8,9 @@ use super::super::responses::{
|
||||
};
|
||||
use super::super::support::admin_provider_ops_json_object_map;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_admin::provider::ops::parse_sub2api_balance_payload;
|
||||
use aether_admin::provider::ops::{
|
||||
parse_sub2api_api_key_usage_payload, parse_sub2api_balance_payload,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
use serde_json::{json, Value};
|
||||
@@ -24,6 +26,23 @@ pub(super) async fn admin_provider_ops_sub2api_balance_payload(
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> serde_json::Value {
|
||||
let start = std::time::Instant::now();
|
||||
if let Some(api_key) = credentials
|
||||
.get("api_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return admin_provider_ops_sub2api_api_key_balance_payload(
|
||||
state,
|
||||
provider_id,
|
||||
base_url,
|
||||
action_config,
|
||||
api_key,
|
||||
proxy_snapshot,
|
||||
start,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
let (access_token, updated_credentials, _frontend_updated_credentials) =
|
||||
match admin_provider_ops_sub2api_exchange_token(
|
||||
state,
|
||||
@@ -191,6 +210,117 @@ pub(super) async fn admin_provider_ops_sub2api_balance_payload(
|
||||
)
|
||||
}
|
||||
|
||||
async fn admin_provider_ops_sub2api_api_key_balance_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
base_url: &str,
|
||||
action_config: &serde_json::Map<String, serde_json::Value>,
|
||||
api_key: &str,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
start: std::time::Instant,
|
||||
) -> serde_json::Value {
|
||||
let usage_endpoint = action_config
|
||||
.get("api_key_usage_endpoint")
|
||||
.or_else(|| action_config.get("usage_endpoint"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("/v1/usage");
|
||||
let usage_url = admin_provider_ops_sub2api_request_url(base_url, usage_endpoint);
|
||||
let auth_value = match reqwest::header::HeaderValue::from_str(&format!("Bearer {api_key}")) {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
return admin_provider_ops_action_error(
|
||||
"parse_error",
|
||||
"query_balance",
|
||||
"API Key 格式无效",
|
||||
Some(start.elapsed().as_millis() as u64),
|
||||
);
|
||||
}
|
||||
};
|
||||
let auth_headers = reqwest::header::HeaderMap::from_iter([
|
||||
(reqwest::header::AUTHORIZATION, auth_value),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
]);
|
||||
let request_id = format!("provider-ops-action:sub2api:usage:{provider_id}");
|
||||
let result = admin_provider_ops_execute_json_request(
|
||||
state,
|
||||
&request_id,
|
||||
reqwest::Method::GET,
|
||||
&usage_url,
|
||||
&auth_headers,
|
||||
None,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await;
|
||||
let response_time_ms = Some(start.elapsed().as_millis() as u64);
|
||||
let (status, response_json) = match result {
|
||||
Ok(result) => result,
|
||||
Err(AdminProviderOpsExecuteJsonError::InvalidJson(message))
|
||||
| Err(AdminProviderOpsExecuteJsonError::Transport(message)) => {
|
||||
return admin_provider_ops_action_error(
|
||||
"network_error",
|
||||
"query_balance",
|
||||
network_error_message(&message),
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
};
|
||||
if matches!(
|
||||
status,
|
||||
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
|
||||
) {
|
||||
return admin_provider_ops_action_error(
|
||||
"auth_failed",
|
||||
"query_balance",
|
||||
"认证失败,请检查 API Key",
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
if status != http::StatusCode::OK {
|
||||
return admin_provider_ops_action_error(
|
||||
"unknown_error",
|
||||
"query_balance",
|
||||
format!(
|
||||
"HTTP {}: {}",
|
||||
status.as_u16(),
|
||||
status.canonical_reason().unwrap_or("Unknown")
|
||||
),
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
|
||||
let data = match parse_sub2api_api_key_usage_payload(action_config, &response_json) {
|
||||
Ok(payload) => payload,
|
||||
Err(message) => {
|
||||
return admin_provider_ops_action_error(
|
||||
if message.contains("无效") {
|
||||
"auth_failed"
|
||||
} else if message == "响应格式无效" {
|
||||
"parse_error"
|
||||
} else {
|
||||
"unknown_error"
|
||||
},
|
||||
"query_balance",
|
||||
message,
|
||||
response_time_ms,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
admin_provider_ops_action_response(
|
||||
"success",
|
||||
"query_balance",
|
||||
data,
|
||||
None,
|
||||
response_time_ms,
|
||||
86400,
|
||||
)
|
||||
}
|
||||
|
||||
fn network_error_message(error: &str) -> String {
|
||||
let normalized = error.trim();
|
||||
let lower = normalized.to_ascii_lowercase();
|
||||
|
||||
@@ -250,6 +250,9 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
payload: AdminProviderOpsSaveConfigRequest,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let architecture_id =
|
||||
admin_provider_ops_pure::normalize_architecture_id(payload.architecture_id.as_str())
|
||||
.to_string();
|
||||
let auth_type = payload.connector.auth_type.trim().to_string();
|
||||
if auth_type.is_empty() || !admin_provider_ops_is_supported_auth_type(auth_type.as_str()) {
|
||||
return Err("connector.auth_type 必须是合法的认证类型".to_string());
|
||||
@@ -257,7 +260,7 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
|
||||
let merged_credentials = admin_provider_ops_merge_credentials(
|
||||
state,
|
||||
payload.architecture_id.as_str(),
|
||||
architecture_id.as_str(),
|
||||
provider,
|
||||
payload.connector.credentials,
|
||||
);
|
||||
@@ -278,7 +281,7 @@ pub(super) fn build_admin_provider_ops_saved_config_value(
|
||||
.collect::<serde_json::Map<String, serde_json::Value>>();
|
||||
|
||||
Ok(json!({
|
||||
"architecture_id": payload.architecture_id,
|
||||
"architecture_id": architecture_id,
|
||||
"base_url": payload.base_url,
|
||||
"connector": {
|
||||
"auth_type": auth_type,
|
||||
@@ -328,14 +331,16 @@ pub(super) fn build_admin_provider_ops_config_payload(
|
||||
});
|
||||
};
|
||||
let connector = admin_provider_ops_connector_object(provider_ops_config);
|
||||
let architecture_id = provider_ops_config
|
||||
.get("architecture_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(admin_provider_ops_pure::normalize_architecture_id)
|
||||
.unwrap_or("generic_api");
|
||||
|
||||
json!({
|
||||
"provider_id": provider_id,
|
||||
"is_configured": true,
|
||||
"architecture_id": provider_ops_config
|
||||
.get("architecture_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("generic_api"),
|
||||
"architecture_id": architecture_id,
|
||||
"base_url": resolve_admin_provider_ops_base_url(
|
||||
provider,
|
||||
endpoints,
|
||||
|
||||
@@ -6,7 +6,8 @@ use crate::handlers::admin::provider::ops::providers::config::persist_admin_prov
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use aether_admin::provider::ops::{
|
||||
admin_provider_ops_frontend_updated_credentials, admin_provider_ops_verify_failure,
|
||||
parse_verify_payload, ADMIN_PROVIDER_OPS_USER_AGENT,
|
||||
admin_provider_ops_verify_success, admin_provider_ops_verify_user_payload,
|
||||
parse_sub2api_api_key_usage_payload, parse_verify_payload, ADMIN_PROVIDER_OPS_USER_AGENT,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
@@ -21,6 +22,21 @@ pub(super) async fn admin_provider_ops_local_sub2api_verify_response(
|
||||
credentials: &Map<String, Value>,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> Value {
|
||||
if let Some(api_key) = credentials
|
||||
.get("api_key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return admin_provider_ops_local_sub2api_api_key_verify_response(
|
||||
state,
|
||||
base_url,
|
||||
api_key,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let (access_token, updated_credentials, frontend_updated_credentials) =
|
||||
match admin_provider_ops_sub2api_exchange_token(
|
||||
state,
|
||||
@@ -93,6 +109,78 @@ pub(super) async fn admin_provider_ops_local_sub2api_verify_response(
|
||||
)
|
||||
}
|
||||
|
||||
async fn admin_provider_ops_local_sub2api_api_key_verify_response(
|
||||
state: &AdminAppState<'_>,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
proxy_snapshot: Option<&ProxySnapshot>,
|
||||
) -> Value {
|
||||
let usage_url = admin_provider_ops_sub2api_request_url(base_url, "/v1/usage");
|
||||
let auth_value = match reqwest::header::HeaderValue::from_str(&format!("Bearer {api_key}")) {
|
||||
Ok(value) => value,
|
||||
Err(_) => return admin_provider_ops_verify_failure("API Key 格式无效"),
|
||||
};
|
||||
let auth_headers = reqwest::header::HeaderMap::from_iter([
|
||||
(reqwest::header::AUTHORIZATION, auth_value),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
]);
|
||||
let auth_headers =
|
||||
admin_provider_ops_headers_with_transport_controls(&auth_headers, None, true);
|
||||
let (status, response_json) = match admin_provider_ops_execute_json_request(
|
||||
state,
|
||||
"provider-ops-verify:sub2api:api-key",
|
||||
reqwest::Method::GET,
|
||||
&usage_url,
|
||||
&auth_headers,
|
||||
None,
|
||||
proxy_snapshot,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(AdminProviderOpsExecuteJsonError::InvalidJson(message))
|
||||
| Err(AdminProviderOpsExecuteJsonError::Transport(message)) => {
|
||||
return admin_provider_ops_verify_failure(
|
||||
admin_provider_ops_verify_execution_error_message(&message),
|
||||
);
|
||||
}
|
||||
};
|
||||
if matches!(
|
||||
status,
|
||||
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
|
||||
) {
|
||||
return admin_provider_ops_verify_failure("认证失败:API Key 无效或已过期");
|
||||
}
|
||||
if status != http::StatusCode::OK {
|
||||
return admin_provider_ops_verify_failure(format!("验证失败:HTTP {}", status.as_u16()));
|
||||
}
|
||||
|
||||
let payload = match parse_sub2api_api_key_usage_payload(&Map::new(), &response_json) {
|
||||
Ok(payload) => payload,
|
||||
Err(message) => return admin_provider_ops_verify_failure(message),
|
||||
};
|
||||
let quota = payload.get("total_available").and_then(Value::as_f64);
|
||||
let extra = payload
|
||||
.get("extra")
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
admin_provider_ops_verify_success(
|
||||
admin_provider_ops_verify_user_payload(
|
||||
Some("Sub2API API Key".to_string()),
|
||||
Some("Sub2API API Key".to_string()),
|
||||
None,
|
||||
quota,
|
||||
Some(extra),
|
||||
),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
// 对齐 Python httpx.AsyncClient(base_url=...) 的行为:
|
||||
// 以 "/" 开头的端点始终相对站点根路径解析,而不是简单字符串拼接。
|
||||
pub(in super::super) fn admin_provider_ops_sub2api_request_url(
|
||||
|
||||
@@ -12,6 +12,13 @@ pub(crate) fn admin_provider_id_for_refresh_quota(request_path: &str) -> Option<
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_provider_id_for_key_balance(request_path: &str) -> Option<String> {
|
||||
request_path
|
||||
.strip_prefix("/api/admin/endpoints/providers/")?
|
||||
.strip_suffix("/key-balance")
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn admin_reveal_key_id(request_path: &str) -> Option<String> {
|
||||
request_path
|
||||
.strip_prefix("/api/admin/endpoints/keys/")?
|
||||
|
||||
@@ -15,9 +15,9 @@ pub(crate) use self::crud::{
|
||||
is_admin_providers_root,
|
||||
};
|
||||
pub(crate) use self::endpoint_keys::{
|
||||
admin_clear_oauth_invalid_key_id, admin_export_key_id, admin_provider_id_for_keys,
|
||||
admin_provider_id_for_refresh_quota, admin_reset_cycle_stats_key_id, admin_reveal_key_id,
|
||||
admin_update_key_id,
|
||||
admin_clear_oauth_invalid_key_id, admin_export_key_id, admin_provider_id_for_key_balance,
|
||||
admin_provider_id_for_keys, admin_provider_id_for_refresh_quota,
|
||||
admin_reset_cycle_stats_key_id, admin_reveal_key_id, admin_update_key_id,
|
||||
};
|
||||
pub(crate) use self::oauth::{
|
||||
admin_provider_oauth_batch_import_provider_id, admin_provider_oauth_batch_import_task_path,
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(crate) async fn build_admin_provider_summary_payload(
|
||||
active_global_model_ids_result,
|
||||
) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.read_provider_quota_snapshot(provider_id),
|
||||
state.list_provider_model_stats(&provider_ids),
|
||||
state.list_active_global_model_ids_by_provider_ids(&provider_ids),
|
||||
@@ -197,7 +197,7 @@ pub(crate) async fn build_admin_providers_summary_payload(
|
||||
} else {
|
||||
let (endpoints_result, keys_result, model_stats_result, active_global_model_refs_result) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
state.list_provider_model_stats(&provider_ids),
|
||||
state.list_active_global_model_ids_by_provider_ids(&provider_ids),
|
||||
);
|
||||
|
||||
@@ -7,7 +7,7 @@ use aether_data_contracts::repository::candidates::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
fn json_truthy(value: &serde_json::Value) -> bool {
|
||||
@@ -27,6 +27,80 @@ fn endpoint_timestamp_or_now(value: Option<u64>, now_unix_secs: u64) -> serde_js
|
||||
.unwrap_or(serde_json::Value::Null)
|
||||
}
|
||||
|
||||
fn finite_json_number(value: Option<&Value>) -> Option<f64> {
|
||||
match value {
|
||||
Some(Value::Number(number)) => number.as_f64().filter(|value| value.is_finite()),
|
||||
Some(Value::String(value)) => value
|
||||
.trim()
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.filter(|value| value.is_finite()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn finite_json_u64(value: Option<&Value>) -> Option<u64> {
|
||||
finite_json_number(value).and_then(|value| {
|
||||
if value >= 0.0 {
|
||||
Some(value as u64)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn latest_key_balance_summary(keys: &[StoredProviderCatalogKey]) -> Value {
|
||||
let mut selected: Option<(u64, &StoredProviderCatalogKey, &Map<String, Value>)> = None;
|
||||
|
||||
for key in keys {
|
||||
let Some(balance) = key
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|metadata| metadata.get("balance_query"))
|
||||
.and_then(Value::as_object)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let Some(updated_at) = finite_json_u64(balance.get("updated_at")) else {
|
||||
continue;
|
||||
};
|
||||
let has_balance_value = ["total_available", "total_used", "total_granted"]
|
||||
.into_iter()
|
||||
.any(|field| finite_json_number(balance.get(field)).is_some());
|
||||
if !has_balance_value {
|
||||
continue;
|
||||
}
|
||||
|
||||
if selected
|
||||
.as_ref()
|
||||
.is_none_or(|(selected_updated_at, _, _)| updated_at > *selected_updated_at)
|
||||
{
|
||||
selected = Some((updated_at, key, balance));
|
||||
}
|
||||
}
|
||||
|
||||
let Some((updated_at, key, balance)) = selected else {
|
||||
return Value::Null;
|
||||
};
|
||||
|
||||
json!({
|
||||
"key_id": key.id.clone(),
|
||||
"key_name": key.name.clone(),
|
||||
"updated_at": updated_at,
|
||||
"architecture_id": balance.get("architecture_id").cloned().unwrap_or(Value::Null),
|
||||
"status": balance.get("status").cloned().unwrap_or_else(|| json!("success")),
|
||||
"executed_at": balance.get("executed_at").cloned().unwrap_or(Value::Null),
|
||||
"response_time_ms": balance.get("response_time_ms").cloned().unwrap_or(Value::Null),
|
||||
"total_available": balance.get("total_available").cloned().unwrap_or(Value::Null),
|
||||
"total_used": balance.get("total_used").cloned().unwrap_or(Value::Null),
|
||||
"total_granted": balance.get("total_granted").cloned().unwrap_or(Value::Null),
|
||||
"currency": balance.get("currency").cloned().unwrap_or_else(|| json!("USD")),
|
||||
"plan_name": balance.get("plan_name").cloned().unwrap_or(Value::Null),
|
||||
"query_config": balance.get("query_config").cloned().unwrap_or(Value::Null),
|
||||
"extra": balance.get("extra").cloned().unwrap_or(Value::Null),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_admin_provider_summary_value(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
@@ -158,6 +232,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
.and_then(|quota| quota.quota_expires_at_unix_secs)
|
||||
.or(provider.quota_expires_at_unix_secs)
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
let key_balance_summary = latest_key_balance_summary(keys);
|
||||
|
||||
json!({
|
||||
"id": provider.id.clone(),
|
||||
@@ -196,6 +271,7 @@ pub(crate) fn build_admin_provider_summary_value(
|
||||
"endpoint_health_details": endpoint_health_details,
|
||||
"ops_configured": ops_configured,
|
||||
"ops_architecture_id": ops_architecture_id,
|
||||
"key_balance_summary": key_balance_summary,
|
||||
"kiro_simulated_cache_enabled": kiro_simulated_cache_enabled,
|
||||
"created_at": endpoint_timestamp_or_now(provider.created_at_unix_ms, now_unix_secs),
|
||||
"updated_at": endpoint_timestamp_or_now(provider.updated_at_unix_secs, now_unix_secs),
|
||||
|
||||
@@ -177,6 +177,21 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.update_provider_catalog_key(key).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
|
||||
&self,
|
||||
key_id: &str,
|
||||
upstream_metadata: Option<&serde_json::Value>,
|
||||
updated_at_unix_secs: Option<u64>,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.app
|
||||
.update_provider_catalog_key_upstream_metadata(
|
||||
key_id,
|
||||
upstream_metadata,
|
||||
updated_at_unix_secs,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_provider_catalog_key(
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
|
||||
@@ -21,6 +21,7 @@ const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] ";
|
||||
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
|
||||
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
|
||||
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
|
||||
const BALANCE_QUERY_SECRET_CIPHERTEXT_KEY: &str = "secret_ciphertext";
|
||||
|
||||
pub(crate) fn provider_catalog_key_supports_format(
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -173,6 +174,31 @@ pub(crate) fn parse_catalog_auth_config_json(
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn sanitized_admin_upstream_metadata(upstream_metadata: Option<&Value>) -> Value {
|
||||
let Some(mut metadata) = upstream_metadata.cloned() else {
|
||||
return Value::Null;
|
||||
};
|
||||
let Some(balance_query) = metadata
|
||||
.as_object_mut()
|
||||
.and_then(|metadata| metadata.get_mut("balance_query"))
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return metadata;
|
||||
};
|
||||
let has_saved_secret = balance_query
|
||||
.remove(BALANCE_QUERY_SECRET_CIPHERTEXT_KEY)
|
||||
.is_some();
|
||||
if has_saved_secret {
|
||||
let query_config = balance_query
|
||||
.entry("query_config".to_string())
|
||||
.or_insert_with(|| Value::Object(Map::new()));
|
||||
if let Some(query_config) = query_config.as_object_mut() {
|
||||
query_config.insert("has_saved_secret".to_string(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
metadata
|
||||
}
|
||||
|
||||
pub(crate) fn default_provider_key_status_snapshot() -> serde_json::Value {
|
||||
json!({
|
||||
"oauth": {
|
||||
@@ -1937,7 +1963,7 @@ pub(crate) fn build_admin_provider_key_response(
|
||||
);
|
||||
payload.insert(
|
||||
"upstream_metadata".to_string(),
|
||||
json!(key.upstream_metadata),
|
||||
sanitized_admin_upstream_metadata(key.upstream_metadata.as_ref()),
|
||||
);
|
||||
payload.insert("proxy".to_string(), json!(key.proxy));
|
||||
payload.insert("fingerprint".to_string(), json!(key.fingerprint));
|
||||
|
||||
@@ -221,6 +221,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("create_endpoint"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("batch_delete_keys"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("refresh_quota"))
|
||||
| (Some("endpoints_manage"), http::Method::POST, Some("query_key_balance"))
|
||||
| (Some("endpoints_manage"), http::Method::PUT, Some("update_key"))
|
||||
| (Some("endpoints_manage"), http::Method::PUT, Some("update_endpoint"))
|
||||
| (Some("modules_manage"), http::Method::PUT, Some("set_enabled"))
|
||||
|
||||
@@ -1490,7 +1490,24 @@ async fn gateway_verifies_admin_provider_ops_locally_for_new_api_with_trusted_ad
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("session=foo")
|
||||
);
|
||||
assert!(headers.contains_key("sec-ch-ua"));
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::USER_AGENT)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("cc-switch/1.0")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(axum::http::header::ACCEPT)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
encoder
|
||||
.write_all(
|
||||
|
||||
13
apps/aether-proxy/install.ps1
Normal file
13
apps/aether-proxy/install.ps1
Normal file
@@ -0,0 +1,13 @@
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
if ($env:AETHER_PROXY_AETHER_URL -and -not $env:AETHER_TUNNEL_AETHER_URL) {
|
||||
$env:AETHER_TUNNEL_AETHER_URL = $env:AETHER_PROXY_AETHER_URL
|
||||
}
|
||||
if ($env:AETHER_PROXY_MANAGEMENT_TOKEN -and -not $env:AETHER_TUNNEL_MANAGEMENT_TOKEN) {
|
||||
$env:AETHER_TUNNEL_MANAGEMENT_TOKEN = $env:AETHER_PROXY_MANAGEMENT_TOKEN
|
||||
}
|
||||
if ($env:AETHER_PROXY_NODE_NAME -and -not $env:AETHER_TUNNEL_NODE_NAME) {
|
||||
$env:AETHER_TUNNEL_NODE_NAME = $env:AETHER_PROXY_NODE_NAME
|
||||
}
|
||||
|
||||
irm 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.ps1' | iex
|
||||
21
apps/aether-proxy/install.sh
Normal file
21
apps/aether-proxy/install.sh
Normal file
@@ -0,0 +1,21 @@
|
||||
#!/bin/sh
|
||||
set -eu
|
||||
|
||||
if [ -n "${AETHER_PROXY_AETHER_URL:-}" ] && [ -z "${AETHER_TUNNEL_AETHER_URL:-}" ]; then
|
||||
export AETHER_TUNNEL_AETHER_URL="${AETHER_PROXY_AETHER_URL}"
|
||||
fi
|
||||
if [ -n "${AETHER_PROXY_MANAGEMENT_TOKEN:-}" ] && [ -z "${AETHER_TUNNEL_MANAGEMENT_TOKEN:-}" ]; then
|
||||
export AETHER_TUNNEL_MANAGEMENT_TOKEN="${AETHER_PROXY_MANAGEMENT_TOKEN}"
|
||||
fi
|
||||
if [ -n "${AETHER_PROXY_NODE_NAME:-}" ] && [ -z "${AETHER_TUNNEL_NODE_NAME:-}" ]; then
|
||||
export AETHER_TUNNEL_NODE_NAME="${AETHER_PROXY_NODE_NAME}"
|
||||
fi
|
||||
|
||||
if command -v curl >/dev/null 2>&1; then
|
||||
curl -fsSL 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh' | sh
|
||||
elif command -v wget >/dev/null 2>&1; then
|
||||
wget -qO- 'https://raw.githubusercontent.com/fawney19/Aether/main/apps/aether-tunnel/install.sh' | sh
|
||||
else
|
||||
printf '%s\n' "[Aether Tunnel] 需要 curl 或 wget 下载安装脚本" >&2
|
||||
exit 1
|
||||
fi
|
||||
@@ -10,6 +10,10 @@ AETHER_TUNNEL_NODE_NAME=jp-proxy-01
|
||||
# Maximum request body buffered for 307/308 replay (supports K/M/G, 0 disables body replay buffering)
|
||||
AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES=5M
|
||||
|
||||
# Optional tunnel TCP address-family restriction (set at most one to true)
|
||||
AETHER_TUNNEL_IPV4_ONLY=false
|
||||
AETHER_TUNNEL_IPV6_ONLY=false
|
||||
|
||||
# Logging
|
||||
AETHER_TUNNEL_LOG_LEVEL=info
|
||||
AETHER_TUNNEL_LOG_DESTINATION=stdout
|
||||
|
||||
@@ -126,6 +126,8 @@ sudo aether-tunnel uninstall
|
||||
| `--tunnel-max-streams` | `AETHER_TUNNEL_MAX_STREAMS` | 自动(硬件估算) | 单连接最大并发 stream 数 |
|
||||
| `--tunnel-ping-interval-ms` | `AETHER_TUNNEL_PING_INTERVAL_MS` | `10000` | WebSocket ping 周期(毫秒) |
|
||||
| `--tunnel-connect-timeout-ms` | `AETHER_TUNNEL_CONNECT_TIMEOUT_MS` | `3000` | tunnel 建连超时(毫秒) |
|
||||
| `--tunnel-ipv4-only` | `AETHER_TUNNEL_IPV4_ONLY` | `false` | 仅使用 IPv4 地址建立直连 WebSocket tunnel;配置 `aether_outbound_proxy_url` 时仅限制代理端点解析 |
|
||||
| `--tunnel-ipv6-only` | `AETHER_TUNNEL_IPV6_ONLY` | `false` | 仅使用 IPv6 地址建立直连 WebSocket tunnel;配置 `aether_outbound_proxy_url` 时仅限制代理端点解析 |
|
||||
| `--tunnel-stale-timeout-ms` | `AETHER_TUNNEL_STALE_TIMEOUT_MS` | `30000` | 无入站数据断连阈值(毫秒) |
|
||||
| `--tunnel-scale-check-interval-ms` | `AETHER_TUNNEL_SCALE_CHECK_INTERVAL_MS` | `1000` | autoscale 采样周期(毫秒) |
|
||||
| `--tunnel-scale-up-threshold-percent` | `AETHER_TUNNEL_SCALE_UP_THRESHOLD_PERCENT` | `50` | 单 tunnel 占用率超过该值时扩容 |
|
||||
@@ -138,6 +140,8 @@ sudo aether-tunnel uninstall
|
||||
|
||||
省略 `tunnel_connections` 时,tunnel 会按设备能力自动计算一个基线值和偏单机上限的扩容上限:默认至少保留 2 条常驻 tunnel,并会更早触发扩容;如果显式设置了 `tunnel_connections` 但没有设置 `tunnel_connections_max`,则保持固定连接池,不自动扩缩。
|
||||
|
||||
`tunnel_ipv4_only` / `tunnel_ipv6_only` 只能二选一。它们只改变 WebSocket tunnel 回连的 TCP 地址选择:直连 Aether 时过滤 Aether 域名的 DNS 结果;配置 `aether_outbound_proxy_url` 时过滤代理服务器端点的 DNS 结果,Host/SNI 仍使用原始 WebSocket URL。该选项不会影响 provider 上游请求;如需限制 provider 上游流量,请在 `upstream_proxy_url` 或系统网络层处理。对于 Cloudflare 等边缘 IP 会变化的域名,优先使用该选项而不是固定 `/etc/hosts`。
|
||||
|
||||
#### 上游 HTTP 请求
|
||||
|
||||
| 参数 | 环境变量 | 默认值 | 说明 |
|
||||
|
||||
@@ -1356,6 +1356,8 @@ mod tests {
|
||||
tunnel_ping_interval_ms: 1_000,
|
||||
tunnel_max_streams: Some(8),
|
||||
tunnel_connect_timeout_ms: 2_000,
|
||||
tunnel_ipv4_only: false,
|
||||
tunnel_ipv6_only: false,
|
||||
tunnel_tcp_keepalive_secs: 30,
|
||||
tunnel_tcp_nodelay: true,
|
||||
tunnel_stale_timeout_ms: 5_000,
|
||||
|
||||
@@ -571,6 +571,30 @@ pub struct Config {
|
||||
)]
|
||||
pub tunnel_connect_timeout_ms: u64,
|
||||
|
||||
/// Force direct WebSocket tunnel TCP connects, or Aether outbound proxy endpoint connects, to IPv4 addresses only.
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_TUNNEL_IPV4_ONLY",
|
||||
default_value_t = false,
|
||||
action = clap::ArgAction::Set,
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1,
|
||||
require_equals = true
|
||||
)]
|
||||
pub tunnel_ipv4_only: bool,
|
||||
|
||||
/// Force direct WebSocket tunnel TCP connects, or Aether outbound proxy endpoint connects, to IPv6 addresses only.
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_TUNNEL_IPV6_ONLY",
|
||||
default_value_t = false,
|
||||
action = clap::ArgAction::Set,
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1,
|
||||
require_equals = true
|
||||
)]
|
||||
pub tunnel_ipv6_only: bool,
|
||||
|
||||
/// WebSocket tunnel TCP keepalive in seconds (0 disables)
|
||||
#[arg(long, env = "AETHER_TUNNEL_TCP_KEEPALIVE", default_value_t = 30)]
|
||||
pub tunnel_tcp_keepalive_secs: u64,
|
||||
@@ -655,6 +679,9 @@ impl Config {
|
||||
if tunnel_connect_timeout.is_zero() {
|
||||
anyhow::bail!("effective tunnel connect timeout must be > 0");
|
||||
}
|
||||
if self.tunnel_ipv4_only && self.tunnel_ipv6_only {
|
||||
anyhow::bail!("tunnel_ipv4_only and tunnel_ipv6_only cannot both be enabled");
|
||||
}
|
||||
let tunnel_ping_interval = self.tunnel_ping_interval()?;
|
||||
if tunnel_ping_interval.is_zero() {
|
||||
anyhow::bail!("effective tunnel ping interval must be > 0");
|
||||
@@ -763,6 +790,16 @@ impl Config {
|
||||
Ok(Duration::from_millis(self.tunnel_connect_timeout_ms))
|
||||
}
|
||||
|
||||
pub fn tunnel_ip_family(&self) -> crate::egress_proxy::IpFamily {
|
||||
if self.tunnel_ipv4_only {
|
||||
crate::egress_proxy::IpFamily::Ipv4Only
|
||||
} else if self.tunnel_ipv6_only {
|
||||
crate::egress_proxy::IpFamily::Ipv6Only
|
||||
} else {
|
||||
crate::egress_proxy::IpFamily::Any
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tunnel_stale_timeout(&self) -> anyhow::Result<Duration> {
|
||||
Ok(Duration::from_millis(self.tunnel_stale_timeout_ms))
|
||||
}
|
||||
@@ -959,6 +996,10 @@ pub struct ConfigFile {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tunnel_connect_timeout_ms: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tunnel_ipv4_only: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tunnel_ipv6_only: Option<bool>,
|
||||
#[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>,
|
||||
@@ -1147,6 +1188,8 @@ impl ConfigFile {
|
||||
TUNNEL_CONNECT_TIMEOUT_MS_ENV,
|
||||
self.tunnel_connect_timeout_ms
|
||||
);
|
||||
set!("AETHER_TUNNEL_IPV4_ONLY", self.tunnel_ipv4_only);
|
||||
set!("AETHER_TUNNEL_IPV6_ONLY", self.tunnel_ipv6_only);
|
||||
set!(
|
||||
"AETHER_TUNNEL_TCP_KEEPALIVE",
|
||||
self.tunnel_tcp_keepalive_secs
|
||||
@@ -1338,6 +1381,20 @@ mod tests {
|
||||
assert_eq!(cfg.allow_private_targets, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_deserializes_tunnel_ip_family_flags() {
|
||||
let cfg: ConfigFile = toml::from_str(
|
||||
r#"
|
||||
tunnel_ipv4_only = true
|
||||
tunnel_ipv6_only = false
|
||||
"#,
|
||||
)
|
||||
.expect("tunnel IP-family TOML");
|
||||
|
||||
assert_eq!(cfg.tunnel_ipv4_only, Some(true));
|
||||
assert_eq!(cfg.tunnel_ipv6_only, Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_file_deserializes_upstream_proxy_url() {
|
||||
let cfg: ConfigFile = toml::from_str("upstream_proxy_url = \"http://proxy.example:8080\"")
|
||||
@@ -1531,6 +1588,131 @@ node_name = "tunnel-test"
|
||||
assert!(config.allow_private_targets);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_defaults_tunnel_ip_family_to_any() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
]);
|
||||
|
||||
assert!(!config.tunnel_ipv4_only);
|
||||
assert!(!config.tunnel_ipv6_only);
|
||||
assert_eq!(
|
||||
config.tunnel_ip_family(),
|
||||
crate::egress_proxy::IpFamily::Any
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_accepts_tunnel_ipv4_only() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-ipv4-only",
|
||||
]);
|
||||
|
||||
assert!(config.tunnel_ipv4_only);
|
||||
assert_eq!(
|
||||
config.tunnel_ip_family(),
|
||||
crate::egress_proxy::IpFamily::Ipv4Only
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_accepts_tunnel_ipv6_only() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-ipv6-only",
|
||||
]);
|
||||
|
||||
assert!(config.tunnel_ipv6_only);
|
||||
assert_eq!(
|
||||
config.tunnel_ip_family(),
|
||||
crate::egress_proxy::IpFamily::Ipv6Only
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_parses_conflicting_tunnel_ip_family_flags_before_validation() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-ipv4-only",
|
||||
"--tunnel-ipv6-only",
|
||||
]);
|
||||
|
||||
assert!(config.tunnel_ipv4_only);
|
||||
assert!(config.tunnel_ipv6_only);
|
||||
let error = config
|
||||
.validate()
|
||||
.expect_err("conflicting tunnel IP-family flags should fail validation");
|
||||
assert!(error.to_string().contains("tunnel_ipv4_only"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_accepts_explicit_false_tunnel_ip_family_flags() {
|
||||
let config = Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
"--tunnel-ipv4-only=false",
|
||||
"--tunnel-ipv6-only=false",
|
||||
]);
|
||||
|
||||
assert!(!config.tunnel_ipv4_only);
|
||||
assert!(!config.tunnel_ipv6_only);
|
||||
config
|
||||
.validate()
|
||||
.expect("explicit false family flags should be valid");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_conflicting_toml_tunnel_ip_family_flags() {
|
||||
let config = Config {
|
||||
tunnel_ipv4_only: true,
|
||||
tunnel_ipv6_only: true,
|
||||
..Config::parse_from([
|
||||
"aether-tunnel",
|
||||
"--aether-url",
|
||||
"https://example.com",
|
||||
"--management-token",
|
||||
"ae_test",
|
||||
"--node-name",
|
||||
"tunnel-test",
|
||||
])
|
||||
};
|
||||
|
||||
let error = config
|
||||
.validate()
|
||||
.expect_err("conflicting TOML-injected tunnel family flags should fail validation");
|
||||
assert!(error.to_string().contains("tunnel_ipv4_only"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_fast_recovery_defaults_use_millisecond_values() {
|
||||
let config = Config::parse_from([
|
||||
|
||||
@@ -8,6 +8,31 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum IpFamily {
|
||||
Any,
|
||||
Ipv4Only,
|
||||
Ipv6Only,
|
||||
}
|
||||
|
||||
impl IpFamily {
|
||||
pub(crate) fn allows(self, addr: SocketAddr) -> bool {
|
||||
match self {
|
||||
Self::Any => true,
|
||||
Self::Ipv4Only => addr.is_ipv4(),
|
||||
Self::Ipv6Only => addr.is_ipv6(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn no_address_message(self, context: &str) -> String {
|
||||
match self {
|
||||
Self::Any => format!("{context} DNS returned no addresses"),
|
||||
Self::Ipv4Only => format!("{context} DNS returned no IPv4 addresses"),
|
||||
Self::Ipv6Only => format!("{context} DNS returned no IPv6 addresses"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum UpstreamProxyScheme {
|
||||
Http,
|
||||
@@ -126,6 +151,7 @@ pub(crate) struct ProxyConnectOptions {
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_nodelay: bool,
|
||||
pub tcp_keepalive: Option<Duration>,
|
||||
pub ip_family: IpFamily,
|
||||
}
|
||||
|
||||
pub(crate) async fn connect_target_via_proxy(
|
||||
@@ -139,6 +165,7 @@ pub(crate) async fn connect_target_via_proxy(
|
||||
options.connect_timeout,
|
||||
options.tcp_nodelay,
|
||||
options.tcp_keepalive,
|
||||
options.ip_family,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -159,6 +186,7 @@ pub(crate) async fn connect_proxy_tcp(
|
||||
connect_timeout: Duration,
|
||||
tcp_nodelay: bool,
|
||||
tcp_keepalive: Option<Duration>,
|
||||
ip_family: IpFamily,
|
||||
) -> io::Result<TcpStream> {
|
||||
let resolved = tokio::time::timeout(
|
||||
connect_timeout,
|
||||
@@ -169,7 +197,7 @@ pub(crate) async fn connect_proxy_tcp(
|
||||
.map_err(|err| io::Error::other(format!("proxy DNS failed: {err}")))?;
|
||||
|
||||
let mut last_error = None;
|
||||
for addr in resolved {
|
||||
for addr in resolved.filter(|addr| ip_family.allows(*addr)) {
|
||||
match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
||||
Ok(Ok(stream)) => {
|
||||
configure_tcp_stream(&stream, tcp_nodelay, tcp_keepalive)?;
|
||||
@@ -185,7 +213,7 @@ pub(crate) async fn connect_proxy_tcp(
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| io::Error::other("proxy DNS returned no addresses")))
|
||||
Err(last_error.unwrap_or_else(|| io::Error::other(ip_family.no_address_message("proxy"))))
|
||||
}
|
||||
|
||||
fn configure_tcp_stream(
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
|
||||
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -10,7 +12,9 @@ use tokio_tungstenite::tungstenite::http;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::egress_proxy::{connect_target_via_proxy, ProxyConnectOptions, UpstreamProxyConfig};
|
||||
use crate::egress_proxy::{
|
||||
connect_target_via_proxy, IpFamily, ProxyConnectOptions, UpstreamProxyConfig,
|
||||
};
|
||||
use crate::state::{AppState, ServerContext};
|
||||
use aether_contracts::tunnel::{CURRENT_TUNNEL_PROTOCOL_VERSION, TUNNEL_PROTOCOL_VERSION_HEADER};
|
||||
|
||||
@@ -325,6 +329,7 @@ async fn connect_tunnel_tcp(
|
||||
tcp_nodelay: state.config.tunnel_tcp_nodelay,
|
||||
tcp_keepalive: (state.config.tunnel_tcp_keepalive_secs > 0)
|
||||
.then(|| Duration::from_secs(state.config.tunnel_tcp_keepalive_secs)),
|
||||
ip_family: state.config.tunnel_ip_family(),
|
||||
},
|
||||
),
|
||||
)
|
||||
@@ -338,15 +343,54 @@ async fn connect_tunnel_tcp(
|
||||
.map_err(anyhow::Error::from);
|
||||
}
|
||||
|
||||
tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
|
||||
let ip_family = state.config.tunnel_ip_family();
|
||||
tokio::time::timeout(
|
||||
connect_timeout,
|
||||
connect_direct_tunnel_tcp(host, port, ip_family),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"tunnel TCP connect timeout ({}ms)",
|
||||
connect_timeout.as_millis()
|
||||
)
|
||||
})?
|
||||
.map_err(anyhow::Error::from)
|
||||
}
|
||||
|
||||
async fn connect_direct_tunnel_tcp(
|
||||
host: &str,
|
||||
port: u16,
|
||||
ip_family: IpFamily,
|
||||
) -> io::Result<TcpStream> {
|
||||
let resolved = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"tunnel TCP connect timeout ({}ms)",
|
||||
connect_timeout.as_millis()
|
||||
)
|
||||
})?
|
||||
.map_err(anyhow::Error::from)
|
||||
.map_err(|err| io::Error::other(format!("tunnel DNS failed: {err}")))?;
|
||||
let addrs = filter_socket_addrs(resolved, ip_family);
|
||||
|
||||
if addrs.is_empty() {
|
||||
return Err(io::Error::other(ip_family.no_address_message("tunnel")));
|
||||
}
|
||||
|
||||
let mut last_error = None;
|
||||
for addr in addrs {
|
||||
match TcpStream::connect(addr).await {
|
||||
Ok(stream) => return Ok(stream),
|
||||
Err(error) => last_error = Some(error),
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| io::Error::other("tunnel DNS returned no addresses")))
|
||||
}
|
||||
|
||||
fn filter_socket_addrs(
|
||||
addrs: impl IntoIterator<Item = SocketAddr>,
|
||||
ip_family: IpFamily,
|
||||
) -> Vec<SocketAddr> {
|
||||
addrs
|
||||
.into_iter()
|
||||
.filter(|addr| ip_family.allows(*addr))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Configure TCP keepalive and NODELAY on an established socket.
|
||||
@@ -392,3 +436,40 @@ fn build_tunnel_url(server: &ServerContext) -> String {
|
||||
};
|
||||
format!("{}/api/internal/proxy-tunnel", ws_base)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::net::{Ipv4Addr, Ipv6Addr};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn mixed_addrs() -> Vec<SocketAddr> {
|
||||
vec![
|
||||
SocketAddr::from((Ipv6Addr::LOCALHOST, 443)),
|
||||
SocketAddr::from((Ipv4Addr::LOCALHOST, 443)),
|
||||
]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_socket_addrs_keeps_all_addresses_by_default() {
|
||||
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Any);
|
||||
|
||||
assert_eq!(addrs.len(), 2);
|
||||
assert!(addrs[0].is_ipv6());
|
||||
assert!(addrs[1].is_ipv4());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_socket_addrs_keeps_only_ipv4_addresses() {
|
||||
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Ipv4Only);
|
||||
|
||||
assert_eq!(addrs, vec![SocketAddr::from((Ipv4Addr::LOCALHOST, 443))]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_socket_addrs_keeps_only_ipv6_addresses() {
|
||||
let addrs = filter_socket_addrs(mixed_addrs(), IpFamily::Ipv6Only);
|
||||
|
||||
assert_eq!(addrs, vec![SocketAddr::from((Ipv6Addr::LOCALHOST, 443))]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -541,6 +541,8 @@ mod tests {
|
||||
tunnel_ping_interval_ms: 1_000,
|
||||
tunnel_max_streams: Some(8),
|
||||
tunnel_connect_timeout_ms: 2_000,
|
||||
tunnel_ipv4_only: false,
|
||||
tunnel_ipv6_only: false,
|
||||
tunnel_tcp_keepalive_secs: 30,
|
||||
tunnel_tcp_nodelay: true,
|
||||
tunnel_stale_timeout_ms: 5_000,
|
||||
|
||||
@@ -2556,6 +2556,8 @@ mod tests {
|
||||
tunnel_ping_interval_ms: 15_000,
|
||||
tunnel_max_streams: Some(8),
|
||||
tunnel_connect_timeout_ms: 15_000,
|
||||
tunnel_ipv4_only: false,
|
||||
tunnel_ipv6_only: false,
|
||||
tunnel_tcp_keepalive_secs: 30,
|
||||
tunnel_tcp_nodelay: true,
|
||||
tunnel_stale_timeout_ms: 45_000,
|
||||
|
||||
@@ -281,6 +281,7 @@ impl Service<Uri> for InstrumentedConnector {
|
||||
connect_timeout: self.connect_timeout,
|
||||
tcp_nodelay: self.tcp_nodelay,
|
||||
tcp_keepalive: self.tcp_keepalive,
|
||||
ip_family: crate::egress_proxy::IpFamily::Any,
|
||||
};
|
||||
let connect_start = std::time::Instant::now();
|
||||
return Box::pin(async move {
|
||||
@@ -347,6 +348,7 @@ async fn connect_via_proxy(
|
||||
options.connect_timeout,
|
||||
options.tcp_nodelay,
|
||||
options.tcp_keepalive,
|
||||
options.ip_family,
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user