mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
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
@@ -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();
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
mod local_denials;
|
||||
mod routing;
|
||||
@@ -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();
|
||||
}
|
||||
Reference in New Issue
Block a user