mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +08:00
fix: harden OAuth identity and cookies and correct quota and JSON display
This commit is contained in:
@@ -8,7 +8,7 @@ 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,
|
||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, GlobalModelReadRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
@@ -20,8 +20,9 @@ use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::super::{
|
||||
build_router_with_state, build_state_with_execution_runtime_override, sample_bound_auth_config,
|
||||
sample_bound_key, sample_endpoint, sample_key, sample_proxy_node, start_server, AppState,
|
||||
build_router_with_state, build_state_with_execution_runtime_override, sample_admin_global_model,
|
||||
sample_bound_auth_config, sample_bound_key, sample_endpoint, sample_key, sample_proxy_node,
|
||||
start_server, AppState,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
@@ -2421,7 +2422,15 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default());
|
||||
let existing_global_model = sample_admin_global_model(
|
||||
"global-claude-sonnet-4",
|
||||
"claude-sonnet-4",
|
||||
"Claude Sonnet 4",
|
||||
);
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::default()
|
||||
.with_admin_global_models(vec![existing_global_model.clone()]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
@@ -2533,7 +2542,16 @@ 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
|
||||
let global_models = global_model_repository
|
||||
.list_admin_global_models(&AdminGlobalModelListQuery {
|
||||
limit: 100,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect("global models should read after quota refresh");
|
||||
assert_eq!(global_models.total, 1);
|
||||
assert_eq!(global_models.items, vec![existing_global_model]);
|
||||
let provider_models = global_model_repository
|
||||
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||
provider_id: "provider-antigravity".to_string(),
|
||||
is_active: None,
|
||||
@@ -2541,15 +2559,8 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
|
||||
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"));
|
||||
.expect("Antigravity provider models should read after quota refresh");
|
||||
assert!(provider_models.is_empty());
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
|
||||
@@ -3916,11 +3916,47 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
|
||||
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,
|
||||
|| {
|
||||
assert_antigravity_oauth_account_uses_google_userinfo_email(
|
||||
"complete",
|
||||
json!({
|
||||
"callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"
|
||||
}),
|
||||
)
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl() {
|
||||
#[test]
|
||||
fn gateway_names_imported_antigravity_oauth_account_from_google_userinfo_email() {
|
||||
run_admin_oauth_test(
|
||||
"gateway_names_imported_antigravity_oauth_account_from_google_userinfo_email",
|
||||
|| {
|
||||
assert_antigravity_oauth_account_uses_google_userinfo_email(
|
||||
"import-refresh-token",
|
||||
json!({"refresh_token": "antigravity-import-refresh-token"}),
|
||||
)
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_names_batch_imported_antigravity_oauth_account_from_google_userinfo_email() {
|
||||
run_admin_oauth_test(
|
||||
"gateway_names_batch_imported_antigravity_oauth_account_from_google_userinfo_email",
|
||||
|| {
|
||||
assert_antigravity_oauth_account_uses_google_userinfo_email(
|
||||
"batch-import",
|
||||
json!({"credentials": "antigravity-import-refresh-token"}),
|
||||
)
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
async fn assert_antigravity_oauth_account_uses_google_userinfo_email(
|
||||
operation: &str,
|
||||
request_body: Value,
|
||||
) {
|
||||
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| {
|
||||
@@ -4025,15 +4061,13 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/complete"
|
||||
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/{operation}"
|
||||
))
|
||||
.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=cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"
|
||||
}))
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
@@ -4041,9 +4075,22 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
|
||||
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);
|
||||
let account_result = if operation == "batch-import" {
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["success"], 1, "payload={payload}");
|
||||
assert_eq!(payload["failed"], 0);
|
||||
assert_eq!(payload["results"][0]["status"], "success");
|
||||
assert_eq!(
|
||||
payload["results"][0]["key_name"],
|
||||
"[email protected]"
|
||||
);
|
||||
&payload["results"][0]
|
||||
} else {
|
||||
assert_eq!(payload["provider_type"], "antigravity");
|
||||
assert_eq!(payload["email"], "[email protected]");
|
||||
&payload
|
||||
};
|
||||
assert_eq!(account_result["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!(
|
||||
@@ -4055,7 +4102,7 @@ async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let key_id = payload["key_id"]
|
||||
let key_id = account_result["key_id"]
|
||||
.as_str()
|
||||
.expect("created key id should be returned")
|
||||
.to_string();
|
||||
|
||||
Reference in New Issue
Block a user