fix: harden OAuth identity and cookies and correct quota and JSON display

This commit is contained in:
elky
2026-09-08 10:51:25 +08:00
parent 7113d04f8a
commit cf8ea19856
24 changed files with 1327 additions and 154 deletions
@@ -1463,7 +1463,7 @@ mod tests {
&auth_config,
Some(0),
),
"antigravity_[email protected]"
"[email protected]"
);
}
@@ -1,3 +1,4 @@
use super::super::helpers::admin_provider_oauth_key_name_from_auth_config;
use super::super::kiro::{
admin_provider_oauth_kiro_refresh_base_url_override, fetch_admin_provider_oauth_kiro_email,
refresh_admin_provider_oauth_kiro_auth_config,
@@ -79,7 +80,7 @@ fn kiro_social_key_name(
.collect::<String>()
})
.unwrap_or_else(|| "unknown".to_string());
format!("kiro_{fallback} ({provider})")
format!("账号_{fallback} ({provider})")
}
fn kiro_social_poll_error_response(error: impl Into<String>) -> Response<Body> {
@@ -1004,10 +1005,11 @@ async fn handle_admin_provider_oauth_windsurf_browser_device_poll(
}
}
} else {
let key_name = email
.as_deref()
.map(|email| format!("windsurf_{email}"))
.unwrap_or_else(|| format!("windsurf_{}", current_unix_secs()));
let key_name = admin_provider_oauth_key_name_from_auth_config(
&provider.provider_type,
&auth_config,
None,
);
match state
.create_provider_oauth_catalog_key(
&provider.id,
@@ -1356,6 +1358,32 @@ mod tests {
use crate::control::GatewayAdminPrincipalContext;
use aether_data::repository::provider_oauth::StoredAdminProviderOAuthDeviceSession;
#[test]
fn kiro_social_key_name_preserves_email_and_auth_method() {
assert_eq!(
super::kiro_social_key_name(
Some(" [email protected] "),
Some("Github"),
Some("refresh-token-1"),
),
"[email protected] (Github)"
);
}
#[test]
fn kiro_social_key_name_without_email_uses_generic_account_prefix() {
for email in [None, Some(""), Some(" ")] {
assert_eq!(
super::kiro_social_key_name(email, Some("Google"), Some("refresh-token-1")),
"账号_154f43 (Google)"
);
assert_eq!(
super::kiro_social_key_name(email, None, None),
"账号_unknown (social)"
);
}
}
fn device_session() -> StoredAdminProviderOAuthDeviceSession {
StoredAdminProviderOAuthDeviceSession {
session_id: "device-session-1".to_string(),
@@ -52,13 +52,12 @@ pub(super) fn admin_provider_oauth_key_name_from_auth_config(
auth_config: &Map<String, Value>,
batch_index: Option<usize>,
) -> String {
let provider_type = provider_type.trim();
if let Some(email) = trimmed_auth_config_string(auth_config, "email") {
return format!("{provider_type}_{email}");
return email;
}
if provider_type.eq_ignore_ascii_case("grok") {
if provider_type.trim().eq_ignore_ascii_case("grok") {
if let Some(user_id) = trimmed_auth_config_string(auth_config, "user_id") {
return format!("grok_{user_id}");
return user_id;
}
}
@@ -68,7 +67,7 @@ pub(super) fn admin_provider_oauth_key_name_from_auth_config(
.map(|duration| duration.as_secs())
.unwrap_or(0);
match batch_index {
Some(index) => format!("{provider_type}_{timestamp}_{index}"),
Some(index) => format!("账号_{timestamp}_{index}"),
None => format!("账号_{timestamp}"),
}
}
@@ -87,6 +86,106 @@ mod tests {
use super::*;
use serde_json::{json, Map};
const PROVIDER_TYPES: &[&str] = &[
"codex",
" Codex ",
"claude_code",
"chatgpt_web",
"gemini_cli",
"antigravity",
"grok",
" Grok ",
"kiro",
"windsurf",
];
#[test]
fn default_key_name_uses_email_without_provider_prefix() {
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(" [email protected] "));
for provider_type in PROVIDER_TYPES {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
"[email protected]"
);
}
}
}
#[test]
fn antigravity_default_key_name_uses_email_without_provider_prefix() {
for email in [" [email protected] ", "[email protected]"] {
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(email));
for provider_type in ["antigravity", " Antigravity "] {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
email.trim()
);
}
}
}
}
#[test]
fn default_key_name_preserves_email_with_provider_prefix() {
for provider_type in PROVIDER_TYPES {
let email = format!("{}[email protected]", provider_type.trim());
let mut auth_config = Map::new();
auth_config.insert("email".to_string(), json!(email));
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
email
);
}
}
}
#[test]
fn default_key_name_without_email_uses_generic_account_name() {
for email in [None, Some(""), Some(" ")] {
let mut auth_config = Map::new();
if let Some(email) = email {
auth_config.insert("email".to_string(), json!(email));
}
for provider_type in PROVIDER_TYPES {
for batch_index in [None, Some(3)] {
let name = admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
);
let suffix = name.strip_prefix("账号_").expect("generic account prefix");
let timestamp = if batch_index.is_some() {
suffix.strip_suffix("_3").expect("batch index suffix")
} else {
suffix
};
assert!(timestamp.parse::<u64>().is_ok());
}
}
}
}
#[test]
fn grok_default_key_name_uses_full_user_id() {
let mut auth_config = Map::new();
@@ -95,10 +194,18 @@ mod tests {
json!("1619039a-0191-4e0a-a490-8f4ad21262c9"),
);
assert_eq!(
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
"grok_1619039a-0191-4e0a-a490-8f4ad21262c9"
);
for provider_type in ["grok", " Grok "] {
for batch_index in [None, Some(3)] {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config(
provider_type,
&auth_config,
batch_index,
),
"1619039a-0191-4e0a-a490-8f4ad21262c9"
);
}
}
}
#[test]
@@ -109,17 +216,22 @@ mod tests {
assert_eq!(
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
"grok_grok@example.com"
"[email protected]"
);
}
#[test]
fn batch_default_key_name_keeps_existing_timestamp_shape() {
fn batch_default_key_name_keeps_distinct_indexes_without_provider_prefix() {
let auth_config = Map::new();
let name = admin_provider_oauth_key_name_from_auth_config("codex", &auth_config, Some(3));
let name = admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, Some(3));
let other_name =
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, Some(4));
assert!(name.starts_with("codex_"));
assert!(name.starts_with("账号_"));
assert!(name.ends_with("_3"));
assert!(other_name.starts_with("账号_"));
assert!(other_name.ends_with("_4"));
assert_ne!(name, other_name);
}
#[test]
@@ -5,7 +5,6 @@ use super::shared::{
quota_key_auto_removed, quota_refresh_success_invalid_state,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::AdminImportProviderModelsRequest;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::{
@@ -24,63 +23,6 @@ use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
fn antigravity_discovered_model_ids(metadata_update: Option<&serde_json::Value>) -> Vec<String> {
metadata_update
.and_then(|value| value.pointer("/antigravity/quota_by_model"))
.and_then(serde_json::Value::as_object)
.into_iter()
.flat_map(|models| models.keys())
.map(String::as_str)
.filter(|model_id| aether_model_fetch::antigravity_model_id_is_routable(model_id))
.map(ToOwned::to_owned)
.collect()
}
async fn sync_antigravity_discovered_models(
state: &AdminAppState<'_>,
provider_id: &str,
metadata_update: Option<&serde_json::Value>,
) {
if !state.has_global_model_data_reader() || !state.has_global_model_data_writer() {
return;
}
let model_ids = antigravity_discovered_model_ids(metadata_update);
if model_ids.is_empty() {
return;
}
let result = state
.build_admin_import_provider_models_payload(
provider_id,
AdminImportProviderModelsRequest {
model_ids,
tiered_pricing: None,
price_per_request: None,
},
)
.await;
match result {
Ok(payload) => {
let errors = payload
.get("errors")
.and_then(serde_json::Value::as_array)
.map(Vec::len)
.unwrap_or(0);
if errors > 0 {
warn!(
provider_id,
errors, "Antigravity discovered-model catalog sync completed with item errors"
);
}
}
Err(error) => warn!(
provider_id,
error = %error,
"Antigravity discovered-model catalog sync failed"
),
}
}
async fn execute_antigravity_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
@@ -380,10 +322,6 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
continue;
}
if status == "success" {
sync_antigravity_discovered_models(state, &provider.id, metadata_update.as_ref()).await;
}
if status == "success" {
success_count += 1;
} else {