mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 09:57:47 +08:00
merge(main): sync latest main into security branch
This commit is contained in:
@@ -1,9 +1,15 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_crypto::{
|
||||
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminProviderModelListQuery, GlobalModelReadRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -2308,6 +2314,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
},
|
||||
"gemini-2.5-pro": {
|
||||
"displayName": "Gemini 2.5 Pro"
|
||||
},
|
||||
"gemini-3.7-flash-tiered": {
|
||||
"displayName": "Gemini 3.7 Flash"
|
||||
},
|
||||
"chat_23310": {
|
||||
"displayName": "Internal Chat"
|
||||
}
|
||||
}
|
||||
}),
|
||||
@@ -2409,6 +2421,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default());
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
@@ -2418,6 +2431,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository.clone())
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
@@ -2519,6 +2533,23 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
.and_then(|value| value.get("remaining_fraction")),
|
||||
Some(&json!(0.25))
|
||||
);
|
||||
let imported_provider_models = global_model_repository
|
||||
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||
provider_id: "provider-antigravity".to_string(),
|
||||
is_active: None,
|
||||
offset: 0,
|
||||
limit: 100,
|
||||
})
|
||||
.await
|
||||
.expect("imported Antigravity provider models should read");
|
||||
let imported_model_names = imported_provider_models
|
||||
.iter()
|
||||
.map(|model| model.provider_model_name.as_str())
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
assert!(imported_model_names.contains("claude-sonnet-4"));
|
||||
assert!(imported_model_names.contains("gemini-2.5-pro"));
|
||||
assert!(imported_model_names.contains("gemini-3.7-flash-tiered"));
|
||||
assert!(!imported_model_names.contains("chat_23310"));
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
|
||||
@@ -2,11 +2,9 @@ use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
|
||||
use aether_data_contracts::repository::global_models::{
|
||||
AdminProviderModelListQuery, GlobalModelReadRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::routing_profiles::StoredRoutingGroup;
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
@@ -843,29 +841,6 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
|
||||
]),
|
||||
);
|
||||
|
||||
let routing_group_repository = Arc::new(InMemoryRoutingGroupRepository::seed(
|
||||
[StoredRoutingGroup {
|
||||
id: "system-default".to_string(),
|
||||
name: "system-default".to_string(),
|
||||
description: Some("test system default routing strategy".to_string()),
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json: json!({
|
||||
"default_policy": {
|
||||
"priority_mode": "global_key",
|
||||
"scheduling_mode": "fixed_order"
|
||||
}
|
||||
}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: Some(1),
|
||||
}],
|
||||
std::iter::empty(),
|
||||
std::iter::empty(),
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
@@ -875,7 +850,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository)
|
||||
.with_routing_group_repository_for_tests(routing_group_repository),
|
||||
.with_system_default_routing_group_for_tests(),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
@@ -898,8 +873,8 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
|
||||
assert_eq!(payload["global_model_name"], "gpt-5");
|
||||
assert_eq!(payload["display_name"], "GPT 5");
|
||||
assert_eq!(payload["global_model_mappings"], json!(["gpt-5-upstream"]));
|
||||
assert_eq!(payload["scheduling_mode"], "fixed_order");
|
||||
assert_eq!(payload["priority_mode"], "global_key");
|
||||
assert_eq!(payload["scheduling_mode"], "cache_affinity");
|
||||
assert_eq!(payload["priority_mode"], "provider");
|
||||
assert_eq!(payload["total_providers"], 2);
|
||||
assert_eq!(payload["active_providers"], 2);
|
||||
|
||||
|
||||
@@ -3912,6 +3912,175 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email() {
|
||||
run_admin_oauth_test(
|
||||
"gateway_names_new_antigravity_oauth_account_from_google_userinfo_email",
|
||||
gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl,
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl() {
|
||||
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 token_hits = Arc::new(Mutex::new(0usize));
|
||||
let token_hits_clone = Arc::clone(&token_hits);
|
||||
let user_info_hits = Arc::new(Mutex::new(0usize));
|
||||
let user_info_hits_clone = Arc::clone(&user_info_hits);
|
||||
let seen_user_info_authorization = Arc::new(Mutex::new(None::<String>));
|
||||
let seen_user_info_authorization_clone = Arc::clone(&seen_user_info_authorization);
|
||||
let google_server = Router::new()
|
||||
.route(
|
||||
"/oauth/token",
|
||||
post(move || {
|
||||
let token_hits_inner = Arc::clone(&token_hits_clone);
|
||||
async move {
|
||||
*token_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"access_token": "antigravity-access-token",
|
||||
"refresh_token": "antigravity-refresh-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"scope": "https://www.googleapis.com/auth/userinfo.email"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/oauth/userinfo",
|
||||
get(move |headers: HeaderMap| {
|
||||
let user_info_hits_inner = Arc::clone(&user_info_hits_clone);
|
||||
let seen_authorization_inner = Arc::clone(&seen_user_info_authorization_clone);
|
||||
async move {
|
||||
*user_info_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
*seen_authorization_inner.lock().expect("mutex should lock") = headers
|
||||
.get(http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
Json(json!({
|
||||
"email": "[email protected]",
|
||||
"verified_email": true,
|
||||
"name": "Antigravity User"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut provider = sample_provider("provider-antigravity", "antigravity", 10);
|
||||
provider.provider_type = "antigravity".to_string();
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-antigravity",
|
||||
"provider-antigravity",
|
||||
"gemini:generate_content",
|
||||
"https://daily-cloudcode-pa.googleapis.com",
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (google_url, google_handle) = start_server(google_server).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),
|
||||
)
|
||||
.with_provider_oauth_state_entry_for_tests(
|
||||
"nonce-antigravity-123",
|
||||
json!({
|
||||
"nonce": "nonce-antigravity-123",
|
||||
"key_id": "",
|
||||
"provider_id": "provider-antigravity",
|
||||
"provider_type": "antigravity",
|
||||
"pkce_verifier": "verifier-antigravity-123",
|
||||
}),
|
||||
)
|
||||
.with_provider_oauth_token_url_for_tests(
|
||||
"antigravity",
|
||||
format!("{google_url}/oauth/token"),
|
||||
)
|
||||
.with_provider_oauth_token_url_for_tests(
|
||||
"antigravity_user_info",
|
||||
format!("{google_url}/oauth/userinfo"),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/complete"
|
||||
))
|
||||
.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!({
|
||||
"callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=nonce-antigravity-123"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["provider_type"], "antigravity");
|
||||
assert_eq!(payload["email"], "[email protected]");
|
||||
assert_eq!(payload["replaced"], false);
|
||||
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
|
||||
assert_eq!(*user_info_hits.lock().expect("mutex should lock"), 1);
|
||||
assert_eq!(
|
||||
seen_user_info_authorization
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.as_deref(),
|
||||
Some("Bearer antigravity-access-token")
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let key_id = payload["key_id"]
|
||||
.as_str()
|
||||
.expect("created key id should be returned")
|
||||
.to_string();
|
||||
let persisted_keys = provider_catalog_repository
|
||||
.list_keys_by_ids(std::slice::from_ref(&key_id))
|
||||
.await
|
||||
.expect("created key should load");
|
||||
let persisted = persisted_keys.first().expect("created key should exist");
|
||||
assert_eq!(persisted.name, "[email protected]");
|
||||
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
persisted
|
||||
.encrypted_auth_config
|
||||
.as_deref()
|
||||
.expect("auth config should be stored"),
|
||||
)
|
||||
.expect("auth config should decrypt");
|
||||
let auth_config: Value =
|
||||
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
|
||||
assert_eq!(auth_config["email"], "[email protected]");
|
||||
assert_eq!(auth_config["refresh_token"], "antigravity-refresh-token");
|
||||
|
||||
gateway_handle.abort();
|
||||
google_handle.abort();
|
||||
upstream_handle.abort();
|
||||
drop(upstream_url);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal() {
|
||||
run_admin_oauth_test(
|
||||
|
||||
@@ -2184,7 +2184,7 @@ async fn gateway_handles_admin_system_config_locally_with_bearer_admin_session()
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let state = AppState::new().expect("gateway should build");
|
||||
let access_token = issue_test_admin_access_token(&state, "device-admin-config").await;
|
||||
let gateway = build_router_with_state(state);
|
||||
@@ -2208,6 +2208,44 @@ async fn gateway_handles_admin_system_config_locally_with_bearer_admin_session()
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_removed_admin_system_provider_priority_mode_with_bearer_admin_session() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/system/configs/provider_priority_mode",
|
||||
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 state = AppState::new().expect("gateway should build");
|
||||
let access_token = issue_test_admin_access_token(&state, "device-admin-config").await;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/system/configs/provider_priority_mode"
|
||||
))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-admin-config")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_sets_admin_system_config_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
Reference in New Issue
Block a user