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

feat(proxy): 实现代理节点批量升级回滚、隧道重定向跟随及远程配置管理

Browse Source 
核心功能：
- 新增代理节点批量升级回滚工作流，支持分批升级、健康探针、跳过/重试/取消等操作
- proxy 隧道流处理器支持 HTTP 重定向跟随（最多 10 跳），区分 307/308 可重播与不可重播请求体
- proxy 协议新增 follow_redirects / http1_only 字段，网关侧同步支持
- 新增代理节点远端配置变更接口（名称、允许端口、调度状态、升级目标等）
- 新增代理节点注册/反注册/心跳的 Admin API，及节点过期清理维护任务
- gateway 隧道 owner-relay 支持流式代理大请求体，新增 5 MiB 默认限制
- 新增 ProxyNodeRegistrationMutation / ProxyNodeRemoteConfigMutation 数据类型
- proxy 配置新增重定向重播预算、心跳间隔等参数，TUI 安装向导同步更新
- 前端 ProxyNodes 页面新增批量升级操作面板及滚动进度展示
This commit is contained in:
fawney19
2026-04-12 16:02:38 +08:00
parent 7c5bb7f383
commit 9703840a36
83 changed files with 11832 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,6 @@
use std::sync::{Arc, Mutex};
use aether_contracts::{ExecutionPlan, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
@@ -10,6 +11,7 @@ use aether_data::repository::oauth_providers::{
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use axum::body::{to_bytes, Body, Bytes};
use axum::response::{IntoResponse, Response};
@@ -19,8 +21,9 @@ use http::{HeaderMap, HeaderValue, StatusCode};
use serde_json::json;
use super::super::{
build_router_with_state, sample_endpoint, sample_key, sample_management_token,
sample_oauth_provider_config, sample_provider, start_server, AppState,
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
sample_key, sample_management_token, sample_oauth_provider_config, sample_provider,
sample_proxy_node, start_server, AppState,
};
use crate::admin_api::{
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
@@ -549,6 +552,152 @@ async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transit
token_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_device_authorize_via_execution_runtime_proxy_node() {
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-kiro"));
assert_eq!(
plan.headers.get("host").map(String::as_str),
Some("oidc.us-east-1.amazonaws.com")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
if plan.request_id == "kiro_device_register" {
assert_eq!(plan.url, "https://oidc.us-east-1.amazonaws.com/client/register");
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"clientId": "kiro-device-client",
"clientSecret": "kiro-device-secret"
}
}
}))
} else {
assert_eq!(plan.request_id, "kiro_device_authorize");
assert_eq!(
plan.url,
"https://oidc.us-east-1.amazonaws.com/device_authorization"
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"deviceCode": "device-code-123",
"userCode": "USER-CODE",
"verificationUri": "https://device.example.com/verify",
"verificationUriComplete": "https://device.example.com/verify?user_code=USER-CODE",
"expiresIn": 600,
"interval": 5
}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let mut manual_node = sample_proxy_node("proxy-node-kiro");
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 (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.attach_proxy_node_repository_for_tests(proxy_node_repository),
)
.with_provider_oauth_device_session_entry_for_tests(
"seed-session",
json!({"status":"seed"}),
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_register",
"https://oidc.us-east-1.amazonaws.com/client/register",
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_authorize",
"https://oidc.us-east-1.amazonaws.com/device_authorization",
);
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-authorize"
))
.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!({
"start_url": "https://view.awsapps.com/start",
"region": "us-east-1",
"proxy_node_id": "proxy-node-kiro",
}))
.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}");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist")
.to_string();
assert_eq!(payload["user_code"], "USER-CODE");
assert_eq!(payload["expires_in"], 600);
assert_eq!(payload["interval"], 5);
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["proxy_node_id"], "proxy-node-kiro");
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_oauth_start_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -1655,6 +1804,134 @@ async fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node() {
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-import"));
assert_eq!(plan.request_id, "provider-oauth:refresh-token");
assert_eq!(plan.method, "POST");
assert_eq!(plan.url, "https://oauth.example/oauth/token");
assert_eq!(
plan.headers.get("content-type").map(String::as_str),
Some("application/x-www-form-urlencoded")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
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();
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 manual_node = sample_proxy_node("proxy-node-codex-import");
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 (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_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",
"proxy_node_id": "proxy-node-codex-import",
"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");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["email"], "alice@example.com");
assert_eq!(payload["replaced"], false);
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,
Some(json!({"node_id":"proxy-node-codex-import","enabled":true}))
);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
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));
@@ -1797,6 +2074,167 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_ad
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_proxy_node() {
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());
assert_eq!(plan.request_id, "kiro_batch_refresh:social");
assert_eq!(plan.url, "https://oauth.example/refreshToken");
assert_eq!(
plan.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-kiro-batch-runtime")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"accessToken": sample_kiro_device_access_token("kiro-runtime@example.com"),
"refreshToken": "kiro-runtime-refresh-token-new",
"expiresIn": 1800,
}
}
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let endpoint = sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
);
let mut existing_key = sample_key(
"key-kiro-batch-runtime",
"provider-kiro",
"kiro:generateAssistantResponse",
"stale-kiro-runtime-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","auth_method":"social","email":"kiro-runtime@example.com","refresh_token":"kiro-runtime-refresh-token-old"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let mut manual_node = sample_proxy_node("proxy-node-kiro-batch-runtime");
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 (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_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_refresh",
"https://oauth.example",
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import"
))
.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!({
"credentials": "kiro-runtime-refresh-token-old",
"proxy_node_id": "proxy-node-kiro-batch-runtime"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("payload should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(payload["results"][0]["key_id"], "key-kiro-batch-runtime");
assert_eq!(payload["results"][0]["replaced"], true);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
drop(plans);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-batch-runtime".to_string()])
.await
.expect("keys should load")
.into_iter()
.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}))
);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["email"], "kiro-runtime@example.com");
assert_eq!(
auth_config["refresh_token"],
"kiro-runtime-refresh-token-new"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_starts_admin_provider_oauth_kiro_batch_import_task_locally_with_trusted_admin_principal(
) {

418
apps/aether-gateway/src/tests/control/admin/provider_ops.rs
View File

@@ -1,6 +1,8 @@
use std::sync::{Arc, Mutex};
use aether_contracts::ExecutionPlan;
use aether_contracts::{
ExecutionPlan, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
@@ -466,7 +468,7 @@ async fn gateway_saves_admin_provider_ops_config_locally_with_trusted_admin_prin
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_eq!(payload["success"], true, "payload={payload}");
assert_eq!(payload["message"], "配置保存成功");
let stored_provider = provider_catalog_repository
@@ -1142,6 +1144,12 @@ async fn gateway_verifies_admin_provider_ops_locally_for_anyrouter_proxy_mode()
assert_eq!(parsed_proxy.password(), Some("supersecret"));
if plan.url.ends_with("/api/user/self") {
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
None
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
@@ -1160,6 +1168,14 @@ async fn gateway_verifies_admin_provider_ops_locally_for_anyrouter_proxy_mode()
}
}))
} else {
assert_eq!(plan.request_id, "provider-ops-acw:anyrouter");
assert_eq!(plan.url, "https://ops.example");
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("false")
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
@@ -1242,9 +1258,159 @@ async fn gateway_verifies_admin_provider_ops_locally_for_anyrouter_proxy_mode()
assert_eq!(payload["data"]["request_count"], 8);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 1);
assert_eq!(plans[0].request_id, "provider-ops-verify:anyrouter");
assert_eq!(plans[0].url, "https://ops.example/api/user/self");
assert_eq!(plans.len(), 2);
assert_eq!(plans[0].request_id, "provider-ops-acw:anyrouter");
assert_eq!(plans[0].url, "https://ops.example");
assert_eq!(plans[1].request_id, "provider-ops-verify:anyrouter");
assert_eq!(plans[1].url, "https://ops.example/api/user/self");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_verifies_admin_provider_ops_sub2api_proxy_mode_via_execution_runtime_http1_only() {
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-sub2api"));
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_HTTP1_ONLY_HEADER)
.map(String::as_str),
Some("true")
);
if plan.url.ends_with("/api/v1/auth/refresh") {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"code": 0,
"data": {
"access_token": "sub2api-access-token",
"refresh_token": "sub2api-refresh-token-new"
}
}
}
}))
} else {
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer sub2api-access-token")
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"code": 0,
"data": {
"username": "sub2api-user",
"email": "sub2api@example.com",
"balance": 8.5,
"points": 1.5,
"status": "active",
"concurrency": 3
}
}
}
}))
}
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let mut manual_node = sample_proxy_node("proxy-node-sub2api");
manual_node.name = "sub2api-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());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_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,
)
.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-sub2api"),
)]),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-ops/providers/provider-openai/verify"
))
.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!({
"architecture_id": "sub2api",
"base_url": "https://sub2api.example",
"connector": {
"auth_type": "session_login",
"config": {},
"credentials": {
"refresh_token": "refresh-token-old",
}
},
"actions": {},
"schedule": {},
}))
.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["success"], true);
assert_eq!(payload["data"]["username"], "sub2api-user");
assert_eq!(
payload["updated_credentials"]["refresh_token"],
"sub2api-refresh-token-new"
);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);
assert_eq!(plans[0].url, "https://sub2api.example/api/v1/auth/refresh");
assert!(
plans[1]
.url
.starts_with("https://sub2api.example/api/v1/auth/me"),
"url={}",
plans[1].url
);
gateway_handle.abort();
execution_runtime_handle.abort();
@@ -1516,6 +1682,109 @@ async fn gateway_verifies_admin_provider_ops_locally_for_new_api_proxy_node_mode
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_verifies_admin_provider_ops_locally_for_new_api_without_proxy_via_execution_runtime(
) {
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());
assert_eq!(plan.request_id, "provider-ops-verify:new_api");
assert_eq!(plan.url, "https://ops.example/api/user/self");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer live-secret-api-key")
);
assert_eq!(
plan.headers.get("new-api-user").map(String::as_str),
Some("42")
);
assert!(plan.proxy.is_none());
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"success": true,
"data": {
"username": "alice",
"display_name": "Alice",
"quota": 42.5,
"used_quota": 12.5,
"request_count": 9
}
}
}
}))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
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,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-ops/providers/provider-openai/verify"
))
.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!({
"architecture_id": "new_api",
"base_url": "https://ops.example",
"connector": {
"auth_type": "api_key",
"config": {},
"credentials": {
"api_key": "live-secret-api-key",
"user_id": "42",
"cookie": "session=foo"
}
},
"actions": {},
"schedule": {},
}))
.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["success"], true);
assert_eq!(payload["data"]["username"], "alice");
assert_eq!(payload["data"]["used_quota"], 12.5);
assert_eq!(payload["data"]["request_count"], 9);
assert_eq!(execution_plans.lock().expect("mutex should lock").len(), 1);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_verifies_admin_provider_ops_locally_for_sub2api_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -2446,6 +2715,145 @@ async fn gateway_handles_admin_provider_ops_balance_locally_for_generic_api_prox
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_ops_balance_locally_without_proxy_via_execution_runtime() {
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());
assert!(plan.proxy.is_none());
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer live-secret-api-key")
);
if plan.url.ends_with("/api/user/checkin") {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"success": true,
"message": "执行层签到成功"
}
}
}))
} else {
assert!(plan.url.ends_with("/api/user/balance"));
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"success": true,
"data": {
"quota": 2500000,
"used_quota": 500000
}
}
}
}))
}
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"provider_ops": {
"architecture_id": "generic_api",
"base_url": "https://ops.example",
"connector": {
"auth_type": "api_key",
"config": {
"auth_method": "bearer"
},
"credentials": {
"api_key": encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"live-secret-api-key",
).expect("api key should encrypt"),
}
}
}
})),
),
],
vec![],
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,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-ops/providers/provider-openai/balance?refresh=false"
))
.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 payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["status"], "success");
assert_eq!(payload["data"]["total_available"], 5.0);
assert_eq!(payload["data"]["total_used"], 1.0);
assert_eq!(payload["data"]["extra"]["checkin_success"], true);
assert_eq!(
payload["data"]["extra"]["checkin_message"],
"执行层签到成功"
);
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);
assert!(plans.iter().all(|plan| plan.proxy.is_none()));
assert!(plans
.iter()
.any(|plan| plan.request_id == "provider-ops-action:probe_checkin"));
assert!(plans.iter().any(|plan| {
plan.request_id == "provider-ops-action:generic_api:query_balance:provider-openai"
}));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_ops_checkin_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

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

@@ -1,18 +1,28 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNodeEvent};
use aether_data::repository::management_tokens::InMemoryManagementTokenRepository;
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeHeartbeatMutation, StoredProxyNodeEvent,
};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use serde_json::json;
use super::super::{build_router_with_state, sample_proxy_node, start_server, AppState};
use super::super::{
build_router_with_state, hash_management_token, sample_management_token, sample_proxy_node,
start_server, AppState,
};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
use crate::maintenance::{
record_proxy_upgrade_traffic_success, skip_proxy_upgrade_rollout_node,
start_proxy_upgrade_rollout,
};
#[tokio::test]
async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal() {
@@ -77,6 +87,7 @@ async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal(
assert_eq!(payload["total"], 1);
assert_eq!(payload["skip"], 0);
assert_eq!(payload["limit"], 10);
assert!(payload["rollout"].is_null());
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
@@ -95,7 +106,502 @@ async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal(
}
#[tokio::test]
async fn gateway_rejects_admin_proxy_nodes_unavailable_routes_locally() {
async fn gateway_reports_active_proxy_upgrade_rollout_in_proxy_node_list() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
alpha.proxy_metadata = Some(json!({ "version": "1.9.0" }));
let mut beta = sample_proxy_node("node-beta");
beta.name = "beta".to_string();
beta.status = "online".to_string();
beta.tunnel_connected = true;
beta.remote_config = None;
beta.proxy_metadata = Some(json!({ "version": "1.9.0" }));
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let rollout = start_proxy_upgrade_rollout(
&data_state,
"2.0.0".to_string(),
2,
120,
Some(crate::maintenance::ProxyUpgradeRolloutProbeConfig {
url: "https://probe.example/health".to_string(),
timeout_secs: 15,
}),
)
.await
.expect("rollout should start");
assert_eq!(rollout.updated, 2);
data_state
.apply_proxy_node_heartbeat(&ProxyNodeHeartbeatMutation {
node_id: "node-alpha".to_string(),
heartbeat_interval: None,
active_connections: None,
total_requests_delta: None,
avg_latency_ms: None,
failed_requests_delta: None,
dns_failures_delta: None,
stream_errors_delta: None,
proxy_metadata: None,
proxy_version: Some("2.0.0".to_string()),
})
.await
.expect("heartbeat should apply");
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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["rollout"]["version"], "2.0.0");
assert_eq!(payload["rollout"]["batch_size"], 2);
assert_eq!(payload["rollout"]["cooldown_secs"], 120);
assert_eq!(
payload["rollout"]["probe"]["url"],
"https://probe.example/health"
);
assert_eq!(payload["rollout"]["probe"]["timeout_secs"], 15);
assert_eq!(
payload["rollout"]["pending_node_ids"],
json!(["node-alpha", "node-beta"])
);
assert_eq!(payload["rollout"]["completed_node_ids"], json!([]));
assert_eq!(payload["rollout"]["conflict_node_ids"], json!([]));
assert_eq!(payload["rollout"]["blocked"], true);
assert!(payload["rollout"]["started_at"].is_string());
assert!(payload["rollout"]["last_dispatched_at"].is_string());
assert!(payload["rollout"]["updated_at"].is_string());
let tracked_nodes = payload["rollout"]["tracked_nodes"]
.as_array()
.expect("tracked_nodes should be array");
assert_eq!(tracked_nodes.len(), 2);
let alpha_status = tracked_nodes
.iter()
.find(|tracked| tracked["node_id"] == "node-alpha")
.expect("alpha status should exist");
assert_eq!(alpha_status["state"], "awaiting_traffic");
assert!(alpha_status["version_confirmed_at"].is_string());
assert!(alpha_status["traffic_confirmed_at"].is_null());
let beta_status = tracked_nodes
.iter()
.find(|tracked| tracked["node_id"] == "node-beta")
.expect("beta status should exist");
assert_eq!(beta_status["state"], "awaiting_version");
assert!(beta_status["version_confirmed_at"].is_null());
assert!(beta_status["traffic_confirmed_at"].is_null());
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_cancels_active_proxy_upgrade_rollout_locally() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
.await
.expect("rollout should start");
assert!(rollout.rollout_active);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/upgrade/cancel"
))
.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["cancelled"], true);
assert_eq!(payload["version"], "2.0.0");
assert_eq!(payload["pending_node_ids"], json!(["node-alpha"]));
assert_eq!(payload["conflict_node_ids"], json!([]));
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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");
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert!(list_payload["rollout"].is_null());
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_clears_proxy_upgrade_rollout_conflicts_locally() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let mut beta = sample_proxy_node("node-beta");
beta.name = "beta".to_string();
beta.status = "online".to_string();
beta.tunnel_connected = true;
beta.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
.await
.expect("rollout should start");
assert_eq!(rollout.updated, 1);
data_state
.update_proxy_node_remote_config(
&aether_data::repository::proxy_nodes::ProxyNodeRemoteConfigMutation {
node_id: "node-beta".to_string(),
node_name: None,
allowed_ports: None,
log_level: None,
heartbeat_interval: None,
scheduling_state: None,
upgrade_to: Some(Some("3.0.0".to_string())),
},
)
.await
.expect("conflict target should update");
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/upgrade/clear-conflicts"
))
.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["cleared"], 1);
assert_eq!(payload["node_ids"], json!(["node-beta"]));
assert_eq!(payload["blocked"], true);
assert_eq!(payload["pending_node_ids"], json!(["node-alpha"]));
let updated_beta = data_state
.find_proxy_node("node-beta")
.await
.expect("node lookup should succeed")
.expect("beta should exist");
let beta_upgrade_to = updated_beta
.remote_config
.as_ref()
.and_then(|value| value.get("upgrade_to"))
.cloned()
.unwrap_or(serde_json::Value::Null);
assert!(beta_upgrade_to.is_null());
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_skips_proxy_upgrade_rollout_node_and_advances_next_wave_locally() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let mut beta = sample_proxy_node("node-beta");
beta.name = "beta".to_string();
beta.status = "online".to_string();
beta.tunnel_connected = true;
beta.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let rollout = start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
.await
.expect("rollout should start");
assert_eq!(rollout.node_ids, vec!["node-alpha"]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/node-alpha/upgrade/skip"
))
.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"], "node-alpha");
assert_eq!(payload["skipped_node_ids"], json!(["node-alpha"]));
assert_eq!(payload["updated"], 1);
assert_eq!(payload["pending_node_ids"], json!(["node-beta"]));
let alpha_after = data_state
.find_proxy_node("node-alpha")
.await
.expect("node lookup should succeed")
.expect("alpha should exist");
let alpha_upgrade_to = alpha_after
.remote_config
.as_ref()
.and_then(|value| value.get("upgrade_to"))
.cloned()
.unwrap_or(serde_json::Value::Null);
assert!(alpha_upgrade_to.is_null());
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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");
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(
list_payload["rollout"]["skipped_node_ids"],
json!(["node-alpha"])
);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_retries_proxy_upgrade_rollout_node_locally() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let mut beta = sample_proxy_node("node-beta");
beta.name = "beta".to_string();
beta.status = "online".to_string();
beta.tunnel_connected = true;
beta.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
.await
.expect("rollout should start");
let _ = skip_proxy_upgrade_rollout_node(&data_state, "node-alpha")
.await
.expect("skip should succeed");
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/node-alpha/upgrade/retry"
))
.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"], "node-alpha");
assert_eq!(payload["skipped_node_ids"], json!([]));
assert_eq!(payload["blocked"], true);
let alpha_after = data_state
.find_proxy_node("node-alpha")
.await
.expect("node lookup should succeed")
.expect("alpha should exist");
let alpha_upgrade_to = alpha_after
.remote_config
.as_ref()
.and_then(|value| value.get("upgrade_to"))
.and_then(serde_json::Value::as_str);
assert_eq!(alpha_upgrade_to, Some("2.0.0"));
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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");
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(list_payload["rollout"]["skipped_node_ids"], json!([]));
let tracked_nodes = list_payload["rollout"]["tracked_nodes"]
.as_array()
.expect("tracked nodes should be array");
let alpha_status = tracked_nodes
.iter()
.find(|tracked| tracked["node_id"] == "node-alpha")
.expect("alpha should be tracked again");
assert_eq!(alpha_status["state"], "awaiting_version");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_restores_skipped_proxy_upgrade_rollout_nodes_locally() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let mut beta = sample_proxy_node("node-beta");
beta.name = "beta".to_string();
beta.status = "online".to_string();
beta.tunnel_connected = true;
beta.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![beta, alpha]));
let data_state = GatewayDataState::with_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
start_proxy_upgrade_rollout(&data_state, "2.0.0".to_string(), 1, 120, None)
.await
.expect("rollout should start");
let _ = skip_proxy_upgrade_rollout_node(&data_state, "node-alpha")
.await
.expect("skip should succeed");
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/proxy-nodes/upgrade/restore-skipped"
))
.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["restored"], 1);
assert_eq!(payload["node_ids"], json!(["node-alpha"]));
assert_eq!(payload["skipped_node_ids"], json!([]));
assert_eq!(payload["updated"], 0);
assert_eq!(payload["blocked"], true);
assert_eq!(payload["pending_node_ids"], json!(["node-beta"]));
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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");
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(list_payload["rollout"]["skipped_node_ids"], json!([]));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_registers_and_unregisters_proxy_nodes_locally_with_management_token_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
@@ -109,35 +615,100 @@ async fn gateway_rejects_admin_proxy_nodes_unavailable_routes_locally() {
}),
);
let raw_token = "ae_proxy_register_test";
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::default());
let state = AppState::new().expect("gateway should build");
let admin_user = state
.create_local_auth_user_with_settings(
Some("proxy-admin@example.com".to_string()),
true,
"admin".to_string(),
"hash".to_string(),
"admin".to_string(),
None,
None,
None,
None,
)
.await
.expect("admin user should be created")
.expect("admin user should exist");
let mut management_token =
sample_management_token("token-proxy-register", &admin_user.id, "proxy-admin", true);
management_token.token.allowed_ips = None;
let management_token_repository =
Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
vec![management_token],
vec![(
hash_management_token(raw_token),
"token-proxy-register".to_string(),
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let state = state.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(management_token_repository)
.attach_proxy_node_repository_for_tests(proxy_node_repository),
);
let token_lookup = state
.get_management_token_with_user_by_hash(&hash_management_token(raw_token))
.await
.expect("token lookup should succeed");
assert!(token_lookup.is_some());
let user_lookup = state
.find_user_auth_by_id(&admin_user.id)
.await
.expect("user lookup should succeed");
assert!(user_lookup.is_some());
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let register_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
.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")
.header(GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(raw_token)
.json(&json!({
"name": "proxy-1",
"ip": "1.1.1.1",
"port": 8080
"port": 0,
"heartbeat_interval": 30,
"tunnel_mode": true
}))
.send()
.await
.expect("request should succeed");
assert_eq!(register_response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(register_response.status(), StatusCode::OK);
let register_payload: serde_json::Value = register_response
.json()
.await
.expect("json body should parse");
let node_id = register_payload["node_id"]
.as_str()
.expect("node_id should be present")
.to_string();
assert_eq!(register_payload["node"]["name"], "proxy-1");
assert_eq!(register_payload["node"]["status"], "offline");
assert_eq!(
register_payload["detail"],
"Admin proxy nodes data unavailable"
register_payload["node"]["registered_by"],
json!(admin_user.id)
);
let unregister_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/unregister"))
.header(GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(raw_token)
.json(&json!({ "node_id": node_id }))
.send()
.await
.expect("request should succeed");
assert_eq!(unregister_response.status(), StatusCode::OK);
let unregister_payload: serde_json::Value = unregister_response
.json()
.await
.expect("json body should parse");
assert_eq!(unregister_payload["message"], "unregistered");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -214,3 +785,336 @@ async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_prin
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_proxy_node_config_and_batches_upgrade_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new()
.route(
"/api/admin/proxy-nodes/node-online/config",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
)
.route(
"/api/admin/proxy-nodes/upgrade",
any(|| async move { (StatusCode::OK, Body::from("unexpected upstream hit")) }),
);
let mut online_node = sample_proxy_node("node-online");
online_node.status = "online".to_string();
online_node.tunnel_connected = true;
let mut online_node_2 = sample_proxy_node("node-zeta");
online_node_2.name = "zeta-online".to_string();
online_node_2.status = "online".to_string();
online_node_2.tunnel_connected = true;
online_node_2.remote_config = None;
let mut offline_node = sample_proxy_node("node-offline");
offline_node.status = "offline".to_string();
offline_node.tunnel_connected = false;
offline_node.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
online_node,
online_node_2,
offline_node,
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let data_state =
GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state.clone()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let config_response = client
.put(format!(
"{gateway_url}/api/admin/proxy-nodes/node-online/config"
))
.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!({
"node_name": "edge-online",
"allowed_ports": [443, 8443],
"log_level": "info",
"heartbeat_interval": 45,
"upgrade_to": null
}))
.send()
.await
.expect("request should succeed");
assert_eq!(config_response.status(), StatusCode::OK);
let config_payload: serde_json::Value = config_response
.json()
.await
.expect("json body should parse");
assert_eq!(config_payload["node_id"], "node-online");
assert_eq!(config_payload["config_version"], 8);
assert_eq!(config_payload["node"]["name"], "edge-online");
assert_eq!(config_payload["remote_config"]["log_level"], "info");
assert!(config_payload["remote_config"].get("upgrade_to").is_none());
let upgrade_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
.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.0.0", "cooldown_secs": 0 }))
.send()
.await
.expect("request should succeed");
assert_eq!(upgrade_response.status(), StatusCode::OK);
let upgrade_payload: serde_json::Value = upgrade_response
.json()
.await
.expect("json body should parse");
assert_eq!(upgrade_payload["version"], "2.0.0");
assert_eq!(upgrade_payload["batch_size"], 1);
assert_eq!(upgrade_payload["updated"], 1);
assert_eq!(upgrade_payload["skipped"], 1);
assert_eq!(upgrade_payload["blocked"], false);
assert_eq!(upgrade_payload["pending_node_ids"], json!(["node-online"]));
assert_eq!(upgrade_payload["node_ids"], json!(["node-online"]));
assert_eq!(upgrade_payload["completed"], 0);
assert_eq!(upgrade_payload["remaining"], 2);
assert_eq!(upgrade_payload["rollout_active"], true);
let blocked_upgrade_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
.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.0.0", "cooldown_secs": 0 }))
.send()
.await
.expect("request should succeed");
assert_eq!(blocked_upgrade_response.status(), StatusCode::OK);
let blocked_upgrade_payload: serde_json::Value = blocked_upgrade_response
.json()
.await
.expect("json body should parse");
assert_eq!(blocked_upgrade_payload["updated"], 0);
assert_eq!(blocked_upgrade_payload["blocked"], true);
assert_eq!(
blocked_upgrade_payload["pending_node_ids"],
json!(["node-online"])
);
let heartbeat_response = client
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-online",
"heartbeat_interval": 45,
"active_connections": 3,
"total_requests": 5,
"avg_latency_ms": 10.0,
"proxy_metadata": { "arch": "arm64" },
"proxy_version": "2.0.0"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(heartbeat_response.status(), StatusCode::OK);
let heartbeat_payload: serde_json::Value = heartbeat_response
.json()
.await
.expect("json body should parse");
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);
assert_eq!(heartbeat_payload["remote_config"]["log_level"], "info");
let second_upgrade_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
.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.0.0", "cooldown_secs": 0 }))
.send()
.await
.expect("request should succeed");
assert_eq!(second_upgrade_response.status(), StatusCode::OK);
let second_upgrade_payload: serde_json::Value = second_upgrade_response
.json()
.await
.expect("json body should parse");
assert_eq!(second_upgrade_payload["batch_size"], 1);
assert_eq!(second_upgrade_payload["updated"], 0);
assert_eq!(second_upgrade_payload["skipped"], 2);
assert_eq!(second_upgrade_payload["blocked"], true);
assert_eq!(
second_upgrade_payload["pending_node_ids"],
json!(["node-online"])
);
assert_eq!(second_upgrade_payload["node_ids"], json!([]));
assert_eq!(second_upgrade_payload["completed"], 0);
assert_eq!(second_upgrade_payload["remaining"], 3);
assert_eq!(second_upgrade_payload["rollout_active"], true);
assert!(
record_proxy_upgrade_traffic_success(&data_state, "node-online")
.await
.expect("traffic confirmation should be recorded")
);
let third_upgrade_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
.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.0.0", "cooldown_secs": 0 }))
.send()
.await
.expect("request should succeed");
assert_eq!(third_upgrade_response.status(), StatusCode::OK);
let third_upgrade_payload: serde_json::Value = third_upgrade_response
.json()
.await
.expect("json body should parse");
assert_eq!(third_upgrade_payload["batch_size"], 1);
assert_eq!(third_upgrade_payload["updated"], 1);
assert_eq!(third_upgrade_payload["skipped"], 1);
assert_eq!(third_upgrade_payload["blocked"], false);
assert_eq!(
third_upgrade_payload["pending_node_ids"],
json!(["node-zeta"])
);
assert_eq!(third_upgrade_payload["node_ids"], json!(["node-zeta"]));
assert_eq!(third_upgrade_payload["completed"], 1);
assert_eq!(third_upgrade_payload["remaining"], 1);
assert_eq!(third_upgrade_payload["rollout_active"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_marks_draining_proxy_nodes_unschedulable_and_rollout_skips_them() {
let mut alpha = sample_proxy_node("node-alpha");
alpha.name = "alpha".to_string();
alpha.status = "online".to_string();
alpha.tunnel_connected = true;
alpha.remote_config = None;
let mut zeta = sample_proxy_node("node-zeta");
zeta.name = "zeta".to_string();
zeta.status = "online".to_string();
zeta.tunnel_connected = true;
zeta.remote_config = None;
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![zeta, alpha]));
let data_state =
GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(&proxy_node_repository))
.with_system_config_values_for_tests(Vec::<(String, serde_json::Value)>::new());
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let config_response = client
.put(format!(
"{gateway_url}/api/admin/proxy-nodes/node-alpha/config"
))
.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!({
"scheduling_state": "draining"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(config_response.status(), StatusCode::OK);
let config_payload: serde_json::Value = config_response
.json()
.await
.expect("json body should parse");
assert_eq!(
config_payload["remote_config"]["scheduling_state"],
"draining"
);
let list_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
let alpha_payload = list_payload["items"]
.as_array()
.expect("items should be array")
.iter()
.find(|item| item["id"] == "node-alpha")
.expect("alpha should exist");
assert_eq!(
alpha_payload["remote_config"]["scheduling_state"],
"draining"
);
let upgrade_response = client
.post(format!("{gateway_url}/api/admin/proxy-nodes/upgrade"))
.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.0.0", "batch_size": 2, "cooldown_secs": 0 }))
.send()
.await
.expect("request should succeed");
assert_eq!(upgrade_response.status(), StatusCode::OK);
let upgrade_payload: serde_json::Value = upgrade_response
.json()
.await
.expect("json body should parse");
assert_eq!(upgrade_payload["updated"], 1);
assert_eq!(upgrade_payload["node_ids"], json!(["node-zeta"]));
assert_eq!(upgrade_payload["pending_node_ids"], json!(["node-zeta"]));
let rollout_response = client
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?skip=0&limit=10"
))
.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");
let rollout_payload: serde_json::Value = rollout_response
.json()
.await
.expect("json body should parse");
assert_eq!(rollout_payload["rollout"]["online_eligible_total"], 1);
gateway_handle.abort();
}

4
apps/aether-gateway/src/tests/control/helpers.rs
View File

@@ -69,6 +69,10 @@ pub(super) fn hash_api_key(value: &str) -> String {
format!("{:x}", hasher.finalize())
}
pub(super) fn hash_management_token(value: &str) -> String {
hash_api_key(value)
}
pub(super) fn test_auth_secret() -> String {
std::env::var("JWT_SECRET_KEY")
.ok()

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

@@ -1,8 +1,12 @@
use std::io;
use std::sync::{Arc, Mutex};
use aether_data::repository::proxy_nodes::ProxyNodeReadRepository;
use axum::body::Body;
use axum::routing::{any, post};
use axum::{extract::Request, Json, Router};
use bytes::Bytes;
use futures_util::stream;
use http::header::HeaderValue;
use http::StatusCode;
use serde_json::json;
@@ -45,6 +49,7 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-123",
"heartbeat_id": 77,
"heartbeat_interval": 45,
"active_connections": 5,
"total_requests": 9,
@@ -61,6 +66,7 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["heartbeat_id"], 77);
assert_eq!(payload["config_version"], 7);
assert_eq!(payload["upgrade_to"], "1.2.3");
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
@@ -105,6 +111,7 @@ async fn gateway_handles_internal_tunnel_node_status_locally_with_loopback() {
"node_id": "node-123",
"connected": true,
"conn_count": 4,
"observed_at_unix_secs": 1_800_000_321u64,
}))
.send()
.await
@@ -114,6 +121,12 @@ async fn gateway_handles_internal_tunnel_node_status_locally_with_loopback() {
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["updated"], json!(true));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let node = repository
.find_proxy_node("node-123")
.await
.expect("lookup should succeed")
.expect("node should exist");
assert_eq!(node.tunnel_connected_at_unix_secs, Some(1_800_000_321));
gateway_handle.abort();
upstream_handle.abort();
@@ -261,6 +274,66 @@ async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
owner_handle.abort();
}
#[tokio::test]
async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
let owner_hits = Arc::new(Mutex::new(0usize));
let owner_hits_clone = Arc::clone(&owner_hits);
let owner = Router::new().route(
"/api/internal/tunnel/relay/node-123",
post(move |body: Body| {
let owner_hits_inner = Arc::clone(&owner_hits_clone);
async move {
*owner_hits_inner.lock().expect("mutex should lock") += 1;
let body = axum::body::to_bytes(body, usize::MAX)
.await
.expect("body should read");
assert_eq!(body, Bytes::from_static(b"relay-stream-envelope"));
(StatusCode::OK, Body::from("stream-ok"))
}
}),
);
let (owner_url, owner_handle) = start_server(owner).await;
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
"tunnel.attachments.node-123".to_string(),
json!({
"gateway_instance_id": "gateway-b",
"relay_base_url": owner_url,
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
)]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_body = reqwest::Body::wrap_stream(stream::iter(vec![
Ok::<Bytes, io::Error>(Bytes::from_static(b"relay-")),
Ok::<Bytes, io::Error>(Bytes::from_static(b"stream-")),
Ok::<Bytes, io::Error>(Bytes::from_static(b"envelope")),
]));
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.body(request_body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
"stream-ok"
);
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1);
gateway_handle.abort();
owner_handle.abort();
}
#[tokio::test]
async fn gateway_does_not_forward_tunnel_relay_twice() {
let owner_hits = Arc::new(Mutex::new(0usize));
@@ -310,3 +383,53 @@ async fn gateway_does_not_forward_tunnel_relay_twice() {
gateway_handle.abort();
owner_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_owner_relay_body_above_configured_limit() {
let owner_hits = Arc::new(Mutex::new(0usize));
let owner_hits_clone = Arc::clone(&owner_hits);
let owner = Router::new().route(
"/api/internal/tunnel/relay/node-123",
post(move |_request: Request| {
let owner_hits_inner = Arc::clone(&owner_hits_clone);
async move {
*owner_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected owner hit"))
}
}),
);
let (owner_url, owner_handle) = start_server(owner).await;
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![
(
"tunnel.attachments.node-123".to_string(),
json!({
"gateway_instance_id": "gateway-b",
"relay_base_url": owner_url,
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
),
("max_request_body_size".to_string(), json!(8)),
]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.body("relay-envelope")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
owner_handle.abort();
}