refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案

将 gateway 内部的 model-fetch、provider-transport、scheduler-core、
usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway
内部模块结构(state/router/cache/data/query 等);移除大量遗留模块
文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关
API 和组件。
This commit is contained in:
fawney19
2026-04-05 20:23:16 +08:00
parent cbc811f6ce
commit 763ff03a7b
777 changed files with 42659 additions and 21469 deletions

View File

@@ -0,0 +1,24 @@
mod adaptive;
mod api_keys;
mod billing;
mod endpoints;
mod gemini_files;
mod health_access;
mod ldap;
mod models;
mod monitoring;
mod oauth;
mod payments;
mod pool;
mod provider_ops;
mod provider_query;
mod provider_strategy;
mod providers;
mod proxy_nodes;
mod security;
mod stats;
mod system;
mod usage;
mod users;
mod video_tasks;
mod wallets;

View File

@@ -0,0 +1,449 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository, StoredProviderCatalogKey,
};
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_key, sample_provider, 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;
fn adaptive_repository_with_keys(
keys: Vec<StoredProviderCatalogKey>,
) -> Arc<InMemoryProviderCatalogReadRepository> {
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
keys,
))
}
async fn assert_adaptive_route_returns_local_503(
data_state: crate::data::GatewayDataState,
method: http::Method,
path: &str,
body: Option<serde_json::Value>,
expected_message: &str,
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
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 mut request = reqwest::Client::new()
.request(method, format!("{gateway_url}{path}"))
.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");
if let Some(body) = body {
request = request.json(&body);
}
let response = request.send().await.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["message"], expected_message);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
const ADAPTIVE_DATA_UNAVAILABLE_MESSAGE: &str = "Admin adaptive data unavailable";
#[tokio::test]
async fn gateway_handles_admin_adaptive_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
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"))
}
}),
);
let mut adaptive_key = sample_key("key-adaptive", "provider-openai", "openai:chat", "sk-test");
adaptive_key.rpm_limit = None;
adaptive_key.learned_rpm_limit = Some(17);
adaptive_key.concurrent_429_count = Some(2);
adaptive_key.rpm_429_count = Some(1);
let mut fixed_key = sample_key("key-fixed", "provider-openai", "openai:chat", "sk-fixed");
fixed_key.rpm_limit = Some(9);
fixed_key.learned_rpm_limit = Some(11);
let repository = adaptive_repository_with_keys(vec![adaptive_key, fixed_key]);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/adaptive/keys"))
.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");
let items = payload.as_array().expect("items should be an array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "key-adaptive");
assert_eq!(items[0]["is_adaptive"], true);
assert_eq!(items[0]["effective_limit"], 17);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_adaptive_summary_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
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"))
}
}),
);
let mut first_key = sample_key("key-a", "provider-openai", "openai:chat", "sk-a");
first_key.rpm_limit = None;
first_key.concurrent_429_count = Some(2);
first_key.rpm_429_count = Some(3);
first_key.adjustment_history = Some(json!([
{"timestamp": "2026-03-28T15:00:00Z", "delta": 2},
{"timestamp": "2026-03-28T14:00:00Z", "delta": -1}
]));
let mut second_key = sample_key("key-b", "provider-openai", "openai:chat", "sk-b");
second_key.rpm_limit = None;
second_key.concurrent_429_count = Some(1);
second_key.rpm_429_count = Some(1);
second_key.adjustment_history = Some(json!([
{"timestamp": "2026-03-28T16:00:00Z", "delta": 1}
]));
let repository = adaptive_repository_with_keys(vec![first_key, second_key]);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/adaptive/summary"))
.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["total_adaptive_keys"], 2);
assert_eq!(payload["total_concurrent_429_errors"], 3);
assert_eq!(payload["total_rpm_429_errors"], 4);
assert_eq!(payload["total_adjustments"], 3);
assert_eq!(payload["recent_adjustments"][0]["key_id"], "key-b");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_adaptive_keys_returns_service_unavailable_when_disabled() {
assert_adaptive_route_returns_local_503(
crate::data::GatewayDataState::disabled(),
http::Method::GET,
"/api/admin/adaptive/keys",
None,
ADAPTIVE_DATA_UNAVAILABLE_MESSAGE,
)
.await;
}
#[tokio::test]
async fn gateway_handles_admin_adaptive_mode_returns_service_unavailable_without_writer() {
let mut adaptive_key = sample_key("key-adaptive", "provider-openai", "openai:chat", "sk-test");
adaptive_key.rpm_limit = None;
let repository = adaptive_repository_with_keys(vec![adaptive_key]);
assert_adaptive_route_returns_local_503(
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
repository,
),
http::Method::PATCH,
"/api/admin/adaptive/keys/key-adaptive/mode",
Some(json!({ "enabled": false, "fixed_limit": 20 })),
ADAPTIVE_DATA_UNAVAILABLE_MESSAGE,
)
.await;
}
#[tokio::test]
async fn gateway_handles_admin_adaptive_stats_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
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"))
}
}),
);
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
key.rpm_limit = None;
key.learned_rpm_limit = Some(12);
key.concurrent_429_count = Some(2);
key.rpm_429_count = Some(1);
key.last_429_at_unix_secs = Some(1_711_111_111);
key.last_429_type = Some("rpm".to_string());
key.adjustment_history = Some(json!([
{"timestamp": "2026-03-28T14:00:00Z", "delta": 2}
]));
key.status_snapshot = Some(json!({
"learning_confidence": 0.9,
"enforcement_active": true,
"observation_count": 8,
"header_observation_count": 3,
"latest_upstream_limit": 15
}));
let repository = adaptive_repository_with_keys(vec![key]);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/adaptive/keys/key-openai/stats"
))
.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["adaptive_mode"], true);
assert_eq!(payload["effective_limit"], 12);
assert_eq!(payload["learning_confidence"], 0.9);
assert_eq!(payload["latest_upstream_limit"], 15);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_toggles_admin_adaptive_mode_locally_with_trusted_admin_principal() {
let repository = adaptive_repository_with_keys({
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
key.rpm_limit = None;
key.learned_rpm_limit = Some(18);
vec![key]
});
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&repository,
)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.patch(format!(
"{gateway_url}/api/admin/adaptive/keys/key-openai/mode"
))
.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!({ "enabled": false, "fixed_limit": 12 }))
.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["is_adaptive"], false);
assert_eq!(payload["rpm_limit"], 12);
let updated = repository
.list_keys_by_ids(&["key-openai".to_string()])
.await
.expect("repository query should succeed");
assert_eq!(updated[0].rpm_limit, Some(12));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_sets_admin_adaptive_limit_locally_with_trusted_admin_principal() {
let repository = adaptive_repository_with_keys({
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
key.rpm_limit = None;
vec![key]
});
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&repository,
)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.patch(format!(
"{gateway_url}/api/admin/adaptive/keys/key-openai/limit?limit=9"
))
.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["rpm_limit"], 9);
assert_eq!(payload["is_adaptive"], false);
let updated = repository
.list_keys_by_ids(&["key-openai".to_string()])
.await
.expect("repository query should succeed");
assert_eq!(updated[0].rpm_limit, Some(9));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_resets_admin_adaptive_learning_locally_with_trusted_admin_principal() {
let repository = adaptive_repository_with_keys({
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
key.rpm_limit = None;
key.learned_rpm_limit = Some(18);
key.concurrent_429_count = Some(2);
key.rpm_429_count = Some(1);
key.last_429_at_unix_secs = Some(1_711_111_111);
key.last_429_type = Some("rpm".to_string());
key.adjustment_history = Some(json!([
{"timestamp": "2026-03-28T14:00:00Z", "delta": 2}
]));
vec![key]
});
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&repository,
)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/adaptive/keys/key-openai/learning"
))
.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["message"], "学习状态已重置");
let updated = repository
.list_keys_by_ids(&["key-openai".to_string()])
.await
.expect("repository query should succeed");
assert_eq!(updated[0].learned_rpm_limit, None);
assert_eq!(updated[0].concurrent_429_count, None);
assert_eq!(updated[0].rpm_429_count, None);
assert_eq!(updated[0].adjustment_history, None);
gateway_handle.abort();
}

View File

@@ -0,0 +1,566 @@
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
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, 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;
fn admin_request(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
builder
.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")
}
async fn start_api_keys_upstream(
path: &'static str,
) -> (String, Arc<Mutex<usize>>, tokio::task::JoinHandle<()>) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
(upstream_url, upstream_hits, upstream_handle)
}
fn sample_standalone_api_key_snapshot(
api_key_id: &str,
user_id: &str,
is_active: bool,
) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"standalone-owner".to_string(),
Some("owner@example.com".to_string()),
"admin".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
api_key_id.to_string(),
Some(format!("key-{api_key_id}")),
is_active,
false,
true,
Some(120),
Some(5),
Some(4_102_444_800),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
fn sample_standalone_export_record(
api_key_id: &str,
user_id: &str,
plaintext_key: &str,
is_active: bool,
) -> StoredAuthApiKeyExportRecord {
StoredAuthApiKeyExportRecord::new(
user_id.to_string(),
api_key_id.to_string(),
format!("hash-{api_key_id}"),
Some(
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, plaintext_key)
.expect("key should encrypt"),
),
Some(format!("key-{api_key_id}")),
Some(json!(["openai"])),
Some(json!(["openai:chat"])),
Some(json!(["gpt-4.1"])),
Some(120),
Some(5),
None,
is_active,
Some(4_102_444_800),
false,
7,
1.25,
true,
)
.expect("export record should build")
}
fn sample_standalone_wallet(api_key_id: &str) -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
format!("wallet-{api_key_id}"),
None,
Some(api_key_id.to_string()),
18.5,
1.5,
"unlimited".to_string(),
"USD".to_string(),
"active".to_string(),
30.0,
2.0,
0.0,
0.0,
1_710_000_000,
)
.expect("wallet should build")
}
fn sample_usage_row(
id: &str,
request_id: &str,
api_key_id: &str,
total_tokens: i32,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some(api_key_id.to_string()),
Some("user-user-1".to_string()),
Some(format!("key-{api_key_id}")),
"openai".to_string(),
"gpt-4.1".to_string(),
Some("gpt-4.1".to_string()),
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
total_tokens / 2,
total_tokens - (total_tokens / 2),
total_tokens,
0.42,
0.42,
Some(200),
None,
None,
Some(240),
Some(80),
"completed".to_string(),
"settled".to_string(),
1_711_000_000,
1_711_000_001,
Some(1_711_000_002),
)
.expect("usage row should build")
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_list_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys").await;
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![
(
None,
sample_standalone_api_key_snapshot("key-1", "user-1", true),
),
(
None,
sample_standalone_api_key_snapshot("key-2", "user-2", false),
),
])
.with_export_records([
sample_standalone_export_record("key-1", "user-1", "sk-key-1-plaintext", true),
sample_standalone_export_record("key-2", "user-2", "sk-key-2-plaintext", false),
]),
);
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_usage_row("usage-1", "req-1", "key-1", 90),
sample_usage_row("usage-2", "req-2", "key-2", 120),
]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_wallet_and_usage_for_tests(
auth_repository,
Arc::new(InMemoryWalletRepository::seed(
Vec::<StoredWalletSnapshot>::new(),
)),
usage_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/api-keys?is_active=true&limit=10&skip=0"
)))
.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["total"], json!(1));
assert_eq!(payload["limit"], json!(10));
assert_eq!(payload["skip"], json!(0));
assert_eq!(payload["api_keys"][0]["id"], json!("key-1"));
assert_eq!(payload["api_keys"][0]["is_standalone"], json!(true));
assert_eq!(
payload["api_keys"][0]["key_display"],
json!("sk-key-1-p...text")
);
assert_eq!(payload["api_keys"][0]["total_requests"], json!(7));
assert_eq!(payload["api_keys"][0]["total_tokens"], json!(90));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_detail_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys/key-1").await;
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
None,
sample_standalone_api_key_snapshot("key-1", "user-1", true),
)])
.with_export_records([sample_standalone_export_record(
"key-1",
"user-1",
"sk-key-1-plaintext",
true,
)]),
);
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
sample_standalone_wallet("key-1"),
]));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![sample_usage_row(
"usage-1", "req-1", "key-1", 150,
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_wallet_and_usage_for_tests(
auth_repository,
wallet_repository,
usage_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys/key-1")),
)
.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["id"], json!("key-1"));
assert_eq!(payload["user_id"], json!("user-1"));
assert_eq!(payload["wallet"]["id"], json!("wallet-key-1"));
assert_eq!(payload["wallet"]["unlimited"], json!(true));
assert_eq!(payload["wallet"]["balance"], json!(20.0));
assert_eq!(payload["key_display"], json!("sk-key-1-p...text"));
assert_eq!(payload["total_tokens"], json!(150));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_full_key_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys/key-1").await;
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
None,
sample_standalone_api_key_snapshot("key-1", "user-1", true),
)])
.with_export_records([sample_standalone_export_record(
"key-1",
"user-1",
"sk-key-1-plaintext",
true,
)]),
);
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
sample_standalone_wallet("key-1"),
]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_and_wallet_for_tests(
auth_repository,
wallet_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/api-keys/key-1?include_key=true"
)))
.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, json!({ "key": "sk-key-1-plaintext" }));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_create_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys").await;
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(Vec::<(
Option<String>,
StoredAuthApiKeySnapshot,
)>::new()));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
Arc::clone(&repository),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response =
admin_request(reqwest::Client::new().post(format!("{gateway_url}/api/admin/api-keys")))
.json(&json!({
"name": "standalone-key",
"rate_limit": 180,
"allowed_providers": ["openai"],
"allowed_api_formats": ["openai:chat"],
"allowed_models": ["gpt-4.1"],
}))
.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["name"], json!("standalone-key"));
assert_eq!(payload["is_standalone"], json!(true));
assert_eq!(payload["rate_limit"], json!(180));
assert_eq!(payload["allowed_providers"], json!(["openai"]));
assert_eq!(payload["allowed_api_formats"], json!(["openai:chat"]));
assert_eq!(payload["allowed_models"], json!(["gpt-4.1"]));
assert_eq!(payload["wallet"], serde_json::Value::Null);
let plaintext = payload["key"]
.as_str()
.expect("plaintext key should exist")
.to_string();
assert!(plaintext.starts_with("sk-"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let list_response =
admin_request(reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys")))
.send()
.await
.expect("list request should succeed");
assert_eq!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("list json should parse");
assert_eq!(list_payload["total"], json!(1));
assert_eq!(list_payload["api_keys"][0]["name"], json!("standalone-key"));
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_update_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
let repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-key-123".to_string()),
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
)])
.with_export_records([sample_standalone_export_record(
"key-123",
"admin-user-123",
"sk-key-123-plaintext",
true,
)]),
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
Arc::clone(&repository),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().put(format!("{gateway_url}/api/admin/api-keys/key-123")),
)
.json(&json!({
"name": "renamed-key",
"rate_limit": 240,
"allowed_providers": ["gemini"],
"allowed_api_formats": ["gemini:chat"],
"allowed_models": ["gemini-2.5-pro"],
}))
.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["id"], json!("key-123"));
assert_eq!(payload["name"], json!("renamed-key"));
assert_eq!(payload["rate_limit"], json!(240));
assert_eq!(payload["allowed_providers"], json!(["gemini"]));
assert_eq!(payload["allowed_api_formats"], json!(["gemini:chat"]));
assert_eq!(payload["allowed_models"], json!(["gemini-2.5-pro"]));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_toggle_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
let repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-key-123".to_string()),
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
)])
.with_export_records([sample_standalone_export_record(
"key-123",
"admin-user-123",
"sk-key-123-plaintext",
true,
)]),
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
Arc::clone(&repository),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().patch(format!("{gateway_url}/api/admin/api-keys/key-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["id"], json!("key-123"));
assert_eq!(payload["is_active"], json!(false));
assert_eq!(payload["message"], json!("API密钥已禁用"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_api_keys_delete_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
let repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-key-123".to_string()),
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
)])
.with_export_records([sample_standalone_export_record(
"key-123",
"admin-user-123",
"sk-key-123-plaintext",
true,
)]),
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
Arc::clone(&repository),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().delete(format!("{gateway_url}/api/admin/api-keys/key-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["message"], json!("API密钥已删除"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let detail_response = admin_request(
reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys/key-123")),
)
.send()
.await
.expect("detail request should succeed");
assert_eq!(detail_response.status(), StatusCode::NOT_FOUND);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,663 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::users::InMemoryUserReadRepository;
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, 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;
async fn assert_admin_billing_get_returns_local(
path: &str,
) -> (serde_json::Value, Arc<Mutex<usize>>) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}{path}"))
.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 status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
gateway_handle.abort();
upstream_handle.abort();
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(status, StatusCode::OK);
(payload, upstream_hits)
}
async fn send_admin_billing_request(
gateway_url: &str,
method: http::Method,
path: &str,
body: Option<serde_json::Value>,
) -> reqwest::Response {
let client = reqwest::Client::new();
let mut request = client
.request(method, format!("{gateway_url}{path}"))
.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");
if let Some(body) = body {
request = request.json(&body);
}
request.send().await.expect("request should succeed")
}
#[tokio::test]
async fn gateway_handles_admin_billing_presets_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/billing/presets",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/billing/presets"))
.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["items"].as_array().map(Vec::len), Some(1));
assert_eq!(payload["items"][0]["name"], "aether-core");
assert_eq!(payload["items"][0]["version"], "1.0");
assert_eq!(payload["items"][0]["collector_count"], 16);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_billing_apply_preset_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/billing/presets/apply",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_billing_collectors_for_tests(Vec::<
crate::AdminBillingCollectorRecord,
>::new()),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let merge_response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/presets/apply",
Some(json!({"preset":"aether-core","mode":"merge"})),
)
.await;
assert_eq!(merge_response.status(), StatusCode::OK);
let merge_payload: serde_json::Value =
merge_response.json().await.expect("json body should parse");
assert_eq!(merge_payload["ok"], json!(true));
assert_eq!(merge_payload["preset"], json!("aether-core"));
assert_eq!(merge_payload["mode"], json!("merge"));
assert_eq!(merge_payload["created"], json!(16));
assert_eq!(merge_payload["updated"], json!(0));
assert_eq!(merge_payload["skipped"], json!(0));
assert_eq!(merge_payload["errors"], json!([]));
let merge_again_response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/presets/apply",
Some(json!({"preset":"default","mode":"merge"})),
)
.await;
assert_eq!(merge_again_response.status(), StatusCode::OK);
let merge_again_payload: serde_json::Value = merge_again_response
.json()
.await
.expect("json body should parse");
assert_eq!(merge_again_payload["ok"], json!(true));
assert_eq!(merge_again_payload["preset"], json!("aether-core"));
assert_eq!(merge_again_payload["created"], json!(0));
assert_eq!(merge_again_payload["updated"], json!(0));
assert_eq!(merge_again_payload["skipped"], json!(16));
let overwrite_response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/presets/apply",
Some(json!({"preset":"aether-core","mode":"overwrite"})),
)
.await;
assert_eq!(overwrite_response.status(), StatusCode::OK);
let overwrite_payload: serde_json::Value = overwrite_response
.json()
.await
.expect("json body should parse");
assert_eq!(overwrite_payload["ok"], json!(true));
assert_eq!(overwrite_payload["preset"], json!("aether-core"));
assert_eq!(overwrite_payload["mode"], json!("overwrite"));
assert_eq!(overwrite_payload["created"], json!(0));
assert_eq!(overwrite_payload["updated"], json!(16));
assert_eq!(overwrite_payload["skipped"], json!(0));
assert_eq!(overwrite_payload["errors"], json!([]));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_conflict_for_admin_billing_apply_preset_when_backend_unavailable() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/billing/presets/apply",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let mut state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_user_reader_for_tests(Arc::new(
InMemoryUserReadRepository::seed_auth_users(Vec::new()),
)),
);
state.admin_billing_collector_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/presets/apply",
Some(json!({"preset":"aether-core","mode":"merge"})),
)
.await;
assert_eq!(response.status(), StatusCode::CONFLICT);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error_code"], "read_only_mode");
assert_eq!(payload["detail"], "当前为只读模式,无法应用计费预设");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_billing_rule_routes_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream_hits_clone_detail = Arc::clone(&upstream_hits);
let upstream = Router::new()
.route(
"/api/admin/billing/rules",
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/billing/rules/rule-1",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone_detail);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_billing_rules_for_tests([
crate::AdminBillingRuleRecord {
id: "rule-1".to_string(),
name: "chat-input".to_string(),
task_type: "chat".to_string(),
global_model_id: Some("global-gpt-5".to_string()),
model_id: Some("model-1".to_string()),
expression: "input_tokens * 0.000001".to_string(),
variables: json!({ "base": 1.5 }),
dimension_mappings: json!({
"input_tokens": { "source": "dimension", "dimension": "input_tokens" }
}),
is_enabled: true,
created_at_unix_secs: 1_710_000_000,
updated_at_unix_secs: 1_710_000_100,
},
crate::AdminBillingRuleRecord {
id: "rule-2".to_string(),
name: "image-output".to_string(),
task_type: "image".to_string(),
global_model_id: None,
model_id: Some("model-2".to_string()),
expression: "images * 0.01".to_string(),
variables: json!({}),
dimension_mappings: json!({
"images": { "source": "dimension", "dimension": "image_count" }
}),
is_enabled: false,
created_at_unix_secs: 1_710_000_000,
updated_at_unix_secs: 1_710_000_050,
},
]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/billing/rules?task_type=chat&is_enabled=true&page=1&page_size=1"
))
.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!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(list_payload["items"].as_array().map(Vec::len), Some(1));
assert_eq!(list_payload["items"][0]["id"], "rule-1");
assert_eq!(list_payload["items"][0]["name"], "chat-input");
assert_eq!(list_payload["items"][0]["task_type"], "chat");
assert_eq!(list_payload["items"][0]["is_enabled"], true);
assert_eq!(list_payload["total"], 1);
assert_eq!(list_payload["page"], 1);
assert_eq!(list_payload["page_size"], 1);
assert_eq!(list_payload["pages"], 1);
let detail_response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/billing/rules/rule-1"))
.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!(detail_response.status(), StatusCode::OK);
let detail_payload: serde_json::Value = detail_response
.json()
.await
.expect("json body should parse");
assert_eq!(detail_payload["id"], "rule-1");
assert_eq!(detail_payload["name"], "chat-input");
assert_eq!(detail_payload["task_type"], "chat");
assert_eq!(detail_payload["global_model_id"], "global-gpt-5");
assert_eq!(detail_payload["model_id"], "model-1");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
let upstream =
Router::new().route(
"/api/admin/billing/rules",
any(|_request: Request| async move {
(StatusCode::OK, Body::from("unexpected upstream hit"))
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let create_response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/rules",
Some(json!({
"name": "rule",
"task_type": "chat",
"model_id": "model-1",
"expression": "input_tokens * 0.000001",
"variables": { "base": 1.5 },
"dimension_mappings": {
"input_tokens": { "source": "dimension", "dimension": "input_tokens" }
},
"is_enabled": true
})),
)
.await;
assert_eq!(create_response.status(), StatusCode::OK);
let create_payload: serde_json::Value = create_response
.json()
.await
.expect("json body should parse");
let rule_id = create_payload["id"]
.as_str()
.expect("rule id should exist")
.to_string();
assert_eq!(create_payload["name"], json!("rule"));
assert_eq!(create_payload["model_id"], json!("model-1"));
let update_response = send_admin_billing_request(
&gateway_url,
http::Method::PUT,
&format!("/api/admin/billing/rules/{rule_id}"),
Some(json!({
"name": "rule-updated",
"task_type": "video",
"global_model_id": "global-model-1",
"expression": "max(output_tokens, 1)",
"variables": { "factor": 2 },
"dimension_mappings": {
"output_tokens": { "source": "dimension", "dimension": "output_tokens" }
},
"is_enabled": false
})),
)
.await;
assert_eq!(update_response.status(), StatusCode::OK);
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(update_payload["id"], json!(rule_id));
assert_eq!(update_payload["name"], json!("rule-updated"));
assert_eq!(update_payload["task_type"], json!("video"));
assert_eq!(update_payload["global_model_id"], json!("global-model-1"));
assert_eq!(update_payload["model_id"], serde_json::Value::Null);
assert_eq!(update_payload["is_enabled"], json!(false));
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_conflict_for_admin_billing_rule_create_when_backend_unavailable() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/billing/rules",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let mut state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_user_reader_for_tests(Arc::new(
InMemoryUserReadRepository::seed_auth_users(Vec::new()),
)),
);
state.admin_billing_rule_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/rules",
Some(json!({
"name": "rule",
"task_type": "chat",
"model_id": "model-1",
"expression": "input_tokens * 0.000001",
"variables": {},
"dimension_mappings": {},
"is_enabled": true
})),
)
.await;
assert_eq!(response.status(), StatusCode::CONFLICT);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error_code"], "read_only_mode");
assert_eq!(payload["detail"], "当前为只读模式,无法创建计费规则");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_billing_collector_routes_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream_hits_clone_list = Arc::clone(&upstream_hits);
let upstream = Router::new()
.route(
"/api/admin/billing/collectors",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone_list);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
)
.route(
"/api/admin/billing/collectors/collector-1",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_billing_collectors_for_tests([
crate::AdminBillingCollectorRecord {
id: "collector-1".to_string(),
api_format: "OPENAI".to_string(),
task_type: "chat".to_string(),
dimension_name: "tokens".to_string(),
source_type: "request".to_string(),
source_path: Some("usage.input_tokens".to_string()),
value_type: "float".to_string(),
transform_expression: None,
default_value: None,
priority: 10,
is_enabled: true,
created_at_unix_secs: 1_710_000_000,
updated_at_unix_secs: 1_710_000_100,
},
crate::AdminBillingCollectorRecord {
id: "collector-2".to_string(),
api_format: "OPENAI".to_string(),
task_type: "chat".to_string(),
dimension_name: "latency".to_string(),
source_type: "computed".to_string(),
source_path: None,
value_type: "int".to_string(),
transform_expression: Some("max(latency_ms, 1)".to_string()),
default_value: Some("1".to_string()),
priority: 5,
is_enabled: false,
created_at_unix_secs: 1_710_000_000,
updated_at_unix_secs: 1_710_000_050,
},
]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let list_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/billing/collectors?api_format=OPENAI&task_type=chat&is_enabled=true&page=1&page_size=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!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(list_payload["items"].as_array().map(Vec::len), Some(1));
assert_eq!(list_payload["items"][0]["id"], "collector-1");
assert_eq!(list_payload["items"][0]["dimension_name"], "tokens");
assert_eq!(list_payload["total"], 1);
assert_eq!(list_payload["page"], 1);
assert_eq!(list_payload["page_size"], 10);
assert_eq!(list_payload["pages"], 1);
let detail_response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/billing/collectors/collector-1"
))
.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!(detail_response.status(), StatusCode::OK);
let detail_payload: serde_json::Value = detail_response
.json()
.await
.expect("json body should parse");
assert_eq!(detail_payload["id"], "collector-1");
assert_eq!(detail_payload["api_format"], "OPENAI");
assert_eq!(detail_payload["source_path"], "usage.input_tokens");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
let upstream =
Router::new().route(
"/api/admin/billing/collectors",
any(|_request: Request| async move {
(StatusCode::OK, Body::from("unexpected upstream hit"))
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let create_response = send_admin_billing_request(
&gateway_url,
http::Method::POST,
"/api/admin/billing/collectors",
Some(json!({
"api_format": "openai",
"task_type": "chat",
"dimension_name": "tokens",
"source_type": "request",
"source_path": "usage.input_tokens",
"value_type": "float",
"priority": 10,
"is_enabled": true
})),
)
.await;
assert_eq!(create_response.status(), StatusCode::OK);
let create_payload: serde_json::Value = create_response
.json()
.await
.expect("json body should parse");
let collector_id = create_payload["id"]
.as_str()
.expect("collector id should exist")
.to_string();
assert_eq!(create_payload["api_format"], json!("OPENAI"));
assert_eq!(create_payload["source_path"], json!("usage.input_tokens"));
let update_response = send_admin_billing_request(
&gateway_url,
http::Method::PUT,
&format!("/api/admin/billing/collectors/{collector_id}"),
Some(json!({
"api_format": "openai",
"task_type": "chat",
"dimension_name": "tokens",
"source_type": "computed",
"source_path": null,
"value_type": "int",
"transform_expression": "max(total_tokens, 1)",
"default_value": "1",
"priority": 12,
"is_enabled": false
})),
)
.await;
assert_eq!(update_response.status(), StatusCode::OK);
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(update_payload["id"], json!(collector_id));
assert_eq!(update_payload["source_type"], json!("computed"));
assert_eq!(update_payload["source_path"], serde_json::Value::Null);
assert_eq!(
update_payload["transform_expression"],
json!("max(total_tokens, 1)")
);
assert_eq!(update_payload["default_value"], json!("1"));
assert_eq!(update_payload["is_enabled"], json!(false));
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,3 @@
mod keys;
mod quota;
mod routes;

View File

@@ -0,0 +1,891 @@
use std::sync::{Arc, Mutex};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use serde_json::json;
use super::super::super::{
build_router_with_state, sample_endpoint, sample_key, sample_provider, 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;
#[tokio::test]
async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/keys",
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"))
}
}),
);
let mut key_a = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key_a.internal_priority = 10;
key_a.request_count = Some(12);
key_a.success_count = Some(9);
key_a.error_count = Some(3);
key_a.total_response_time_ms = Some(1800);
key_a.created_at_unix_secs = Some(1_711_000_000);
key_a.updated_at_unix_secs = Some(1_711_000_100);
key_a.last_used_at_unix_secs = Some(1_711_000_120);
key_a.note = Some("primary key".to_string());
key_a.status_snapshot = Some(json!({
"oauth": {"code": "none", "requires_reauth": false, "expiring_soon": false},
"account": {"code": "ok", "blocked": false, "recoverable": false},
"quota": {"code": "unknown", "exhausted": false}
}));
let mut key_b = sample_key(
"key-openai-b",
"provider-openai",
"openai:chat",
"sk-test-b",
);
key_b.internal_priority = 20;
key_b.request_count = Some(4);
key_b.success_count = Some(4);
key_b.error_count = Some(0);
key_b.total_response_time_ms = Some(400);
key_b.created_at_unix_secs = Some(1_711_100_000);
key_b.updated_at_unix_secs = Some(1_711_100_100);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![key_a, key_b],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_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/endpoints/providers/provider-openai/keys?skip=0&limit=50"
))
.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");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 2);
assert_eq!(items[0]["id"], "key-openai-a");
assert_eq!(items[0]["internal_priority"], 10);
assert_eq!(items[0]["request_count"], 12);
assert_eq!(items[0]["success_count"], 9);
assert_eq!(items[0]["error_count"], 3);
assert_eq!(items[0]["note"], "primary key");
assert_eq!(items[0]["api_key_masked"], "sk-test-a***");
assert_eq!(items[1]["id"], "key-openai-b");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/keys",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys"
))
.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!({
"api_formats": ["openai:chat"],
"api_key": "sk-created-openai",
"name": "created key",
"internal_priority": 15,
"capabilities": {"cache_1h": true},
"note": "created from rust"
}))
.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["provider_id"], "provider-openai");
assert_eq!(payload["name"], "created key");
assert_eq!(payload["internal_priority"], 15);
assert_eq!(payload["api_formats"], json!(["openai:chat"]));
assert_eq!(payload["api_key_masked"], "sk-creat***enai");
assert_eq!(payload["request_count"], 0);
assert_eq!(payload["success_count"], 0);
assert_eq!(payload["error_count"], 0);
assert_eq!(payload["note"], "created from rust");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("keys should read");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].name, "created key");
assert_eq!(keys[0].auth_type, "api_key");
assert_eq!(keys[0].internal_priority, 15);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reveals_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a/reveal",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/reveal"
))
.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["auth_type"], "api_key");
assert_eq!(payload["api_key"], "sk-test-a");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_exports_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-kiro-a/export",
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"))
}
}),
);
let mut key = sample_key(
"key-kiro-a",
"provider-kiro",
"claude:cli",
"oauth-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","auth_method":"idc","refresh_token":"rt-kiro-123"}"#,
)
.expect("auth config ciphertext should build"),
);
key.upstream_metadata = Some(json!({"kiro": {"email": "alice@example.com"}}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-kiro", "kiro", 10)],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/keys/key-kiro-a/export"
))
.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["provider_type"], "kiro");
assert_eq!(payload["auth_method"], "idc");
assert_eq!(payload["refresh_token"], "rt-kiro-123");
assert_eq!(payload["email"], "alice@example.com");
assert_eq!(payload["name"], "default");
assert!(payload.get("exported_at").is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_clears_admin_provider_key_oauth_invalid_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid",
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"))
}
}),
);
let mut key = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some("token expired".to_string());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid"
))
.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["message"], "已清除 OAuth 失效标记");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-openai-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].oauth_invalid_at_unix_secs, None);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_noops_admin_provider_key_oauth_invalid_clear_when_marker_absent() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid"
))
.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["message"], "该 Key 当前无失效标记,无需清除");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a",
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"))
}
}),
);
let mut key = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key.learned_rpm_limit = Some(88);
key.allowed_models = Some(json!(["gpt-4.1"]));
key.fingerprint = Some(json!({"user_agent": "old-ua"}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.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!({
"api_key": "sk-updated-openai",
"name": "updated key",
"internal_priority": 15,
"rpm_limit": null,
"allowed_models": [],
"note": "updated from rust",
"is_active": false,
"fingerprint": {"user_agent": "new-ua"}
}))
.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["name"], "updated key");
assert_eq!(payload["internal_priority"], 15);
assert_eq!(payload["rpm_limit"], serde_json::Value::Null);
assert_eq!(payload["learned_rpm_limit"], serde_json::Value::Null);
assert_eq!(payload["allowed_models"], json!([]));
assert_eq!(payload["note"], "updated from rust");
assert_eq!(payload["is_active"], false);
assert_eq!(payload["fingerprint"]["user_agent"], "new-ua");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-openai-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].name, "updated key");
assert_eq!(reloaded[0].internal_priority, 15);
assert_eq!(reloaded[0].rpm_limit, None);
assert_eq!(reloaded[0].learned_rpm_limit, None);
assert_eq!(reloaded[0].allowed_models, None);
assert_eq!(reloaded[0].note.as_deref(), Some("updated from rust"));
assert!(!reloaded[0].is_active);
let decrypted = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
&reloaded[0].encrypted_api_key,
)
.expect("ciphertext should decrypt");
assert_eq!(decrypted, "sk-updated-openai");
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_provider_key_update_when_api_key_duplicates_existing_key() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![
sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
),
sample_key(
"key-openai-b",
"provider-openai",
"openai:chat",
"sk-test-b",
),
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.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!({
"api_key": "sk-test-b"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["detail"]
.as_str()
.expect("detail should be string")
.contains("该 API Key 已存在于当前 Provider 中"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_deletes_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-openai-a",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
))
.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["message"], "Key key-openai-a 已删除");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-openai-a".to_string()])
.await
.expect("keys should read");
assert!(reloaded.is_empty());
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_batch_deletes_admin_provider_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/batch-delete",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![
sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
),
sample_key(
"key-openai-b",
"provider-openai",
"openai:chat",
"sk-test-b",
),
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/keys/batch-delete"
))
.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!({
"ids": ["key-openai-a", "key-missing"]
}))
.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_count"], 1);
assert_eq!(payload["failed_count"], 1);
assert_eq!(payload["failed"][0]["id"], "key-missing");
assert_eq!(payload["failed"][0]["error"], "not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].id, "key-openai-b");
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/grouped-by-format",
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"))
}
}),
);
let mut key_a = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key_a.internal_priority = 10;
key_a.request_count = Some(12);
key_a.success_count = Some(9);
key_a.created_at_unix_secs = Some(1_711_000_000);
key_a.updated_at_unix_secs = Some(1_711_000_100);
key_a.capabilities = Some(json!({"cache_1h": true, "gemini_files": false}));
key_a.global_priority_by_format = Some(json!({"openai:chat": 3}));
key_a.health_by_format = Some(json!({"openai:chat": {"health_score": 0.8}}));
key_a.circuit_breaker_by_format = Some(json!({"openai:chat": {"open": false}}));
let mut key_b = sample_key("key-claude-a", "provider-claude", "claude:chat", "sk-ant-a");
key_b.internal_priority = 20;
key_b.request_count = Some(2);
key_b.success_count = Some(1);
key_b.created_at_unix_secs = Some(1_711_100_000);
key_b.updated_at_unix_secs = Some(1_711_100_100);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10),
sample_provider("provider-claude", "claude", 20)
.with_transport_fields(false, false, true, None, None, None, None, None, None),
],
vec![
sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
),
sample_endpoint(
"endpoint-claude-chat",
"provider-claude",
"claude:chat",
"https://api.claude.example",
),
],
vec![key_a, key_b],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/keys/grouped-by-format"
))
.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["openai:chat"][0]["id"], "key-openai-a");
assert_eq!(payload["openai:chat"][0]["provider_name"], "openai");
assert_eq!(
payload["openai:chat"][0]["endpoint_base_url"],
"https://api.openai.example"
);
assert_eq!(payload["openai:chat"][0]["capabilities"], json!(["1h缓存"]));
assert_eq!(payload["claude:chat"][0]["provider_active"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,644 @@
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use axum::body::{to_bytes, Body};
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::json;
use super::super::super::{
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
sample_key, start_server,
};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_admin_principal() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
provider_api_format: String,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
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"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 12.5,
"reset_after_seconds": 18000,
"reset_at": 1_900_000_000u64,
"window_minutes": 300
},
"secondary_window": {
"used_percent": 55.0,
"reset_after_seconds": 604800,
"reset_at": 1_900_500_000u64,
"window_minutes": 10080
}
},
"credits": {
"has_credits": true,
"balance": 42.0,
"unlimited": false
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-codex".to_string(),
"codex".to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:cli",
"https://chatgpt.com/backend-api",
)],
vec![sample_key(
"key-codex-a",
"provider-codex",
"openai:cli",
"sk-codex-123",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/refresh-quota"
))
.header(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["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-codex-123"
);
assert_eq!(
seen_execution_runtime_request.provider_api_format,
"openai:cli"
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("plan_type")),
Some(&json!("plus"))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("primary_used_percent")),
Some(&json!(55.0))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("secondary_used_percent")),
Some(&json!(12.5))
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_admin_principal() {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
provider_api_format: String,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-kiro/refresh-quota",
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"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO+"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 5.0,
"usageLimitWithPrecision": 20.0,
"nextDateReset": 1_900_000_000u64
}],
"desktopUserInfo": {
"email": "dev@example.com"
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"access_token":"kiro-access-token",
"api_region":"us-west-2",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"kiro_version":"1.2.3"
}"#,
)
.expect("auth config ciphertext should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-kiro-a".to_string(),
"provider-kiro".to_string(),
"default".to_string(),
"bearer".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["claude:cli"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-kiro".to_string(),
"kiro".to_string(),
Some("https://example.com".to_string()),
"kiro".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-kiro-cli",
"provider-kiro",
"claude:cli",
"https://q.us-west-2.amazonaws.com",
)],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-kiro/refresh-quota"
))
.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["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert!(seen_execution_runtime_request
.url
.starts_with("https://q.us-west-2.amazonaws.com/getUsageLimits?"),);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer kiro-access-token"
);
assert_eq!(
seen_execution_runtime_request.provider_api_format,
"kiro:usage"
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("kiro"))
.and_then(|value| value.get("subscription_title")),
Some(&json!("KIRO PRO+"))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("kiro"))
.and_then(|value| value.get("remaining")),
Some(&json!(15.0))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("kiro"))
.and_then(|value| value.get("email")),
Some(&json!("dev@example.com"))
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
provider_api_format: String,
request_body: Option<serde_json::Value>,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-antigravity/refresh-quota",
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"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
request_body: plan.body.json_body.clone(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"models": {
"claude-sonnet-4": {
"displayName": "Claude Sonnet 4",
"quotaInfo": {
"remainingFraction": 0.25,
"resetTime": "2026-03-27T00:00:00Z"
}
},
"gemini-2.5-pro": {
"displayName": "Gemini 2.5 Pro"
}
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"project_id":"project-ant-123",
"client_version":"1.18.4",
"session_id":"session-ant-1"
}"#,
)
.expect("auth config ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-antigravity-a".to_string(),
"provider-antigravity".to_string(),
"default".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["gemini:chat"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "ya29.ant-token")
.expect("api key ciphertext should build"),
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-antigravity".to_string(),
"antigravity".to_string(),
Some("https://example.com".to_string()),
"antigravity".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-antigravity-chat",
"provider-antigravity",
"gemini:chat",
"https://daily-cloudcode-pa.googleapis.com",
)],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-antigravity/refresh-quota"
))
.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["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer ya29.ant-token"
);
assert_eq!(
seen_execution_runtime_request.provider_api_format,
"antigravity:fetch_available_models"
);
assert_eq!(
seen_execution_runtime_request.request_body,
Some(json!({ "project": "project-ant-123" }))
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-antigravity-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("antigravity"))
.and_then(|value| value.get("models"))
.and_then(|value| value.get("claude-sonnet-4"))
.and_then(|value| value.get("remaining_fraction")),
Some(&json!(0.25))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("antigravity"))
.and_then(|value| value.get("models"))
.and_then(|value| value.get("claude-sonnet-4"))
.and_then(|value| value.get("used_percent")),
Some(&json!(75.0))
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,609 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use serde_json::json;
use super::super::super::{
build_router_with_state, sample_endpoint, sample_key, sample_provider, 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;
const ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL: &str = "Admin endpoint data unavailable";
#[tokio::test]
async fn gateway_returns_service_unavailable_for_admin_provider_endpoints_when_catalog_reader_unavailable(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/endpoints",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints?skip=0&limit=50"
))
.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::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/endpoints",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![
sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100)),
sample_endpoint(
"endpoint-openai-embed",
"provider-openai",
"openai:embedding",
"https://api.openai.example",
)
.with_timestamps(Some(1_710_000_000), Some(1_710_000_100)),
],
vec![
sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
),
sample_key(
"key-openai-b",
"provider-openai",
"openai:chat",
"sk-test-b",
),
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_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/endpoints/providers/provider-openai/endpoints?skip=0&limit=50"
))
.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");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 2);
assert_eq!(items[0]["id"], "endpoint-openai-chat");
assert_eq!(items[0]["provider_name"], "openai");
assert_eq!(items[0]["api_format"], "openai:chat");
assert_eq!(items[0]["total_keys"], 2);
assert_eq!(items[0]["active_keys"], 2);
assert_eq!(items[1]["id"], "endpoint-openai-embed");
assert_eq!(items[1]["total_keys"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_service_unavailable_for_admin_provider_endpoint_create_when_catalog_writer_unavailable(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/endpoints",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_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/endpoints/providers/provider-openai/endpoints"
))
.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!({
"provider_id": "provider-openai",
"api_format": "openai:chat",
"base_url": "https://api.openai.example/"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-openai/endpoints",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints"
))
.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!({
"provider_id": "provider-openai",
"api_format": "openai:chat",
"base_url": "https://api.openai.example/",
"custom_path": "/v1/chat/completions",
"max_retries": 5,
"config": {"foo": "bar"},
"proxy": {"url": "http://proxy.internal", "password": "secret"}
}))
.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["provider_id"], "provider-openai");
assert_eq!(payload["provider_name"], "openai");
assert_eq!(payload["api_format"], "openai:chat");
assert_eq!(payload["base_url"], "https://api.openai.example");
assert_eq!(payload["custom_path"], "/v1/chat/completions");
assert_eq!(payload["max_retries"], 5);
assert_eq!(payload["total_keys"], 0);
assert_eq!(payload["active_keys"], 0);
assert_eq!(payload["proxy"]["url"], "http://proxy.internal");
assert_eq!(payload["proxy"]["password"], "***");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].provider_id, "provider-openai");
assert_eq!(endpoints[0].api_format, "openai:chat");
assert_eq!(endpoints[0].base_url, "https://api.openai.example");
assert_eq!(endpoints[0].max_retries, Some(5));
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/endpoint-openai-chat",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)
.with_transport_fields(
"https://api.openai.example".to_string(),
None,
None,
Some(2),
Some("/v1/chat/completions".to_string()),
Some(json!({"foo":"old"})),
None,
Some(json!({"url":"http://proxy.internal","password":"secret"})),
)
.expect("endpoint transport should build")],
vec![sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/endpoint-openai-chat"
))
.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!({
"base_url": "https://updated.openai.example/",
"custom_path": "/v1/responses",
"max_retries": 5,
"is_active": false,
"config": {"foo": "new"},
"proxy": {"url": "http://proxy-2.internal"}
}))
.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["id"], "endpoint-openai-chat");
assert_eq!(payload["provider_id"], "provider-openai");
assert_eq!(payload["api_format"], "openai:chat");
assert_eq!(payload["base_url"], "https://updated.openai.example");
assert_eq!(payload["custom_path"], "/v1/responses");
assert_eq!(payload["max_retries"], 5);
assert_eq!(payload["is_active"], false);
assert_eq!(payload["total_keys"], 1);
assert_eq!(payload["active_keys"], 1);
assert_eq!(payload["proxy"]["url"], "http://proxy-2.internal");
assert_eq!(payload["proxy"]["password"], "***");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let endpoints = provider_catalog_repository
.list_endpoints_by_ids(&["endpoint-openai-chat".to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].base_url, "https://updated.openai.example");
assert_eq!(endpoints[0].custom_path.as_deref(), Some("/v1/responses"));
assert_eq!(endpoints[0].max_retries, Some(5));
assert!(!endpoints[0].is_active);
assert_eq!(endpoints[0].config, Some(json!({"foo":"new"})));
assert_eq!(
endpoints[0].proxy,
Some(json!({"url":"http://proxy-2.internal","password":"secret"}))
);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_deletes_admin_provider_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/endpoint-openai-chat",
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"))
}
}),
);
let mut key_a = sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
);
key_a.api_formats = Some(json!(["openai:chat", "openai:embedding"]));
let mut key_b = sample_key(
"key-openai-b",
"provider-openai",
"openai:chat",
"sk-test-b",
);
key_b.api_formats = Some(json!(["openai:chat"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![key_a, key_b],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/endpoints/endpoint-openai-chat"
))
.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["message"], "Endpoint endpoint-openai-chat 已删除");
assert_eq!(payload["affected_keys_count"], 2);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let endpoints = provider_catalog_repository
.list_endpoints_by_ids(&["endpoint-openai-chat".to_string()])
.await
.expect("endpoints should read");
assert!(endpoints.is_empty());
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])
.await
.expect("keys should read");
let key_a = keys
.iter()
.find(|key| key.id == "key-openai-a")
.expect("key a should exist");
let key_b = keys
.iter()
.find(|key| key.id == "key-openai-b")
.expect("key b should exist");
assert_eq!(key_a.api_formats, Some(json!(["openai:embedding"])));
assert_eq!(key_b.api_formats, Some(json!([])));
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_get_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/endpoint-openai-chat",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![sample_endpoint(
"endpoint-openai-chat",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100))],
vec![sample_key(
"key-openai-a",
"provider-openai",
"openai:chat",
"sk-test-a",
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_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/endpoints/endpoint-openai-chat"
))
.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["id"], "endpoint-openai-chat");
assert_eq!(payload["provider_id"], "provider-openai");
assert_eq!(payload["provider_name"], "openai");
assert_eq!(payload["api_format"], "openai:chat");
assert_eq!(payload["total_keys"], 1);
assert_eq!(payload["active_keys"], 1);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_default_body_rules_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/defaults/openai:cli/body-rules",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/defaults/openai:cli/body-rules?provider_type=codex"
))
.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["api_format"], "openai:cli");
let rules = payload["body_rules"]
.as_array()
.expect("body_rules should be an array");
assert_eq!(rules.len(), 5);
assert_eq!(rules[0]["action"], "drop");
assert_eq!(rules[0]["path"], "max_output_tokens");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,349 @@
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::gemini_file_mappings::{
GeminiFileMappingReadRepository, InMemoryGeminiFileMappingRepository,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use axum::body::{to_bytes, Body};
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
use http::StatusCode;
use serde_json::json;
use std::sync::{Arc, Mutex};
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, sample_provider,
start_server,
};
use crate::constants::{
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
#[derive(Debug, Clone)]
struct SeenAdminGeminiUploadExecution {
key_id: String,
method: String,
url: String,
auth_header_value: String,
content_type: String,
body_bytes_b64: String,
}
fn sample_admin_gemini_provider() -> StoredProviderCatalogProvider {
sample_provider("provider-gemini-admin-1", "gemini", 10).with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_admin_gemini_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-gemini-admin-1".to_string(),
"provider-gemini-admin-1".to_string(),
"gemini:chat".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_admin_gemini_key(id: &str, name: &str, secret: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
"provider-gemini-admin-1".to_string(),
name.to_string(),
"api_key".to_string(),
Some(json!({"gemini_files": true})),
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["gemini:chat"])),
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)
.expect("api key should encrypt"),
None,
None,
Some(json!({"gemini:chat": 1})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
fn build_admin_gemini_upload_multipart(
boundary: &str,
file_name: &str,
mime_type: &str,
body: &[u8],
) -> Vec<u8> {
let mut payload = Vec::new();
payload.extend_from_slice(
format!(
"--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"{file_name}\"\r\nContent-Type: {mime_type}\r\n\r\n"
)
.as_bytes(),
);
payload.extend_from_slice(body);
payload.extend_from_slice(format!("\r\n--{boundary}--\r\n").as_bytes());
payload
}
#[tokio::test]
async fn gateway_uploads_admin_gemini_file_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/gemini-files/upload",
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"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(Vec::<SeenAdminGeminiUploadExecution>::new()));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
.expect("execution runtime payload should parse");
let key_id = payload
.get("key_id")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string();
seen_execution_runtime_inner
.lock()
.expect("mutex should lock")
.push(SeenAdminGeminiUploadExecution {
key_id: key_id.clone(),
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
auth_header_value: payload
.get("headers")
.and_then(|value| value.get("x-goog-api-key"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
content_type: payload
.get("headers")
.and_then(|value| value.get("content-type"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
body_bytes_b64: payload
.get("body")
.and_then(|value| value.get("body_bytes_b64"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
match key_id.as_str() {
"key-gemini-admin-ok" => Json(json!({
"request_id": "trace-admin-gemini-upload-123:ok",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"file": {
"name": "files/admin-upload-ok"
}
}
}
})),
"key-gemini-admin-fail" => Json(json!({
"request_id": "trace-admin-gemini-upload-123:fail",
"status_code": 429,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "quota exceeded"
}
}
}
})),
other => panic!("unexpected key id: {other}"),
}
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_admin_gemini_provider()],
vec![sample_admin_gemini_endpoint()],
vec![
sample_admin_gemini_key("key-gemini-admin-ok", "ok-key", "sk-admin-gemini-ok"),
sample_admin_gemini_key("key-gemini-admin-fail", "fail-key", "sk-admin-gemini-fail"),
],
));
let gemini_file_mapping_repository = Arc::new(InMemoryGeminiFileMappingRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let data_state =
GatewayDataState::with_request_candidate_and_gemini_file_mapping_repository_for_tests(
request_candidate_repository,
Arc::clone(&gemini_file_mapping_repository),
)
.attach_provider_catalog_repository_for_tests(Arc::clone(&provider_catalog_repository))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let boundary = "----AetherAdminGeminiUploadBoundary";
let upload_body = build_admin_gemini_upload_multipart(
boundary,
"admin-upload.txt",
"text/plain",
b"admin-upload-body",
);
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/gemini-files/upload?key_ids=key-gemini-admin-ok,key-gemini-admin-fail"
))
.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(
http::header::CONTENT_TYPE,
format!("multipart/form-data; boundary={boundary}"),
)
.header(TRACE_ID_HEADER, "trace-admin-gemini-upload-123")
.body(upload_body)
.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["display_name"], "admin-upload.txt");
assert_eq!(payload["mime_type"], "text/plain");
assert_eq!(payload["size_bytes"], json!(17));
assert_eq!(payload["success_count"], json!(1));
assert_eq!(payload["fail_count"], json!(1));
assert_eq!(payload["results"].as_array().map(Vec::len), Some(2));
assert_eq!(payload["results"][0]["key_id"], "key-gemini-admin-ok");
assert_eq!(payload["results"][0]["key_name"], "ok-key");
assert_eq!(payload["results"][0]["success"], json!(true));
assert_eq!(payload["results"][0]["file_name"], "files/admin-upload-ok");
assert_eq!(payload["results"][0]["error"], serde_json::Value::Null);
assert_eq!(payload["results"][1]["key_id"], "key-gemini-admin-fail");
assert_eq!(payload["results"][1]["key_name"], "fail-key");
assert_eq!(payload["results"][1]["success"], json!(false));
assert_eq!(payload["results"][1]["file_name"], serde_json::Value::Null);
assert_eq!(payload["results"][1]["error"], "quota exceeded");
let seen_requests = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(seen_requests.len(), 2);
for request in &seen_requests {
assert_eq!(request.method, "POST");
assert_eq!(
request.url,
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable"
);
assert_eq!(request.content_type, "text/plain");
assert_eq!(
BASE64_STANDARD
.decode(&request.body_bytes_b64)
.expect("execution runtime body should decode"),
b"admin-upload-body"
);
}
let seen_by_key_id = seen_requests
.iter()
.map(|request| (request.key_id.as_str(), request))
.collect::<std::collections::BTreeMap<_, _>>();
assert_eq!(
seen_by_key_id
.get("key-gemini-admin-ok")
.expect("ok request should exist")
.auth_header_value,
"sk-admin-gemini-ok"
);
assert_eq!(
seen_by_key_id
.get("key-gemini-admin-fail")
.expect("fail request should exist")
.auth_header_value,
"sk-admin-gemini-fail"
);
let stored_mapping = gemini_file_mapping_repository
.find_by_file_name("files/admin-upload-ok")
.await
.expect("mapping lookup should succeed")
.expect("mapping should exist");
assert_eq!(stored_mapping.key_id, "key-gemini-admin-ok");
assert_eq!(stored_mapping.user_id, None);
assert_eq!(
stored_mapping.display_name.as_deref(),
Some("admin-upload.txt")
);
assert_eq!(stored_mapping.mime_type.as_deref(), Some("text/plain"));
assert!(gemini_file_mapping_repository
.find_by_file_name("files/admin-upload-fail")
.await
.expect("mapping lookup should succeed")
.is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,225 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::auth_modules::{
InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
};
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_ldap_module_config, 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;
#[tokio::test]
async fn gateway_handles_admin_ldap_config_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/ldap/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"))
}
}),
);
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::seed(
Vec::<StoredOAuthProviderModuleConfig>::new(),
Some(sample_ldap_module_config()),
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
auth_module_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/ldap/config"))
.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 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["server_url"], "ldaps://ldap.example.com");
assert_eq!(payload["bind_dn"], "cn=admin,dc=example,dc=com");
assert_eq!(payload["base_dn"], "dc=example,dc=com");
assert_eq!(payload["has_bind_password"], json!(true));
assert_eq!(payload["user_search_filter"], "(uid={username})");
assert_eq!(payload["username_attr"], "uid");
assert_eq!(payload["email_attr"], "mail");
assert_eq!(payload["display_name_attr"], "displayName");
assert_eq!(payload["is_enabled"], json!(true));
assert_eq!(payload["is_exclusive"], json!(false));
assert_eq!(payload["use_starttls"], json!(true));
assert_eq!(payload["connect_timeout"], json!(10));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_ldap_config_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/ldap/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"))
}
}),
);
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::seed(
Vec::<StoredOAuthProviderModuleConfig>::new(),
Some(sample_ldap_module_config()),
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
auth_module_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let update_response = client
.put(format!("{gateway_url}/api/admin/ldap/config"))
.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!({
"server_url": "mockldap://ldap.internal.example.com",
"bind_dn": "cn=svc,dc=example,dc=com",
"bind_password": "secret123",
"base_dn": "ou=people,dc=example,dc=com",
"user_search_filter": "(uid={username})",
"username_attr": "uid",
"email_attr": "mail",
"display_name_attr": "cn",
"is_enabled": true,
"is_exclusive": false,
"use_starttls": false,
"connect_timeout": 20
}))
.send()
.await
.expect("update request should succeed");
let status = update_response.status();
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={update_payload}");
assert_eq!(update_payload["message"], "LDAP配置更新成功");
let get_response = client
.get(format!("{gateway_url}/api/admin/ldap/config"))
.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 status = get_response.status();
let payload: serde_json::Value = get_response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(
payload["server_url"],
"mockldap://ldap.internal.example.com"
);
assert_eq!(payload["bind_dn"], "cn=svc,dc=example,dc=com");
assert_eq!(payload["base_dn"], "ou=people,dc=example,dc=com");
assert_eq!(payload["display_name_attr"], "cn");
assert_eq!(payload["has_bind_password"], json!(true));
assert_eq!(payload["connect_timeout"], json!(20));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_tests_admin_ldap_connection_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/ldap/test",
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"))
}
}),
);
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
auth_module_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/ldap/test"))
.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!({
"server_url": "mockldap://ldap.example.com",
"bind_dn": "cn=admin,dc=example,dc=com",
"bind_password": "secret123",
"base_dn": "dc=example,dc=com",
"use_starttls": false,
"connect_timeout": 10
}))
.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["success"], json!(true));
assert_eq!(payload["message"], "连接成功");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,92 @@
use std::sync::{Arc, Mutex};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use super::super::super::{build_router_with_state, start_server, AppState};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
#[tokio::test]
async fn gateway_handles_admin_external_models_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/models/external",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/models/external"))
.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::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(
payload["detail"],
"External models catalog requires Rust admin backend"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_clears_admin_external_models_cache_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/models/external/cache",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!("{gateway_url}/api/admin/models/external/cache"))
.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"], false);
assert_eq!(payload["message"], "Redis 未启用");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,3 @@
mod external;
mod global;
mod provider;

View File

@@ -0,0 +1,639 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::global_models::{
AdminProviderModelListQuery, GlobalModelReadRepository, InMemoryGlobalModelReadRepository,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use serde_json::json;
use super::super::super::{
build_router_with_state, sample_admin_global_model, sample_admin_provider_model,
sample_provider, 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;
#[tokio::test]
async fn gateway_handles_admin_provider_models_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/models",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_provider_models(vec![
sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
),
]),
);
let quota_repository = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_global_model_and_quota_readers_for_tests(
provider_catalog_repository,
global_model_repository,
quota_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/providers/provider-openai/models?skip=0&limit=20"
))
.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");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "model-openai-gpt5");
assert_eq!(items[0]["provider_id"], "provider-openai");
assert_eq!(items[0]["global_model_id"], "global-gpt-5");
assert_eq!(items[0]["provider_model_name"], "gpt-5-upstream");
assert_eq!(items[0]["global_model_name"], "gpt-5");
assert_eq!(items[0]["global_model_display_name"], "GPT 5");
assert_eq!(items[0]["effective_input_price"], 3.0);
assert_eq!(items[0]["effective_output_price"], 15.0);
assert_eq!(items[0]["effective_supports_streaming"], true);
assert_eq!(items[0]["created_at"], "2024-03-21T05:46:40Z");
assert_eq!(items[0]["updated_at"], "2024-03-21T05:48:20Z");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_model_detail_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/models/model-openai-gpt5",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_provider_models(vec![
sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
),
]),
);
let quota_repository = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_global_model_and_quota_readers_for_tests(
provider_catalog_repository,
global_model_repository,
quota_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
))
.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["id"], "model-openai-gpt5");
assert_eq!(
payload["provider_model_mappings"],
json!([{"name": "gpt-5-upstream-alias", "priority": 1}])
);
assert_eq!(
payload["effective_config"],
json!({
"billing": {"currency": "USD", "mode": "local"},
"provider_hint": "gpt-5-upstream",
"streaming": true,
"vision": false
})
);
assert_eq!(payload["effective_supports_vision"], true);
assert_eq!(payload["effective_supports_image_generation"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_provider_model_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/models",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/providers/provider-openai/models"
))
.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!({
"provider_model_name": "gpt-5-upstream",
"global_model_id": "global-gpt-5",
"supports_vision": true,
"config": {"provider_hint": "gpt-5-upstream"}
}))
.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["provider_id"], "provider-openai");
assert_eq!(payload["global_model_id"], "global-gpt-5");
assert_eq!(payload["provider_model_name"], "gpt-5-upstream");
assert_eq!(payload["effective_supports_vision"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let created = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: "provider-openai".to_string(),
is_active: None,
offset: 0,
limit: 20,
})
.await
.expect("models should read");
assert_eq!(created.len(), 1);
assert_eq!(created[0].provider_model_name, "gpt-5-upstream");
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_and_deletes_admin_provider_model_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/models/model-openai-gpt5",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
sample_admin_global_model("global-gpt-5-mini", "gpt-5-mini", "GPT 5 mini"),
])
.with_admin_provider_models(vec![sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
)]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let update_response = reqwest::Client::new()
.patch(format!(
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
))
.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!({
"provider_model_name": "gpt-5-mini-upstream",
"global_model_id": "global-gpt-5-mini",
"supports_streaming": false,
"is_available": false
}))
.send()
.await
.expect("request should succeed");
assert_eq!(update_response.status(), StatusCode::OK);
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(update_payload["provider_model_name"], "gpt-5-mini-upstream");
assert_eq!(update_payload["global_model_id"], "global-gpt-5-mini");
assert_eq!(update_payload["is_available"], false);
let delete_response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
))
.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!(delete_response.status(), StatusCode::OK);
let delete_payload: serde_json::Value = delete_response
.json()
.await
.expect("json body should parse");
assert_eq!(
delete_payload["message"],
"Model 'gpt-5-mini-upstream' deleted successfully"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: "provider-openai".to_string(),
is_active: None,
offset: 0,
limit: 20,
})
.await
.expect("models should read");
assert!(models.is_empty());
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_batch_creates_admin_provider_models_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/models/batch",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/providers/provider-openai/models/batch"
))
.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!([
{"provider_model_name": "gpt-5-upstream", "global_model_id": "global-gpt-5"},
{"provider_model_name": "gpt-4.1-upstream", "global_model_id": "global-gpt-4.1"}
]))
.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.as_array().expect("payload array").len(), 2);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_available_source_models_locally_with_trusted_admin_principal(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/providers/provider-openai/available-source-models",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-5",
"gpt-5",
"GPT 5",
)])
.with_admin_provider_models(vec![
sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
),
sample_admin_provider_model(
"model-openai-gpt5-b",
"provider-openai",
"global-gpt-5",
"gpt-5-alt",
),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/providers/provider-openai/available-source-models"
))
.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["total"], 1);
assert_eq!(payload["models"][0]["global_model_name"], "gpt-5");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_assigns_and_imports_admin_provider_models_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new()
.route(
"/api/admin/providers/provider-openai/assign-global-models",
any({
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}
}),
)
.route(
"/api/admin/providers/provider-openai/import-from-upstream",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
])
.with_admin_provider_models(vec![sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
)]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let assign_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/providers/provider-openai/assign-global-models"
))
.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!({"global_model_ids": ["global-gpt-5", "global-gpt-4.1"]}))
.send()
.await
.expect("request should succeed");
assert_eq!(assign_response.status(), StatusCode::OK);
let assign_payload: serde_json::Value = assign_response
.json()
.await
.expect("json body should parse");
assert_eq!(
assign_payload["success"]
.as_array()
.expect("success array")
.len(),
1
);
assert_eq!(
assign_payload["errors"]
.as_array()
.expect("errors array")
.len(),
1
);
let import_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/providers/provider-openai/import-from-upstream"
))
.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!({
"model_ids": ["gpt-5-upstream", "brand-new-model"],
"price_per_request": 0.1
}))
.send()
.await
.expect("request should succeed");
assert_eq!(import_response.status(), StatusCode::OK);
let import_payload: serde_json::Value = import_response
.json()
.await
.expect("json body should parse");
assert_eq!(
import_payload["success"]
.as_array()
.expect("success array")
.len(),
2
);
assert_eq!(import_payload["success"][1]["created_global_model"], true);
assert_eq!(*upstream_hits_clone.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,752 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::wallet::StoredWalletSnapshot;
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, start_server, AppState};
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
fn admin_request(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
builder
.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")
}
async fn start_payments_upstream(
path: &'static str,
) -> (String, Arc<Mutex<usize>>, tokio::task::JoinHandle<()>) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
(upstream_url, upstream_hits, upstream_handle)
}
fn sample_wallet(wallet_id: &str, user_id: &str) -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
wallet_id.to_string(),
Some(user_id.to_string()),
None,
12.5,
2.5,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
30.0,
10.0,
3.0,
1.5,
1_710_000_000,
)
.expect("wallet should build")
}
#[allow(clippy::too_many_arguments)]
fn sample_payment_order(
order_id: &str,
wallet_id: &str,
user_id: &str,
amount_usd: f64,
payment_method: &str,
status: &str,
created_at_unix_secs: u64,
expires_at_unix_secs: Option<u64>,
) -> crate::AdminWalletPaymentOrderRecord {
crate::AdminWalletPaymentOrderRecord {
id: order_id.to_string(),
order_no: format!("po-{order_id}"),
wallet_id: wallet_id.to_string(),
user_id: Some(user_id.to_string()),
amount_usd,
pay_amount: None,
pay_currency: None,
exchange_rate: None,
refunded_amount_usd: 0.0,
refundable_amount_usd: amount_usd,
payment_method: payment_method.to_string(),
gateway_order_id: None,
status: status.to_string(),
gateway_response: None,
created_at_unix_secs,
paid_at_unix_secs: None,
credited_at_unix_secs: None,
expires_at_unix_secs,
}
}
fn sample_payment_callback(
callback_id: &str,
payment_order_id: Option<&str>,
payment_method: &str,
callback_key: &str,
status: &str,
created_at_unix_secs: u64,
) -> crate::state::AdminPaymentCallbackRecord {
crate::state::AdminPaymentCallbackRecord {
id: callback_id.to_string(),
payment_order_id: payment_order_id.map(str::to_string),
payment_method: payment_method.to_string(),
callback_key: callback_key.to_string(),
order_no: payment_order_id.map(|value| format!("po-{value}")),
gateway_order_id: Some(format!("gw-{callback_id}")),
payload_hash: Some(format!("hash-{callback_id}")),
signature_valid: status != "failed",
status: status.to_string(),
payload: Some(json!({ "source": callback_id })),
error_message: (status == "failed").then(|| "signature mismatch".to_string()),
created_at_unix_secs,
processed_at_unix_secs: Some(created_at_unix_secs.saturating_add(60)),
}
}
#[tokio::test]
async fn gateway_handles_admin_payments_list_orders_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_wallet_payment_orders_for_tests([
sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"alipay",
"pending",
1_710_000_000,
Some(4_102_444_800),
),
sample_payment_order(
"order-2",
"wallet-2",
"user-2",
8.0,
"manual",
"failed",
1_709_000_000,
None,
),
]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/payments/orders?status=pending&payment_method=alipay&limit=20&offset=0"
)))
.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");
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "order-1");
assert_eq!(items[0]["wallet_id"], "wallet-1");
assert_eq!(items[0]["payment_method"], "alipay");
assert_eq!(items[0]["status"], "pending");
assert_eq!(payload["total"], 1);
assert_eq!(payload["limit"], 20);
assert_eq!(payload["offset"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_get_order_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-1").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(
reqwest::Client::new().get(format!("{gateway_url}/api/admin/payments/orders/order-1")),
)
.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["order"]["id"], "order-1");
assert_eq!(payload["order"]["payment_method"], "wechat");
assert_eq!(payload["order"]["amount_usd"], 12.5);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_expire_order_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-1/expire").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/expire"
)))
.json(&json!({}))
.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["expired"], true);
assert_eq!(payload["order"]["status"], "expired");
assert_eq!(
payload["order"]["gateway_response"]["expire_reason"],
"admin_mark_expired"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_credit_order_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-1/credit").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/credit"
)))
.json(&json!({
"gateway_order_id": "gateway-order-1",
"pay_amount": 91.25,
"pay_currency": "cny",
"exchange_rate": 7.3,
"gateway_response": {
"channel": "manual-review"
}
}))
.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["credited"], true);
assert_eq!(payload["order"]["status"], "credited");
assert_eq!(payload["order"]["gateway_order_id"], "gateway-order-1");
assert_eq!(payload["order"]["pay_amount"], 91.25);
assert_eq!(payload["order"]["pay_currency"], "CNY");
assert_eq!(payload["order"]["exchange_rate"], 7.3);
assert_eq!(
payload["order"]["gateway_response"]["channel"],
"manual-review"
);
assert_eq!(payload["order"]["gateway_response"]["manual_credit"], true);
assert_eq!(
payload["order"]["gateway_response"]["credited_by"],
"admin-user-123"
);
assert!(payload["order"]["paid_at"].as_str().is_some());
assert!(payload["order"]["credited_at"].as_str().is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_fail_order_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-1/fail").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/fail"
)))
.json(&json!({}))
.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["order"]["status"], "failed");
assert_eq!(
payload["order"]["gateway_response"]["failure_reason"],
"admin_mark_failed"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_callbacks_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/callbacks").await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_payment_callbacks_for_tests([
sample_payment_callback(
"callback-1",
Some("order-1"),
"alipay",
"callback-key-1",
"processed",
1_710_100_000,
),
sample_payment_callback(
"callback-2",
Some("order-2"),
"wechat",
"callback-key-2",
"failed",
1_710_000_000,
),
]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/payments/callbacks?payment_method=alipay&limit=10&offset=0"
)))
.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");
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "callback-1");
assert_eq!(items[0]["payment_order_id"], "order-1");
assert_eq!(items[0]["payment_method"], "alipay");
assert_eq!(items[0]["callback_key"], "callback-key-1");
assert_eq!(items[0]["signature_valid"], true);
assert_eq!(items[0]["status"], "processed");
assert_eq!(items[0]["payload"]["source"], "callback-1");
assert!(items[0]["processed_at"].as_str().is_some());
assert_eq!(payload["total"], 1);
assert_eq!(payload["limit"], 10);
assert_eq!(payload["offset"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_trailing_slash_routes_locally_with_trusted_admin_principal()
{
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new()
.route(
"/api/admin/payments/orders/order-1",
any({
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}
}),
)
.route(
"/api/admin/payments/orders/order-1/expire",
any({
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}
}),
)
.route(
"/api/admin/payments/orders/order-1/credit",
any({
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}
}),
)
.route(
"/api/admin/payments/orders/order-1/fail",
any({
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let detail = admin_request(
reqwest::Client::new().get(format!("{gateway_url}/api/admin/payments/orders/order-1/")),
)
.send()
.await
.expect("detail request should succeed");
assert_eq!(detail.status(), StatusCode::OK);
let detail_payload: serde_json::Value = detail.json().await.expect("json body should parse");
assert_eq!(detail_payload["order"]["id"], "order-1");
let expire = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/expire/"
)))
.json(&json!({}))
.send()
.await
.expect("expire request should succeed");
assert_eq!(expire.status(), StatusCode::OK);
let expire_payload: serde_json::Value = expire.json().await.expect("json body should parse");
assert_eq!(expire_payload["order"]["status"], "expired");
let mut state = AppState::new().expect("gateway should build");
state = state
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle_credit) = start_server(gateway).await;
let credit = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/credit/"
)))
.json(&json!({
"gateway_order_id": "gateway-order-1",
"pay_amount": 91.25,
"pay_currency": "cny",
"exchange_rate": 7.3
}))
.send()
.await
.expect("credit request should succeed");
assert_eq!(credit.status(), StatusCode::OK);
let credit_payload: serde_json::Value = credit.json().await.expect("json body should parse");
assert_eq!(credit_payload["order"]["status"], "credited");
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
"order-1",
"wallet-1",
"user-1",
12.5,
"wechat",
"pending",
1_710_000_000,
Some(4_102_444_800),
)]),
);
let (gateway_url, gateway_handle_fail) = start_server(gateway).await;
let fail = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-1/fail/"
)))
.json(&json!({}))
.send()
.await
.expect("fail request should succeed");
assert_eq!(fail.status(), StatusCode::OK);
let fail_payload: serde_json::Value = fail.json().await.expect("json body should parse");
assert_eq!(fail_payload["order"]["status"], "failed");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
gateway_handle_credit.abort();
gateway_handle_fail.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_list_orders_locally_without_payment_backend() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders").await;
let mut state = AppState::new().expect("gateway should build");
state.admin_wallet_payment_order_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/payments/orders?limit=10&offset=0"
)))
.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["items"], json!([]));
assert_eq!(payload["total"], 0);
assert_eq!(payload["limit"], 10);
assert_eq!(payload["offset"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_get_order_locally_without_payment_backend() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-missing").await;
let mut state = AppState::new().expect("gateway should build");
state.admin_wallet_payment_order_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/payments/orders/order-missing"
)))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "Payment order read backend unavailable");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_expire_order_with_backend_unavailable_detail() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/orders/order-missing/expire").await;
let mut state = AppState::new().expect("gateway should build");
state.admin_wallet_payment_order_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().post(format!(
"{gateway_url}/api/admin/payments/orders/order-missing/expire"
)))
.json(&json!({}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "Payment order write backend unavailable");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_payments_callbacks_locally_without_payment_backend() {
let (upstream_url, upstream_hits, upstream_handle) =
start_payments_upstream("/api/admin/payments/callbacks").await;
let mut state = AppState::new().expect("gateway should build");
state.admin_payment_callback_store = None;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/payments/callbacks?payment_method=alipay&limit=10&offset=5"
)))
.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["items"], json!([]));
assert_eq!(payload["total"], 0);
assert_eq!(payload["limit"], 10);
assert_eq!(payload["offset"], 5);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_payments_empty_order_identifier_locally_with_trusted_admin_principal(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().fallback(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"))
}
}));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let expire = admin_request(
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//expire")),
)
.json(&json!({}))
.send()
.await
.expect("expire request should succeed");
assert_eq!(expire.status(), StatusCode::NOT_FOUND);
let expire_payload: serde_json::Value = expire.json().await.expect("json body should parse");
assert_eq!(expire_payload["detail"], "Payment order not found");
let credit = admin_request(
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//credit")),
)
.json(&json!({}))
.send()
.await
.expect("credit request should succeed");
assert_eq!(credit.status(), StatusCode::NOT_FOUND);
let credit_payload: serde_json::Value = credit.json().await.expect("json body should parse");
assert_eq!(credit_payload["detail"], "Payment order not found");
let fail = admin_request(
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//fail")),
)
.json(&json!({}))
.send()
.await
.expect("fail request should succeed");
assert_eq!(fail.status(), StatusCode::NOT_FOUND);
let fail_payload: serde_json::Value = fail.json().await.expect("json body should parse");
assert_eq!(fail_payload["detail"], "Payment order not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,344 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
};
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_admin_global_model, sample_admin_provider_model, sample_key,
sample_provider, 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;
async fn assert_admin_provider_query_route(
path: &str,
request_payload: serde_json::Value,
expected_status: StatusCode,
expected_payload_assertions: impl FnOnce(&serde_json::Value),
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}{path}"))
.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(&request_payload)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), expected_status);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
expected_payload_assertions(&payload);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-query/models",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "OpenAI", 10)],
vec![StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat".to_string(),
"provider-openai".to_string(),
"openai:chat".to_string(),
Some("chat".to_string()),
Some("primary".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com/v1".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")],
vec![
{
let mut key = sample_key(
"key-openai-allowed",
"provider-openai",
"openai:chat",
"sk-test",
);
key.allowed_models = Some(json!(["gpt-5"]));
key
},
sample_key(
"key-openai-all",
"provider-openai",
"openai:chat",
"sk-test-2",
),
],
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
])
.with_admin_provider_models(vec![
{
let mut model = sample_admin_provider_model(
"provider-model-gpt-5",
"provider-openai",
"global-gpt-5",
"gpt-5",
);
model.global_model_name = Some("gpt-5".to_string());
model.global_model_display_name = Some("GPT 5".to_string());
model
},
{
let mut model = sample_admin_provider_model(
"provider-model-gpt-4.1",
"provider-openai",
"global-gpt-4.1",
"gpt-4.1",
);
model.global_model_name = Some("gpt-4.1".to_string());
model.global_model_display_name = Some("GPT 4.1".to_string());
model
},
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.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!({
"provider_id": "provider-openai",
"api_key_id": "key-openai-allowed"
}))
.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"], json!(true));
assert_eq!(payload["provider"]["id"], "provider-openai");
assert_eq!(payload["provider"]["name"], "OpenAI");
assert_eq!(payload["provider"]["display_name"], "OpenAI");
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
assert_eq!(payload["data"]["from_cache"], json!(true));
assert_eq!(payload["data"]["keys_total"], json!(2));
let models = payload["data"]["models"]
.as_array()
.expect("models should be an array");
assert_eq!(models.len(), 2);
let model_ids: Vec<_> = models
.iter()
.map(|model| {
(
model["id"].as_str().expect("id should be present"),
model["display_name"]
.as_str()
.expect("display_name should be present"),
)
})
.collect();
assert_eq!(model_ids, vec![("gpt-4.1", "GPT 4.1"), ("gpt-5", "GPT 5")]);
for model in models {
assert_eq!(model["owned_by"], "OpenAI");
assert_eq!(model["api_format"], "openai:chat");
assert_eq!(model["api_formats"], json!(["openai:chat"]));
}
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_locally_with_trusted_admin_principal() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model",
json!({ "provider_id": "provider-openai", "model": "gpt-4.1" }),
StatusCode::OK,
|payload| {
assert_eq!(payload["success"], json!(false));
assert_eq!(payload["tested"], json!(false));
assert!(payload["provider_id"].as_str().is_some());
},
)
.await;
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_failover_locally_with_trusted_admin_principal(
) {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model-failover",
json!({
"provider_id": "provider-openai",
"failover_models": ["gpt-4.1", "gpt-4o-mini"]
}),
StatusCode::OK,
|payload| {
assert_eq!(payload["success"], json!(false));
assert_eq!(payload["tested"], json!(false));
assert!(payload["provider_id"].as_str().is_some());
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_invalid_json_body() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-query/models",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/models"))
.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(reqwest::header::CONTENT_TYPE, "application/json")
.body("{")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], json!("Invalid JSON request body"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_without_provider_id() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model",
json!({ "model": "gpt-4.1" }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("provider_id is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_without_model() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model",
json!({ "provider_id": "provider-openai" }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("model is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_failover_without_provider_id() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model-failover",
json!({ "failover_models": ["gpt-4.1"] }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(payload["detail"], json!("provider_id is required"));
},
)
.await;
}
#[tokio::test]
async fn gateway_rejects_admin_provider_query_test_model_failover_without_models() {
assert_admin_provider_query_route(
"/api/admin/provider-query/test-model-failover",
json!({ "provider_id": "provider-openai", "failover_models": [] }),
StatusCode::BAD_REQUEST,
|payload| {
assert_eq!(
payload["detail"],
json!("failover_models should not be empty")
);
},
)
.await;
}

View File

@@ -0,0 +1,617 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredProviderUsageWindow};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::{Method, StatusCode};
use serde_json::json;
use super::super::{build_router_with_state, sample_provider, 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;
async fn assert_provider_strategy_route_returns_local_503(
data_state: GatewayDataState,
method: Method,
path: &str,
body: Option<serde_json::Value>,
expected_message: &str,
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
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 mut request = reqwest::Client::new()
.request(method, format!("{gateway_url}{path}"))
.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");
if let Some(body) = body {
request = request.json(&body);
}
let response = request.send().await.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["error"]["message"], expected_message);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
const PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE: &str = "Admin provider strategy data unavailable";
const PROVIDER_STRATEGY_STATS_DATA_UNAVAILABLE_MESSAGE: &str =
"Admin provider strategy stats data unavailable";
#[tokio::test]
async fn gateway_handles_admin_provider_strategy_list_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/strategies",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-strategy/strategies"
))
.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["total"], 1);
assert_eq!(payload["strategies"][0]["name"], "sticky_priority");
assert_eq!(payload["strategies"][0]["priority"], 110);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_strategy_billing_locally_returns_service_unavailable_when_disabled(
) {
assert_provider_strategy_route_returns_local_503(
crate::data::GatewayDataState::disabled(),
Method::PUT,
"/api/admin/provider-strategy/providers/provider-openai/billing",
Some(json!({
"billing_type": "monthly_quota",
"monthly_quota_usd": 100.0,
"quota_reset_day": 30,
"quota_last_reset_at": "2024-03-21T00:00:00Z",
"quota_expires_at": "2024-04-21T00:00:00Z",
"rpm_limit": 20,
"provider_priority": 5
})),
PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE,
)
.await;
}
#[tokio::test]
async fn gateway_handles_admin_provider_strategy_stats_locally_returns_service_unavailable_without_usage_reader(
) {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let catalog_reader: Arc<dyn ProviderCatalogReadRepository> =
provider_catalog_repository.clone();
assert_provider_strategy_route_returns_local_503(
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
catalog_reader,
),
Method::GET,
"/api/admin/provider-strategy/providers/provider-openai/stats?hours=1",
None,
PROVIDER_STRATEGY_STATS_DATA_UNAVAILABLE_MESSAGE,
)
.await;
}
#[tokio::test]
async fn gateway_handles_admin_provider_strategy_quota_locally_returns_service_unavailable_without_writer(
) {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![],
vec![],
));
let catalog_reader: Arc<dyn ProviderCatalogReadRepository> =
provider_catalog_repository.clone();
assert_provider_strategy_route_returns_local_503(
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
catalog_reader,
),
Method::DELETE,
"/api/admin/provider-strategy/providers/provider-openai/quota",
None,
PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE,
)
.await;
}
async fn gateway_resets_admin_provider_strategy_quota_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/providers/provider-openai/quota",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 10).with_billing_fields(
Some("monthly_quota".to_string()),
Some(100.0),
Some(12.5),
Some(30),
Some(1_711_000_000),
Some(1_711_000_000 + 30 * 24 * 60 * 60),
),
],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/quota"
))
.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["provider_name"], "openai");
assert_eq!(payload["previous_used"], 12.5);
assert_eq!(payload["current_used"], 0.0);
let updated = provider_catalog_repository
.list_providers_by_ids(&["provider-openai".to_string()])
.await
.expect("provider query should succeed");
assert_eq!(updated[0].monthly_used_usd, Some(0.0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_admin_provider_strategy_billing_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/providers/provider-openai/billing",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 7).with_billing_fields(
Some("pay_as_you_go".to_string()),
None,
Some(12.5),
Some(9),
Some(1_711_000_000),
None,
),
],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/billing"
))
.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!({
"billing_type": "monthly_quota",
"monthly_quota_usd": 50.0,
"quota_last_reset_at": "2026-03-01T00:00:00Z",
"quota_expires_at": "2026-04-01T00:00:00Z"
}))
.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["provider"]["id"], "provider-openai");
assert_eq!(payload["provider"]["billing_type"], "monthly_quota");
assert_eq!(payload["provider"]["provider_priority"], 100);
let updated = provider_catalog_repository
.list_providers_by_ids(&["provider-openai".to_string()])
.await
.expect("provider query should succeed");
assert_eq!(updated[0].billing_type.as_deref(), Some("monthly_quota"));
assert_eq!(updated[0].monthly_quota_usd, Some(50.0));
assert_eq!(updated[0].monthly_used_usd, Some(12.5));
assert_eq!(updated[0].quota_reset_day, Some(30));
assert_eq!(updated[0].provider_priority, 100);
assert_eq!(
updated[0].quota_last_reset_at_unix_secs,
Some(1_772_323_200)
);
assert_eq!(updated[0].quota_expires_at_unix_secs, Some(1_775_001_600));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_syncs_admin_provider_strategy_monthly_usage_from_reset_window_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/providers/provider-openai/billing",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 7).with_billing_fields(
Some("pay_as_you_go".to_string()),
None,
Some(12.5),
Some(9),
Some(1_711_000_000),
None,
),
],
vec![],
vec![],
));
let usage_repository = Arc::new(
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_772_236_799,
8,
7,
1,
100.0,
9.99,
)
.expect("window should build"),
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_772_323_200,
5,
5,
0,
110.0,
1.25,
)
.expect("window should build"),
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_772_409_600,
6,
5,
1,
150.0,
0.75,
)
.expect("window should build"),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
Arc::clone(&provider_catalog_repository),
Arc::clone(&usage_repository),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/billing"
))
.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!({
"billing_type": "monthly_quota",
"monthly_quota_usd": 50.0,
"quota_last_reset_at": "2026-03-01T00:00:00Z",
"quota_expires_at": "2026-04-01T00:00:00Z"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let updated = provider_catalog_repository
.list_providers_by_ids(&["provider-openai".to_string()])
.await
.expect("provider query should succeed");
assert_eq!(updated[0].monthly_used_usd, Some(2.0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_reads_admin_provider_strategy_stats_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/providers/provider-openai/stats",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 7).with_billing_fields(
Some("monthly_quota".to_string()),
Some(80.0),
Some(12.5),
Some(30),
Some(1_711_000_000),
Some(1_774_966_400),
),
],
vec![],
vec![],
));
let usage_repository = Arc::new(
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_700_000_000,
10,
9,
1,
120.0,
1.25,
)
.expect("window should build"),
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_700_003_600,
6,
5,
1,
180.0,
0.75,
)
.expect("window should build"),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
Arc::clone(&provider_catalog_repository),
Arc::clone(&usage_repository),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/stats?hours=999999"
))
.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["provider_id"], "provider-openai");
assert_eq!(payload["provider_name"], "openai");
assert_eq!(payload["billing_info"]["billing_type"], "monthly_quota");
assert_eq!(payload["billing_info"]["monthly_quota_usd"], 80.0);
assert_eq!(payload["billing_info"]["monthly_used_usd"], 12.5);
assert_eq!(payload["billing_info"]["quota_remaining_usd"], 67.5);
assert_eq!(payload["usage_stats"]["total_requests"], 16);
assert_eq!(payload["usage_stats"]["successful_requests"], 14);
assert_eq!(payload["usage_stats"]["failed_requests"], 2);
assert_eq!(payload["usage_stats"]["success_rate"], 0.875);
assert_eq!(payload["usage_stats"]["avg_response_time_ms"], 150.0);
assert_eq!(payload["usage_stats"]["total_cost_usd"], 2.0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rounds_admin_provider_strategy_stats_total_cost_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-strategy/providers/provider-openai/stats",
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"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "openai", 7).with_billing_fields(
Some("monthly_quota".to_string()),
Some(80.0),
Some(12.5),
Some(30),
Some(1_711_000_000),
Some(1_774_966_400),
),
],
vec![],
vec![],
));
let usage_repository = Arc::new(
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
StoredProviderUsageWindow::new(
"provider-openai".to_string(),
1_700_000_000,
2,
2,
0,
120.0,
1.23456,
)
.expect("window should build"),
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
Arc::clone(&provider_catalog_repository),
Arc::clone(&usage_repository),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/stats?hours=999999"
))
.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["usage_stats"]["total_cost_usd"], 1.2346);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,216 @@
use std::sync::{Arc, Mutex};
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, 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 crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::data::GatewayDataState;
#[tokio::test]
async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal() {
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",
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"))
}
}),
);
let mut manual_node = sample_proxy_node("proxy-node-manual");
manual_node.name = "alpha-manual".to_string();
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
manual_node.proxy_username = Some("alice".to_string());
manual_node.proxy_password = Some("supersecret".to_string());
manual_node.last_heartbeat_at_unix_secs = None;
manual_node.tunnel_connected_at_unix_secs = None;
let mut tunnel_node = sample_proxy_node("proxy-node-tunnel");
tunnel_node.name = "zeta-tunnel".to_string();
tunnel_node.status = "offline".to_string();
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
tunnel_node,
manual_node,
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
proxy_node_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes?status=online&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["total"], 1);
assert_eq!(payload["skip"], 0);
assert_eq!(payload["limit"], 10);
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], "proxy-node-manual");
assert_eq!(items[0]["name"], "alpha-manual");
assert_eq!(items[0]["status"], "online");
assert_eq!(items[0]["is_manual"], true);
assert_eq!(items[0]["proxy_url"], "http://proxy.example:8080");
assert_eq!(items[0]["proxy_username"], "alice");
assert_eq!(items[0]["proxy_password"], "su****et");
assert!(items[0]["created_at"].is_string());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_admin_proxy_nodes_unavailable_routes_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/register",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
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")
.json(&json!({
"name": "proxy-1",
"ip": "1.1.1.1",
"port": 8080
}))
.send()
.await
.expect("request should succeed");
assert_eq!(register_response.status(), StatusCode::SERVICE_UNAVAILABLE);
let register_payload: serde_json::Value = register_response
.json()
.await
.expect("json body should parse");
assert_eq!(
register_payload["detail"],
"Admin proxy nodes data unavailable"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_principal() {
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-1/events",
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"))
}
}),
);
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed_with_events(
vec![sample_proxy_node("node-1")],
vec![
StoredProxyNodeEvent {
id: 1,
node_id: "node-1".to_string(),
event_type: "connected".to_string(),
detail: Some("older".to_string()),
created_at_unix_secs: Some(1_710_000_000),
},
StoredProxyNodeEvent {
id: 2,
node_id: "node-1".to_string(),
event_type: "disconnected".to_string(),
detail: Some("newer".to_string()),
created_at_unix_secs: Some(1_710_000_100),
},
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
proxy_node_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/proxy-nodes/node-1/events?limit=1"
))
.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");
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["id"], 2);
assert_eq!(items[0]["event_type"], "disconnected");
assert_eq!(items[0]["detail"], "newer");
assert_eq!(items[0]["created_at"], "2024-03-09T16:01:40Z");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,406 @@
use std::sync::{Arc, Mutex};
use axum::body::{Body, Bytes};
use axum::routing::any;
use axum::{extract::Request, Router};
use http::{HeaderMap, HeaderValue, StatusCode};
use serde_json::json;
use super::super::{build_router_with_state, start_server, AppState};
use crate::audit::AdminAuditEvent;
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::control::resolve_public_request_context;
use crate::handlers::admin::maybe_build_local_admin_security_response;
async fn send_admin_security_request(
gateway: Router,
method: reqwest::Method,
path: &str,
body: Option<serde_json::Value>,
) -> (StatusCode, serde_json::Value, usize) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
path,
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let mut request = client
.request(method, format!("{gateway_url}{path}"))
.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");
if let Some(body) = body {
request = request.json(&body);
}
let response = request.send().await.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let upstream_count = *upstream_hits.lock().expect("mutex should lock");
gateway_handle.abort();
upstream_handle.abort();
(status, payload, upstream_count)
}
fn trusted_admin_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
headers.insert(
TRUSTED_ADMIN_USER_ID_HEADER,
HeaderValue::from_static("admin-user-123"),
);
headers.insert(
TRUSTED_ADMIN_USER_ROLE_HEADER,
HeaderValue::from_static("admin"),
);
headers.insert(
TRUSTED_ADMIN_SESSION_ID_HEADER,
HeaderValue::from_static("session-123"),
);
headers.insert(
TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER,
HeaderValue::from_static("management-token-123"),
);
headers
}
async fn local_admin_security_response(
state: &AppState,
method: http::Method,
uri: &str,
body: Option<serde_json::Value>,
) -> axum::response::Response<Body> {
let headers = trusted_admin_headers();
let request_context = resolve_public_request_context(
state,
&method,
&uri.parse().expect("uri should parse"),
&headers,
"trace-123",
)
.await
.expect("request context should resolve");
let body_bytes = body.map(|value| Bytes::from(value.to_string()));
maybe_build_local_admin_security_response(state, &request_context, body_bytes.as_ref())
.await
.expect("local security response should build")
.expect("security route should resolve locally")
}
#[tokio::test]
async fn gateway_handles_admin_security_blacklist_add_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::POST,
"/api/admin/security/ip/blacklist",
Some(json!({ "ip_address": "1.2.3.4", "reason": "manual", "ttl": 60 })),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["success"], true);
assert_eq!(payload["message"], "IP 1.2.3.4 已加入黑名单");
assert_eq!(payload["reason"], "manual");
assert_eq!(payload["ttl"], 60);
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn local_admin_security_blacklist_add_attaches_explicit_audit() {
let state = AppState::new().expect("gateway should build");
let response = local_admin_security_response(
&state,
http::Method::POST,
"/api/admin/security/ip/blacklist",
Some(json!({ "ip_address": "1.2.3.4", "reason": "manual", "ttl": 60 })),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("blacklist add should attach audit");
assert_eq!(audit.event_name, "admin_security_blacklist_added");
assert_eq!(audit.action, "add_security_blacklist_entry");
assert_eq!(audit.target_type, "security_blacklist_entry");
assert_eq!(audit.target_id, "1.2.3.4");
}
#[tokio::test]
async fn gateway_handles_admin_security_blacklist_remove_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_security_blacklist_for_tests([(
"1.2.3.4".to_string(),
"manual".to_string(),
)]),
);
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::DELETE,
"/api/admin/security/ip/blacklist/1.2.3.4",
None,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["success"], true);
assert_eq!(payload["message"], "IP 1.2.3.4 已从黑名单移除");
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn gateway_rejects_admin_security_blacklist_remove_without_ip_address() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::DELETE,
"/api/admin/security/ip/blacklist/",
None,
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(payload["detail"], "缺少 ip_address");
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn gateway_handles_admin_security_blacklist_stats_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_security_blacklist_for_tests([
("1.2.3.4".to_string(), "manual".to_string()),
("5.6.7.8".to_string(), "abuse".to_string()),
]),
);
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::GET,
"/api/admin/security/ip/blacklist/stats",
None,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["available"], true);
assert_eq!(payload["total"], 2);
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn local_admin_security_blacklist_list_attaches_explicit_audit() {
let state = AppState::new()
.expect("gateway should build")
.with_admin_security_blacklist_for_tests([
("5.6.7.8".to_string(), "abuse".to_string()),
("1.2.3.4".to_string(), "manual".to_string()),
]);
let response = local_admin_security_response(
&state,
http::Method::GET,
"/api/admin/security/ip/blacklist",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("blacklist list should attach audit");
assert_eq!(audit.event_name, "admin_security_blacklist_viewed");
assert_eq!(audit.action, "view_security_blacklist");
assert_eq!(audit.target_type, "security_blacklist");
assert_eq!(audit.target_id, "global");
}
#[tokio::test]
async fn gateway_handles_admin_security_whitelist_add_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::POST,
"/api/admin/security/ip/whitelist",
Some(json!({ "ip_address": "1.2.3.4" })),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["success"], true);
assert_eq!(payload["message"], "IP 1.2.3.4 已加入白名单");
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn local_admin_security_whitelist_add_attaches_explicit_audit() {
let state = AppState::new().expect("gateway should build");
let response = local_admin_security_response(
&state,
http::Method::POST,
"/api/admin/security/ip/whitelist",
Some(json!({ "ip_address": "1.2.3.4" })),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("whitelist add should attach audit");
assert_eq!(audit.event_name, "admin_security_whitelist_added");
assert_eq!(audit.action, "add_security_whitelist_entry");
assert_eq!(audit.target_type, "security_whitelist_entry");
assert_eq!(audit.target_id, "1.2.3.4");
}
#[tokio::test]
async fn gateway_handles_admin_security_whitelist_remove_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_security_whitelist_for_tests(["1.2.3.4".to_string()]),
);
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::DELETE,
"/api/admin/security/ip/whitelist/1.2.3.4",
None,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["success"], true);
assert_eq!(payload["message"], "IP 1.2.3.4 已从白名单移除");
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn gateway_rejects_admin_security_whitelist_remove_without_ip_address() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::DELETE,
"/api/admin/security/ip/whitelist/",
None,
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(payload["detail"], "缺少 ip_address");
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn gateway_handles_admin_security_whitelist_list_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_security_whitelist_for_tests([
"10.0.0.0/24".to_string(),
"1.2.3.4".to_string(),
]),
);
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::GET,
"/api/admin/security/ip/whitelist",
None,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["whitelist"], json!(["1.2.3.4", "10.0.0.0/24"]));
assert_eq!(payload["total"], 2);
assert_eq!(upstream_count, 0);
}
#[tokio::test]
async fn local_admin_security_whitelist_list_attaches_explicit_audit() {
let state = AppState::new()
.expect("gateway should build")
.with_admin_security_whitelist_for_tests([
"10.0.0.0/24".to_string(),
"1.2.3.4".to_string(),
]);
let response = local_admin_security_response(
&state,
http::Method::GET,
"/api/admin/security/ip/whitelist",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("whitelist list should attach audit");
assert_eq!(audit.event_name, "admin_security_whitelist_viewed");
assert_eq!(audit.action, "view_security_whitelist");
assert_eq!(audit.target_type, "security_whitelist");
assert_eq!(audit.target_id, "global");
}
#[tokio::test]
async fn gateway_handles_admin_security_blacklist_list_locally_with_trusted_admin_principal() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_admin_security_blacklist_for_tests([
("5.6.7.8".to_string(), "abuse".to_string()),
("1.2.3.4".to_string(), "manual".to_string()),
]),
);
let (status, payload, upstream_count) = send_admin_security_request(
gateway,
reqwest::Method::GET,
"/api/admin/security/ip/blacklist",
None,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(payload["total"], 2);
let items = payload["items"].as_array().expect("items array exists");
assert_eq!(items.len(), 2);
assert_eq!(items[0]["ip_address"], "1.2.3.4");
assert_eq!(items[0]["reason"], "manual");
assert_eq!(items[1]["ip_address"], "5.6.7.8");
assert_eq!(items[1]["reason"], "abuse");
assert_eq!(upstream_count, 0);
}

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,944 @@
use std::sync::{Arc, Mutex};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
};
use aether_data::repository::video_tasks::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::response::Response;
use axum::routing::{any, get, post};
use axum::{extract::Request, Json, Router};
use http::{HeaderMap, HeaderValue, StatusCode};
use serde_json::json;
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
sample_key, sample_provider, start_server, AppState,
};
use crate::audit::AdminAuditEvent;
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::control::resolve_public_request_context;
use crate::data::GatewayDataState;
use crate::handlers::admin::maybe_build_local_admin_video_tasks_response;
fn trusted_admin_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
headers.insert(
TRUSTED_ADMIN_USER_ID_HEADER,
HeaderValue::from_static("admin-user-123"),
);
headers.insert(
TRUSTED_ADMIN_USER_ROLE_HEADER,
HeaderValue::from_static("admin"),
);
headers.insert(
TRUSTED_ADMIN_SESSION_ID_HEADER,
HeaderValue::from_static("session-123"),
);
headers.insert(
TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER,
HeaderValue::from_static("management-token-123"),
);
headers
}
async fn local_admin_video_tasks_response(
state: &AppState,
method: http::Method,
uri: &str,
_body: Option<serde_json::Value>,
) -> axum::response::Response<Body> {
let headers = trusted_admin_headers();
let request_context = resolve_public_request_context(
state,
&method,
&uri.parse().expect("uri should parse"),
&headers,
"trace-123",
)
.await
.expect("request context should resolve");
maybe_build_local_admin_video_tasks_response(state, &request_context)
.await
.expect("local video tasks response should build")
.expect("video tasks route should resolve locally")
}
fn sample_admin_video_task(
id: &str,
status: VideoTaskStatus,
created_at_unix_secs: u64,
user_id: &str,
username: &str,
provider_id: &str,
model: &str,
prompt: &str,
) -> UpsertVideoTask {
UpsertVideoTask {
id: id.to_string(),
short_id: Some(format!("short-{id}")),
request_id: format!("request-{id}"),
user_id: Some(user_id.to_string()),
api_key_id: Some(format!("api-key-{id}")),
username: Some(username.to_string()),
api_key_name: Some("primary".to_string()),
external_task_id: Some(format!("ext-{id}")),
provider_id: Some(provider_id.to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
client_api_format: Some("openai:video".to_string()),
provider_api_format: Some("openai:video".to_string()),
format_converted: false,
model: Some(model.to_string()),
prompt: Some(prompt.to_string()),
original_request_body: Some(json!({ "prompt": prompt })),
duration_seconds: Some(4),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("1280x720".to_string()),
status,
progress_percent: if matches!(status, VideoTaskStatus::Completed) {
100
} else {
50
},
progress_message: Some("ok".to_string()),
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: None,
poll_count: 1,
max_poll_count: 360,
created_at_unix_secs,
submitted_at_unix_secs: Some(created_at_unix_secs),
completed_at_unix_secs: if matches!(status, VideoTaskStatus::Completed) {
Some(created_at_unix_secs + 30)
} else {
None
},
updated_at_unix_secs: created_at_unix_secs + 5,
error_code: None,
error_message: None,
video_url: Some(format!("https://example.com/{id}.mp4")),
request_metadata: None,
}
}
#[tokio::test]
async fn gateway_handles_admin_video_tasks_list_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/video-tasks",
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"))
}
}),
);
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-completed",
VideoTaskStatus::Completed,
1_710_000_100,
"user-1",
"alice",
"provider-openai",
"gpt-video",
&"x".repeat(120),
))
.await
.expect("task should upsert");
repository
.upsert(sample_admin_video_task(
"task-processing",
VideoTaskStatus::Processing,
1_710_000_200,
"user-2",
"bob",
"provider-anthropic",
"claude-video",
"short prompt",
))
.await
.expect("task should upsert");
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-openai", "OpenAI", 10),
sample_provider("provider-anthropic", "Anthropic", 20),
],
vec![],
vec![],
));
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
Arc::clone(&repository),
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY,
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
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/video-tasks?status=completed&page=1&page_size=20"
))
.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["total"], json!(1));
assert_eq!(payload["page"], json!(1));
assert_eq!(payload["page_size"], json!(20));
assert_eq!(payload["pages"], json!(1));
assert_eq!(payload["items"].as_array().map(Vec::len), Some(1));
assert_eq!(payload["items"][0]["id"], "task-completed");
assert_eq!(payload["items"][0]["username"], "alice");
assert_eq!(payload["items"][0]["provider_name"], "OpenAI");
assert_eq!(payload["items"][0]["status"], "completed");
assert!(payload["items"][0]["prompt"]
.as_str()
.is_some_and(|value| value.ends_with("...")));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_video_tasks_stats_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/video-tasks/stats",
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"))
}
}),
);
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-processing",
VideoTaskStatus::Processing,
1_710_000_000,
"user-1",
"alice",
"provider-openai",
"gpt-video",
"prompt one",
))
.await
.expect("task should upsert");
repository
.upsert(sample_admin_video_task(
"task-completed",
VideoTaskStatus::Completed,
1_710_000_100,
"user-2",
"bob",
"provider-openai",
"gpt-video",
"prompt two",
))
.await
.expect("task should upsert");
repository
.upsert(sample_admin_video_task(
"task-failed",
VideoTaskStatus::Failed,
1_710_000_200,
"user-2",
"bob",
"provider-openai",
"veo-2",
"prompt three",
))
.await
.expect("task should upsert");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_video_task_data_repository_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/video-tasks/stats"))
.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["total"], json!(3));
assert_eq!(payload["today_count"], json!(0));
assert_eq!(payload["active_users"], json!(2));
assert_eq!(payload["processing_count"], json!(1));
assert_eq!(payload["by_status"]["processing"], json!(1));
assert_eq!(payload["by_status"]["completed"], json!(1));
assert_eq!(payload["by_status"]["failed"], json!(1));
assert_eq!(payload["by_model"]["gpt-video"], json!(2));
assert_eq!(payload["by_model"]["veo-2"], json!(1));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_video_task_detail_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/video-tasks/task-detail",
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"))
}
}),
);
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-detail",
VideoTaskStatus::Completed,
1_710_000_300,
"user-9",
"charlie",
"provider-openai",
"gpt-video",
"detail prompt",
))
.await
.expect("task should upsert");
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "OpenAI", 10)],
vec![sample_endpoint(
"endpoint-1",
"provider-openai",
"openai:video",
"https://api.openai.example",
)],
vec![],
));
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
Arc::clone(&repository),
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY,
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
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/video-tasks/task-detail"))
.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["id"], "task-detail");
assert_eq!(payload["username"], "charlie");
assert_eq!(payload["provider_name"], "OpenAI");
assert_eq!(payload["endpoint"]["id"], "endpoint-1");
assert_eq!(payload["endpoint"]["api_format"], "openai:video");
assert_eq!(payload["status"], "completed");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn local_admin_video_task_detail_attaches_explicit_audit() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-detail-audit",
VideoTaskStatus::Completed,
1_710_000_350,
"user-3",
"dana",
"provider-openai",
"gpt-video",
"detail audit prompt",
))
.await
.expect("task should upsert");
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "OpenAI", 10)],
vec![sample_endpoint(
"endpoint-1",
"provider-openai",
"openai:video",
"https://api.openai.example",
)],
vec![],
));
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
Arc::clone(&repository),
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let response = local_admin_video_tasks_response(
&state,
http::Method::GET,
"/api/admin/video-tasks/task-detail-audit",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("video task detail should attach audit");
assert_eq!(audit.event_name, "admin_video_task_detail_viewed");
assert_eq!(audit.action, "view_video_task_detail");
assert_eq!(audit.target_type, "video_task");
assert_eq!(audit.target_id, "task-detail-audit");
}
#[tokio::test]
async fn gateway_cancels_admin_video_task_locally_with_trusted_admin_principal() {
#[derive(Debug, Clone, PartialEq, Eq)]
struct SeenExecutionRuntimeSyncRequest {
method: String,
url: String,
authorization: String,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/video-tasks/task-openai-cancel/cancel",
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"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).expect("payload should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
method: payload
.get("method")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
url: payload
.get("url")
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
});
Json(json!({
"request_id": "cancel-openai-task-1",
"status_code": 204,
"headers": {},
"telemetry": {
"elapsed_ms": 11
}
}))
}
}),
);
let repository = Arc::new(InMemoryVideoTaskRepository::default());
let mut task = sample_admin_video_task(
"task-openai-cancel",
VideoTaskStatus::Processing,
1_710_000_400,
"user-1",
"alice",
"provider-openai",
"sora-2",
"cancel prompt",
);
task.external_task_id = Some("ext-video-task-123".to_string());
task.request_metadata = Some(json!({
"rust_local_snapshot": {
"OpenAi": {
"local_task_id": "task-openai-cancel",
"upstream_task_id": "ext-video-task-123",
"created_at_unix_secs": 1710000400,
"user_id": "user-1",
"api_key_id": "api-key-task-openai-cancel",
"model": "sora-2",
"prompt": "cancel prompt",
"size": "1280x720",
"seconds": "4",
"remixed_from_video_id": null,
"status": "Processing",
"progress_percent": 50,
"completed_at_unix_secs": null,
"expires_at_unix_secs": null,
"error_code": null,
"error_message": null,
"video_url": null,
"persistence": {
"request_id": "request-task-openai-cancel",
"username": "alice",
"api_key_name": "primary",
"client_api_format": "openai:video",
"provider_api_format": "openai:video",
"original_request_body": {
"prompt": "cancel prompt"
},
"format_converted": false
},
"transport": {
"upstream_base_url": "https://api.openai.example",
"provider_name": "openai-video",
"provider_id": "provider-openai",
"endpoint_id": "endpoint-1",
"key_id": "provider-key-1",
"headers": {
"authorization": "Bearer sk-upstream-openai-video",
"content-type": "application/json"
},
"content_type": "application/json",
"model_name": "sora-2-upstream",
"proxy": null,
"tls_profile": null,
"timeouts": null
}
}
}
}));
repository
.upsert(task)
.await
.expect("upsert should succeed");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_video_task_data_repository_for_tests(Arc::clone(&repository)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/video-tasks/task-openai-cancel/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);
assert_eq!(
response
.json::<serde_json::Value>()
.await
.expect("body should parse"),
json!({
"id": "task-openai-cancel",
"status": "cancelled",
"message": "Task cancelled successfully",
})
);
let detail = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/video-tasks/task-openai-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("detail request should succeed");
assert_eq!(detail.status(), StatusCode::OK);
let detail_json: serde_json::Value = detail.json().await.expect("detail should parse");
assert_eq!(detail_json["status"], "cancelled");
assert_eq!(
detail_json["request_metadata"]["rust_local_snapshot"]["OpenAi"]["status"],
"Cancelled"
);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime should be called");
assert_eq!(seen_execution_runtime_request.method, "DELETE");
assert_eq!(
seen_execution_runtime_request.url,
"https://api.openai.example/v1/videos/ext-video-task-123"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-upstream-openai-video"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn local_admin_video_task_cancel_attaches_explicit_audit() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
let mut task = sample_admin_video_task(
"task-cancel-audit",
VideoTaskStatus::Processing,
1_710_000_450,
"user-4",
"erin",
"provider-openai",
"gpt-video",
"cancel audit prompt",
);
task.client_api_format = None;
task.provider_api_format = None;
task.request_metadata = Some(json!({}));
repository.upsert(task).await.expect("task should upsert");
let state = AppState::new()
.expect("gateway state should build")
.with_video_task_data_repository_for_tests(repository);
let response = local_admin_video_tasks_response(
&state,
http::Method::POST,
"/api/admin/video-tasks/task-cancel-audit/cancel",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("video task cancel should attach audit");
assert_eq!(audit.event_name, "admin_video_task_cancelled");
assert_eq!(audit.action, "cancel_video_task");
assert_eq!(audit.target_type, "video_task");
assert_eq!(audit.target_id, "task-cancel-audit");
}
#[tokio::test]
async fn gateway_redirects_admin_video_task_video_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/video-tasks/task-redirect/video",
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"))
}
}),
);
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-redirect",
VideoTaskStatus::Completed,
1_710_000_500,
"user-1",
"alice",
"provider-openai",
"gpt-video",
"redirect prompt",
))
.await
.expect("task should upsert");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_video_task_data_repository_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("client should build");
let response = client
.get(format!(
"{gateway_url}/api/admin/video-tasks/task-redirect/video"
))
.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::TEMPORARY_REDIRECT);
assert_eq!(
response
.headers()
.get(http::header::LOCATION)
.and_then(|value| value.to_str().ok()),
Some("https://example.com/task-redirect.mp4")
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn local_admin_video_task_video_redirect_attaches_explicit_audit() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(sample_admin_video_task(
"task-video-audit",
VideoTaskStatus::Completed,
1_710_000_550,
"user-5",
"frank",
"provider-openai",
"gpt-video",
"video audit prompt",
))
.await
.expect("task should upsert");
let state = AppState::new()
.expect("gateway state should build")
.with_video_task_data_repository_for_tests(repository);
let response = local_admin_video_tasks_response(
&state,
http::Method::GET,
"/api/admin/video-tasks/task-video-audit/video",
None,
)
.await;
assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
let audit = response
.extensions()
.get::<AdminAuditEvent>()
.cloned()
.expect("video task video should attach audit");
assert_eq!(audit.event_name, "admin_video_task_video_viewed");
assert_eq!(audit.action, "view_video_task_video");
assert_eq!(audit.target_type, "video_task_video");
assert_eq!(audit.target_id, "task-video-audit");
}
#[tokio::test]
async fn gateway_proxies_admin_video_task_video_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let seen_api_key = Arc::new(Mutex::new(None::<String>));
let seen_api_key_clone = Arc::clone(&seen_api_key);
let upstream = Router::new()
.route(
"/api/admin/video-tasks/task-proxy/video",
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(
"/generativelanguage.googleapis.com/v1beta/files/video-task-123:download",
any(move |request: Request| {
let seen_api_key_inner = Arc::clone(&seen_api_key_clone);
async move {
*seen_api_key_inner.lock().expect("mutex should lock") = request
.headers()
.get("x-goog-api-key")
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from("proxied-video-bytes"))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("video/mp4"),
);
response
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let repository = Arc::new(InMemoryVideoTaskRepository::default());
let mut task = sample_admin_video_task(
"task-proxy",
VideoTaskStatus::Completed,
1_710_000_600,
"user-1",
"alice",
"provider-gemini",
"veo-3",
"proxy prompt",
);
task.provider_id = Some("provider-gemini".to_string());
task.endpoint_id = Some("endpoint-gemini".to_string());
task.key_id = Some("key-gemini".to_string());
task.client_api_format = Some("gemini:video".to_string());
task.provider_api_format = Some("gemini:video".to_string());
task.video_url = Some(format!(
"{upstream_url}/generativelanguage.googleapis.com/v1beta/files/video-task-123:download"
));
repository.upsert(task).await.expect("task should upsert");
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-gemini", "Gemini", 10)],
vec![sample_endpoint(
"endpoint-gemini",
"provider-gemini",
"gemini:video",
"https://generativelanguage.googleapis.com",
)],
vec![sample_key(
"key-gemini",
"provider-gemini",
"gemini:video",
"gemini-upstream-secret",
)],
));
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
Arc::clone(&repository),
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY,
);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state 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/video-tasks/task-proxy/video"
))
.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);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("video/mp4")
);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_DISPOSITION)
.and_then(|value| value.to_str().ok()),
Some("inline; filename=\"video_task-proxy.mp4\"")
);
assert_eq!(
response.bytes().await.expect("body should read"),
Bytes::from_static(b"proxied-video-bytes")
);
assert_eq!(
seen_api_key.lock().expect("mutex should lock").as_deref(),
Some("gemini-upstream-secret")
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,569 @@
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use chrono::{DateTime, Utc};
use hmac::Mac;
use serde_json::{json, Map, Value};
use sha2::{Digest, Sha256};
use super::{
RequestCandidateStatus, StoredAdminGlobalModel, StoredAdminProviderModel,
StoredAuthApiKeySnapshot, StoredLdapModuleConfig, StoredManagementToken,
StoredManagementTokenUserSummary, StoredManagementTokenWithUser, StoredOAuthProviderConfig,
StoredOAuthProviderModuleConfig, StoredProviderActiveGlobalModel,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
StoredProviderModelStats, StoredProviderQuotaSnapshot, StoredProxyNode,
StoredPublicGlobalModel, StoredRequestCandidate,
};
use crate::AppState;
pub(super) fn sample_currently_usable_auth_snapshot(
api_key_id: &str,
user_id: &str,
) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-5"])),
)
.expect("auth snapshot should build")
}
pub(super) fn sample_expired_auth_snapshot(
api_key_id: &str,
user_id: &str,
) -> StoredAuthApiKeySnapshot {
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
snapshot.api_key_expires_at_unix_secs = Some(1);
snapshot
}
pub(super) fn sample_locked_auth_snapshot(
api_key_id: &str,
user_id: &str,
) -> StoredAuthApiKeySnapshot {
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
snapshot.api_key_is_locked = true;
snapshot
}
pub(super) fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
pub(super) fn test_auth_secret() -> String {
std::env::var("JWT_SECRET_KEY")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "aether-rust-dev-jwt-secret".to_string())
}
pub(super) fn build_test_auth_token(
token_type: &str,
mut payload: serde_json::Map<String, serde_json::Value>,
expires_at: chrono::DateTime<chrono::Utc>,
) -> String {
use base64::Engine as _;
let header = json!({ "alg": "HS256", "typ": "JWT" });
payload.insert("exp".to_string(), json!(expires_at.timestamp()));
payload.insert("type".to_string(), json!(token_type));
let header_segment = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
serde_json::to_vec(&header)
.expect("jwt header should serialize")
.as_slice(),
);
let payload_segment = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
serde_json::to_vec(&payload)
.expect("jwt payload should serialize")
.as_slice(),
);
let signing_input = format!("{header_segment}.{payload_segment}");
let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(test_auth_secret().as_bytes())
.expect("jwt secret should build");
mac.update(signing_input.as_bytes());
let signature = mac.finalize().into_bytes();
format!(
"{header_segment}.{payload_segment}.{}",
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(signature.as_slice())
)
}
pub(super) async fn issue_test_admin_access_token(
state: &AppState,
client_device_id: &str,
) -> String {
let user = state
.create_local_auth_user_with_settings(
Some("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 now = chrono::Utc::now();
let session_id = "session-admin-token".to_string();
let refresh_token = "refresh-session-admin-token".to_string();
let session = crate::data::state::StoredUserSessionRecord::new(
session_id.clone(),
user.id.clone(),
client_device_id.to_string(),
None,
crate::data::state::StoredUserSessionRecord::hash_refresh_token(
&refresh_token,
),
None,
None,
Some(now),
Some(now + chrono::Duration::days(7)),
None,
None,
Some("127.0.0.1".to_string()),
Some("admin-test".to_string()),
Some(now),
Some(now),
)
.expect("session should build");
state
.create_user_session(session)
.await
.expect("session should persist")
.expect("session should exist");
build_test_auth_token(
"access",
serde_json::Map::from_iter([
("user_id".to_string(), json!(user.id)),
("role".to_string(), json!("admin")),
(
"created_at".to_string(),
json!(user.created_at.map(|value| value.to_rfc3339())),
),
("session_id".to_string(), json!(session_id)),
]),
now + chrono::Duration::hours(12),
)
}
pub(super) fn sample_provider(
id: &str,
name: &str,
priority: i32,
) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
id.to_string(),
name.to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_routing_fields(priority)
}
pub(super) fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
StoredProxyNode::new(
node_id.to_string(),
"proxy-node-1".to_string(),
"127.0.0.1".to_string(),
0,
false,
"offline".to_string(),
30,
0,
0,
0,
0,
0,
true,
false,
7,
)
.expect("proxy node should build")
.with_runtime_fields(
Some("test".to_string()),
Some("system".to_string()),
Some(1_710_000_000),
None,
None,
None,
None,
Some(1_710_000_010),
Some(json!({
"upgrade_to": "1.2.3",
"allowed_ports": [443],
})),
Some(1_709_000_000),
Some(1_710_000_100),
)
}
pub(super) fn sample_provider_quota(provider_id: &str) -> StoredProviderQuotaSnapshot {
StoredProviderQuotaSnapshot::new(
provider_id.to_string(),
"monthly_quota".to_string(),
Some(100.0),
12.5,
Some(30),
Some(1_711_000_000),
Some(1_711_000_000 + 30 * 24 * 60 * 60),
true,
)
.expect("provider quota should build")
}
pub(super) fn sample_provider_model_stats(
provider_id: &str,
total_models: i64,
active_models: i64,
) -> StoredProviderModelStats {
StoredProviderModelStats::new(provider_id.to_string(), total_models, active_models)
.expect("provider model stats should build")
}
pub(super) fn sample_provider_active_global_model(
provider_id: &str,
global_model_id: &str,
) -> StoredProviderActiveGlobalModel {
StoredProviderActiveGlobalModel::new(provider_id.to_string(), global_model_id.to_string())
.expect("provider active global model should build")
}
pub(super) fn sample_management_token(
token_id: &str,
user_id: &str,
username: &str,
is_active: bool,
) -> StoredManagementTokenWithUser {
let token = StoredManagementToken::new(
token_id.to_string(),
user_id.to_string(),
format!("{username}-token"),
)
.expect("management token should build")
.with_display_fields(
Some(format!("{username} token")),
Some("ae_test".to_string()),
Some(json!(["127.0.0.1"])),
)
.with_runtime_fields(
Some(4_102_444_800),
Some(1_711_000_000),
Some("127.0.0.1".to_string()),
7,
is_active,
)
.with_timestamps(Some(1_710_000_000), Some(1_711_000_100));
let user = StoredManagementTokenUserSummary::new(
user_id.to_string(),
Some(format!("{username}@example.com")),
username.to_string(),
"admin".to_string(),
)
.expect("management token user should build");
StoredManagementTokenWithUser::new(token, user)
}
pub(super) fn sample_admin_provider_model(
id: &str,
provider_id: &str,
global_model_id: &str,
provider_model_name: &str,
) -> StoredAdminProviderModel {
StoredAdminProviderModel::new(
id.to_string(),
provider_id.to_string(),
global_model_id.to_string(),
provider_model_name.to_string(),
Some(json!([{"name": format!("{provider_model_name}-alias"), "priority": 1}])),
Some(0.02),
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 3.0,
"output_price_per_1m": 15.0,
}]
})),
Some(true),
Some(true),
None,
Some(false),
Some(false),
true,
true,
Some(json!({"billing": {"mode": "local"}, "provider_hint": provider_model_name})),
Some(1_711_000_000),
Some(1_711_000_100),
Some("gpt-5".to_string()),
Some("GPT 5".to_string()),
Some(0.03),
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 4.0,
"output_price_per_1m": 20.0,
}]
})),
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
)
.expect("admin provider model should build")
}
pub(super) fn sample_admin_global_model(
id: &str,
name: &str,
display_name: &str,
) -> StoredAdminGlobalModel {
StoredAdminGlobalModel::new(
id.to_string(),
name.to_string(),
display_name.to_string(),
true,
Some(0.03),
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 4.0,
"output_price_per_1m": 20.0,
}]
})),
Some(json!(["streaming", "vision"])),
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
Some(1_711_000_000),
Some(1_711_000_100),
)
.expect("admin global model should build")
}
pub(super) fn sample_public_global_model_with_mappings(
id: &str,
name: &str,
display_name: &str,
mappings: &[&str],
) -> StoredPublicGlobalModel {
StoredPublicGlobalModel::new(
id.to_string(),
name.to_string(),
Some(display_name.to_string()),
true,
None,
None,
None,
Some(json!({ "model_mappings": mappings })),
0,
)
.expect("public global model should build")
}
pub(super) fn sample_oauth_module_provider(
provider_type: &str,
display_name: &str,
) -> StoredOAuthProviderModuleConfig {
StoredOAuthProviderModuleConfig::new(
provider_type.to_string(),
display_name.to_string(),
"client-id".to_string(),
Some("encrypted-secret".to_string()),
"https://example.com/oauth/callback".to_string(),
)
.expect("oauth module provider should build")
}
pub(super) fn sample_ldap_module_config() -> StoredLdapModuleConfig {
StoredLdapModuleConfig {
server_url: "ldaps://ldap.example.com".to_string(),
bind_dn: "cn=admin,dc=example,dc=com".to_string(),
bind_password_encrypted: Some("encrypted-password".to_string()),
base_dn: "dc=example,dc=com".to_string(),
user_search_filter: Some("(uid={username})".to_string()),
username_attr: Some("uid".to_string()),
email_attr: Some("mail".to_string()),
display_name_attr: Some("displayName".to_string()),
is_enabled: true,
is_exclusive: false,
use_starttls: true,
connect_timeout: Some(10),
}
}
pub(super) fn sample_oauth_provider_config(provider_type: &str) -> StoredOAuthProviderConfig {
StoredOAuthProviderConfig::new(
provider_type.to_string(),
"Linux Do".to_string(),
"client-id".to_string(),
"https://backend.example.com/oauth/callback".to_string(),
"https://frontend.example.com/auth/callback".to_string(),
)
.expect("oauth provider config should build")
.with_config_fields(
Some(
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "secret-value")
.expect("secret should encrypt"),
),
Some("https://connect.linux.do/oauth2/authorize".to_string()),
Some("https://connect.linux.do/oauth2/token".to_string()),
Some("https://connect.linux.do/api/user".to_string()),
Some(vec!["openid".to_string()]),
Some(json!({"email": "email"})),
Some(json!({"team": true})),
true,
)
}
pub(super) fn sample_endpoint(
id: &str,
provider_id: &str,
api_format: &str,
base_url: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
id.to_string(),
provider_id.to_string(),
api_format.to_string(),
None,
None,
true,
)
.expect("endpoint should build")
.with_health_score(0.9)
.with_transport_fields(
base_url.to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
pub(super) fn sample_key(
id: &str,
provider_id: &str,
api_format: &str,
secret: &str,
) -> StoredProviderCatalogKey {
let encrypted_api_key = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)
.expect("api key ciphertext should build");
StoredProviderCatalogKey::new(
id.to_string(),
provider_id.to_string(),
"default".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!([api_format])),
encrypted_api_key,
None,
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build")
}
pub(super) fn sample_request_candidate(
id: &str,
request_id: &str,
endpoint_id: &str,
status: RequestCandidateStatus,
created_at_unix_secs: i64,
finished_at_unix_secs: Option<i64>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some(endpoint_id.to_string()),
Some("key-1".to_string()),
status,
None,
false,
Some(200),
matches!(status, RequestCandidateStatus::Failed).then_some("rate_limit".to_string()),
None,
Some(120),
Some(1),
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
finished_at_unix_secs,
)
.expect("request candidate should build")
}
pub(super) fn sample_recent_key_rpm_candidate(
id: &str,
request_id: &str,
endpoint_id: &str,
key_id: &str,
now_unix_secs: i64,
concurrent_requests: i32,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some(endpoint_id.to_string()),
Some(key_id.to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(120),
Some(concurrent_requests),
None,
None,
now_unix_secs - 10,
Some(now_unix_secs - 10),
Some(now_unix_secs - 8),
)
.expect("request candidate should build")
}

View File

@@ -0,0 +1,312 @@
use std::sync::{Arc, Mutex};
use axum::body::Body;
use axum::routing::{any, post};
use axum::{extract::Request, Json, Router};
use http::header::HeaderValue;
use http::StatusCode;
use serde_json::json;
use super::{
build_router_with_state, sample_proxy_node, start_server, AppState, GatewayDataState,
InMemoryProxyNodeRepository, TRACE_ID_HEADER,
};
#[tokio::test]
async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/internal/tunnel/heartbeat",
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"))
}
}),
);
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-123",
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
Arc::clone(&repository),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
.json(&json!({
"node_id": "node-123",
"heartbeat_interval": 45,
"active_connections": 5,
"total_requests": 9,
"avg_latency_ms": 12.5,
"failed_requests": 1,
"dns_failures": 2,
"stream_errors": 3,
"proxy_metadata": {"arch": "arm64"},
"proxy_version": "2.0.0",
}))
.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["config_version"], 7);
assert_eq!(payload["upgrade_to"], "1.2.3");
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_internal_tunnel_node_status_locally_with_loopback() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/internal/tunnel/node-status",
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"))
}
}),
);
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
"node-123",
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
Arc::clone(&repository),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/node-status"))
.json(&json!({
"node_id": "node-123",
"connected": true,
"conn_count": 4,
}))
.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["updated"], json!(true));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_owns_proxy_tunnel_path_without_proxying_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/internal/proxy-tunnel",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/internal/proxy-tunnel"))
.header("x-node-id", "node-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_internal_tunnel_relay_locally_without_proxying_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/internal/tunnel/relay/node-123",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
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(Vec::<u8>::new())
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_forwards_tunnel_relay_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 |headers: axum::http::HeaderMap, body: Body| {
let owner_hits_inner = Arc::clone(&owner_hits_clone);
async move {
*owner_hits_inner.lock().expect("mutex should lock") += 1;
assert_eq!(
headers
.get(aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gateway-a")
);
let body = axum::body::to_bytes(body, usize::MAX)
.await
.expect("body should read");
let mut response = axum::http::Response::builder()
.status(StatusCode::OK)
.body(Body::from(body))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/octet-stream"),
);
response
}
}),
);
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 response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.header(TRACE_ID_HEADER, "trace-owner-forward")
.body("relay-envelope")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok()),
Some("trace-owner-forward")
);
assert_eq!(
response.text().await.expect("body should read"),
"relay-envelope"
);
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));
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,
}),
)]);
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"))
.header(
aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER,
"gateway-z",
)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
owner_handle.abort();
}

View File

@@ -0,0 +1,61 @@
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
InMemoryAuthModuleReadRepository, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig,
};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::global_models::{
GlobalModelReadRepository, InMemoryGlobalModelReadRepository, StoredAdminGlobalModel,
StoredAdminProviderModel, StoredProviderActiveGlobalModel, StoredProviderModelStats,
StoredPublicGlobalModel,
};
use aether_data::repository::management_tokens::{
InMemoryManagementTokenRepository, ManagementTokenReadRepository, StoredManagementToken,
StoredManagementTokenUserSummary, StoredManagementTokenWithUser,
};
use aether_data::repository::oauth_providers::{
InMemoryOAuthProviderRepository, OAuthProviderReadRepository, StoredOAuthProviderConfig,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use aether_data::repository::quota::{
InMemoryProviderQuotaRepository, StoredProviderQuotaSnapshot,
};
use aether_data::repository::wallet::InMemoryWalletRepository;
use axum::body::{to_bytes, Body, Bytes};
use axum::response::Response;
use axum::routing::{any, post};
use axum::{extract::Request, Json, Router};
use http::header::{HeaderName, HeaderValue};
use http::{HeaderMap, StatusCode};
use serde_json::json;
mod admin;
mod helpers;
mod internal;
mod proxy;
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
Infallible, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
use crate::constants::*;
use crate::data::GatewayDataState;
use helpers::*;

View File

@@ -0,0 +1,801 @@
use std::sync::{Arc, Mutex};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::json;
use super::super::{
build_router_with_state, hash_api_key, sample_currently_usable_auth_snapshot,
sample_expired_auth_snapshot, sample_locked_auth_snapshot, start_server, AppState,
GatewayDataState, InMemoryAuthApiKeySnapshotRepository, InMemoryWalletRepository,
};
use crate::constants::{
CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_AUTH_DENIED,
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER,
TRUSTED_AUTH_API_KEY_ID_HEADER, TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER,
};
#[tokio::test]
async fn gateway_locally_denies_explicit_trusted_balance_failure_without_hitting_control_or_upstream(
) {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"auth_context": {
"user_id": "user-from-control",
"api_key_id": "key-from-control",
"balance_remaining": 99.0,
"access_allowed": true
}
}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_currently_usable_auth_snapshot("key-123", "user-123"),
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-control-balance-denied-1")
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
.header(TRUSTED_AUTH_BALANCE_HEADER, "0")
.header(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "false")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("ai_public")
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "balance_exceeded");
assert_eq!(payload["error"]["message"], "余额不足(剩余: $0.00");
assert_eq!(payload["error"]["details"]["balance_type"], "USD");
assert_eq!(payload["error"]["details"]["remaining"], 0.0);
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_invalid_trusted_snapshot_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"auth_context": {
"user_id": "user-from-control",
"api_key_id": "key-from-control",
"access_allowed": true
}
}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_expired_auth_snapshot("key-123", "user-123"),
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-control-invalid-trusted-1")
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "无效的API密钥");
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_missing_wallet_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_currently_usable_auth_snapshot("key-123", "user-123"),
)]));
let wallet_repository = Arc::new(InMemoryWalletRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let data_state =
GatewayDataState::with_auth_and_wallet_for_tests(auth_repository, wallet_repository);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-control-wallet-missing-1")
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "钱包不可用");
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-other")),
sample_currently_usable_auth_snapshot("key-123", "user-123"),
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(http::header::AUTHORIZATION, "Bearer sk-missing")
.header(TRACE_ID_HEADER, "trace-control-invalid-bearer-1")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "无效的API密钥");
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_admin_proxy_without_admin_principal_and_without_hitting_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/health/summary",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/endpoints/health/summary"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["detail"], "admin authentication required");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_unclassified_oauth_public_route_as_local_not_found_without_hitting_upstream(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/oauth/linuxdo/unsupported",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/oauth/linuxdo/unsupported"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "Route not found");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_rejects_unclassified_admin_route_without_hitting_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/unsupported",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/unsupported"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(
payload["detail"],
"admin proxy route not implemented in rust frontdoor"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_disallowed_claude_api_format_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/messages",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut snapshot = sample_currently_usable_auth_snapshot("key-claude-123", "user-claude-123");
snapshot.api_key_allowed_providers = Some(vec!["claude".to_string()]);
snapshot.user_allowed_providers = Some(vec!["claude".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-claude-auth-123")),
snapshot,
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages"))
.header(http::header::CONTENT_TYPE, "application/json")
.header("x-api-key", "sk-claude-auth-123")
.header(TRACE_ID_HEADER, "trace-control-claude-format-1")
.body("{\"model\":\"claude-3-7-sonnet\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(
payload["error"]["message"],
"当前密钥不允许访问 claude:chat 格式"
);
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_disallowed_provider_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/messages",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut snapshot =
sample_currently_usable_auth_snapshot("key-claude-allow-123", "user-claude-allow-123");
snapshot.api_key_allowed_api_formats = None;
snapshot.user_allowed_api_formats = None;
snapshot.api_key_allowed_providers = Some(vec!["openai".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-claude-provider-123")),
snapshot,
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/messages"))
.header(http::header::CONTENT_TYPE, "application/json")
.header("x-api-key", "sk-claude-provider-123")
.header(TRACE_ID_HEADER, "trace-control-provider-guard-1")
.body("{\"model\":\"claude-3-7-sonnet\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(
payload["error"]["message"],
"当前密钥不允许访问 claude 提供商"
);
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_disallowed_gemini_model_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1beta/models/{*path}",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut snapshot = sample_currently_usable_auth_snapshot("key-gemini-123", "user-gemini-123");
snapshot.api_key_allowed_providers = Some(vec!["gemini".to_string()]);
snapshot.user_allowed_providers = Some(vec!["gemini".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["gemini:chat".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["gemini-1.5-pro".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("gemini-client-key-123")),
snapshot,
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/models/gemini-2.5-pro:generateContent"
))
.header(http::header::CONTENT_TYPE, "application/json")
.header("x-goog-api-key", "gemini-client-key-123")
.header(TRACE_ID_HEADER, "trace-control-gemini-model-1")
.body("{\"contents\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(
payload["error"]["message"],
"当前密钥不允许访问模型 gemini-2.5-pro"
);
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_locked_trusted_snapshot_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_locked_auth_snapshot("key-locked-123", "user-locked-123"),
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(TRACE_ID_HEADER, "trace-control-locked-trusted-1")
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-locked-123")
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-locked-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(
payload["error"]["message"],
"该密钥已被管理员锁定,请联系管理员"
);
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_locally_denies_disallowed_openai_model_without_hitting_control_or_upstream() {
let auth_context_hits = Arc::new(Mutex::new(0usize));
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
let public_hits = Arc::new(Mutex::new(0usize));
let public_hits_clone = Arc::clone(&public_hits);
let upstream = Router::new()
.route(
"/api/internal/gateway/auth-context",
any(move |_request: Request| {
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
async move {
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({"auth_context": null}))
}
}),
)
.route(
"/v1/chat/completions",
any(move |_request: Request| {
let public_hits_inner = Arc::clone(&public_hits_clone);
async move {
*public_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut snapshot =
sample_currently_usable_auth_snapshot("key-openai-model-123", "user-openai-model-123");
snapshot.api_key_allowed_models = Some(vec!["gpt-4.1".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-openai-model-guard-123")),
snapshot,
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-openai-model-guard-123",
)
.header(TRACE_ID_HEADER, "trace-control-openai-model-guard-1")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("response json should parse");
assert_eq!(payload["error"]["type"], "http_error");
assert_eq!(payload["error"]["message"], "当前密钥不允许访问模型 gpt-5");
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}

View File

@@ -0,0 +1,2 @@
mod local_denials;
mod routing;

View File

@@ -0,0 +1,152 @@
use std::sync::{Arc, Mutex};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use super::super::{
build_router, build_router_with_state, hash_api_key, sample_currently_usable_auth_snapshot,
start_server, AppState, InMemoryAuthApiKeySnapshotRepository,
};
use crate::constants::{
CONTROL_ENDPOINT_SIGNATURE_HEADER, CONTROL_EXECUTION_RUNTIME_HEADER,
CONTROL_ROUTE_CLASS_HEADER, CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER,
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
};
#[tokio::test]
async fn gateway_rejects_spoofed_admin_principal_headers_without_gateway_marker_locally() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/health/api-formats",
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"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router().expect("gateway should build");
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/health/api-formats"
))
.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::UNAUTHORIZED);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["detail"], "admin authentication required");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_models_support_routes_locally_with_public_support_control_headers() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/v1/models",
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,
[(GATEWAY_HEADER, "fallback-probe")],
Body::from("{\"object\":\"list\",\"data\":[]}"),
)
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-test-models")),
sample_currently_usable_auth_snapshot("key-models-1", "user-models-1"),
)]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(auth_repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/v1/models?limit=20"))
.header(http::header::AUTHORIZATION, "Bearer sk-test-models")
.header(TRACE_ID_HEADER, "trace-models-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok()),
Some("public_support")
);
assert_eq!(
response
.headers()
.get(CONTROL_EXECUTION_RUNTIME_HEADER)
.and_then(|value| value.to_str().ok()),
Some("false")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_FAMILY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("models")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("list")
);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:chat")
);
assert_eq!(
response
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok()),
Some("trace-models-123")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["detail"],
"public support route not implemented in rust frontdoor"
);
assert_eq!(payload["route_family"], "models");
assert_eq!(payload["route_kind"], "list");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}