mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
Fix OAuth token import and table filters
This commit is contained in:
+105
-38
@@ -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()));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user