Fix OAuth token import and table filters

This commit is contained in:
fawney19
2026-05-01 02:14:49 +08:00
parent 9570e5c2c1
commit 4fc7cecf30
54 changed files with 3160 additions and 727 deletions
@@ -1,3 +1,4 @@
use super::super::token_import::build_codex_access_token_import_auth_config;
use super::kiro_import::execute_admin_provider_oauth_kiro_batch_import;
use super::parse::{
apply_admin_provider_oauth_batch_import_hints, extract_admin_provider_oauth_batch_error_detail,
@@ -22,12 +23,19 @@ use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
};
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
use crate::GatewayError;
use aether_admin::provider::oauth::parse_admin_provider_oauth_kiro_batch_import_entries;
use serde_json::json;
use aether_contracts::ProxySnapshot;
use serde_json::{json, Map, Value};
use std::time::{SystemTime, UNIX_EPOCH};
struct AdminProviderOAuthResolvedBatchImport {
access_token: String,
auth_config: Map<String, Value>,
expires_at: Option<u64>,
}
pub(super) fn estimate_admin_provider_oauth_batch_import_total(
provider_type: &str,
raw_credentials: &str,
@@ -70,6 +78,90 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
}
}
async fn resolve_admin_provider_oauth_batch_import_tokens(
state: &AdminAppState<'_>,
template: AdminProviderOAuthTemplate,
provider_type: &str,
entry: &AdminProviderOAuthBatchImportEntry,
request_proxy: Option<ProxySnapshot>,
) -> Result<AdminProviderOAuthResolvedBatchImport, String> {
let refresh_token = entry
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
let access_token = entry
.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(refresh_token) = refresh_token {
let token_payload = match exchange_admin_provider_oauth_refresh_token(
state,
template,
refresh_token,
request_proxy.clone(),
)
.await
{
Ok(payload) => payload,
Err(response) => {
let detail = extract_admin_provider_oauth_batch_error_detail(response).await;
if provider_type.eq_ignore_ascii_case("codex") {
if let Some(access_token) = access_token {
let (auth_config, expires_at) = build_codex_access_token_import_auth_config(
access_token,
Some(refresh_token),
entry.expires_at,
Some(detail.as_str()),
);
return Ok(AdminProviderOAuthResolvedBatchImport {
access_token: access_token.to_string(),
auth_config,
expires_at,
});
}
}
return Err(format!("Token 验证失败: {detail}"));
}
};
let (mut auth_config, access_token, returned_refresh_token, expires_at) =
build_provider_oauth_auth_config_from_token_payload(provider_type, &token_payload);
let Some(access_token) = access_token else {
return Err("Token 刷新返回缺少 access_token".to_string());
};
let refresh_token = returned_refresh_token
.or_else(|| Some(refresh_token.to_string()))
.filter(|value| !value.trim().is_empty());
if let Some(refresh_token) = refresh_token.as_ref() {
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
}
return Ok(AdminProviderOAuthResolvedBatchImport {
access_token,
auth_config,
expires_at,
});
}
if let Some(access_token) = access_token {
if !provider_type.eq_ignore_ascii_case("codex") {
return Err("Access Token 导入仅支持 Codex Provider".to_string());
}
let (auth_config, expires_at) =
build_codex_access_token_import_auth_config(access_token, None, entry.expires_at, None);
return Ok(AdminProviderOAuthResolvedBatchImport {
access_token: access_token.to_string(),
auth_config,
expires_at,
});
}
Err("Refresh Token 或 Access Token 不能为空".to_string())
}
pub(super) async fn execute_admin_provider_oauth_batch_import(
state: &AdminAppState<'_>,
provider_id: &str,
@@ -145,24 +237,22 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
let mut failed = 0usize;
for (index, entry) in entries.iter().enumerate() {
let token_payload = match exchange_admin_provider_oauth_refresh_token(
let resolved_import = match resolve_admin_provider_oauth_batch_import_tokens(
state,
template,
entry.refresh_token.as_str(),
provider_type,
entry,
request_proxy.clone(),
)
.await
{
Ok(payload) => payload,
Err(response) => {
Ok(value) => value,
Err(error) => {
failed += 1;
results.push(json!({
"index": index,
"status": "error",
"error": format!(
"Token 验证失败: {}",
extract_admin_provider_oauth_batch_error_detail(response).await
),
"error": error,
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
@@ -176,34 +266,11 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
continue;
}
};
let (mut auth_config, access_token, returned_refresh_token, expires_at) =
build_provider_oauth_auth_config_from_token_payload(provider_type, &token_payload);
let Some(access_token) = access_token else {
failed += 1;
results.push(json!({
"index": index,
"status": "error",
"error": "Token 刷新返回缺少 access_token",
"replaced": false,
}));
maybe_report_admin_provider_oauth_batch_import_progress(
&mut progress,
entries.len(),
success,
failed,
&results,
)
.await;
continue;
};
let refresh_token = returned_refresh_token
.or_else(|| Some(entry.refresh_token.clone()))
.filter(|value| !value.trim().is_empty());
if let Some(refresh_token) = refresh_token.as_ref() {
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
}
let AdminProviderOAuthResolvedBatchImport {
access_token,
mut auth_config,
expires_at,
} = resolved_import;
apply_admin_provider_oauth_batch_import_hints(provider_type, entry, &mut auth_config);
let duplicate =
@@ -1,5 +1,6 @@
use super::super::token_import::{import_tokens_from_raw_token, normalize_single_import_tokens};
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::state::current_unix_secs;
use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_u64_value};
use axum::{
body::{to_bytes, Body, Bytes},
http,
@@ -18,7 +19,9 @@ pub(super) struct AdminProviderOAuthBatchImportRequest {
#[derive(Debug, Clone)]
pub(super) struct AdminProviderOAuthBatchImportEntry {
pub refresh_token: String,
pub refresh_token: Option<String>,
pub access_token: Option<String>,
pub expires_at: Option<u64>,
pub account_id: Option<String>,
pub account_user_id: Option<String>,
pub plan_type: Option<String>,
@@ -73,8 +76,11 @@ fn extract_admin_provider_oauth_batch_import_entry(
if refresh_token.is_empty() {
None
} else {
let (refresh_token, access_token) = import_tokens_from_raw_token(refresh_token);
Some(AdminProviderOAuthBatchImportEntry {
refresh_token: refresh_token.to_string(),
refresh_token,
access_token,
expires_at: None,
account_id: None,
account_user_id: None,
plan_type: None,
@@ -88,7 +94,19 @@ fn extract_admin_provider_oauth_batch_import_entry(
object
.get("refresh_token")
.or_else(|| object.get("refreshToken")),
)?;
);
let access_token = coerce_admin_provider_oauth_import_str(
object
.get("access_token")
.or_else(|| object.get("accessToken")),
);
let (refresh_token, access_token) =
normalize_single_import_tokens(refresh_token.as_deref(), access_token.as_deref());
if refresh_token.is_none() && access_token.is_none() {
return None;
}
let expires_at =
json_u64_value(object.get("expires_at").or_else(|| object.get("expiresAt")));
let account_id = coerce_admin_provider_oauth_import_str(
object
.get("account_id")
@@ -121,6 +139,8 @@ fn extract_admin_provider_oauth_batch_import_entry(
let email = coerce_admin_provider_oauth_import_str(object.get("email"));
Some(AdminProviderOAuthBatchImportEntry {
refresh_token,
access_token,
expires_at,
account_id,
account_user_id,
plan_type,
@@ -163,13 +183,18 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
raw.lines()
.map(str::trim)
.filter(|line| !line.is_empty() && !line.starts_with('#'))
.map(|refresh_token| AdminProviderOAuthBatchImportEntry {
refresh_token: refresh_token.to_string(),
account_id: None,
account_user_id: None,
plan_type: None,
user_id: None,
email: None,
.map(|token| {
let (refresh_token, access_token) = import_tokens_from_raw_token(token);
AdminProviderOAuthBatchImportEntry {
refresh_token,
access_token,
expires_at: None,
account_id: None,
account_user_id: None,
plan_type: None,
user_id: None,
email: None,
}
})
.collect()
}
@@ -291,3 +316,47 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
"updated_at": updated_at,
})
}
#[cfg(test)]
mod tests {
use super::parse_admin_provider_oauth_batch_import_entries;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use serde_json::json;
fn unsigned_jwt(payload: serde_json::Value) -> String {
let header = json!({"alg": "none", "typ": "JWT"});
let encode = |value: serde_json::Value| {
URL_SAFE_NO_PAD.encode(serde_json::to_vec(&value).expect("jwt json should serialize"))
};
format!("{}.{}.signature", encode(header), encode(payload))
}
#[test]
fn parses_access_token_only_entry() {
let entries = parse_admin_provider_oauth_batch_import_entries(
r#"[{"accessToken":"at_1","expiresAt":2100000000,"accountId":"acc-1","email":"[email protected]"}]"#,
);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].refresh_token, None);
assert_eq!(entries[0].access_token.as_deref(), Some("at_1"));
assert_eq!(entries[0].expires_at, Some(2_100_000_000));
assert_eq!(entries[0].account_id.as_deref(), Some("acc-1"));
assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
}
#[test]
fn parses_plain_jwt_line_as_access_token() {
let token = unsigned_jwt(json!({
"iss": "https://auth.openai.com",
"aud": ["https://api.openai.com/v1"],
"exp": 2_000_000_000u64,
}));
let entries = parse_admin_provider_oauth_batch_import_entries(&token);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].refresh_token, None);
assert_eq!(entries[0].access_token.as_deref(), Some(token.as_str()));
}
}