mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-05 00:47:48 +08:00
feat(gateway): harden provider request execution
Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
This commit is contained in:
@@ -1604,7 +1604,7 @@ async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_b
|
||||
(
|
||||
"provider-vertex-ai-reconcile",
|
||||
"vertex_ai",
|
||||
3usize,
|
||||
2usize,
|
||||
"gemini:generate_content",
|
||||
"https://aiplatform.googleapis.com",
|
||||
"Vertex AI 暂不支持自动刷新额度",
|
||||
|
||||
@@ -48,6 +48,45 @@ async fn gateway_blocks_blacklisted_ip_before_routing() {
|
||||
assert_eq!(payload["error"]["message"], "当前 IP 已被禁止访问");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_shapes_blacklist_rejections_for_claude_routes_before_routing() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_blacklist_for_tests([(
|
||||
"127.0.0.1".to_string(),
|
||||
"blocked".to_string(),
|
||||
)]),
|
||||
);
|
||||
|
||||
for path in ["/v1/messages", "/v1/messages/count_tokens"] {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri(path)
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.body(Body::from(r#"{"model":"claude-sonnet-4","messages":[]}"#))
|
||||
.expect("request should build");
|
||||
|
||||
let response = send_request(gateway.clone(), request).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN, "path: {path}");
|
||||
let payload = response
|
||||
.into_body()
|
||||
.collect()
|
||||
.await
|
||||
.expect("body should collect")
|
||||
.to_bytes();
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&payload).expect("response should be json");
|
||||
assert_eq!(payload["type"], "error", "path: {path}");
|
||||
assert_eq!(payload["error"]["type"], "permission_error", "path: {path}");
|
||||
assert_eq!(
|
||||
payload["error"]["message"], "当前 IP 已被禁止访问",
|
||||
"path: {path}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_blocks_forwarded_ip_from_trusted_proxy() {
|
||||
let gateway = build_router_with_state(
|
||||
|
||||
@@ -1758,6 +1758,144 @@ async fn gateway_reports_field_path_for_invalid_admin_system_config_import_shape
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import() {
|
||||
run_admin_system_import_test(
|
||||
"gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import",
|
||||
gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_rejects_invalid_anthropic_profiles_during_admin_system_config_import_impl() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(build_empty_admin_system_data_state()),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
for config_scope in ["provider", "endpoint"] {
|
||||
let mut payload = sample_system_import_payload();
|
||||
let invalid_config = json!({
|
||||
"anthropic": {"compatibility_profile": "claude_cod_typo"}
|
||||
});
|
||||
if config_scope == "provider" {
|
||||
payload["providers"][0]["config"] = invalid_config;
|
||||
} else {
|
||||
payload["providers"][0]["endpoints"][0]["config"] = invalid_config;
|
||||
}
|
||||
|
||||
let response = client
|
||||
.post(format!("{gateway_url}/api/admin/system/config/import"))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.expect("invalid Anthropic profile import should complete locally");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let body: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(
|
||||
body["detail"], "无效的 Anthropic compatibility profile",
|
||||
"unexpected {config_scope} validation response: {body}"
|
||||
);
|
||||
}
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import() {
|
||||
run_admin_system_import_test(
|
||||
"gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import",
|
||||
gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_does_not_restore_retired_vertex_claude_endpoint_from_system_import_impl() {
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::seed(Vec::<
|
||||
StoredPublicGlobalModel,
|
||||
>::new()));
|
||||
let data_state = build_admin_system_data_state_with_repositories(
|
||||
Arc::clone(&provider_catalog_repository),
|
||||
Arc::clone(&global_model_repository),
|
||||
);
|
||||
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 payload = sample_system_import_payload();
|
||||
payload["providers"][0]["name"] = json!("legacy-vertex-backup");
|
||||
payload["providers"][0]["provider_type"] = json!("vertex_ai");
|
||||
payload["providers"][0]["endpoints"] = json!([
|
||||
{
|
||||
"api_format": "gemini:generate_content",
|
||||
"base_url": "https://aiplatform.googleapis.com",
|
||||
"max_retries": 2,
|
||||
"is_active": true
|
||||
},
|
||||
{
|
||||
"api_format": "claude:messages",
|
||||
"base_url": "https://aiplatform.googleapis.com",
|
||||
"max_retries": 2,
|
||||
"is_active": true
|
||||
}
|
||||
]);
|
||||
payload["providers"][0]["api_keys"] = json!([]);
|
||||
payload["providers"][0]["models"] = json!([]);
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/admin/system/config/import"))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.expect("legacy Vertex import should complete");
|
||||
|
||||
let status = response.status();
|
||||
let response_body: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={response_body}");
|
||||
assert_eq!(response_body["stats"]["endpoints"]["created"], json!(1));
|
||||
assert_eq!(response_body["stats"]["endpoints"]["skipped"], json!(1));
|
||||
assert!(response_body["stats"]["errors"]
|
||||
.as_array()
|
||||
.is_some_and(|errors| errors.iter().any(|error| {
|
||||
error
|
||||
.as_str()
|
||||
.is_some_and(|error| error.contains("claude:messages"))
|
||||
})));
|
||||
|
||||
let providers = provider_catalog_repository
|
||||
.list_providers(false)
|
||||
.await
|
||||
.expect("providers should load");
|
||||
assert_eq!(providers.len(), 1);
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_provider_ids(std::slice::from_ref(&providers[0].id))
|
||||
.await
|
||||
.expect("endpoints should load");
|
||||
assert_eq!(endpoints.len(), 1, "unexpected endpoints: {endpoints:?}");
|
||||
assert_eq!(endpoints[0].api_format, "gemini:generate_content");
|
||||
assert!(endpoints[0].is_active);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_imports_admin_system_config_with_numeric_string_prices() {
|
||||
run_admin_system_import_test(
|
||||
|
||||
@@ -322,6 +322,7 @@ async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_o
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert!(payload.get("type").is_none());
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "无效的API密钥");
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
@@ -331,6 +332,57 @@ async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_o
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_claude_routes_use_anthropic_authentication_error_for_invalid_api_key() {
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-other-claude-key")),
|
||||
sample_currently_usable_auth_snapshot("key-claude-other", "user-claude-other"),
|
||||
)]));
|
||||
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 client = reqwest::Client::new();
|
||||
|
||||
for (path, trace_id) in [
|
||||
("/v1/messages", "trace-control-claude-invalid-key-messages"),
|
||||
(
|
||||
"/v1/messages/count_tokens",
|
||||
"trace-control-claude-invalid-key-count-tokens",
|
||||
),
|
||||
] {
|
||||
let response = client
|
||||
.post(format!("{gateway_url}{path}"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-missing-claude-key")
|
||||
.header(TRACE_ID_HEADER, trace_id)
|
||||
.body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete locally");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED, "path: {path}");
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED),
|
||||
"path: {path}"
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["type"], "error", "path: {path}");
|
||||
assert_eq!(
|
||||
payload["error"]["type"], "authentication_error",
|
||||
"path: {path}"
|
||||
);
|
||||
}
|
||||
|
||||
gateway_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));
|
||||
@@ -498,7 +550,8 @@ async fn gateway_locally_denies_disallowed_claude_api_format_without_hitting_con
|
||||
|
||||
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["type"], "error");
|
||||
assert_eq!(payload["error"]["type"], "permission_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 claude:messages 格式"
|
||||
@@ -581,7 +634,8 @@ async fn gateway_locally_denies_disallowed_provider_without_hitting_control_or_u
|
||||
|
||||
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["type"], "error");
|
||||
assert_eq!(payload["error"]["type"], "permission_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问 claude 提供商"
|
||||
@@ -817,3 +871,52 @@ async fn gateway_locally_denies_disallowed_openai_model_without_hitting_control_
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_disallowed_claude_model_with_anthropic_permission_error() {
|
||||
let mut snapshot =
|
||||
sample_currently_usable_auth_snapshot("key-claude-model-123", "user-claude-model-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!["claude:messages".to_string()]);
|
||||
snapshot.user_allowed_api_formats = Some(vec!["claude:messages".to_string()]);
|
||||
snapshot.api_key_allowed_models = Some(vec!["claude-haiku-4-5".to_string()]);
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-claude-model-guard-123")),
|
||||
snapshot,
|
||||
)]));
|
||||
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-model-guard-123")
|
||||
.header(TRACE_ID_HEADER, "trace-control-claude-model-guard-1")
|
||||
.body("{\"model\":\"claude-sonnet-4-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should complete locally");
|
||||
|
||||
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["type"], "error");
|
||||
assert_eq!(payload["error"]["type"], "permission_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前用户、用户组或密钥的访问控制策略不允许访问模型 claude-sonnet-4-5"
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user