i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 09:57:47 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
664c063a062726bb9e833e74cb29859f941fc43b
Aether/apps/aether-gateway/src/tests/control/admin/oauth.rs
T

9204 lines
358 KiB
Rust
Raw Blame History

use std::sync::{Arc, Mutex};
use aether_contracts::{
ExecutionPlan, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::management_tokens::{
InMemoryManagementTokenRepository, ManagementTokenReadRepository,
};
use aether_data::repository::oauth_providers::{
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
};
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, PoolMemberHardState, PoolMemberIdentity, PoolScoreReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::response::{IntoResponse, Response};
use axum::routing::{any, delete, get, patch, post, put};
use axum::{extract::Request, Json, Router};
use http::{HeaderMap, HeaderValue, StatusCode};
use serde_json::json;
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, hash_management_token,
sample_endpoint, sample_key, sample_management_token, sample_oauth_provider_config,
sample_provider, sample_proxy_node, start_server, AppState,
};
use crate::admin_api::{
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
};
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
use crate::audit::AdminAuditEvent;
use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use crate::control::resolve_public_request_context;
use crate::data::GatewayDataState;
const ADMIN_OAUTH_TEST_STACK_BYTES: usize = 16 * 1024 * 1024;
fn run_admin_oauth_test<F, Fut>(test_name: &'static str, make_future: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + 'static,
{
let handle = std::thread::Builder::new()
.name(test_name.to_string())
.stack_size(ADMIN_OAUTH_TEST_STACK_BYTES)
.spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("test runtime should build");
runtime.block_on(make_future());
})
.expect("admin oauth test thread should spawn");
if let Err(payload) = handle.join() {
std::panic::resume_unwind(payload);
}
}
fn trusted_admin_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
headers.insert(
TRUSTED_ADMIN_USER_ID_HEADER,
HeaderValue::from_static("admin-user-123"),
);
headers.insert(
TRUSTED_ADMIN_USER_ROLE_HEADER,
HeaderValue::from_static("admin"),
);
headers.insert(
TRUSTED_ADMIN_SESSION_ID_HEADER,
HeaderValue::from_static("session-123"),
);
headers.insert(
TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER,
HeaderValue::from_static("management-token-123"),
);
headers
}
async fn local_admin_provider_oauth_response(
state: &AppState,
method: http::Method,
uri: &str,
body: Option<serde_json::Value>,
) -> Response<Body> {
let headers = trusted_admin_headers();
let request_context = resolve_public_request_context(
state,
&method,
&uri.parse().expect("uri should parse"),
&headers,
"trace-123",
)
.await
.expect("request context should resolve");
let body_bytes = body.map(|value| Bytes::from(value.to_string()));
maybe_build_local_admin_provider_oauth_response(
&AdminAppState::new(state),
&AdminRequestContext::new(&request_context),
body_bytes.as_ref(),
)
.await
.expect("local provider oauth response should build")
.expect("provider oauth route should resolve locally")
}
fn sample_kiro_device_access_token(email: &str) -> String {
use base64::Engine as _;
let header =
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
json!({
"email": email,
"sub": "kiro-user-123",
})
.to_string(),
);
format!("{header}.{payload}.sig")
}
fn sample_kiro_device_access_token_without_email() -> String {
use base64::Engine as _;
let header =
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
json!({
"sub": "kiro-user-123",
})
.to_string(),
);
format!("{header}.{payload}.sig")
}
fn sample_codex_access_token_with_profile_email(email: &str, account_id: &str) -> String {
use base64::Engine as _;
let header =
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
json!({
"iss": "https://auth.openai.com",
"aud": ["https://api.openai.com/v1"],
"exp": 2_000_000_000u64,
"https://api.openai.com/profile": {
"email": email,
"email_verified": true,
},
"https://api.openai.com/auth": {
"chatgpt_account_id": account_id,
},
})
.to_string(),
);
format!("{header}.{payload}.sig")
}
fn codex_import_token_execution_result(request_id: &str) -> serde_json::Value {
json!({
"request_id": request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "imported-codex-access-token",
"refresh_token": "imported-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
})
}
fn codex_quota_execution_result(request_id: &str) -> serde_json::Value {
json!({
"request_id": request_id,
"status_code": 200,
"headers": {},
"body": {
"json_body": {
"user": {
"id": "user-codex-123",
"email": "[email protected]"
},
"account": {
"id": "acct-codex-123",
"plan_type": "plus"
},
"plan": {
"type": "Plus",
"title": "ChatGPT Plus"
}
}
}
})
}
fn windsurf_register_user_execution_result(request_id: &str) -> serde_json::Value {
json!({
"request_id": request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"sessionToken": "devin-session-token$registered",
"name": "Windsurf User",
"apiServerUrl": "https://server.codeium.com"
}
}
})
}
fn assert_single_provider_oauth_refresh_token_plan<'a>(
plans: &'a [ExecutionPlan],
) -> &'a ExecutionPlan {
let token_plans = plans
.iter()
.filter(|plan| plan.request_id == "provider-oauth:refresh-token")
.collect::<Vec<_>>();
assert_eq!(
token_plans.len(),
1,
"expected exactly one provider-oauth refresh-token plan, got {:?}",
plans
.iter()
.map(|plan| plan.request_id.as_str())
.collect::<Vec<_>>()
);
assert!(
plans.iter().all(|plan| {
plan.request_id == "provider-oauth:refresh-token"
|| plan.request_id.starts_with("codex-quota:")
}),
"unexpected execution plans: {:?}",
plans
.iter()
.map(|plan| plan.request_id.as_str())
.collect::<Vec<_>>()
);
token_plans[0]
}
#[test]
fn gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_supported_types_locally_with_trusted_admin_principal_impl(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-oauth/supported-types",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-oauth/supported-types"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("items should be array");
assert_eq!(items.len(), 6);
assert_eq!(items[0]["provider_type"], "claude_code");
assert_eq!(items[1]["provider_type"], "codex");
assert_eq!(items[2]["provider_type"], "chatgpt_web");
assert_eq!(items[3]["provider_type"], "gemini_cli");
assert_eq!(items[4]["provider_type"], "antigravity");
assert_eq!(items[5]["provider_type"], "windsurf");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_authorize_for_windsurf_browser() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_authorize_for_windsurf_browser",
gateway_handles_admin_provider_oauth_device_authorize_for_windsurf_browser_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_authorize_for_windsurf_browser_impl() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-authorize",
Some(json!({
"auth_type": "browser",
"login_option": "github",
"proxy_node_id": "proxy-node-windsurf"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist");
assert_eq!(payload["auth_type"], "browser");
assert_eq!(payload["login_option"], "github");
assert_eq!(payload["redirect_uri"], "show-auth-token");
assert_eq!(payload["callback_required"], true);
let authorization_url = payload["verification_uri_complete"]
.as_str()
.expect("authorization url should exist");
let parsed = url::Url::parse(authorization_url).expect("authorization url should parse");
let params = parsed
.query_pairs()
.into_owned()
.collect::<std::collections::BTreeMap<_, _>>();
assert_eq!(
parsed.as_str().split('?').next(),
Some("https://windsurf.com/windsurf/signin")
);
assert_eq!(
params.get("response_type").map(String::as_str),
Some("token")
);
assert_eq!(params.get("state").map(String::as_str), Some(session_id));
assert_eq!(
params.get("redirect_uri").map(String::as_str),
Some("show-auth-token")
);
assert_eq!(
params.get("redirect_parameters_type").map(String::as_str),
Some("query")
);
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["provider_id"], "provider-windsurf");
assert_eq!(stored["auth_type"], "browser");
assert_eq!(stored["social_provider"], "github");
assert_eq!(stored["redirect_uri"], "show-auth-token");
assert_eq!(stored["proxy_node_id"], "proxy-node-windsurf");
assert_eq!(stored["status"], "pending");
}
#[test]
fn gateway_rejects_generic_oauth_start_for_windsurf_provider() {
run_admin_oauth_test(
"gateway_rejects_generic_oauth_start_for_windsurf_provider",
gateway_rejects_generic_oauth_start_for_windsurf_provider_impl,
);
}
async fn gateway_rejects_generic_oauth_start_for_windsurf_provider_impl() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/start",
None,
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert!(
payload["detail"]
.as_str()
.is_some_and(|detail| detail.contains("浏览器登录")),
"payload={payload}"
);
}
#[test]
fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token",
gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_poll_for_windsurf_one_time_token_impl() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:windsurf-register:new" {
return Json(windsurf_register_user_execution_result(&plan.request_id));
}
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}),
);
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut proxy_node = sample_proxy_node("proxy-node-windsurf");
proxy_node.status = "online".to_string();
proxy_node.is_manual = true;
proxy_node.tunnel_mode = false;
proxy_node.tunnel_connected = false;
proxy_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![proxy_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY)
.attach_proxy_node_repository_for_tests(proxy_node_repository),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-windsurf",
json!({
"provider_id": "provider-windsurf",
"region": "",
"client_id": "",
"client_secret": "",
"device_code": "",
"auth_type": "browser",
"social_provider": "google",
"code_verifier": null,
"redirect_uri": "show-auth-token",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": "proxy-node-windsurf",
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
);
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"token": "ott$browser-token"
})),
)
.await;
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["replaced"], false);
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-windsurf")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "authorized");
let key_id = stored["key_id"]
.as_str()
.expect("key_id should be stored")
.to_string();
assert_eq!(payload["key_id"], key_id);
let persisted = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert_eq!(persisted.auth_type, "oauth");
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-windsurf", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "devin-session-token$registered");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "windsurf");
assert_eq!(auth_config["auth_method"], "browser");
assert_eq!(auth_config["register_source"], "new");
assert_eq!(auth_config["social_provider"], "google");
{
let plans = execution_plans.lock().expect("mutex should lock");
let register_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:windsurf-register:new")
.expect("register plan should execute");
assert_eq!(register_plan.method, "POST");
assert_eq!(
register_plan.content_type.as_deref(),
Some("application/proto")
);
assert!(register_plan.body.json_body.is_none());
let encoded_body = register_plan
.body
.body_bytes_b64
.as_deref()
.expect("register body should be bytes");
use base64::Engine as _;
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(encoded_body)
.expect("register body should decode");
let mut expected_body = vec![0x0a, "ott$browser-token".len() as u8];
expected_body.extend_from_slice(b"ott$browser-token");
assert_eq!(body_bytes, expected_body);
assert_eq!(
register_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-windsurf")
);
}
execution_runtime_handle.abort();
}
#[test]
fn gateway_rejects_windsurf_callback_state_mismatch_and_missing_token() {
run_admin_oauth_test(
"gateway_rejects_windsurf_callback_state_mismatch_and_missing_token",
gateway_rejects_windsurf_callback_state_mismatch_and_missing_token_impl,
);
}
async fn gateway_rejects_windsurf_callback_state_mismatch_and_missing_token_impl() {
let mut provider = sample_provider("provider-windsurf", "windsurf", 10);
provider.provider_type = "windsurf".to_string();
let endpoint = sample_endpoint(
"endpoint-windsurf-chat",
"provider-windsurf",
"openai:chat",
"https://server.codeium.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_device_session_entry_for_tests(
"session-windsurf",
json!({
"provider_id": "provider-windsurf",
"region": "",
"client_id": "",
"client_secret": "",
"device_code": "",
"auth_type": "browser",
"social_provider": "google",
"code_verifier": null,
"redirect_uri": "show-auth-token",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
);
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?token=ott$wrong-state&state=wrong-state"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["status"], "error");
assert!(payload["error"]
.as_str()
.is_some_and(|error| error.contains("state")));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-windsurf/device-poll",
Some(json!({
"session_id": "session-windsurf",
"callback_url": "https://windsurf.com/show-auth-token?state=session-windsurf"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["status"], "error");
assert!(payload["error"]
.as_str()
.is_some_and(|error| error.contains("token")));
}
#[test]
fn gateway_handles_admin_provider_oauth_device_authorize_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_authorize_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_device_authorize_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_authorize_locally_with_trusted_admin_principal_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let register_hits = Arc::new(Mutex::new(0usize));
let register_hits_clone = Arc::clone(&register_hits);
let authorize_hits = Arc::new(Mutex::new(0usize));
let authorize_hits_clone = Arc::clone(&authorize_hits);
let oidc_server = Router::new()
.route(
"/client/register",
post(move |_request: Request| {
let register_hits_inner = Arc::clone(&register_hits_clone);
async move {
*register_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"clientId": "kiro-device-client",
"clientSecret": "kiro-device-secret",
}))
}
}),
)
.route(
"/device_authorization",
post(move |_request: Request| {
let authorize_hits_inner = Arc::clone(&authorize_hits_clone);
async move {
*authorize_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"deviceCode": "device-code-123",
"userCode": "USER-CODE",
"verificationUri": "https://device.example.com/verify",
"verificationUriComplete": "https://device.example.com/verify?user_code=USER-CODE",
"expiresIn": 600,
"interval": 5,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (oidc_url, oidc_handle) = start_server(oidc_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_device_session_entry_for_tests(
"seed-session",
json!({"status":"seed"}),
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_register",
format!("{oidc_url}/client/register"),
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_authorize",
format!("{oidc_url}/device_authorization"),
);
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-authorize"
))
.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!({
"start_url": "https://view.awsapps.com/start",
"region": "us-east-1",
"proxy_node_id": "proxy-node-kiro",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist")
.to_string();
assert_eq!(payload["user_code"], "USER-CODE");
assert_eq!(
payload["verification_uri_complete"],
"https://device.example.com/verify?user_code=USER-CODE"
);
assert_eq!(payload["expires_in"], 600);
assert_eq!(payload["interval"], 5);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*register_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*authorize_hits.lock().expect("mutex should lock"), 1);
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["provider_id"], "provider-kiro");
assert_eq!(stored["region"], "us-east-1");
assert_eq!(stored["client_id"], "kiro-device-client");
assert_eq!(stored["client_secret"], "kiro-device-secret");
assert_eq!(stored["device_code"], "device-code-123");
assert_eq!(stored["proxy_node_id"], "proxy-node-kiro");
assert_eq!(stored["status"], "pending");
gateway_handle.abort();
oidc_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_authorize_for_kiro_google_social() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_authorize_for_kiro_google_social",
gateway_handles_admin_provider_oauth_device_authorize_for_kiro_google_social_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_authorize_for_kiro_google_social_impl() {
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_token_url_for_tests(
"kiro_social_portal",
"https://portal.example.com/signin",
);
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/device-authorize",
Some(json!({
"auth_type": "google"
})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist");
assert_eq!(payload["auth_type"], "google");
assert_eq!(payload["callback_required"], true);
assert_eq!(payload["redirect_uri"], "http://localhost:49153");
let authorization_url = payload["verification_uri_complete"]
.as_str()
.expect("authorization url should exist");
let parsed = url::Url::parse(authorization_url).expect("authorization url should parse");
let params = parsed
.query_pairs()
.into_owned()
.collect::<std::collections::BTreeMap<_, _>>();
assert_eq!(
parsed.as_str().split('?').next(),
Some("https://portal.example.com/signin")
);
assert_eq!(
params.get("redirect_uri").map(String::as_str),
Some("http://localhost:49153")
);
assert_eq!(params.get("state").map(String::as_str), Some(session_id));
assert_eq!(
params.get("code_challenge_method").map(String::as_str),
Some("S256")
);
assert_eq!(
params.get("redirect_from").map(String::as_str),
Some("KiroIDE")
);
assert_eq!(
params.get("login_option").map(String::as_str),
Some("google")
);
assert!(params
.get("code_challenge")
.is_some_and(|value| !value.is_empty()));
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["provider_id"], "provider-kiro");
assert_eq!(stored["auth_type"], "social");
assert_eq!(stored["social_provider"], "Google");
assert_eq!(stored["redirect_uri"], "http://localhost:49153");
assert!(stored["code_verifier"]
.as_str()
.is_some_and(|value| !value.is_empty()));
assert!(stored["machine_id"]
.as_str()
.is_some_and(|value| !value.is_empty()));
assert_eq!(stored["status"], "pending");
}
#[test]
fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_poll_locally_with_trusted_admin_principal_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let access_token = sample_kiro_device_access_token("[email protected]");
let expected_access_token = access_token.clone();
let token_server = Router::new().route(
"/token",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let access_token_inner = access_token.clone();
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"accessToken": access_token_inner,
"refreshToken": "kiro-device-refresh-token",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-123",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "kiro-device-client",
"client_secret": "kiro-device-secret",
"device_code": "device-code-123",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": "proxy-node-kiro",
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string());
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-poll"
))
.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!({
"session_id": "session-123"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 2);
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-123")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "authorized");
assert_eq!(stored["email"], "[email protected]");
assert_eq!(stored["replaced"], false);
let key_id = stored["key_id"]
.as_str()
.expect("key_id should be stored")
.to_string();
assert_eq!(payload["key_id"], key_id);
let persisted = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert_eq!(persisted.auth_type, "oauth");
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-kiro", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, expected_access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["auth_method"], "idc");
assert_eq!(auth_config["refresh_token"], "kiro-device-refresh-token");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["client_id"], "kiro-device-client");
assert_eq!(auth_config["client_secret"], "kiro-device-secret");
assert_eq!(auth_config["region"], "us-east-1");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_poll_for_kiro_social_callback() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_poll_for_kiro_social_callback",
gateway_handles_admin_provider_oauth_device_poll_for_kiro_social_callback_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_poll_for_kiro_social_callback_impl() {
let token_requests = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let token_requests_clone = Arc::clone(&token_requests);
let access_token = sample_kiro_device_access_token("[email protected]");
let expected_access_token = access_token.clone();
let token_server = Router::new().route(
"/oauth/token",
post(move |request: Request| {
let token_requests_inner = Arc::clone(&token_requests_clone);
let access_token_inner = access_token.clone();
async move {
let user_agent = request
.headers()
.get("user-agent")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let raw_body = String::from_utf8(
to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read")
.to_vec(),
)
.expect("body should be utf8");
token_requests_inner
.lock()
.expect("mutex should lock")
.push((user_agent, raw_body));
Json(json!({
"accessToken": access_token_inner,
"refreshToken": "kiro-social-refresh-token",
"profileArn": "arn:aws:kiro:profile/social",
"idToken": "id-token-123",
"tokenType": "Bearer",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-social",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "",
"client_secret": "",
"device_code": "",
"auth_type": "social",
"social_provider": "Github",
"code_verifier": "verifier-123",
"redirect_uri": "http://localhost:49153",
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_token",
format!("{token_url}/oauth/token"),
);
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-poll"
))
.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!({
"session_id": "session-social",
"callback_url": "http://localhost:49153/signin/callback?login_option=github&code=social-code-123&state=session-social"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
{
let requests = token_requests.lock().expect("mutex should lock");
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].0,
"KiroIDE-0.6.18-123e4567-e89b-12d3-a456-426614174000"
);
assert!(requests[0].1.contains("\"code\":\"social-code-123\""));
assert!(requests[0].1.contains("\"code_verifier\":\"verifier-123\""));
assert!(requests[0].1.contains(
"\"redirect_uri\":\"http://localhost:49153/signin/callback?login_option=github\""
));
}
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-social")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "authorized");
assert_eq!(stored["email"], "[email protected]");
let key_id = stored["key_id"]
.as_str()
.expect("key_id should be stored")
.to_string();
assert_eq!(payload["key_id"], key_id);
let persisted = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert_eq!(persisted.name, "[email protected] (Github)");
assert_eq!(persisted.auth_type, "oauth");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, expected_access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["provider"], "Github");
assert_eq!(auth_config["auth_method"], "social");
assert_eq!(auth_config["refresh_token"], "kiro-social-refresh-token");
assert_eq!(auth_config["profile_arn"], "arn:aws:kiro:profile/social");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(
auth_config["machine_id"],
"123e4567-e89b-12d3-a456-426614174000"
);
assert_eq!(auth_config["kiro_version"], "0.6.18");
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_keeps_admin_provider_oauth_device_poll_pending_for_authorization_pending_error() {
run_admin_oauth_test(
"gateway_keeps_admin_provider_oauth_device_poll_pending_for_authorization_pending_error",
gateway_keeps_admin_provider_oauth_device_poll_pending_for_authorization_pending_error_impl,
);
}
async fn gateway_keeps_admin_provider_oauth_device_poll_pending_for_authorization_pending_error_impl(
) {
let token_server = Router::new().route(
"/token",
post(move |_request: Request| async move {
(
StatusCode::BAD_REQUEST,
Json(json!({
"error": "authorization_pending",
"error_description": "waiting for user confirmation",
})),
)
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_device_session_entry_for_tests(
"session-pending",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "kiro-device-client",
"client_secret": "kiro-device-secret",
"device_code": "device-code-123",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/device-poll",
Some(json!({ "session_id": "session-pending" })),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
assert!(
response.extensions().get::<AdminAuditEvent>().is_none(),
"pending state should not attach terminal audit"
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["status"], "pending");
assert_eq!(payload["replaced"], false);
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-pending")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "pending");
assert_eq!(stored["error_msg"], serde_json::Value::Null);
token_handle.abort();
}
#[test]
fn gateway_revalidates_kiro_device_poll_via_idc_refresh_and_backfills_email() {
run_admin_oauth_test(
"gateway_revalidates_kiro_device_poll_via_idc_refresh_and_backfills_email",
gateway_revalidates_kiro_device_poll_via_idc_refresh_and_backfills_email_impl,
);
}
async fn gateway_revalidates_kiro_device_poll_via_idc_refresh_and_backfills_email_impl() {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_requests = Arc::new(Mutex::new(Vec::<String>::new()));
let token_requests_clone = Arc::clone(&token_requests);
let initial_access_token = sample_kiro_device_access_token_without_email();
let refreshed_access_token = sample_kiro_device_access_token_without_email();
let expected_refreshed_access_token = refreshed_access_token.clone();
let token_server = Router::new().route(
"/token",
post(move |request: Request| {
let token_requests_inner = Arc::clone(&token_requests_clone);
let initial_access_token_inner = initial_access_token.clone();
let refreshed_access_token_inner = refreshed_access_token.clone();
async move {
let raw_body = String::from_utf8(
to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read")
.to_vec(),
)
.expect("body should be utf8");
token_requests_inner
.lock()
.expect("mutex should lock")
.push(raw_body.clone());
if raw_body.contains("urn:ietf:params:oauth:grant-type:device_code") {
return Json(json!({
"accessToken": initial_access_token_inner,
"refreshToken": "kiro-device-refresh-token-initial",
"expiresIn": 1800,
}))
.into_response();
}
if raw_body.contains("\"grantType\":\"refresh_token\"") {
return Json(json!({
"accessToken": refreshed_access_token_inner,
"refreshToken": "kiro-device-refresh-token-rotated",
"expiresIn": 2400,
}))
.into_response();
}
(
StatusCode::BAD_REQUEST,
Json(json!({
"error": "unexpected_request",
"body": raw_body,
})),
)
.into_response()
}
}),
);
let usage_hits = Arc::new(Mutex::new(0usize));
let usage_hits_clone = Arc::clone(&usage_hits);
let usage_server = Router::new().route(
"/getUsageLimits",
get(move |_request: Request| {
let usage_hits_inner = Arc::clone(&usage_hits_clone);
async move {
*usage_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO+"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 5.0,
"usageLimitWithPrecision": 20.0,
"nextDateReset": 1_900_000_000u64
}],
"desktopUserInfo": {
"email": "[email protected]"
}
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let (usage_url, usage_handle) = start_server(usage_server).await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-refresh-email",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "kiro-device-client",
"client_secret": "kiro-device-secret",
"device_code": "device-code-123",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": "proxy-node-kiro",
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string())
.with_provider_oauth_token_url_for_tests(
"kiro_device_email",
format!("{usage_url}/getUsageLimits"),
);
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-poll"
))
.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!({
"session_id": "session-refresh-email"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["status"], "authorized");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*usage_hits.lock().expect("mutex should lock"), 1);
{
let requests = token_requests.lock().expect("mutex should lock");
assert_eq!(requests.len(), 2);
assert!(
requests
.iter()
.any(|body| body.contains("urn:ietf:params:oauth:grant-type:device_code")),
"requests={requests:?}"
);
assert!(
requests
.iter()
.any(|body| body.contains("\"grantType\":\"refresh_token\"")),
"requests={requests:?}"
);
}
let stored = state
.load_provider_oauth_device_session_for_tests("device_auth_session:session-refresh-email")
.expect("device session should persist");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["status"], "authorized");
assert_eq!(stored["email"], "[email protected]");
assert_eq!(stored["replaced"], false);
let key_id = stored["key_id"]
.as_str()
.expect("key_id should be stored")
.to_string();
assert_eq!(payload["key_id"], key_id);
let persisted = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert_eq!(persisted.auth_type, "oauth");
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-kiro", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, expected_refreshed_access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["auth_method"], "idc");
assert_eq!(
auth_config["refresh_token"],
"kiro-device-refresh-token-rotated"
);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["client_id"], "kiro-device-client");
assert_eq!(auth_config["client_secret"], "kiro-device-secret");
assert_eq!(auth_config["region"], "us-east-1");
gateway_handle.abort();
usage_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state(
) {
run_admin_oauth_test(
"local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state",
local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state_impl,
);
}
async fn local_admin_provider_oauth_device_poll_attaches_audit_only_when_transition_reaches_terminal_state_impl(
) {
let access_token = sample_kiro_device_access_token("[email protected]");
let token_server = Router::new().route(
"/token",
post(move |_request: Request| {
let access_token_inner = access_token.clone();
async move {
Json(json!({
"accessToken": access_token_inner,
"refreshToken": "kiro-device-refresh-token",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let terminal_state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_device_session_entry_for_tests(
"session-terminal",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "kiro-device-client",
"client_secret": "kiro-device-secret",
"device_code": "device-code-123",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "pending",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": null,
"email": null,
"replaced": false,
"error_msg": null,
}),
)
.with_provider_oauth_token_url_for_tests("kiro_device_poll", format!("{token_url}/token"))
.with_provider_oauth_token_url_for_tests("kiro_idc_refresh", token_url.to_string());
let terminal_response = local_admin_provider_oauth_response(
&terminal_state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/device-poll",
Some(json!({ "session_id": "session-terminal" })),
)
.await;
assert_eq!(terminal_response.status(), StatusCode::OK);
let terminal_audit = terminal_response
.extensions()
.get::<AdminAuditEvent>()
.expect("terminal transition should attach audit");
assert_eq!(
terminal_audit.event_name,
"admin_provider_oauth_device_authorization_completed"
);
assert_eq!(
terminal_audit.action,
"poll_provider_oauth_device_authorization_terminal_state"
);
assert_eq!(terminal_audit.target_type, "provider_oauth_device_session");
assert_eq!(terminal_audit.target_id, "session-terminal");
let non_terminal_state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
))
.with_provider_oauth_device_session_entry_for_tests(
"session-existing-authorized",
json!({
"provider_id": "provider-kiro",
"region": "us-east-1",
"client_id": "kiro-device-client",
"client_secret": "kiro-device-secret",
"device_code": "device-code-123",
"interval": 5,
"expires_at_unix_secs": 4_102_444_800u64,
"status": "authorized",
"proxy_node_id": null,
"created_at_unix_ms": 1_711_000_000u64,
"key_id": "key-existing",
"email": "[email protected]",
"replaced": false,
"error_msg": null,
}),
);
let non_terminal_response = local_admin_provider_oauth_response(
&non_terminal_state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/device-poll",
Some(json!({ "session_id": "session-existing-authorized" })),
)
.await;
assert_eq!(non_terminal_response.status(), StatusCode::OK);
assert!(
non_terminal_response
.extensions()
.get::<AdminAuditEvent>()
.is_none(),
"existing terminal session should not attach a new audit event"
);
token_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_device_authorize_via_execution_runtime_proxy_node() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_device_authorize_via_execution_runtime_proxy_node",
gateway_handles_admin_provider_oauth_device_authorize_via_execution_runtime_proxy_node_impl,
);
}
async fn gateway_handles_admin_provider_oauth_device_authorize_via_execution_runtime_proxy_node_impl(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-kiro"));
assert_eq!(
plan.headers.get("host").map(String::as_str),
Some("oidc.us-east-1.amazonaws.com")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
if plan.request_id == "kiro_device_register" {
assert_eq!(plan.url, "https://oidc.us-east-1.amazonaws.com/client/register");
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"clientId": "kiro-device-client",
"clientSecret": "kiro-device-secret"
}
}
}))
} else {
assert_eq!(plan.request_id, "kiro_device_authorize");
assert_eq!(
plan.url,
"https://oidc.us-east-1.amazonaws.com/device_authorization"
);
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"deviceCode": "device-code-123",
"userCode": "USER-CODE",
"verificationUri": "https://device.example.com/verify",
"verificationUriComplete": "https://device.example.com/verify?user_code=USER-CODE",
"expiresIn": 600,
"interval": 5
}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let mut manual_node = sample_proxy_node("proxy-node-kiro");
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.attach_proxy_node_repository_for_tests(proxy_node_repository),
)
.with_provider_oauth_device_session_entry_for_tests(
"seed-session",
json!({"status":"seed"}),
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_register",
"https://oidc.us-east-1.amazonaws.com/client/register",
)
.with_provider_oauth_token_url_for_tests(
"kiro_device_authorize",
"https://oidc.us-east-1.amazonaws.com/device_authorization",
);
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/device-authorize"
))
.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!({
"start_url": "https://view.awsapps.com/start",
"region": "us-east-1",
"proxy_node_id": "proxy-node-kiro",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
let session_id = payload["session_id"]
.as_str()
.expect("session_id should exist")
.to_string();
assert_eq!(payload["user_code"], "USER-CODE");
assert_eq!(payload["expires_in"], 600);
assert_eq!(payload["interval"], 5);
let stored = state
.load_provider_oauth_device_session_for_tests(&format!("device_auth_session:{session_id}"))
.expect("device session should be stored");
let stored: serde_json::Value =
serde_json::from_str(&stored).expect("device session json should parse");
assert_eq!(stored["proxy_node_id"], "proxy-node-kiro");
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_start_key_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_start_key_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_start_key_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_start_key_locally_with_trusted_admin_principal_impl()
{
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-oauth/keys/key-codex-oauth/start",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:chat",
"oauth-access-token",
);
key.auth_type = "oauth".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth/start"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(
payload["redirect_uri"],
"http://localhost:1455/auth/callback"
);
assert!(payload["authorization_url"]
.as_str()
.is_some_and(|url| url.contains("state=")));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_start_provider_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_start_provider_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_start_provider_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_start_provider_locally_with_trusted_admin_principal_impl(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-oauth/providers/provider-codex/start",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/start"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(
payload["redirect_uri"],
"http://localhost:1455/auth/callback"
);
assert!(payload["authorization_url"]
.as_str()
.is_some_and(|url| url.contains("state=")));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_batch_import_task_status_locally_with_trusted_admin_principal(
) {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_batch_import_task_status_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_batch_import_task_status_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_batch_import_task_status_locally_with_trusted_admin_principal_impl(
) {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/task-123",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_provider_oauth_batch_task_entry_for_tests(
"task-123",
json!({
"task_id": "task-123",
"provider_id": "provider-codex",
"provider_type": "codex",
"status": "completed",
"total": 2,
"processed": 2,
"success": 1,
"failed": 1,
"created_count": 0,
"replaced_count": 1,
"progress_percent": 100,
"message": "导入完成：成功 1，失败 1",
"error": null,
"error_samples": [{
"index": 1,
"status": "error",
"error": "Token 验证失败: invalid_grant",
"replaced": false
}],
"created_at": 1700000001u64,
"started_at": 1700000002u64,
"finished_at": 1700000003u64,
"updated_at": 1700000004u64,
}),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/task-123"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["task_id"], "task-123");
assert_eq!(payload["provider_id"], "provider-codex");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["status"], "completed");
assert_eq!(payload["total"], 2);
assert_eq!(payload["processed"], 2);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 1);
assert_eq!(payload["created_count"], 0);
assert_eq!(payload["replaced_count"], 1);
assert_eq!(payload["progress_percent"], 100);
assert_eq!(payload["error_samples"].as_array().map(Vec::len), Some(1));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn local_admin_provider_oauth_batch_task_status_attaches_audit_only_for_terminal_states() {
run_admin_oauth_test(
"local_admin_provider_oauth_batch_task_status_attaches_audit_only_for_terminal_states",
local_admin_provider_oauth_batch_task_status_attaches_audit_only_for_terminal_states_impl,
);
}
async fn local_admin_provider_oauth_batch_task_status_attaches_audit_only_for_terminal_states_impl()
{
let completed_state = AppState::new()
.expect("gateway should build")
.with_provider_oauth_batch_task_entry_for_tests(
"task-completed",
json!({
"task_id": "task-completed",
"provider_id": "provider-codex",
"provider_type": "codex",
"status": "completed",
"total": 1,
"processed": 1,
"success": 1,
"failed": 0,
"progress_percent": 100,
"message": "导入完成",
"error": null,
"error_samples": [],
"created_at": 1700000001u64,
"started_at": 1700000002u64,
"finished_at": 1700000003u64,
"updated_at": 1700000004u64,
}),
);
let completed_response = local_admin_provider_oauth_response(
&completed_state,
http::Method::GET,
"/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/task-completed",
None,
)
.await;
assert_eq!(completed_response.status(), StatusCode::OK);
let completed_audit = completed_response
.extensions()
.get::<AdminAuditEvent>()
.expect("completed task should attach audit");
assert_eq!(
completed_audit.event_name,
"admin_provider_oauth_batch_task_completed_viewed"
);
assert_eq!(
completed_audit.action,
"view_provider_oauth_batch_task_terminal_state"
);
assert_eq!(completed_audit.target_type, "provider_oauth_batch_task");
assert_eq!(completed_audit.target_id, "provider-codex:task-completed");
let processing_state = AppState::new()
.expect("gateway should build")
.with_provider_oauth_batch_task_entry_for_tests(
"task-processing",
json!({
"task_id": "task-processing",
"provider_id": "provider-codex",
"provider_type": "codex",
"status": "processing",
"total": 3,
"processed": 1,
"success": 1,
"failed": 0,
"progress_percent": 33,
"message": "处理中",
"error": null,
"error_samples": [],
"created_at": 1700000001u64,
"started_at": 1700000002u64,
"finished_at": null,
"updated_at": 1700000003u64,
}),
);
let processing_response = local_admin_provider_oauth_response(
&processing_state,
http::Method::GET,
"/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/task-processing",
None,
)
.await;
assert_eq!(processing_response.status(), StatusCode::OK);
assert!(
processing_response
.extensions()
.get::<AdminAuditEvent>()
.is_none(),
"processing task status should not attach audit"
);
}
#[test]
fn gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_principal",
gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_batch_imports_admin_provider_oauth_locally_with_trusted_admin_principal_impl() {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(Vec::<String>::new()));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
let raw_body = String::from_utf8(
to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read")
.to_vec(),
)
.expect("body should be utf8");
token_hits_inner
.lock()
.expect("mutex should lock")
.push(raw_body.clone());
if raw_body.contains("refresh_token=batch-refresh-success") {
Json(json!({
"access_token": "batch-imported-codex-access-token",
"refresh_token": "batch-imported-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-batch-123",
"account_user_id": "acct-user-batch-123",
"plan_type": "plus",
"user_id": "user-batch-123",
}))
.into_response()
} else {
(
StatusCode::BAD_REQUEST,
Json(json!({
"error": "invalid_grant",
"error_description": "refresh token invalid",
})),
)
.into_response()
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let mut existing_key = sample_key(
"key-codex-batch-duplicate",
"provider-codex",
"openai:chat",
"stale-batch-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","account_id":"acct-batch-123","account_user_id":"acct-user-batch-123","plan_type":"plus","refresh_token":"old-refresh-token"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import"
))
.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!({
"credentials": "batch-refresh-success\nbatch-refresh-error\n",
"proxy_node_id": "proxy-node-batch-import"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["total"], 2);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 1);
let results = payload["results"]
.as_array()
.expect("results should be array");
assert_eq!(results.len(), 2);
assert_eq!(results[0]["status"], "success");
assert_eq!(results[0]["key_id"], "key-codex-batch-duplicate");
assert_eq!(results[0]["replaced"], true);
assert_eq!(results[1]["status"], "error");
assert!(
results[1]["error"]
.as_str()
.expect("error should be string")
.contains("Token 验证失败"),
"payload={payload}"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(token_hits.lock().expect("mutex should lock").len(), 2);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-batch-duplicate".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-batch-import", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "batch-imported-codex-access-token");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
run_admin_oauth_test(
"gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints",
gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints_impl,
);
}
async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::BAD_REQUEST,
Json(json!({"error": "unexpected refresh exchange"})),
)
}
}),
);
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
provider.provider_type = "chatgpt_web".to_string();
provider.config = Some(json!({"pool_advanced": {}}));
let endpoint = sample_endpoint(
"endpoint-chatgpt-web-image",
"provider-chatgpt-web",
"openai:image",
"https://chatgpt.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"chatgpt_web",
format!("{token_url}/oauth/token"),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let credentials = json!([
{
"accessToken": "chatgpt-web-batch-access-token",
"expiresAt": 2_100_000_000u64,
"email": "[email protected]",
"accountId": "acct-pool-image",
"accountUserId": "user-pool-image__acct-pool-image",
"planType": "plus",
"userId": "user-pool-image"
}
])
.to_string();
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-chatgpt-web/batch-import"
))
.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!({
"credentials": credentials,
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-chatgpt-web".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert_eq!(persisted.expires_at_unix_secs, Some(2_100_000_000));
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "chatgpt-web-batch-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "chatgpt_web");
assert_eq!(auth_config["access_token_import_temporary"], true);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-pool-image");
assert_eq!(
auth_config["account_user_id"],
"user-pool-image__acct-pool-image"
);
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["user_id"], "user-pool-image");
let score_scope = provider_key_pool_score_scope();
let score_identity =
PoolMemberIdentity::provider_api_key("provider-chatgpt-web", persisted.id.clone());
let scores = pool_score_repository
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
})
.await
.expect("pool score should load");
assert_eq!(scores.len(), 1);
assert_eq!(scores[0].member_id, persisted.id);
assert_eq!(scores[0].hard_state, PoolMemberHardState::Unknown);
assert!(scores[0].score > 0.0);
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_starts_admin_provider_oauth_batch_import_task_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_starts_admin_provider_oauth_batch_import_task_locally_with_trusted_admin_principal",
gateway_starts_admin_provider_oauth_batch_import_task_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_starts_admin_provider_oauth_batch_import_task_locally_with_trusted_admin_principal_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"access_token": "task-imported-codex-access-token",
"refresh_token": "task-imported-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-task-123",
"account_user_id": "acct-user-task-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let submit_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks"
))
.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!({
"credentials": "task-batch-refresh-success"
}))
.send()
.await
.expect("submit request should succeed");
assert_eq!(submit_response.status(), StatusCode::OK);
let submit_payload: serde_json::Value = submit_response
.json()
.await
.expect("submit payload should parse");
assert_eq!(submit_payload["status"], "submitted");
assert_eq!(submit_payload["total"], 1);
let task_id = submit_payload["task_id"]
.as_str()
.expect("task id should exist")
.to_string();
let mut status_payload = serde_json::Value::Null;
for _ in 0..40 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/{task_id}"
))
.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")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
status_payload = response.json().await.expect("status payload should parse");
if status_payload["status"] == "completed" {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
assert_eq!(
status_payload["status"], "completed",
"payload={status_payload}"
);
assert_eq!(status_payload["total"], 1);
assert_eq!(status_payload["processed"], 1);
assert_eq!(status_payload["success"], 1);
assert_eq!(status_payload["failed"], 0);
assert_eq!(status_payload["created_count"], 1);
assert_eq!(status_payload["replaced_count"], 0);
assert_eq!(status_payload["progress_percent"], 100);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_updates_admin_provider_oauth_batch_import_task_progress() {
run_admin_oauth_test(
"gateway_updates_admin_provider_oauth_batch_import_task_progress",
gateway_updates_admin_provider_oauth_batch_import_task_progress_impl,
);
}
async fn gateway_updates_admin_provider_oauth_batch_import_task_progress_impl() {
let upstream = Router::new().fallback(any(|| async {
(StatusCode::OK, Body::from("quota refresh body"))
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
let hit = {
let mut guard = token_hits_inner.lock().expect("mutex should lock");
*guard += 1;
*guard
};
if hit == 2 {
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
Json(json!({
"access_token": format!("progress-codex-access-token-{hit}"),
"refresh_token": format!("progress-codex-refresh-token-{hit}"),
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": format!("progress-{hit}@example.com"),
"account_id": format!("acct-progress-{hit}"),
"account_user_id": format!("acct-user-progress-{hit}"),
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let submit_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks"
))
.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!({
"credentials": "progress-refresh-one\nprogress-refresh-two"
}))
.send()
.await
.expect("submit request should succeed");
assert_eq!(submit_response.status(), StatusCode::OK);
let submit_payload: serde_json::Value = submit_response
.json()
.await
.expect("submit payload should parse");
let task_id = submit_payload["task_id"]
.as_str()
.expect("task id should exist")
.to_string();
let mut progress_payload = serde_json::Value::Null;
for _ in 0..50 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/{task_id}"
))
.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")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
progress_payload = response.json().await.expect("status payload should parse");
if progress_payload["status"] == "processing" && progress_payload["processed"] == 1 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
assert_eq!(
progress_payload["status"], "processing",
"payload={progress_payload}"
);
assert_eq!(progress_payload["total"], 2);
assert_eq!(progress_payload["processed"], 1);
assert_eq!(progress_payload["success"], 1);
assert_eq!(progress_payload["failed"], 0);
assert_eq!(progress_payload["created_count"], 1);
assert_eq!(progress_payload["progress_percent"], 50);
let mut completed_payload = progress_payload;
for _ in 0..50 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/batch-import/tasks/{task_id}"
))
.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")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
completed_payload = response.json().await.expect("status payload should parse");
if completed_payload["status"] == "completed" {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
}
assert_eq!(
completed_payload["status"], "completed",
"payload={completed_payload}"
);
assert_eq!(completed_payload["processed"], 2);
assert_eq!(completed_payload["success"], 2);
assert_eq!(completed_payload["progress_percent"], 100);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 2);
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_principal",
gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_principal_impl() {
#[derive(Debug, Clone)]
struct SeenTokenRequest {
content_type: String,
body: String,
}
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let seen_token = Arc::new(Mutex::new(None::<SeenTokenRequest>));
let seen_token_clone = Arc::clone(&seen_token);
let token_server = Router::new().route(
"/oauth/token",
any(move |request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let seen_token_inner = Arc::clone(&seen_token_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
*seen_token_inner.lock().expect("mutex should lock") = Some(SeenTokenRequest {
content_type: parts
.headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(raw_body.to_vec())
.expect("token request body should be utf8"),
});
Json(json!({
"access_token": "new-codex-access-token",
"refresh_token": "new-codex-refresh-token",
"token_type": "Bearer",
"expiresAt": 4_102_444_800u64,
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth",
"provider-codex",
"openai:chat",
"__placeholder__",
);
key.auth_type = "oauth".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_state_entry_for_tests(
"nonce-codex-123",
json!({
"nonce": "nonce-codex-123",
"key_id": "key-codex-oauth",
"provider_id": "provider-codex",
"provider_type": "codex",
"pkce_verifier": "verifier-codex-123",
}),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth/complete"
))
.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:1455/auth/callback?code=code-codex-123&state=nonce-codex-123"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["expires_at"], 4_102_444_800u64);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["account_state_recheck_attempted"], false);
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_token = seen_token
.lock()
.expect("mutex should lock")
.clone()
.expect("token request should be recorded");
assert_eq!(seen_token.content_type, "application/x-www-form-urlencoded");
assert!(seen_token.body.contains("grant_type=authorization_code"));
assert!(seen_token
.body
.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
assert!(seen_token.body.contains("code=code-codex-123"));
assert!(seen_token.body.contains("code_verifier=verifier-codex-123"));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert_eq!(persisted.expires_at_unix_secs, Some(4_102_444_800));
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "new-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(auth_config["refresh_token"], "new-codex-refresh-token");
assert_eq!(auth_config["expires_at"], 4_102_444_800u64);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_completes_admin_provider_oauth_provider_locally_with_trusted_admin_principal",
gateway_completes_admin_provider_oauth_provider_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_admin_principal_impl()
{
#[derive(Debug, Clone)]
struct SeenTokenRequest {
content_type: String,
body: String,
}
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let seen_token = Arc::new(Mutex::new(None::<SeenTokenRequest>));
let seen_token_clone = Arc::clone(&seen_token);
let token_server = Router::new().route(
"/oauth/token",
any(move |request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let seen_token_inner = Arc::clone(&seen_token_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
*seen_token_inner.lock().expect("mutex should lock") = Some(SeenTokenRequest {
content_type: parts
.headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(raw_body.to_vec())
.expect("token request body should be utf8"),
});
Json(json!({
"access_token": "provider-codex-access-token",
"refresh_token": "provider-codex-refresh-token",
"token_type": "Bearer",
"expires_at": 4_102_444_800u64,
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let mut existing_key = sample_key(
"key-codex-inactive-duplicate",
"provider-codex",
"openai:chat",
"stale-codex-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","refresh_token":"old-refresh-token"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_state_entry_for_tests(
"nonce-provider-codex-123",
json!({
"nonce": "nonce-provider-codex-123",
"key_id": "",
"provider_id": "provider-codex",
"provider_type": "codex",
"pkce_verifier": "verifier-provider-codex-123",
}),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/complete"
))
.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:1455/auth/callback?code=provider-code-123&state=nonce-provider-codex-123",
"proxy_node_id": "proxy-node-codex-oauth",
"name": "should-not-override-inactive-name"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["key_id"], "key-codex-inactive-duplicate");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["expires_at"], 4_102_444_800u64);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_token = seen_token
.lock()
.expect("mutex should lock")
.clone()
.expect("token request should be recorded");
assert_eq!(seen_token.content_type, "application/x-www-form-urlencoded");
assert!(seen_token.body.contains("grant_type=authorization_code"));
assert!(seen_token.body.contains("code=provider-code-123"));
assert!(seen_token
.body
.contains("code_verifier=verifier-provider-codex-123"));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-inactive-duplicate".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(persisted.expires_at_unix_secs, Some(4_102_444_800));
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-oauth", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "provider-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(auth_config["refresh_token"], "provider-codex-refresh-token");
assert_eq!(auth_config["expires_at"], 4_102_444_800u64);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal",
gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal_impl(
) {
#[derive(Debug, Clone)]
struct SeenTokenRequest {
content_type: String,
body: String,
}
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let seen_token = Arc::new(Mutex::new(None::<SeenTokenRequest>));
let seen_token_clone = Arc::clone(&seen_token);
let token_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let seen_token_inner = Arc::clone(&seen_token_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
*seen_token_inner.lock().expect("mutex should lock") = Some(SeenTokenRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
Json(json!({
"access_token": "imported-codex-access-token",
"refresh_token": "imported-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let mut existing_key = sample_key(
"key-codex-import-duplicate",
"provider-codex",
"openai:chat",
"stale-imported-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","refresh_token":"old-refresh-token"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"proxy_node_id": "proxy-node-codex-import",
"name": "should-not-override-inactive-name"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["key_id"], "key-codex-import-duplicate");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_token = seen_token
.lock()
.expect("mutex should lock")
.clone()
.expect("token request should be recorded");
assert_eq!(seen_token.content_type, "application/x-www-form-urlencoded");
assert!(seen_token.body.contains("grant_type=refresh_token"));
assert!(seen_token
.body
.contains("refresh_token=provider-import-refresh-token"));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-import-duplicate".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "imported-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(auth_config["refresh_token"], "imported-codex-refresh-token");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_imports_codex_access_token_without_refresh_token_as_temporary_account() {
run_admin_oauth_test(
"gateway_imports_codex_access_token_without_refresh_token_as_temporary_account",
gateway_imports_codex_access_token_without_refresh_token_as_temporary_account_impl,
);
}
async fn gateway_imports_codex_access_token_without_refresh_token_as_temporary_account_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move || {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::BAD_REQUEST,
Json(json!({"error": "unexpected refresh exchange"})),
)
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let access_token =
sample_codex_access_token_with_profile_email("[email protected]", "acct-profile-123");
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"access_token": access_token,
"name": "temporary-codex-access-token",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], false);
assert_eq!(payload["temporary"], true);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(*token_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert_eq!(persisted.expires_at_unix_secs, Some(2_000_000_000));
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(auth_config["access_token_import_temporary"], true);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-profile-123");
assert!(auth_config.get("refresh_token").is_none());
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_imports_codex_header_authorization_without_overwriting_payload_access_token() {
run_admin_oauth_test(
"gateway_imports_codex_header_authorization_without_overwriting_payload_access_token",
gateway_imports_codex_header_authorization_without_overwriting_payload_access_token_impl,
);
}
async fn gateway_imports_codex_header_authorization_without_overwriting_payload_access_token_impl()
{
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let access_token =
sample_codex_access_token_with_profile_email("[email protected]", "acct-profile-123");
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"access_token": access_token,
"headers": {
"authorization": "Bearer imported-session-token",
"chatgpt-account-id": "acct-header"
},
"name": "temporary-codex-header-auth",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["access_token"], access_token);
assert_eq!(
auth_config["headers"]["authorization"],
"Bearer imported-session-token"
);
assert_eq!(auth_config["headers"]["chatgpt-account-id"], "acct-header");
gateway_handle.abort();
}
#[test]
fn gateway_imports_chatgpt_web_access_token_without_refresh_token_as_temporary_account() {
run_admin_oauth_test(
"gateway_imports_chatgpt_web_access_token_without_refresh_token_as_temporary_account",
gateway_imports_chatgpt_web_access_token_without_refresh_token_as_temporary_account_impl,
);
}
async fn gateway_imports_chatgpt_web_access_token_without_refresh_token_as_temporary_account_impl()
{
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move || {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::BAD_REQUEST,
Json(json!({"error": "unexpected refresh exchange"})),
)
}
}),
);
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
provider.provider_type = "chatgpt_web".to_string();
let endpoint = sample_endpoint(
"endpoint-chatgpt-web-image",
"provider-chatgpt-web",
"openai:image",
"https://chatgpt.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"chatgpt_web",
format!("{token_url}/oauth/token"),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let access_token =
sample_codex_access_token_with_profile_email("[email protected]", "acct-image-123");
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-chatgpt-web/import-refresh-token"
))
.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!({
"access_token": access_token,
"name": "temporary-chatgpt-web-access-token",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "chatgpt_web");
assert_eq!(payload["has_refresh_token"], false);
assert_eq!(payload["temporary"], true);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(*token_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-chatgpt-web".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert_eq!(persisted.expires_at_unix_secs, Some(2_000_000_000));
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "chatgpt_web");
assert_eq!(auth_config["access_token_import_temporary"], true);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-image-123");
assert!(auth_config.get("refresh_token").is_none());
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp() {
run_admin_oauth_test(
"gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp",
gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp_impl,
);
}
async fn gateway_imports_codex_access_token_with_payload_expires_at_when_token_has_no_exp_impl() {
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-responses",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"access_token": "opaque-codex-access-token",
"expiresAt": 2_100_000_000u64,
"name": "temporary-codex-opaque-access-token",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["expires_at"], 2_100_000_000u64);
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert_eq!(persisted.expires_at_unix_secs, Some(2_100_000_000));
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["expires_at"], 2_100_000_000u64);
assert_eq!(auth_config["access_token_import_temporary"], true);
gateway_handle.abort();
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_over_active_expired_duplicate() {
run_admin_oauth_test(
"gateway_imports_admin_provider_oauth_refresh_token_over_active_expired_duplicate",
gateway_imports_admin_provider_oauth_refresh_token_over_active_expired_duplicate_impl,
);
}
async fn gateway_imports_admin_provider_oauth_refresh_token_over_active_expired_duplicate_impl() {
#[derive(Debug, Clone)]
struct SeenTokenRequest {
content_type: String,
body: String,
}
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let seen_token = Arc::new(Mutex::new(None::<SeenTokenRequest>));
let seen_token_clone = Arc::clone(&seen_token);
let token_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let seen_token_inner = Arc::clone(&seen_token_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
*seen_token_inner.lock().expect("mutex should lock") = Some(SeenTokenRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
Json(json!({
"access_token": "imported-expired-codex-access-token",
"refresh_token": "imported-expired-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let mut existing_key = sample_key(
"key-codex-import-expired-duplicate",
"provider-codex",
"openai:chat",
"stale-imported-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = true;
existing_key.expires_at_unix_secs = Some(1);
existing_key.oauth_invalid_at_unix_secs = Some(1_700_000_000);
existing_key.oauth_invalid_reason =
Some("[REFRESH_FAILED] refresh_token 无效、已过期或已撤销，请重新登录授权".to_string());
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","refresh_token":"old-refresh-token","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"proxy_node_id": "proxy-node-codex-import",
"name": "should-not-override-active-expired-name"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["key_id"], "key-codex-import-expired-duplicate");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_token = seen_token
.lock()
.expect("mutex should lock")
.clone()
.expect("token request should be recorded");
assert_eq!(seen_token.content_type, "application/x-www-form-urlencoded");
assert!(seen_token.body.contains("grant_type=refresh_token"));
assert!(seen_token
.body
.contains("refresh_token=provider-import-refresh-token"));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-import-expired-duplicate".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
assert!(persisted.is_active);
assert_eq!(
persisted.proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
assert_eq!(persisted.oauth_invalid_at_unix_secs, None);
assert_eq!(persisted.oauth_invalid_reason, None);
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "imported-expired-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"imported-expired-codex-refresh-token"
);
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
gateway_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_import_invalidate_cached_oauth_entry_before_followup_resolution() {
run_admin_oauth_test(
"gateway_import_invalidate_cached_oauth_entry_before_followup_resolution",
gateway_import_invalidate_cached_oauth_entry_before_followup_resolution_impl,
);
}
async fn gateway_import_invalidate_cached_oauth_entry_before_followup_resolution_impl() {
let token_server = Router::new().route(
"/oauth/token",
post(move |body: Bytes| async move {
let body_text = String::from_utf8(body.to_vec()).unwrap_or_default();
if body_text.contains("refresh_token=old-refresh-token") {
Json(json!({
"access_token": "cached-old-codex-access-token",
"refresh_token": "cached-old-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
} else {
assert!(
body_text.contains("refresh_token=provider-import-refresh-token"),
"unexpected token request body: {body_text}"
);
Json(json!({
"access_token": "imported-fresh-codex-access-token",
"refresh_token": "imported-fresh-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus",
}))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let mut existing_key = sample_key(
"key-codex-import-cache-duplicate",
"provider-codex",
"openai:chat",
"stale-imported-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.expires_at_unix_secs = Some(1);
existing_key.oauth_invalid_at_unix_secs = Some(1_700_000_000);
existing_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","refresh_token":"old-refresh-token","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let app_state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", format!("{token_url}/oauth/token"))
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let stale_transport = app_state
.read_provider_transport_snapshot(
"provider-codex",
"endpoint-codex-chat",
"key-codex-import-cache-duplicate",
)
.await
.expect("transport should load")
.expect("transport should exist");
let cached_entry = app_state
.force_local_oauth_refresh_entry(&stale_transport)
.await
.expect("initial refresh should succeed")
.expect("initial refresh should return cached entry");
assert_eq!(
cached_entry.auth_header_value,
"Bearer cached-old-codex-access-token"
);
let mut replaceable_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-import-cache-duplicate".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("key should exist");
replaceable_key.oauth_invalid_at_unix_secs = Some(1_700_000_000);
replaceable_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
provider_catalog_repository
.update_key(&replaceable_key)
.await
.expect("key should update");
let gateway = build_router_with_state(app_state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"name": "should-not-override-cache-duplicate-name"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["key_id"], "key-codex-import-cache-duplicate");
assert_eq!(payload["replaced"], true);
let fresh_transport = app_state
.read_provider_transport_snapshot(
"provider-codex",
"endpoint-codex-chat",
"key-codex-import-cache-duplicate",
)
.await
.expect("transport should load")
.expect("transport should exist");
let resolved = app_state
.resolve_local_oauth_request_auth(&fresh_transport)
.await
.expect("oauth auth should resolve")
.expect("oauth auth should exist");
match resolved {
crate::provider_transport::LocalResolvedOAuthRequestAuth::Header { value, .. } => {
assert_eq!(value, "Bearer imported-fresh-codex-access-token");
}
crate::provider_transport::LocalResolvedOAuthRequestAuth::Kiro(_) => {
panic!("codex should resolve to header auth")
}
}
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_rejects_kiro_single_refresh_token_import_with_clear_error() {
run_admin_oauth_test(
"gateway_rejects_kiro_single_refresh_token_import_with_clear_error",
gateway_rejects_kiro_single_refresh_token_import_with_clear_error_impl,
);
}
async fn gateway_rejects_kiro_single_refresh_token_import_with_clear_error_impl() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![{
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
provider
}],
vec![sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
)],
Vec::new(),
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_provider_oauth_response(
&state,
http::Method::POST,
"/api/admin/provider-oauth/providers/provider-kiro/import-refresh-token",
Some(json!({
"refresh_token": "kiro-refresh-token"
})),
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
assert_eq!(
payload["detail"],
json!("Kiro 不支持单条 Refresh Token 导入，请使用批量导入或设备授权。")
);
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node() {
run_admin_oauth_test(
"gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node",
gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node_impl,
);
}
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_proxy_node_impl()
{
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-codex-import"));
let timeouts = plan.timeouts.as_ref().expect("timeouts should exist");
assert_eq!(timeouts.connect_ms, Some(60_000));
assert_eq!(timeouts.read_ms, Some(60_000));
assert_eq!(timeouts.write_ms, Some(60_000));
assert_eq!(timeouts.pool_ms, Some(60_000));
assert_eq!(timeouts.total_ms, Some(60_000));
if plan.request_id == "provider-oauth:refresh-token" {
assert_eq!(plan.method, "POST");
assert_eq!(plan.url, "https://oauth.example/oauth/token");
assert_eq!(
plan.headers.get("content-type").map(String::as_str),
Some("application/x-www-form-urlencoded")
);
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
Json(codex_import_token_execution_result(&plan.request_id))
} else {
assert!(
plan.request_id.starts_with("codex-quota:"),
"unexpected execution plan: {}",
plan.request_id
);
Json(codex_quota_execution_result(&plan.request_id))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut manual_node = sample_proxy_node("proxy-node-codex-import");
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"proxy_node_id": "proxy-node-codex-import",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
assert_eq!(
keys[0].proxy,
Some(json!({"node_id": "proxy-node-codex-import", "enabled": true}))
);
let plans = execution_plans.lock().expect("mutex should lock");
let token_plan = assert_single_provider_oauth_refresh_token_plan(&plans);
assert_eq!(
token_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-codex-import")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_provider_proxy_before_system_proxy(
) {
run_admin_oauth_test(
"gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_provider_proxy_before_system_proxy",
gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_provider_proxy_before_system_proxy_impl,
);
}
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_provider_proxy_before_system_proxy_impl(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-codex-provider"));
if plan.request_id == "provider-oauth:refresh-token" {
Json(codex_import_token_execution_result(&plan.request_id))
} else {
assert!(
plan.request_id.starts_with("codex-quota:"),
"unexpected execution plan: {}",
plan.request_id
);
Json(codex_quota_execution_result(&plan.request_id))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-codex-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut provider_node = sample_proxy_node("proxy-node-codex-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let mut system_node = sample_proxy_node("proxy-node-codex-system");
system_node.status = "online".to_string();
system_node.is_manual = true;
system_node.tunnel_mode = false;
system_node.tunnel_connected = false;
system_node.proxy_url = Some("http://proxy-system.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
provider_node,
system_node,
]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-codex-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "codex");
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].proxy, None);
let plans = execution_plans.lock().expect("mutex should lock");
let token_plan = assert_single_provider_oauth_refresh_token_plan(&plans);
assert_eq!(
token_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-codex-provider")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_system_proxy() {
run_admin_oauth_test(
"gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_system_proxy",
gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_system_proxy_impl,
);
}
async fn gateway_imports_admin_provider_oauth_refresh_token_via_execution_runtime_system_proxy_impl(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
let proxy = plan.proxy.as_ref().expect("proxy snapshot should exist");
assert_eq!(proxy.node_id.as_deref(), Some("proxy-node-codex-system"));
if plan.request_id == "provider-oauth:refresh-token" {
Json(codex_import_token_execution_result(&plan.request_id))
} else {
assert!(
plan.request_id.starts_with("codex-quota:"),
"unexpected execution plan: {}",
plan.request_id
);
Json(codex_quota_execution_result(&plan.request_id))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let mut manual_node = sample_proxy_node("proxy-node-codex-system");
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-codex-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "codex");
let plans = execution_plans.lock().expect("mutex should lock");
let token_plan = assert_single_provider_oauth_refresh_token_plan(&plans);
assert_eq!(
token_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-codex-system")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_import_refresh_token_surfaces_execution_runtime_error_detail() {
run_admin_oauth_test(
"gateway_import_refresh_token_surfaces_execution_runtime_error_detail",
gateway_import_refresh_token_surfaces_execution_runtime_error_detail_impl,
);
}
async fn gateway_import_refresh_token_surfaces_execution_runtime_error_detail_impl() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|| async { StatusCode::INTERNAL_SERVER_ERROR }),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.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!({
"refresh_token": "provider-import-refresh-token",
"name": "codex-import"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::BAD_REQUEST, "payload={payload}");
assert!(
payload["detail"]
.as_str()
.expect("detail should be string")
.contains("execution runtime returned HTTP 500"),
"payload={payload}"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_admin_principal",
gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_batch_imports_admin_provider_oauth_kiro_locally_with_trusted_admin_principal_impl()
{
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/refreshToken",
post(move |_request: Request| {
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"accessToken": sample_kiro_device_access_token("[email protected]"),
"refreshToken": "kiro-batch-refresh-token-new",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let endpoint = sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
);
let mut existing_key = sample_key(
"key-kiro-batch-duplicate",
"provider-kiro",
"kiro:generateAssistantResponse",
"stale-kiro-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","auth_method":"social","email":"[email protected]","refresh_token":"kiro-batch-refresh-old"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (refresh_url, refresh_handle) = start_server(refresh_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_refresh",
refresh_url.to_string(),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import"
))
.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!({
"credentials": "kiro-batch-refresh-token",
"proxy_node_id": "proxy-node-kiro-batch"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("payload should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
let results = payload["results"]
.as_array()
.expect("results should be array");
assert_eq!(results.len(), 1);
assert_eq!(results[0]["status"], "success");
assert_eq!(results[0]["key_id"], "key-kiro-batch-duplicate");
assert_eq!(results[0]["auth_method"], "social");
assert_eq!(results[0]["replaced"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-batch-duplicate".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch", "enabled": true}))
);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["auth_method"], "social");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["refresh_token"], "kiro-batch-refresh-token-new");
gateway_handle.abort();
refresh_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_batch_imports_admin_provider_oauth_kiro_over_active_expired_duplicate() {
run_admin_oauth_test(
"gateway_batch_imports_admin_provider_oauth_kiro_over_active_expired_duplicate",
gateway_batch_imports_admin_provider_oauth_kiro_over_active_expired_duplicate_impl,
);
}
async fn gateway_batch_imports_admin_provider_oauth_kiro_over_active_expired_duplicate_impl() {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/refreshToken",
post(move |_request: Request| {
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"accessToken": sample_kiro_device_access_token("[email protected]"),
"refreshToken": "kiro-batch-refresh-token-replaced",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let endpoint = sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
);
let mut existing_key = sample_key(
"key-kiro-batch-expired-duplicate",
"provider-kiro",
"kiro:generateAssistantResponse",
"stale-kiro-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = true;
existing_key.oauth_invalid_at_unix_secs = Some(1_700_000_000);
existing_key.oauth_invalid_reason = Some("Kiro Token 无效或已过期".to_string());
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","auth_method":"social","email":"[email protected]","refresh_token":"kiro-batch-refresh-old"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (refresh_url, refresh_handle) = start_server(refresh_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_refresh",
refresh_url.to_string(),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import"
))
.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!({
"credentials": "kiro-batch-refresh-token",
"proxy_node_id": "proxy-node-kiro-batch"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("payload should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
let results = payload["results"]
.as_array()
.expect("results should be array");
assert_eq!(results.len(), 1);
assert_eq!(results[0]["status"], "success");
assert_eq!(results[0]["key_id"], "key-kiro-batch-expired-duplicate");
assert_eq!(results[0]["auth_method"], "social");
assert_eq!(results[0]["replaced"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-batch-expired-duplicate".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch", "enabled": true}))
);
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
assert_eq!(stored_key.oauth_invalid_reason, None);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["auth_method"], "social");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(
auth_config["refresh_token"],
"kiro-batch-refresh-token-replaced"
);
gateway_handle.abort();
refresh_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_proxy_node() {
run_admin_oauth_test(
"gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_proxy_node",
gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_proxy_node_impl,
);
}
async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_proxy_node_impl() {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
assert_eq!(
plan.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-kiro-batch-runtime")
);
if plan.request_id == "provider-oauth:kiro-social-refresh" {
assert_eq!(plan.url, "https://oauth.example/refreshToken");
assert_eq!(
plan.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
return Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"accessToken": sample_kiro_device_access_token("[email protected]"),
"refreshToken": "kiro-runtime-refresh-token-new",
"expiresIn": 1800,
}
}
}));
}
assert_eq!(plan.request_id, "kiro-quota:key-kiro-batch-runtime");
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO+"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 5.0,
"usageLimitWithPrecision": 20.0,
"nextDateReset": 1_900_000_000u64
}],
"desktopUserInfo": {
"email": "[email protected]"
}
}
}
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let endpoint = sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
);
let mut existing_key = sample_key(
"key-kiro-batch-runtime",
"provider-kiro",
"kiro:generateAssistantResponse",
"stale-kiro-runtime-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.is_active = false;
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"kiro","auth_method":"social","email":"[email protected]","refresh_token":"kiro-runtime-refresh-token-old"}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![existing_key],
));
let mut manual_node = sample_proxy_node("proxy-node-kiro-batch-runtime");
manual_node.status = "online".to_string();
manual_node.is_manual = true;
manual_node.tunnel_mode = false;
manual_node.tunnel_connected = false;
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![manual_node]));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_refresh",
"https://oauth.example",
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import"
))
.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!({
"credentials": "kiro-runtime-refresh-token-old",
"proxy_node_id": "proxy-node-kiro-batch-runtime"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("payload should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(payload["results"][0]["key_id"], "key-kiro-batch-runtime");
assert_eq!(payload["results"][0]["replaced"], true);
for _ in 0..40 {
let plan_count = execution_plans.lock().expect("mutex should lock").len();
if plan_count == 2 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
{
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);
}
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-batch-runtime".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("persisted key should exist");
assert!(stored_key.is_active);
assert_eq!(
stored_key.proxy,
Some(json!({"node_id": "proxy-node-kiro-batch-runtime", "enabled": true}))
);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(
auth_config["refresh_token"],
"kiro-runtime-refresh-token-new"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_starts_admin_provider_oauth_kiro_batch_import_task_locally_with_trusted_admin_principal()
{
run_admin_oauth_test(
"gateway_starts_admin_provider_oauth_kiro_batch_import_task_locally_with_trusted_admin_principal",
gateway_starts_admin_provider_oauth_kiro_batch_import_task_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_starts_admin_provider_oauth_kiro_batch_import_task_locally_with_trusted_admin_principal_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/refreshToken",
post(move |_request: Request| {
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"accessToken": sample_kiro_device_access_token("[email protected]"),
"refreshToken": "kiro-task-refresh-token-new",
"expiresIn": 1800,
}))
}
}),
);
let mut provider = sample_provider("provider-kiro", "kiro", 10);
provider.provider_type = "kiro".to_string();
let endpoint = sample_endpoint(
"endpoint-kiro-chat",
"provider-kiro",
"kiro:generateAssistantResponse",
"https://service.kiro.dev",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (refresh_url, refresh_handle) = start_server(refresh_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_token_url_for_tests(
"kiro_social_refresh",
refresh_url.to_string(),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let submit_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import/tasks"
))
.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!({
"credentials": "kiro-task-refresh-token"
}))
.send()
.await
.expect("submit request should succeed");
assert_eq!(submit_response.status(), StatusCode::OK);
let submit_payload: serde_json::Value = submit_response
.json()
.await
.expect("submit payload should parse");
assert_eq!(submit_payload["status"], "submitted");
let task_id = submit_payload["task_id"]
.as_str()
.expect("task id should exist")
.to_string();
let mut status_payload = serde_json::Value::Null;
for _ in 0..40 {
let response = client
.get(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-kiro/batch-import/tasks/{task_id}"
))
.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")
.send()
.await
.expect("status request should succeed");
assert_eq!(response.status(), StatusCode::OK);
status_payload = response.json().await.expect("status payload should parse");
if status_payload["status"] == "completed" {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
assert_eq!(
status_payload["status"], "completed",
"payload={status_payload}"
);
assert_eq!(status_payload["total"], 1);
assert_eq!(status_payload["processed"], 1);
assert_eq!(status_payload["success"], 1);
assert_eq!(status_payload["failed"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-kiro".to_string()])
.await
.expect("keys should load");
assert_eq!(keys.len(), 1);
gateway_handle.abort();
refresh_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_marks_lazy_codex_oauth_refresh_failures_as_invalid() {
run_admin_oauth_test(
"gateway_marks_lazy_codex_oauth_refresh_failures_as_invalid",
gateway_marks_lazy_codex_oauth_refresh_failures_as_invalid_impl,
);
}
async fn gateway_marks_lazy_codex_oauth_refresh_failures_as_invalid_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::UNAUTHORIZED,
Json(json!({
"error": {
"message": "Your refresh token has already been used to generate a new access token. Please try signing in again.",
"type": "invalid_request_error",
"param": serde_json::Value::Null,
"code": "refresh_token_reused"
}
})),
)
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-lazy",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let transport = state
.read_provider_transport_snapshot(
"provider-codex",
"endpoint-codex-cli",
"key-codex-oauth-lazy",
)
.await
.expect("transport snapshot should load")
.expect("transport snapshot should exist");
let resolved = state
.resolve_local_oauth_request_auth(&transport)
.await
.expect("refresh-token reuse should degrade into oauth-unavailable");
assert_eq!(resolved, None);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-lazy".to_string()])
.await
.expect("keys should list")
.into_iter()
.next()
.expect("oauth key should exist");
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权")
);
let oauth_snapshot = stored_key
.status_snapshot
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|snapshot| snapshot.get("oauth"))
.and_then(serde_json::Value::as_object)
.expect("oauth status snapshot should exist");
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("invalid")
);
assert_eq!(
oauth_snapshot
.get("source")
.and_then(serde_json::Value::as_str),
Some("oauth_refresh")
);
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(true)
);
token_handle.abort();
}
#[test]
fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_principal",
gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_principal_impl() {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
}))
}
}),
);
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 401,
headers: std::collections::BTreeMap::new(),
body: None,
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.oauth_invalid_at_unix_secs = Some(1_700_000_000);
key.oauth_invalid_reason = Some("[REFRESH_FAILED] stale token".to_string());
key.status_snapshot = Some(json!({
"oauth": {
"code": "expired",
"label": "已过期",
"reason": "Access Token 已过期，等待自动续期",
"expires_at": 1u64,
"invalid_at": 1_700_000_000u64,
"source": "expires_at",
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": null,
"reason": null,
"blocked": false,
"source": null,
"recoverable": false
},
"quota": {
"code": "unknown",
"label": null,
"reason": null,
"exhausted": false,
"usage_ratio": null,
"updated_at": null,
"reset_seconds": null,
"plan_type": null
}
}));
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["email"], "[email protected]");
let account_state_recheck_attempted = payload["account_state_recheck_attempted"]
.as_bool()
.expect("account_state_recheck_attempted should be bool");
if account_state_recheck_attempted {
let account_state_recheck_error = payload["account_state_recheck_error"]
.as_str()
.expect("account_state_recheck_error should be string when attempted");
assert!(
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error == "wham/usage API 返回状态码 403"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
} else {
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
}
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer refreshed-codex-access-token"
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-refresh".to_string()])
.await
.expect("keys should list")
.into_iter()
.next()
.expect("refreshed key should exist");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("refreshed api key should decrypt");
assert_eq!(decrypted_api_key, "refreshed-codex-access-token");
if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 401"
{
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] Codex Token 已过期 (401)")
);
} else if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 403"
{
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert!(stored_key
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.contains("(403)")));
} else {
assert_eq!(stored_key.oauth_invalid_at_unix_secs, None);
assert_eq!(stored_key.oauth_invalid_reason, None);
}
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("refreshed auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(
auth_config["refresh_token"],
serde_json::Value::String("refreshed-codex-refresh-token".to_string())
);
assert_eq!(
auth_config["email"],
serde_json::Value::String("[email protected]".to_string())
);
let status_snapshot = stored_key
.status_snapshot
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("status snapshot should exist");
let oauth_snapshot = status_snapshot
.get("oauth")
.and_then(serde_json::Value::as_object)
.expect("oauth snapshot should exist");
assert_eq!(
oauth_snapshot.get("expires_at"),
auth_config.get("expires_at")
);
if stored_key
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with("[OAUTH_EXPIRED]"))
{
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("expired")
);
assert_eq!(
oauth_snapshot
.get("label")
.and_then(serde_json::Value::as_str),
Some("已过期")
);
assert_eq!(
oauth_snapshot.get("reason"),
Some(&json!("Codex Token 已过期 (401)"))
);
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(false)
);
assert_eq!(
oauth_snapshot
.get("expiring_soon")
.and_then(serde_json::Value::as_bool),
Some(false)
);
} else if stored_key.oauth_invalid_reason.is_some() {
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("invalid")
);
assert_eq!(
oauth_snapshot
.get("label")
.and_then(serde_json::Value::as_str),
Some("已失效")
);
assert_eq!(
oauth_snapshot
.get("reason")
.and_then(serde_json::Value::as_str)
.is_some_and(|reason| !reason.trim().is_empty()),
true
);
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(true)
);
assert_eq!(
oauth_snapshot
.get("expiring_soon")
.and_then(serde_json::Value::as_bool),
Some(false)
);
} else {
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("expiring")
);
assert_eq!(
oauth_snapshot
.get("label")
.and_then(serde_json::Value::as_str),
Some("即将过期")
);
assert_eq!(oauth_snapshot.get("reason"), Some(&serde_json::Value::Null));
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(false)
);
assert_eq!(
oauth_snapshot
.get("expiring_soon")
.and_then(serde_json::Value::as_bool),
Some(true)
);
}
gateway_handle.abort();
execution_runtime_handle.abort();
token_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint() {
run_admin_oauth_test(
"gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint",
gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_impl,
);
}
async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"expires_in": 3600,
"token_type": "Bearer"
}))
}
}),
);
let seen_endpoint_id = Arc::new(Mutex::new(None::<String>));
let seen_endpoint_id_clone = Arc::clone(&seen_endpoint_id);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_endpoint_id_inner = Arc::clone(&seen_endpoint_id_clone);
async move {
let plan: ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_endpoint_id_inner.lock().expect("mutex should lock") =
Some(plan.endpoint_id.clone());
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: std::collections::BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 12.5,
"window_minutes": 300
},
"secondary_window": {
"used_percent": 55.0,
"window_minutes": 10080
}
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut provider = sample_provider("provider-codex-missing-endpoint", "codex", 10);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-missing-endpoint-refresh",
"provider-codex-missing-endpoint",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-missing-endpoint-refresh/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["account_state_recheck_attempted"], true);
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-codex-missing-endpoint".to_string()])
.await
.expect("endpoints should read");
let responses_endpoint = endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("openai responses endpoint should be reconciled");
assert!(endpoints
.iter()
.any(|endpoint| endpoint.api_format == "openai:search"));
assert_eq!(
responses_endpoint.base_url,
"https://chatgpt.com/backend-api/codex"
);
assert_eq!(
*seen_endpoint_id.lock().expect("mutex should lock"),
Some(responses_endpoint.id.clone())
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-missing-endpoint-refresh".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("refreshed key should exist");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_api_key
.as_deref()
.expect("api key should exist"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "refreshed-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"refreshed-codex-refresh-token"
);
gateway_handle.abort();
execution_runtime_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint() {
run_admin_oauth_test(
"gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint",
gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint_impl,
);
}
async fn gateway_manual_kiro_oauth_refresh_reconciles_missing_fixed_endpoint_impl() {
run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(None, None, true).await;
}
#[test]
fn gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance() {
run_admin_oauth_test(
"gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance",
gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance_impl,
);
}
async fn gateway_manual_kiro_oauth_refresh_uses_disabled_fixed_endpoint_for_maintenance_impl() {
let mut endpoint = sample_endpoint(
"endpoint-kiro-disabled-maintenance",
"provider-kiro-oauth-refresh",
"claude:messages",
"https://q.{region}.amazonaws.com",
);
endpoint.is_active = false;
run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(
Some(endpoint),
Some("endpoint-kiro-disabled-maintenance"),
false,
)
.await;
}
async fn run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(
initial_endpoint: Option<StoredProviderCatalogEndpoint>,
expected_endpoint_id: Option<&str>,
expected_endpoint_active: bool,
) {
let refreshed_access_token = sample_kiro_device_access_token("[email protected]");
let expected_access_token = refreshed_access_token.clone();
let refreshed_refresh_token = "s".repeat(120);
let expected_refresh_token = refreshed_refresh_token.clone();
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/refreshToken",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let access_token_inner = refreshed_access_token.clone();
let refresh_token_inner = refreshed_refresh_token.clone();
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"accessToken": access_token_inner,
"refreshToken": refresh_token_inner,
"expiresIn": 3600,
"profileArn": "arn:aws:kiro:profile/manual-refresh"
}))
}
}),
);
let seen_endpoint_id = Arc::new(Mutex::new(None::<String>));
let seen_endpoint_id_clone = Arc::clone(&seen_endpoint_id);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_endpoint_id_inner = Arc::clone(&seen_endpoint_id_clone);
async move {
let plan: ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_endpoint_id_inner.lock().expect("mutex should lock") =
Some(plan.endpoint_id.clone());
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: std::collections::BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 2.0,
"usageLimitWithPrecision": 10.0,
"nextDateReset": 1_900_000_000u64
}],
"desktopUserInfo": {
"email": "[email protected]"
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut provider = sample_provider("provider-kiro-oauth-refresh", "kiro", 10);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-oauth-refresh",
"provider-kiro-oauth-refresh",
"claude:messages",
"stale-kiro-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"provider_type": "kiro",
"auth_method": "social",
"refresh_token": "r".repeat(120),
"machine_id": "123e4567-e89b-12d3-a456-426614174000",
"kiro_version": "1.2.3",
"expires_at": 1u64
})
.to_string(),
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
initial_endpoint.into_iter().collect(),
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::kiro::KiroOAuthRefreshAdapter::default()
.with_refresh_base_urls(Some(token_url), None),
)
as Arc<dyn crate::provider_transport::oauth_refresh::LocalOAuthRefreshAdapter>,
]);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-kiro-oauth-refresh/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "kiro");
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["account_state_recheck_attempted"], true);
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
assert!(
*token_hits.lock().expect("mutex should lock") >= 1,
"Kiro refresh endpoint should be called"
);
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-kiro-oauth-refresh".to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].api_format, "claude:messages");
assert_eq!(endpoints[0].base_url, "https://q.{region}.amazonaws.com");
assert_eq!(endpoints[0].is_active, expected_endpoint_active);
if let Some(expected_endpoint_id) = expected_endpoint_id {
assert_eq!(endpoints[0].id, expected_endpoint_id);
}
assert_eq!(
*seen_endpoint_id.lock().expect("mutex should lock"),
Some(endpoints[0].id.clone())
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-kiro-oauth-refresh".to_string()])
.await
.expect("keys should load")
.into_iter()
.next()
.expect("refreshed key should exist");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_api_key
.as_deref()
.expect("api key should exist"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, expected_access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["provider_type"], "kiro");
assert_eq!(auth_config["refresh_token"], expected_refresh_token);
gateway_handle.abort();
execution_runtime_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload() {
run_admin_oauth_test(
"gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload",
gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload_impl,
);
}
async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::UNAUTHORIZED,
Json(json!({
"error": {
"message": "Your refresh token has already been used to generate a new access token. Please try signing in again.",
"type": "invalid_request_error",
"param": serde_json::Value::Null,
"code": "refresh_token_reused"
}
})),
)
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-invalid",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(4_102_444_800);
key.status_snapshot = Some(json!({
"oauth": {
"code": "valid",
"label": "有效",
"reason": serde_json::Value::Null,
"expires_at": 4_102_444_800u64,
"invalid_at": serde_json::Value::Null,
"source": "expires_at",
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"blocked": false,
"source": serde_json::Value::Null,
"recoverable": false
},
"quota": {
"code": "unknown",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"exhausted": false,
"usage_ratio": serde_json::Value::Null,
"updated_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"plan_type": serde_json::Value::Null
}
}));
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"used-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":4102444800}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let refresh_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-invalid/refresh"
))
.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")
.send()
.await
.expect("refresh request should succeed");
assert_eq!(refresh_response.status(), StatusCode::BAD_REQUEST);
let refresh_payload: serde_json::Value = refresh_response
.json()
.await
.expect("refresh payload should parse");
assert_eq!(
refresh_payload["detail"],
json!("Token 刷新失败：refresh_token 已被使用并轮换，请重新登录授权")
);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let pool_response = client
.get(format!(
"{gateway_url}/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all"
))
.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")
.send()
.await
.expect("pool request should succeed");
assert_eq!(pool_response.status(), StatusCode::OK);
let pool_payload: serde_json::Value = pool_response
.json()
.await
.expect("pool payload should parse");
let keys = pool_payload["keys"]
.as_array()
.expect("keys should be array");
assert_eq!(keys.len(), 1);
assert_eq!(
keys[0]["oauth_invalid_reason"],
json!(
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权"
)
);
assert_eq!(
keys[0]["status_snapshot"]["oauth"]["code"],
json!("reauth_required")
);
assert_eq!(
keys[0]["status_snapshot"]["oauth"]["reason"],
json!("Token 续期失败 (401): refresh_token 已被使用并轮换，请重新登录授权")
);
assert_eq!(
keys[0]["status_snapshot"]["oauth"]["source"],
json!("oauth_refresh")
);
assert_eq!(
keys[0]["status_snapshot"]["oauth"]["requires_reauth"],
json!(true)
);
assert_eq!(
keys[0]["status_snapshot"]["oauth"]["usable_until_expiry"],
json!(true)
);
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_expiry() {
run_admin_oauth_test(
"gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_expiry",
gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_expiry_impl,
);
}
async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_expiry_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
"/oauth/token",
post(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::UNAUTHORIZED,
Json(json!({
"error": {
"message": "Could not validate your refresh token. Please try signing in again.",
"type": "invalid_request_error",
"param": serde_json::Value::Null,
"code": "refresh_token_expired"
}
})),
)
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"auto_remove_banned_keys": true
}
})),
);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-expired",
"provider-codex",
"openai:responses",
"expired-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"expired-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let refresh_response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-expired/refresh"
))
.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")
.send()
.await
.expect("refresh request should succeed");
assert_eq!(refresh_response.status(), StatusCode::OK);
let refresh_payload: serde_json::Value = refresh_response
.json()
.await
.expect("refresh payload should parse");
assert_eq!(refresh_payload["status"], json!("auto_removed"));
assert_eq!(refresh_payload["message"], json!("已自动删除"));
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-refresh-expired".to_string()])
.await
.expect("keys should read");
assert!(keys.is_empty());
gateway_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy(
) {
run_admin_oauth_test(
"gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy",
gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy_impl,
);
}
async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy_impl(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:local-refresh-token" {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-provider",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let mut provider_node = sample_proxy_node("proxy-node-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let mut system_node = sample_proxy_node("proxy-node-system");
system_node.status = "online".to_string();
system_node.is_manual = true;
system_node.tunnel_mode = false;
system_node.tunnel_connected = false;
system_node.proxy_url = Some("http://proxy-system.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
provider_node,
system_node,
]));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-provider/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let plans = execution_plans.lock().expect("mutex should lock");
let refresh_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:local-refresh-token")
.expect("local refresh plan should exist");
assert_eq!(
refresh_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-provider")
);
assert_eq!(
refresh_plan
.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_refreshes_admin_provider_oauth_key_tunnel_proxy_with_direct_refresh_controls() {
run_admin_oauth_test(
"gateway_refreshes_admin_provider_oauth_key_tunnel_proxy_with_direct_refresh_controls",
gateway_refreshes_admin_provider_oauth_key_tunnel_proxy_with_direct_refresh_controls_impl,
);
}
async fn gateway_refreshes_admin_provider_oauth_key_tunnel_proxy_with_direct_refresh_controls_impl()
{
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:local-refresh-token" {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({
"mode": "tunnel",
"node_id": "proxy-node-tunnel",
"enabled": true,
"tunnel_base_url": "http://gateway-owner.internal"
}));
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh-tunnel",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-tunnel/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let plans = execution_plans.lock().expect("mutex should lock");
let refresh_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:local-refresh-token")
.expect("local refresh plan should exist");
assert_eq!(
refresh_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.mode.as_deref()),
Some("tunnel")
);
assert_eq!(
refresh_plan
.headers
.get(EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER)
.map(String::as_str),
Some("true")
);
assert_eq!(
refresh_plan
.headers
.get(EXECUTION_REQUEST_HTTP1_ONLY_HEADER)
.map(String::as_str),
Some("true")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_consecutive_manual_oauth_refresh_uses_rotated_refresh_token() {
run_admin_oauth_test(
"gateway_consecutive_manual_oauth_refresh_uses_rotated_refresh_token",
gateway_consecutive_manual_oauth_refresh_uses_rotated_refresh_token_impl,
);
}
async fn gateway_consecutive_manual_oauth_refresh_uses_rotated_refresh_token_impl() {
let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new()));
let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let refresh_request_bodies_inner = Arc::clone(&refresh_request_bodies_clone);
async move {
if plan.request_id == "provider-oauth:local-refresh-token" {
use base64::Engine as _;
let body_text = plan
.body
.body_bytes_b64
.as_deref()
.and_then(|body| {
base64::engine::general_purpose::STANDARD.decode(body).ok()
})
.and_then(|body| String::from_utf8(body).ok())
.unwrap_or_default();
refresh_request_bodies_inner
.lock()
.expect("mutex should lock")
.push(body_text.clone());
let (access_token, refresh_token) =
if body_text.contains("refresh_token=old-codex-refresh-token") {
(
"refreshed-codex-access-token",
"rotated-codex-refresh-token",
)
} else if body_text.contains("refresh_token=rotated-codex-refresh-token") {
(
"refreshed-codex-access-token-2",
"rotated-codex-refresh-token-2",
)
} else {
return Json(json!({
"request_id": plan.request_id,
"status_code": 401,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "unexpected refresh token"
}
}
}
}));
};
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": access_token,
"refresh_token": refresh_token,
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-consecutive-refresh",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for _ in 0..2 {
let response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-consecutive-refresh/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
}
let bodies = refresh_request_bodies
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(bodies.len(), 2);
assert!(
bodies[0].contains("refresh_token=old-codex-refresh-token"),
"unexpected first refresh body: {}",
bodies[0]
);
assert!(
bodies[1].contains("refresh_token=rotated-codex-refresh-token"),
"unexpected second refresh body: {}",
bodies[1]
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-consecutive-refresh".to_string()])
.await
.expect("keys should list")
.into_iter()
.next()
.expect("refreshed key should exist");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(
auth_config["refresh_token"],
"rotated-codex-refresh-token-2"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait() {
run_admin_oauth_test(
"gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait",
gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait_impl,
);
}
async fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_after_lock_wait_impl() {
let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new()));
let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let refresh_request_bodies_inner = Arc::clone(&refresh_request_bodies_clone);
async move {
if plan.request_id == "provider-oauth:local-refresh-token" {
use base64::Engine as _;
let body_text = plan
.body
.body_bytes_b64
.as_deref()
.and_then(|body| {
base64::engine::general_purpose::STANDARD.decode(body).ok()
})
.and_then(|body| String::from_utf8(body).ok())
.unwrap_or_default();
refresh_request_bodies_inner
.lock()
.expect("mutex should lock")
.push(body_text.clone());
let (access_token, refresh_token, delay_ms) =
if body_text.contains("refresh_token=old-codex-refresh-token") {
(
"refreshed-codex-access-token",
"rotated-codex-refresh-token",
200u64,
)
} else if body_text.contains("refresh_token=rotated-codex-refresh-token") {
(
"refreshed-codex-access-token-2",
"rotated-codex-refresh-token-2",
0u64,
)
} else {
return Json(json!({
"request_id": plan.request_id,
"status_code": 401,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "Could not validate your refresh token. Please try signing in again."
}
}
}
}));
};
if delay_ms > 0 {
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
}
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": access_token,
"refresh_token": refresh_token,
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-concurrent-refresh",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.proxy = Some(json!({"node_id":"proxy-node-key","enabled":true}));
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let mut key_node = sample_proxy_node("proxy-node-key");
key_node.status = "online".to_string();
key_node.is_manual = true;
key_node.tunnel_mode = false;
key_node.tunnel_connected = false;
key_node.proxy_url = Some("http://proxy-key.example:8080".to_string());
let mut provider_node = sample_proxy_node("proxy-node-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
key_node,
provider_node,
]));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let refresh_url = format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-concurrent-refresh/refresh"
);
let request_a = client
.post(refresh_url.clone())
.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")
.send();
let request_b = client
.post(refresh_url)
.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")
.send();
let (response_a, response_b) = tokio::join!(request_a, request_b);
let response_a = response_a.expect("first request should succeed");
let response_b = response_b.expect("second request should succeed");
assert_eq!(response_a.status(), StatusCode::OK);
assert_eq!(response_b.status(), StatusCode::OK);
let bodies = refresh_request_bodies
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(bodies.len(), 2);
assert!(
bodies[0].contains("refresh_token=old-codex-refresh-token"),
"unexpected first refresh body: {}",
bodies[0]
);
assert!(
bodies[1].contains("refresh_token=rotated-codex-refresh-token"),
"unexpected second refresh body: {}",
bodies[1]
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-concurrent-refresh".to_string()])
.await
.expect("keys should list")
.into_iter()
.next()
.expect("refreshed key should exist");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
stored_key
.encrypted_auth_config
.as_deref()
.expect("auth config should exist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(
auth_config["refresh_token"],
"rotated-codex-refresh-token-2"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over_stale_runtime_cache() {
run_admin_oauth_test(
"gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over_stale_runtime_cache",
gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over_stale_runtime_cache_impl,
);
}
async fn gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over_stale_runtime_cache_impl(
) {
let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new()));
let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let refresh_request_bodies_inner = Arc::clone(&refresh_request_bodies_clone);
async move {
if plan.request_id == "provider-oauth:local-refresh-token" {
use base64::Engine as _;
let body_text = plan
.body
.body_bytes_b64
.as_deref()
.and_then(|body| {
base64::engine::general_purpose::STANDARD.decode(body).ok()
})
.and_then(|body| String::from_utf8(body).ok())
.unwrap_or_default();
refresh_request_bodies_inner
.lock()
.expect("mutex should lock")
.push(body_text.clone());
if body_text.contains("refresh_token=old-codex-refresh-token") {
return Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "cached-codex-access-token",
"refresh_token": "cached-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}));
}
if body_text.contains("refresh_token=fresh-codex-refresh-token") {
return Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "fresh-codex-access-token-2",
"refresh_token": "fresh-codex-refresh-token-2",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}));
}
Json(json!({
"request_id": plan.request_id,
"status_code": 401,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"error": {
"message": "Could not validate your refresh token. Please try signing in again."
}
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-stale-cache",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":1700000001}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let app_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
let stale_transport = app_state
.read_provider_transport_snapshot(
"provider-codex",
"endpoint-codex-cli",
"key-codex-oauth-stale-cache",
)
.await
.expect("transport should load")
.expect("transport should exist");
let cached_entry = app_state
.force_local_oauth_refresh_entry(&stale_transport)
.await
.expect("initial refresh should succeed")
.expect("initial refresh should return cached entry");
assert_eq!(
cached_entry.auth_header_value,
"Bearer cached-codex-access-token"
);
let mut updated_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-stale-cache".to_string()])
.await
.expect("keys should list")
.into_iter()
.next()
.expect("key should exist");
updated_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#,
)
.expect("updated auth config ciphertext should build"),
);
provider_catalog_repository
.update_key(&updated_key)
.await
.expect("key should update");
let gateway = build_router_with_state(app_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-stale-cache/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let bodies = refresh_request_bodies
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(bodies.len(), 2);
assert!(
bodies[0].contains("refresh_token=old-codex-refresh-token"),
"unexpected first refresh body: {}",
bodies[0]
);
assert!(
bodies[1].contains("refresh_token=fresh-codex-refresh-token"),
"unexpected second refresh body: {}",
bodies[1]
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy(
) {
run_admin_oauth_test(
"gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy",
gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy_impl,
);
}
async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy_impl(
) {
let execution_plans = Arc::new(Mutex::new(Vec::<ExecutionPlan>::new()));
let execution_plans_clone = Arc::clone(&execution_plans);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let execution_plans_inner = Arc::clone(&execution_plans_clone);
async move {
execution_plans_inner
.lock()
.expect("mutex should lock")
.push(plan.clone());
if plan.request_id == "provider-oauth:local-refresh-token" {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"access_token": "refreshed-codex-access-token",
"refresh_token": "refreshed-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 1800,
"scope": "openid email profile offline_access",
"email": "[email protected]",
"account_id": "acct-codex-123",
"plan_type": "plus"
}
}
}))
} else {
Json(json!({
"request_id": plan.request_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {}
}
}))
}
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.proxy = Some(json!({"node_id":"proxy-node-provider","enabled":true}));
let endpoint = sample_endpoint(
"endpoint-codex-cli",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
);
let mut key = sample_key(
"key-codex-oauth-refresh",
"provider-codex",
"openai:responses",
"stale-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.proxy = Some(json!({"node_id":"proxy-node-key","enabled":true}));
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"[email protected]","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let mut key_node = sample_proxy_node("proxy-node-key");
key_node.status = "online".to_string();
key_node.is_manual = true;
key_node.tunnel_mode = false;
key_node.tunnel_connected = false;
key_node.proxy_url = Some("http://proxy-key.example:8080".to_string());
let mut provider_node = sample_proxy_node("proxy-node-provider");
provider_node.status = "online".to_string();
provider_node.is_manual = true;
provider_node.tunnel_mode = false;
provider_node.tunnel_connected = false;
provider_node.proxy_url = Some("http://proxy-provider.example:8080".to_string());
let mut system_node = sample_proxy_node("proxy-node-system");
system_node.status = "online".to_string();
system_node.is_manual = true;
system_node.tunnel_mode = false;
system_node.tunnel_connected = false;
system_node.proxy_url = Some("http://proxy-system.example:8080".to_string());
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
key_node,
provider_node,
system_node,
]));
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
),
]);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.attach_proxy_node_repository_for_tests(proxy_node_repository)
.with_system_config_values_for_tests(vec![(
"system_proxy_node_id".to_string(),
json!("proxy-node-system"),
)])
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let plans = execution_plans.lock().expect("mutex should lock");
let refresh_plan = plans
.iter()
.find(|plan| plan.request_id == "provider-oauth:local-refresh-token")
.expect("local refresh plan should exist");
assert_eq!(
refresh_plan
.proxy
.as_ref()
.and_then(|proxy| proxy.node_id.as_deref()),
Some("proxy-node-key")
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_handles_admin_provider_oauth_unavailable_routes_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_provider_oauth_unavailable_routes_locally_with_trusted_admin_principal",
gateway_handles_admin_provider_oauth_unavailable_routes_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_provider_oauth_unavailable_routes_locally_with_trusted_admin_principal_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for path in [
"/api/admin/provider-oauth/providers/provider-123/import-refresh-token",
"/api/admin/provider-oauth/providers/provider-123/batch-import",
"/api/admin/provider-oauth/providers/provider-123/batch-import/tasks",
"/api/admin/provider-oauth/providers/provider-123/device-authorize",
"/api/admin/provider-oauth/providers/provider-123/device-poll",
] {
let response = client
.post(format!("{gateway_url}{path}"))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "Admin provider OAuth data unavailable");
}
let refresh_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-123/refresh"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(refresh_response.status(), StatusCode::NOT_FOUND);
let refresh_payload: serde_json::Value = refresh_response
.json()
.await
.expect("json body should parse");
assert_eq!(refresh_payload["detail"], "Key 不存在");
let complete_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/keys/key-123/complete"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(complete_response.status(), StatusCode::BAD_REQUEST);
let complete_payload: serde_json::Value = complete_response
.json()
.await
.expect("json body should parse");
assert_eq!(complete_payload["detail"], "请求体必须是合法的 JSON 对象");
let provider_complete_response = client
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-123/complete"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(provider_complete_response.status(), StatusCode::BAD_REQUEST);
let provider_complete_payload: serde_json::Value = provider_complete_response
.json()
.await
.expect("json body should parse");
assert_eq!(
provider_complete_payload["detail"],
"请求体必须是合法的 JSON 对象"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_oauth_supported_types_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_oauth_supported_types_locally_with_trusted_admin_principal",
gateway_handles_admin_oauth_supported_types_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_oauth_supported_types_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/supported-types",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/oauth/supported-types"))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("items should be array");
assert_eq!(items.len(), 2);
assert_eq!(items[0]["provider_type"], "linuxdo");
assert_eq!(items[0]["display_name"], "Linux Do");
assert_eq!(items[1]["provider_type"], "custom_oidc");
assert_eq!(items[1]["display_name"], "Custom OIDC");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_oauth_provider_list_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_oauth_provider_list_locally_with_trusted_admin_principal",
gateway_handles_admin_oauth_provider_list_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_oauth_provider_list_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::seed(vec![
sample_oauth_provider_config("linuxdo"),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/oauth/providers"))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["provider_type"], "linuxdo");
assert_eq!(items[0]["has_secret"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_upserts_admin_oauth_provider_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_upserts_admin_oauth_provider_locally_with_trusted_admin_principal",
gateway_upserts_admin_oauth_provider_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_upserts_admin_oauth_provider_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/linuxdo",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository.clone(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!("{gateway_url}/api/admin/oauth/providers/linuxdo"))
.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!({
"display_name": "Linux Do",
"client_id": "client-id",
"client_secret": "secret-value",
"authorization_url_override": "https://connect.linux.do/oauth2/authorize",
"token_url_override": "https://connect.linux.do/oauth2/token",
"userinfo_url_override": "https://connect.linux.do/api/user",
"scopes": ["openid", "profile"],
"redirect_uri": "https://backend.example.com/oauth/callback",
"frontend_callback_url": "https://frontend.example.com/auth/callback",
"attribute_mapping": {"email": "email"},
"extra_config": {"team": true},
"is_enabled": true,
"force": false
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "linuxdo");
assert_eq!(payload["has_secret"], true);
assert_eq!(payload["is_enabled"], true);
let stored = repository
.get_oauth_provider_config("linuxdo")
.await
.expect("lookup should succeed")
.expect("provider should exist");
assert!(stored.client_secret_encrypted.is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_rejects_custom_oidc_without_allowed_domains() {
run_admin_oauth_test(
"gateway_rejects_custom_oidc_without_allowed_domains",
gateway_rejects_custom_oidc_without_allowed_domains_impl,
);
}
async fn gateway_rejects_custom_oidc_without_allowed_domains_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/custom_oidc",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/oauth/providers/custom_oidc"
))
.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!({
"display_name": "Custom OIDC",
"client_id": "custom-client",
"authorization_url_override": "https://idp.example.com/oauth/authorize",
"token_url_override": "https://idp.example.com/oauth/token",
"userinfo_url_override": "https://idp.example.com/oauth/userinfo",
"scopes": ["openid", "profile", "email"],
"redirect_uri": "https://backend.example.com/oauth/callback",
"frontend_callback_url": "https://frontend.example.com/auth/callback",
"attribute_mapping": {"sub": "sub", "email": "email"},
"extra_config": {},
"is_enabled": true
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(
payload["error"]["message"],
"custom_oidc 必须在 extra_config.allowed_domains 配置域名白名单"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_upserts_custom_oidc_with_allowed_domains() {
run_admin_oauth_test(
"gateway_upserts_custom_oidc_with_allowed_domains",
gateway_upserts_custom_oidc_with_allowed_domains_impl,
);
}
async fn gateway_upserts_custom_oidc_with_allowed_domains_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/custom_oidc",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository.clone(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/oauth/providers/custom_oidc"
))
.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!({
"display_name": "Custom OIDC",
"client_id": "custom-client",
"authorization_url_override": "https://idp.example.com/oauth/authorize",
"token_url_override": "https://idp.example.com/oauth/token",
"userinfo_url_override": "https://idp.example.com/oauth/userinfo",
"scopes": ["openid", "profile", "email"],
"redirect_uri": "https://backend.example.com/oauth/callback",
"frontend_callback_url": "https://frontend.example.com/auth/callback",
"attribute_mapping": {"sub": "id", "email": "profile.email"},
"extra_config": {"allowed_domains": ["idp.example.com"]},
"is_enabled": true
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "custom_oidc");
assert_eq!(payload["is_enabled"], true);
let stored = repository
.get_oauth_provider_config("custom_oidc")
.await
.expect("lookup should succeed")
.expect("provider should exist");
assert_eq!(
stored.authorization_url_override.as_deref(),
Some("https://idp.example.com/oauth/authorize")
);
assert_eq!(
stored
.extra_config
.as_ref()
.and_then(|value| {
value
.get("allowed_domains")
.and_then(serde_json::Value::as_array)
})
.map(Vec::len),
Some(1)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_upserts_multiple_custom_oidc_configs() {
run_admin_oauth_test(
"gateway_upserts_multiple_custom_oidc_configs",
gateway_upserts_multiple_custom_oidc_configs_impl,
);
}
async fn gateway_upserts_multiple_custom_oidc_configs_impl() {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let repository = Arc::new(InMemoryOAuthProviderRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository.clone(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
for (provider_type, display_name, host) in [
("custom_oidc_work", "Work OIDC", "work-idp.example.com"),
(
"custom_oidc_personal",
"Personal OIDC",
"personal-idp.example.com",
),
] {
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/oauth/providers/{provider_type}"
))
.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!({
"display_name": display_name,
"client_id": format!("{provider_type}-client"),
"authorization_url_override": format!("https://{host}/oauth/authorize"),
"token_url_override": format!("https://{host}/oauth/token"),
"userinfo_url_override": format!("https://{host}/oauth/userinfo"),
"scopes": ["openid", "profile", "email"],
"redirect_uri": format!("https://backend.example.com/api/oauth/{provider_type}/callback"),
"frontend_callback_url": "https://frontend.example.com/auth/callback",
"attribute_mapping": {"sub": "id", "email": "profile.email"},
"extra_config": {"allowed_domains": [host]},
"is_enabled": true
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], provider_type);
assert_eq!(payload["display_name"], display_name);
}
let stored = repository
.list_oauth_provider_configs()
.await
.expect("list should succeed");
assert_eq!(stored.len(), 2);
assert!(stored
.iter()
.any(|provider| provider.provider_type == "custom_oidc_work"));
assert!(stored
.iter()
.any(|provider| provider.provider_type == "custom_oidc_personal"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_tests_admin_oauth_linuxdo_endpoints_locally_with_configured_secret() {
run_admin_oauth_test(
"gateway_tests_admin_oauth_linuxdo_endpoints_locally_with_configured_secret",
gateway_tests_admin_oauth_linuxdo_endpoints_locally_with_configured_secret_impl,
);
}
async fn gateway_tests_admin_oauth_linuxdo_endpoints_locally_with_configured_secret_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/linuxdo/test",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let authorization_hits = Arc::new(Mutex::new(0usize));
let authorization_hits_clone = Arc::clone(&authorization_hits);
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let oauth_endpoints = Router::new()
.route(
"/oauth2/authorize",
any(move |_request: Request| {
let authorization_hits_inner = Arc::clone(&authorization_hits_clone);
async move {
*authorization_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::BAD_REQUEST, Body::from("missing query"))
}
}),
)
.route(
"/oauth2/token",
any(move |_request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
(
StatusCode::METHOD_NOT_ALLOWED,
Body::from("method not allowed"),
)
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::seed(vec![
sample_oauth_provider_config("linuxdo"),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (oauth_url, oauth_handle) = start_server(oauth_endpoints).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/oauth/providers/linuxdo/test"
))
.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!({
"client_id": "client-id",
"authorization_url_override": format!("{oauth_url}/oauth2/authorize"),
"token_url_override": format!("{oauth_url}/oauth2/token"),
"redirect_uri": "http://localhost:8084/api/oauth/linuxdo/callback"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["authorization_url_reachable"], true);
assert_eq!(payload["token_url_reachable"], true);
assert_eq!(payload["secret_status"], "configured");
assert_eq!(*authorization_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
oauth_handle.abort();
let _ = upstream_url;
}
#[test]
fn gateway_tests_admin_oauth_linuxdo_reports_invalid_endpoint_urls() {
run_admin_oauth_test(
"gateway_tests_admin_oauth_linuxdo_reports_invalid_endpoint_urls",
gateway_tests_admin_oauth_linuxdo_reports_invalid_endpoint_urls_impl,
);
}
async fn gateway_tests_admin_oauth_linuxdo_reports_invalid_endpoint_urls_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/linuxdo/test",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/oauth/providers/linuxdo/test"
))
.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!({
"client_id": "client-id",
"authorization_url_override": "not-a-url",
"token_url_override": "not-a-url",
"redirect_uri": "http://localhost:8084/api/oauth/linuxdo/callback"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["authorization_url_reachable"], false);
assert_eq!(payload["token_url_reachable"], false);
assert_eq!(payload["secret_status"], "not_provided");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
let _ = upstream_url;
}
#[test]
fn gateway_deletes_admin_oauth_provider_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_deletes_admin_oauth_provider_locally_with_trusted_admin_principal",
gateway_deletes_admin_oauth_provider_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_deletes_admin_oauth_provider_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/oauth/providers/linuxdo",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryOAuthProviderRepository::seed(vec![
sample_oauth_provider_config("linuxdo"),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_oauth_provider_repository_for_tests(
repository.clone(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!("{gateway_url}/api/admin/oauth/providers/linuxdo"))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "删除成功");
assert!(repository
.get_oauth_provider_config("linuxdo")
.await
.expect("lookup should succeed")
.is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_handles_admin_management_token_detail_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_handles_admin_management_token_detail_locally_with_trusted_admin_principal",
gateway_handles_admin_management_token_detail_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_handles_admin_management_token_detail_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/management-tokens/{token_id}",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryManagementTokenRepository::seed(vec![
sample_management_token("mt-admin-1", "user-1", "alice", true),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/management-tokens/mt-admin-1"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["id"], "mt-admin-1");
assert_eq!(payload["user"]["email"], "[email protected]");
assert_eq!(payload["usage_count"], 7);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_creates_updates_and_regenerates_admin_management_token_locally_with_permissions() {
run_admin_oauth_test(
"gateway_creates_updates_and_regenerates_admin_management_token_locally_with_permissions",
gateway_creates_updates_and_regenerates_admin_management_token_locally_with_permissions_impl,
);
}
async fn gateway_creates_updates_and_regenerates_admin_management_token_locally_with_permissions_impl(
) {
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| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}));
let repository = Arc::new(InMemoryManagementTokenRepository::default());
let state = AppState::new().expect("gateway should build");
let admin_user = state
.create_local_auth_user_with_settings(
Some("[email protected]".to_string()),
true,
"admin".to_string(),
"hash".to_string(),
"admin".to_string(),
None,
None,
None,
None,
)
.await
.expect("admin user should be created")
.expect("admin user should exist");
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(state.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(repository.clone()),
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let create_response = client
.post(format!("{gateway_url}/api/admin/management-tokens"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, admin_user.id.as_str())
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"name": "admin-token",
"description": "admin token",
"allowed_ips": ["127.0.0.1"],
"permissions": ["admin:usage:read", "admin:pool:write"],
"expires_at": "2099-01-01T00:00:00Z",
}))
.send()
.await
.expect("request should succeed");
let create_status = create_response.status();
let create_body = create_response.text().await.expect("body should read");
assert_eq!(create_status, StatusCode::CREATED, "{create_body}");
let create_payload: serde_json::Value =
serde_json::from_str(&create_body).expect("json body should parse");
let token_id = create_payload["data"]["id"]
.as_str()
.expect("token id should exist")
.to_string();
assert_eq!(create_payload["message"], "Management Token 创建成功");
assert_eq!(create_payload["data"]["user"]["id"], json!(admin_user.id));
assert_eq!(
create_payload["data"]["permissions"],
json!(["admin:pool:write", "admin:usage:read"])
);
let created_secret = create_payload["token"].as_str().unwrap_or_default();
assert!(created_secret.starts_with("ae-"));
let update_response = client
.put(format!(
"{gateway_url}/api/admin/management-tokens/{token_id}"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, admin_user.id.as_str())
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"name": "admin-token-updated",
"description": null,
"allowed_ips": null,
"permissions": ["admin:usage:read"],
"expires_at": null,
}))
.send()
.await
.expect("request should succeed");
assert_eq!(update_response.status(), StatusCode::OK);
let update_payload: serde_json::Value = update_response
.json()
.await
.expect("json body should parse");
assert_eq!(update_payload["message"], "更新成功");
assert_eq!(update_payload["data"]["name"], "admin-token-updated");
assert_eq!(
update_payload["data"]["permissions"],
json!(["admin:usage:read"])
);
assert_eq!(
update_payload["data"]["allowed_ips"],
serde_json::Value::Null
);
let regenerate_response = client
.post(format!(
"{gateway_url}/api/admin/management-tokens/{token_id}/regenerate"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, admin_user.id.as_str())
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(regenerate_response.status(), StatusCode::OK);
let regenerate_payload: serde_json::Value = regenerate_response
.json()
.await
.expect("json body should parse");
assert_eq!(regenerate_payload["message"], "Token 已重新生成");
assert!(regenerate_payload["token"]
.as_str()
.unwrap_or_default()
.starts_with("ae-"));
assert_eq!(
repository
.get_management_token_with_user(&token_id)
.await
.expect("lookup should succeed")
.expect("token should remain")
.token
.permissions,
Some(json!(["admin:usage:read"]))
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
drop(upstream_url);
}
#[test]
fn gateway_allows_management_token_with_pool_write_for_provider_oauth_batch_import() {
run_admin_oauth_test(
"gateway_allows_management_token_with_pool_write_for_provider_oauth_batch_import",
gateway_allows_management_token_with_pool_write_for_provider_oauth_batch_import_impl,
);
}
async fn gateway_allows_management_token_with_pool_write_for_provider_oauth_batch_import_impl() {
let raw_token = "ae-provider-oauth-batch-pool-write";
let state = AppState::new().expect("gateway should build");
let admin_user = state
.create_local_auth_user_with_settings(
Some("[email protected]".to_string()),
true,
"admin".to_string(),
"hash".to_string(),
"admin".to_string(),
None,
None,
None,
None,
)
.await
.expect("admin user should be created")
.expect("admin user should exist");
let mut management_token = sample_management_token(
"token-provider-oauth-batch-pool",
&admin_user.id,
"provider-oauth-pool",
true,
);
management_token.token.allowed_ips = None;
management_token.token.permissions = Some(json!(["admin:pool:read", "admin:pool:write"]));
let management_token_repository =
Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
vec![management_token],
vec![(
hash_management_token(raw_token),
"token-provider-oauth-batch-pool".to_string(),
)],
));
let gateway = build_router_with_state(state.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(management_token_repository),
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-123/batch-import"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(raw_token)
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE, "{payload}");
assert_eq!(payload["detail"], "Admin provider OAuth data unavailable");
gateway_handle.abort();
}
#[test]
fn gateway_rejects_management_token_without_pool_write_for_provider_oauth_batch_import() {
run_admin_oauth_test(
"gateway_rejects_management_token_without_pool_write_for_provider_oauth_batch_import",
gateway_rejects_management_token_without_pool_write_for_provider_oauth_batch_import_impl,
);
}
async fn gateway_rejects_management_token_without_pool_write_for_provider_oauth_batch_import_impl()
{
let raw_token = "ae-provider-oauth-batch-pool-denied";
let state = AppState::new().expect("gateway should build");
let admin_user = state
.create_local_auth_user_with_settings(
Some("[email protected]".to_string()),
true,
"admin".to_string(),
"hash".to_string(),
"admin".to_string(),
None,
None,
None,
None,
)
.await
.expect("admin user should be created")
.expect("admin user should exist");
let mut management_token = sample_management_token(
"token-provider-oauth-batch-denied",
&admin_user.id,
"provider-oauth-denied",
true,
);
management_token.token.allowed_ips = None;
management_token.token.permissions = Some(json!(["admin:usage:read"]));
let management_token_repository =
Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
vec![management_token],
vec![(
hash_management_token(raw_token),
"token-provider-oauth-batch-denied".to_string(),
)],
));
let gateway = build_router_with_state(state.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(management_token_repository),
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-123/batch-import"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(raw_token)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "management token permission denied");
assert_eq!(payload["required_permission"], "admin:pool:write");
assert_eq!(payload["route_family"], "provider_oauth_manage");
gateway_handle.abort();
}
#[test]
fn gateway_deletes_admin_management_token_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_deletes_admin_management_token_locally_with_trusted_admin_principal",
gateway_deletes_admin_management_token_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_deletes_admin_management_token_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/management-tokens/{token_id}",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryManagementTokenRepository::seed(vec![
sample_management_token("mt-admin-1", "user-1", "alice", true),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.delete(format!(
"{gateway_url}/api/admin/management-tokens/mt-admin-1"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "删除成功");
assert_eq!(
repository
.get_management_token_with_user("mt-admin-1")
.await
.expect("lookup should succeed"),
None
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_toggles_admin_management_token_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
"gateway_toggles_admin_management_token_locally_with_trusted_admin_principal",
gateway_toggles_admin_management_token_locally_with_trusted_admin_principal_impl,
);
}
async fn gateway_toggles_admin_management_token_locally_with_trusted_admin_principal_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/management-tokens/{token_id}/status",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let repository = Arc::new(InMemoryManagementTokenRepository::seed(vec![
sample_management_token("mt-admin-1", "user-1", "alice", true),
]));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(repository.clone()),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.patch(format!(
"{gateway_url}/api/admin/management-tokens/mt-admin-1/status"
))
.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")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["message"], "Token 已禁用");
assert_eq!(payload["data"]["is_active"], false);
assert_eq!(
repository
.get_management_token_with_user("mt-admin-1")
.await
.expect("lookup should succeed")
.expect("token should remain")
.token
.is_active,
false
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[test]
fn gateway_rejects_partial_management_token_for_admin_management_token_routes() {
run_admin_oauth_test(
"gateway_rejects_partial_management_token_for_admin_management_token_routes",
gateway_rejects_partial_management_token_for_admin_management_token_routes_impl,
);
}
async fn gateway_rejects_partial_management_token_for_admin_management_token_routes_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/management-tokens",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let raw_token = "ae-management-partial-access";
let state = AppState::new().expect("gateway should build");
let admin_user = state
.create_local_auth_user_with_settings(
Some("[email protected]".to_string()),
true,
"admin".to_string(),
"hash".to_string(),
"admin".to_string(),
None,
None,
None,
None,
)
.await
.expect("admin user should be created")
.expect("admin user should exist");
let mut management_token = sample_management_token(
"mt-admin-partial",
&admin_user.id,
"management-partial",
true,
);
management_token.token.allowed_ips = None;
management_token.token.permissions = Some(json!(["admin:usage:read"]));
let repository = Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
vec![management_token],
vec![(
hash_management_token(raw_token),
"mt-admin-partial".to_string(),
)],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(state.with_data_state_for_tests(
GatewayDataState::with_management_token_repository_for_tests(repository),
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/admin/management-tokens"))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.bearer_auth(raw_token)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["detail"], "management token permission denied");
assert_eq!(
payload["required_permission"],
"admin:management_tokens:read"
);
assert_eq!(payload["route_family"], "management_tokens_manage");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
Reference in New Issue View Git Blame Copy Permalink