mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
fix: harden OAuth identity and cookies and correct quota and JSON display
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user