mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-15 00:06:24 +00:00
feat: add admin PKCE auth flow
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
Codex
File diff suppressed because one or more lines are too long
@@ -2981,7 +2981,7 @@ The Rust project has no legacy compatibility requirement, but the in-scope behav
|
||||
- [x] **HTTP auth boundary:** restore a real `src/http/auth.rs` module instead of scattered header checks. It validates `Bearer cpr_...` client API keys for `/v1/*`, validates admin session cookies for `/admin/*`, rejects cross-use in both directions, and returns the documented OpenAI/admin response families.
|
||||
- [x] **Account SQLite repository:** implement `src/accounts/repository.rs` for encrypted access/refresh tokens, account metadata, quota JSON, usage counters, labels, status transitions, pagination, batch-safe mutations, metadata-only auth status reads, and account-wide delete for logout. Do not add TypeScript JSON migration or legacy import shims.
|
||||
- [ ] **Admin account routes:** rebuild `/auth/status`, `/auth/accounts`, import/export, batch delete/status, health-check, per-account refresh, label, delete, reset usage, quota, Cookie CRUD, and quota warnings as `/admin/*` endpoints with admin envelopes and cursor pagination where a list is returned. Encrypted import, manual token add, CLI import, paginated sanitized list, native/Sub2API export, Sub2API import marking, label/status/delete, batch delete/status, Cookie CRUD, health/probe, refresh, quota, reset usage, admin auth status, and admin logout are complete. Quota warnings remain.
|
||||
- [ ] **Login routes:** implement OAuth PKCE login-start/code-relay/callback, device-login/device-poll, CLI auth import, manual access-token import, and logout as account-login flows. CLI auth import, validated manual token import, device login/poll, admin-session-gated auth status, and logout are complete. OAuth PKCE remains. Keep this separate from admin password login and client API key creation.
|
||||
- [x] **Login routes:** implement OAuth PKCE login-start/code-relay/callback, device-login/device-poll, CLI auth import, manual access-token import, and logout as account-login flows. CLI auth import, validated manual token import, device login/poll, OAuth PKCE login-start/code-relay/callback, admin-session-gated auth status, and logout are complete. Keep this separate from admin password login and client API key creation.
|
||||
- [x] **Client API key management:** implement admin CRUD for local client API keys: list/create/delete/batch-delete/label/status/export/import. Create/list/status/delete/batch-delete/label/export/import and `/v1/*` SQLite verification are complete. Remove provider/model binding semantics from the TypeScript version; Rust client API keys authorize only local `/v1/*` access.
|
||||
- [ ] **Settings:** implement admin settings for server/runtime fields that remain in scope: default model, reasoning effort, service tier, model aliases, refresh enabled/margin/concurrency, max concurrent per account, request interval, rotation strategy, tier priority, quota refresh/thresholds/skip exhausted, logs state/capacity/body capture, and usage history retention. Read-only `/admin/settings` is complete; persistent mutation and config-store writeback remain. Remove proxy URL, Ollama, third-party provider, OpenAI official key, and Electron/self-update settings.
|
||||
- [ ] **Diagnostics:** expand health/admin diagnostics to include authenticated state, pool summary, capacity summary, transport/fingerprint status, paths, runtime metadata, and test-connection checks. Keep production-local gating for sensitive debug endpoints.
|
||||
@@ -2999,7 +2999,7 @@ This pass re-read the TypeScript routes and core modules under `/home/zyy/Codes/
|
||||
- [x] **OpenAI Chat Completions route wiring:** Rust mounts `/v1/chat/completions` to a dedicated Chat handler that parses Chat `messages`, translates to Codex Responses requests, and returns OpenAI Chat completion JSON or SSE chunks.
|
||||
- [ ] **Responses request semantics:** Rust now retains `service_tier` normalization (`fast` -> `priority` upstream), `prompt_cache_key`, `client_metadata`, `include`, `text.format`, `tool_choice`, `parallel_tool_calls`, Codex context headers, session/window/thread metadata, and reasoning include defaults. Remaining parity gaps include tuple-schema reconversion edge cases, sanitized reasoning replay behavior, and real-client verification for installation/fingerprint metadata.
|
||||
- [ ] **Responses stream/error parity:** Rust reconstructs non-stream output from `response.output_item.done` and `response.output_text.delta`, maps non-stream SSE `error` and `response.failed` to OpenAI-compatible upstream errors, and keeps streaming passthrough while recording terminal failure lifecycle logs. Remaining gaps include upstream premature-close classification, empty-response retry, stream close `response.failed` injection, tuple reconversion, reasoning replay metadata, function-call metadata, image-generation usage, and detailed OpenAI-compatible error mapping.
|
||||
- [ ] **Account login flows:** Rust has admin login, native/Sub2API account import, Codex CLI auth import, validated manual token import, OpenAI OAuth device login/poll, admin-session-gated account logout, and `/auth/status`-equivalent `/admin/auth/status`. Still missing account-login flows for OAuth PKCE login-start/code-relay/callback. These must be rebuilt as admin/session-gated account-login endpoints, not by treating local client API keys as passwords.
|
||||
- [x] **Account login flows:** Rust has admin login, native/Sub2API account import, Codex CLI auth import, validated manual token import, OpenAI OAuth device login/poll, OAuth PKCE login-start/code-relay/callback, admin-session-gated account logout, and `/auth/status`-equivalent `/admin/auth/status`. These are rebuilt as admin/session-gated account-login endpoints, not by treating local client API keys as passwords.
|
||||
- [ ] **Admin account management:** Rust has paginated sanitized list, native/Sub2API import, native/Sub2API export, manual token add with upstream refresh-token exchange and JWT metadata/expires extraction, device-code token import, label/status/delete, batch delete/status, Cookie get/set/delete, per-account refresh/probe, reset usage, quota fetch/update, auth status, and logout. Missing quota warnings, persistence-health reporting, and durable quota-window restore.
|
||||
- [x] **Client API key management:** Rust creates/lists/status-toggles/deletes/batch-deletes/labels/exports/imports local `cpr_` keys and verifies them for `/v1/*`. Do not port the TypeScript provider/model binding or provider model-fetch APIs.
|
||||
- [ ] **Settings:** Rust exposes read-only `/admin/settings` for retained fields. Missing admin-session-gated mutation/writeback with validation for default model, default reasoning effort, service tier, aliases, refresh settings, max concurrent accounts, request interval, rotation strategy, tier priority, quota refresh/threshold/skip settings, log state/capacity/body capture, and usage retention. Do not re-add proxy URL, provider, Ollama, proxy API key, or Electron/self-update settings.
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
|
||||
**Tech Stack:** Rust, Axum, Reqwest/rustls, serde, sqlx/SQLite, wiremock, tokio, futures, optional WebSocket transport dependency.
|
||||
|
||||
**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Commit `350ffbe` completed Task 4: `previous_response_id` requests use WebSocket `response.create` and return SSE-compatible output. Commit `3ecd5b6` partially completed Task 5 for non-streaming `/v1/responses`: 429/402/403 account-state classification and fallback retry across imported accounts. Commit `3960a35` extends the same fallback path to HTTP SSE Responses setup and Chat Completions. Commit `6d18d2c` completes Task 5 policy for WebSocket-backed Responses: handshake upstream errors are classified before request body send, non-history WebSocket requests may fallback accounts, successful handshakes capture rate-limit headers, and `previous_response_id` requests remain account-affine for streaming and non-streaming clients. Commit `1738f24` adds Task 6 cache groundwork: persisted backend model snapshots and cached `/v1/models*` reads. Commit `0b6cc7e` adds Task 7 import compatibility for Sub2API OpenAI OAuth exports and marks the detected source format while ignoring proxy-only fields. Commit `f75db60` adds `/admin/refresh-models` backed by imported accounts. Commit `4b19846` syncs refreshed model-plan allowlists into the runtime account pool. Commit `db6cb98` adds admin account label/status/delete mutations with runtime pool synchronization. Commit `16eb4ed` adds batch delete/status mutations. Commit `8af350d` adds local client API key status/delete mutations. Commit `203270f` adds encrypted account Cookie get/set/delete. Commit `960bbb2` adds local client API key label and batch-delete mutations. Commit `72438c6` adds native and Sub2API account export. Commit `16938e6` adds manual account creation for Codex OAuth tokens. Commit `869b6f8` narrows `POST /admin/accounts` to TS-like `token`/`refreshToken` import semantics for OpenAI GPT/Codex accounts. Commit `e1f35e0` adds deterministic Codex CLI `auth.json` import through the same validated JWT-claim path. Commit `93cfee3` adds admin health-check, per-account refresh, reset-usage, and quota routes for imported Codex accounts. Commit `bfe561d` maps HTTP SSE terminal failures to OpenAI-compatible upstream errors for non-streaming clients while preserving streaming passthrough. Commit `cf5ff0b` adds Rust-local client API key metadata export/import with target-local key rotation. Commit `a32a53c` adds admin-session-gated `/admin/auth/status` and `/admin/auth/logout`, returning sanitized account/pool state and clearing SQLite accounts plus the runtime pool without token exposure. The current device login pass adds admin-session-gated `/admin/auth/device-login` and `/admin/auth/device-poll/{device_code}` backed by the OpenAI OAuth device-code flow and shared validated account import path. Remaining major work continues with OAuth PKCE login-start/code-relay/callback.
|
||||
**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Commit `350ffbe` completed Task 4: `previous_response_id` requests use WebSocket `response.create` and return SSE-compatible output. Commit `3ecd5b6` partially completed Task 5 for non-streaming `/v1/responses`: 429/402/403 account-state classification and fallback retry across imported accounts. Commit `3960a35` extends the same fallback path to HTTP SSE Responses setup and Chat Completions. Commit `6d18d2c` completes Task 5 policy for WebSocket-backed Responses: handshake upstream errors are classified before request body send, non-history WebSocket requests may fallback accounts, successful handshakes capture rate-limit headers, and `previous_response_id` requests remain account-affine for streaming and non-streaming clients. Commit `1738f24` adds Task 6 cache groundwork: persisted backend model snapshots and cached `/v1/models*` reads. Commit `0b6cc7e` adds Task 7 import compatibility for Sub2API OpenAI OAuth exports and marks the detected source format while ignoring proxy-only fields. Commit `f75db60` adds `/admin/refresh-models` backed by imported accounts. Commit `4b19846` syncs refreshed model-plan allowlists into the runtime account pool. Commit `db6cb98` adds admin account label/status/delete mutations with runtime pool synchronization. Commit `16eb4ed` adds batch delete/status mutations. Commit `8af350d` adds local client API key status/delete mutations. Commit `203270f` adds encrypted account Cookie get/set/delete. Commit `960bbb2` adds local client API key label and batch-delete mutations. Commit `72438c6` adds native and Sub2API account export. Commit `16938e6` adds manual account creation for Codex OAuth tokens. Commit `869b6f8` narrows `POST /admin/accounts` to TS-like `token`/`refreshToken` import semantics for OpenAI GPT/Codex accounts. Commit `e1f35e0` adds deterministic Codex CLI `auth.json` import through the same validated JWT-claim path. Commit `93cfee3` adds admin health-check, per-account refresh, reset-usage, and quota routes for imported Codex accounts. Commit `bfe561d` maps HTTP SSE terminal failures to OpenAI-compatible upstream errors for non-streaming clients while preserving streaming passthrough. Commit `cf5ff0b` adds Rust-local client API key metadata export/import with target-local key rotation. Commit `a32a53c` adds admin-session-gated `/admin/auth/status` and `/admin/auth/logout`, returning sanitized account/pool state and clearing SQLite accounts plus the runtime pool without token exposure. Commit `1d6c301` adds admin-session-gated `/admin/auth/device-login` and `/admin/auth/device-poll/{device_code}` backed by the OpenAI OAuth device-code flow and shared validated account import path. The PKCE auth-flow phase adds `/admin/auth/login-start`, `/admin/auth/code-relay`, `/admin/auth/callback`, and the registered `/auth/callback` redirect path, with in-memory PKCE sessions, OpenAI authorization-code exchange, and shared validated account import.
|
||||
|
||||
**Scoped HTTP SSE error update:** Non-WebSocket `/v1/responses` HTTP SSE collection now detects upstream `event: error` and `event: response.failed` terminal frames. Non-streaming clients receive a `502` OpenAI-compatible `upstream_error` body with the upstream message, while streaming clients keep passthrough SSE frames and record the lifecycle event as an upstream SSE failure.
|
||||
|
||||
@@ -29,12 +29,12 @@ This audit is the correction point against the TypeScript reference at `/home/zy
|
||||
| Local client API key label/batch-delete / `960bbb2` | TS label/batch-delete applies to third-party provider key pool entries. | Rust label/batch-delete applies to local HMAC-hashed `cpr_` client keys. | Rust local extension, possible drift if strict parity is required. | Reclassify in docs as local admin utility. Remove only if strict original parity excludes local client-key management beyond create/list/status/delete. |
|
||||
| API key export/import variants | TS can export/import third-party provider keys for reimport. | Rust local `cpr_` key plaintext is shown only at creation and then HMAC-hashed; provider key pools are not implemented. Local `/admin/api-keys/export` exports only metadata, and `/admin/api-keys/import` rotates metadata into new target-local `cpr_` keys with plaintext returned once. | Rust local extension, intentional provider scope cut. | Keep local metadata export/import as client auth utility only. Do not port TS provider key import/export, provider/model binding, proxy pools, or non-OpenAI provider behavior. |
|
||||
| Model refresh and plan allowlist sync / `f75db60`, `4b19846` | TS has provider/key model catalog behavior mixed with account/provider management. | Rust probes Codex backend models via imported accounts and syncs plan allowlists into the account pool. | Partial parity, scoped Rust adaptation. | Keep; avoid third-party provider catalog behavior. Finish full Codex model metadata parity later if needed. |
|
||||
| OAuth/device/CLI account login/import variants | Original scope includes ChatGPT/Codex OAuth login, device login, CLI token import, manual validated token import, logout/status, quota, and health checks. | Rust now supports validated manual `token`/`refreshToken` import, deterministic admin CLI import from Codex `auth.json`, OpenAI OAuth device login/poll, admin quota fetch, health check, per-account refresh, reset-usage, sanitized admin auth status, and admin logout for imported accounts. Rust still lacks OAuth PKCE login-start/code-relay/callback. | Partial OpenAI/Codex parity. | Continue with OAuth PKCE login variants without adding proxy pools or non-OpenAI providers. |
|
||||
| OAuth/device/CLI account login/import variants | Original scope includes ChatGPT/Codex OAuth login, device login, CLI token import, manual validated token import, logout/status, quota, and health checks. | Rust now supports validated manual `token`/`refreshToken` import, deterministic admin CLI import from Codex `auth.json`, OpenAI OAuth device login/poll, admin-session-gated PKCE login-start/code-relay/callback, admin quota fetch, health check, per-account refresh, reset-usage, sanitized admin auth status, and admin logout for imported accounts. | Scoped OpenAI/Codex parity complete for account-login operations. | Keep admin-session-gated Rust routes and do not add proxy pools or non-OpenAI providers. |
|
||||
| Proxy/provider metadata and routes | TS project includes per-account proxy assignment, proxy health, and non-OpenAI provider routes. | Rust imports ignore proxy fields and has no outbound proxy/provider pools. | Intentional scope cut. | Keep omitted. Do not add proxy routing, proxy UI/API, Anthropic/Gemini/Ollama/custom provider support, or Electron/frontend flows. |
|
||||
|
||||
Immediate priority after this audit:
|
||||
|
||||
1. Finish in-scope account operation parity: OAuth PKCE login-start/code-relay/callback.
|
||||
1. Finish remaining in-scope account operation gaps such as quota warnings and durable quota cooldown persistence.
|
||||
2. Run final verification for the completed Task 5 fallback paths, including WebSocket history affinity.
|
||||
3. Keep local client API key utilities documented as Rust-local auth, not TypeScript provider-key parity.
|
||||
|
||||
@@ -188,9 +188,9 @@ Immediate priority after this audit:
|
||||
- [x] Write failing tests for OpenAI OAuth device login/poll.
|
||||
- [x] Implement `/admin/auth/device-login` and `/admin/auth/device-poll/{device_code}` with the OpenAI device-code grant and shared validated account import path.
|
||||
- [x] Run: `cargo test --test admin_accounts_route_test admin_auth_device`; `cargo test --test oauth_refresh_test`; `cargo test --test admin_accounts_route_test`
|
||||
- [ ] Write failing tests for OAuth PKCE login-start/code-relay/callback.
|
||||
- [ ] Implement remaining OAuth PKCE scoped operations needed to operate Codex-backed OpenAI GPT routes.
|
||||
- [ ] Run: `cargo test --test admin_accounts_route_test`
|
||||
- [x] Write failing tests for OAuth PKCE login-start/code-relay/callback.
|
||||
- [x] Implement `/admin/auth/login-start`, `/admin/auth/code-relay`, `/admin/auth/callback`, and the registered `/auth/callback` redirect path with in-memory PKCE sessions and shared validated account import.
|
||||
- [x] Run: `cargo test --test admin_accounts_route_test admin_auth_`; `cargo test --test oauth_refresh_test`; `cargo test --test admin_accounts_route_test`
|
||||
|
||||
### Task 8: Final Verification And Docs
|
||||
|
||||
|
||||
+14
-8
@@ -8,14 +8,15 @@ use tower_http::trace::TraceLayer;
|
||||
use crate::{
|
||||
http::{
|
||||
admin::{
|
||||
account_quota, accounts, api_keys, auth_device_login, auth_device_poll, auth_logout,
|
||||
auth_status, batch_delete_accounts, batch_delete_api_keys, batch_update_account_status,
|
||||
create_account, create_api_key, delete_account, delete_account_cookies, delete_api_key,
|
||||
export_accounts, export_api_keys, get_account_cookies, health_check_accounts,
|
||||
import_accounts, import_api_keys, import_cli_auth, login, logs, refresh_account,
|
||||
refresh_models, reset_account_usage, set_account_cookies, settings,
|
||||
update_account_label, update_account_status, update_api_key_label,
|
||||
update_api_key_status, usage_stats, usage_stats_summary,
|
||||
account_quota, accounts, api_keys, auth_callback, auth_code_relay, auth_device_login,
|
||||
auth_device_poll, auth_login_start, auth_logout, auth_status, batch_delete_accounts,
|
||||
batch_delete_api_keys, batch_update_account_status, create_account, create_api_key,
|
||||
delete_account, delete_account_cookies, delete_api_key, export_accounts,
|
||||
export_api_keys, get_account_cookies, health_check_accounts, import_accounts,
|
||||
import_api_keys, import_cli_auth, login, logs, refresh_account, refresh_models,
|
||||
reset_account_usage, set_account_cookies, settings, update_account_label,
|
||||
update_account_status, update_api_key_label, update_api_key_status, usage_stats,
|
||||
usage_stats_summary,
|
||||
},
|
||||
health::health,
|
||||
middleware::attach_request_id,
|
||||
@@ -40,6 +41,11 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.route("/admin/login", post(login))
|
||||
.route("/admin/auth/status", get(auth_status))
|
||||
.route("/admin/auth/logout", post(auth_logout))
|
||||
.route("/admin/auth/login-start", post(auth_login_start))
|
||||
.route("/admin/auth/code-relay", post(auth_code_relay))
|
||||
.route("/admin/auth/callback", get(auth_callback))
|
||||
// OpenAI OAuth 客户端注册的是 /auth/callback；handler 内仍要求管理员会话。
|
||||
.route("/auth/callback", get(auth_callback))
|
||||
.route("/admin/auth/device-login", post(auth_device_login))
|
||||
.route(
|
||||
"/admin/auth/device-poll/{device_code}",
|
||||
|
||||
@@ -1,6 +1,12 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use rand::Rng;
|
||||
use reqwest::{Client, StatusCode};
|
||||
use serde::Deserialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::auth::{
|
||||
@@ -12,6 +18,8 @@ const CODEX_OAUTH_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
|
||||
const OPENAI_OAUTH_AUTHORIZE_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
|
||||
const OPENAI_OAUTH_DEVICE_CODE_ENDPOINT: &str = "https://auth.openai.com/oauth/device/code";
|
||||
const OPENAI_OAUTH_TOKEN_ENDPOINT: &str = "https://auth.openai.com/oauth/token";
|
||||
const OAUTH_CALLBACK_PORT: u16 = 1455;
|
||||
const OAUTH_CALLBACK_PATH: &str = "/auth/callback";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OAuthConfig {
|
||||
@@ -66,11 +74,115 @@ impl OAuthError {
|
||||
|
||||
#[async_trait]
|
||||
pub trait OAuthClient: Send + Sync + 'static {
|
||||
async fn exchange_code(
|
||||
&self,
|
||||
code: &str,
|
||||
code_verifier: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<TokenPair, OAuthError>;
|
||||
|
||||
async fn request_device_code(&self) -> Result<DeviceCode, OAuthError>;
|
||||
|
||||
async fn poll_device_token(&self, device_code: &str) -> Result<TokenPair, OAuthError>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PkceLogin {
|
||||
pub auth_url: String,
|
||||
pub state: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PkceSession {
|
||||
pub code_verifier: String,
|
||||
pub redirect_uri: String,
|
||||
pub return_host: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct PendingPkceSession {
|
||||
session: PkceSession,
|
||||
created_at: DateTime<Utc>,
|
||||
exchanging: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct PkceSessionStore {
|
||||
pending: BTreeMap<String, PendingPkceSession>,
|
||||
completed: BTreeMap<String, DateTime<Utc>>,
|
||||
ttl: Duration,
|
||||
}
|
||||
|
||||
impl Default for PkceSessionStore {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
pending: BTreeMap::new(),
|
||||
completed: BTreeMap::new(),
|
||||
ttl: Duration::minutes(5),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PkceSessionStore {
|
||||
pub fn start_login(&mut self, return_host: &str, config: &OAuthConfig) -> PkceLogin {
|
||||
self.cleanup(Utc::now());
|
||||
let (code_verifier, code_challenge) = generate_pkce_pair();
|
||||
let state = random_state();
|
||||
let redirect_uri = format!("http://localhost:{OAUTH_CALLBACK_PORT}{OAUTH_CALLBACK_PATH}");
|
||||
let auth_url = build_auth_url(config, &redirect_uri, &state, &code_challenge);
|
||||
self.pending.insert(
|
||||
state.clone(),
|
||||
PendingPkceSession {
|
||||
session: PkceSession {
|
||||
code_verifier,
|
||||
redirect_uri,
|
||||
return_host: return_host.to_string(),
|
||||
},
|
||||
created_at: Utc::now(),
|
||||
exchanging: false,
|
||||
},
|
||||
);
|
||||
PkceLogin { auth_url, state }
|
||||
}
|
||||
|
||||
pub fn try_acquire(&mut self, state: &str) -> Option<PkceSession> {
|
||||
self.cleanup(Utc::now());
|
||||
let pending = self.pending.get_mut(state)?;
|
||||
if pending.exchanging {
|
||||
return None;
|
||||
}
|
||||
pending.exchanging = true;
|
||||
Some(pending.session.clone())
|
||||
}
|
||||
|
||||
pub fn release(&mut self, state: &str) {
|
||||
if let Some(pending) = self.pending.get_mut(state) {
|
||||
pending.exchanging = false;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn complete(&mut self, state: &str) {
|
||||
self.pending.remove(state);
|
||||
self.completed.insert(state.to_string(), Utc::now());
|
||||
}
|
||||
|
||||
pub fn is_completed_or_exchanging(&mut self, state: &str) -> bool {
|
||||
self.cleanup(Utc::now());
|
||||
self.completed.contains_key(state)
|
||||
|| self
|
||||
.pending
|
||||
.get(state)
|
||||
.is_some_and(|pending| pending.exchanging)
|
||||
}
|
||||
|
||||
fn cleanup(&mut self, now: DateTime<Utc>) {
|
||||
self.pending
|
||||
.retain(|_, pending| now - pending.created_at <= self.ttl);
|
||||
self.completed
|
||||
.retain(|_, completed_at| now - *completed_at <= self.ttl);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OpenAiOAuthRefresher {
|
||||
client: Client,
|
||||
@@ -101,6 +213,35 @@ struct OAuthErrorResponse {
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthClient for OpenAiOAuthRefresher {
|
||||
async fn exchange_code(
|
||||
&self,
|
||||
code: &str,
|
||||
code_verifier: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<TokenPair, OAuthError> {
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.config.token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "authorization_code"),
|
||||
("client_id", self.config.client_id.as_str()),
|
||||
("code", code),
|
||||
("redirect_uri", redirect_uri),
|
||||
("code_verifier", code_verifier),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| OAuthError::Transport)?;
|
||||
let status = response.status();
|
||||
let body = response.text().await.map_err(|_| OAuthError::Transport)?;
|
||||
if !status.is_success() {
|
||||
return Err(OAuthError::Rejected(format!(
|
||||
"Token exchange failed ({status}): {body}"
|
||||
)));
|
||||
}
|
||||
parse_token_pair(&body).map_err(|_| OAuthError::Transport)
|
||||
}
|
||||
|
||||
async fn request_device_code(&self) -> Result<DeviceCode, OAuthError> {
|
||||
let response = self
|
||||
.client
|
||||
@@ -195,6 +336,60 @@ fn classify_device_poll_failure(status: StatusCode, body: &str) -> OAuthError {
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_pkce_pair() -> (String, String) {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rng().fill_bytes(&mut bytes);
|
||||
let code_verifier = URL_SAFE_NO_PAD.encode(bytes);
|
||||
let digest = Sha256::digest(code_verifier.as_bytes());
|
||||
let code_challenge = URL_SAFE_NO_PAD.encode(digest);
|
||||
(code_verifier, code_challenge)
|
||||
}
|
||||
|
||||
fn random_state() -> String {
|
||||
let mut bytes = [0u8; 16];
|
||||
rand::rng().fill_bytes(&mut bytes);
|
||||
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
|
||||
}
|
||||
|
||||
fn build_auth_url(
|
||||
config: &OAuthConfig,
|
||||
redirect_uri: &str,
|
||||
state: &str,
|
||||
code_challenge: &str,
|
||||
) -> String {
|
||||
let params = [
|
||||
("response_type", "code"),
|
||||
("client_id", config.client_id.as_str()),
|
||||
("redirect_uri", redirect_uri),
|
||||
("scope", "openid profile email offline_access"),
|
||||
("code_challenge", code_challenge),
|
||||
("code_challenge_method", "S256"),
|
||||
("id_token_add_organizations", "true"),
|
||||
("codex_cli_simplified_flow", "true"),
|
||||
("state", state),
|
||||
("originator", "codex_cli_rs"),
|
||||
];
|
||||
let query = params
|
||||
.iter()
|
||||
.map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(value)))
|
||||
.collect::<Vec<_>>()
|
||||
.join("&");
|
||||
format!("{}?{query}", config.auth_endpoint)
|
||||
}
|
||||
|
||||
fn percent_encode(value: &str) -> String {
|
||||
let mut encoded = String::new();
|
||||
for byte in value.bytes() {
|
||||
match byte {
|
||||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
|
||||
encoded.push(char::from(byte));
|
||||
}
|
||||
_ => encoded.push_str(&format!("%{byte:02X}")),
|
||||
}
|
||||
}
|
||||
encoded
|
||||
}
|
||||
|
||||
fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure {
|
||||
let lower = body.to_ascii_lowercase();
|
||||
if lower.contains("quota") {
|
||||
|
||||
+336
-2
@@ -4,7 +4,7 @@ use axum::{
|
||||
body::Bytes,
|
||||
extract::{Path, Query, State},
|
||||
http::{
|
||||
header::{HeaderValue, SET_COOKIE},
|
||||
header::{HeaderValue, HOST, LOCATION, SET_COOKIE},
|
||||
HeaderMap, StatusCode,
|
||||
},
|
||||
response::{IntoResponse, Response},
|
||||
@@ -33,7 +33,7 @@ use crate::{
|
||||
auth::{
|
||||
admin_session::verify_admin_password,
|
||||
cli_import::{default_codex_home, read_cli_auth_from_home},
|
||||
oauth::{DeviceCode, OAuthError},
|
||||
oauth::{DeviceCode, OAuthClient, OAuthConfig, OAuthError, PkceSession},
|
||||
refresh::RefreshFailure,
|
||||
token::TokenPair,
|
||||
},
|
||||
@@ -264,6 +264,34 @@ pub struct AdminAuthDevicePollData {
|
||||
pub code: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminAuthLoginStartData {
|
||||
pub auth_url: String,
|
||||
pub state: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminAuthCodeRelayRequest {
|
||||
pub callback_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminAuthCallbackQuery {
|
||||
pub code: Option<String>,
|
||||
pub state: Option<String>,
|
||||
pub error: Option<String>,
|
||||
pub error_description: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminAuthCodeRelayData {
|
||||
pub success: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminSettingsData {
|
||||
@@ -1149,6 +1177,196 @@ pub async fn auth_device_poll(
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn auth_login_start(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
let return_host = request_host(&headers, state.config());
|
||||
let login = state
|
||||
.oauth_sessions()
|
||||
.lock()
|
||||
.await
|
||||
.start_login(&return_host, &OAuthConfig::codex_default());
|
||||
AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
AdminAuthLoginStartData {
|
||||
auth_url: login.auth_url,
|
||||
state: login.state,
|
||||
},
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub async fn auth_code_relay(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<AdminAuthCodeRelayRequest>,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
|
||||
let Ok((code, oauth_state)) = parse_callback_url(&payload.callback_url) else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(40001, "Invalid OAuth callback URL", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let Some(repo) = state.account_repository() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(
|
||||
50001,
|
||||
"Account repository is not initialized",
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let Some(oauth_client) = state.oauth_client() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
let session = match acquire_pkce_session(&state, &oauth_state, &request_id).await {
|
||||
Ok(Some(session)) => session,
|
||||
Ok(None) => {
|
||||
return AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAuthCodeRelayData { success: true }, request_id),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
Err(response) => return response,
|
||||
};
|
||||
match exchange_pkce_session(
|
||||
&state,
|
||||
&repo,
|
||||
oauth_client.as_ref(),
|
||||
&oauth_state,
|
||||
&code,
|
||||
session,
|
||||
&request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAuthCodeRelayData { success: true }, request_id),
|
||||
)
|
||||
.into_response(),
|
||||
Err(response) => response,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn auth_callback(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<AdminAuthCallbackQuery>,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
if query.error.is_some() {
|
||||
return AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(
|
||||
40001,
|
||||
query
|
||||
.error_description
|
||||
.unwrap_or_else(|| "OAuth error".to_string()),
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let (Some(code), Some(oauth_state)) = (query.code.as_deref(), query.state.as_deref()) else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(40001, "Missing OAuth code or state", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let Some(repo) = state.account_repository() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(
|
||||
50001,
|
||||
"Account repository is not initialized",
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let Some(oauth_client) = state.oauth_client() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let session = match acquire_pkce_session(&state, oauth_state, &request_id).await {
|
||||
Ok(Some(session)) => session,
|
||||
Ok(None) => return redirect_to_host(&request_host(&headers, state.config()), request_id),
|
||||
Err(response) => return response,
|
||||
};
|
||||
let return_host = session.return_host.clone();
|
||||
match exchange_pkce_session(
|
||||
&state,
|
||||
&repo,
|
||||
oauth_client.as_ref(),
|
||||
oauth_state,
|
||||
code,
|
||||
session,
|
||||
&request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => redirect_to_host(&return_host, request_id),
|
||||
Err(response) => response,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn logs(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
@@ -4206,6 +4424,122 @@ fn oauth_error_response(error: OAuthError, request_id: String, message: &'static
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn request_host(headers: &HeaderMap, config: &AppConfig) -> String {
|
||||
headers
|
||||
.get(HOST)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| format!("localhost:{}", config.server.port))
|
||||
}
|
||||
|
||||
fn parse_callback_url(callback_url: &str) -> Result<(String, String), ()> {
|
||||
let callback_url = callback_url.trim();
|
||||
if callback_url.is_empty() {
|
||||
return Err(());
|
||||
}
|
||||
let url = Url::parse(callback_url).map_err(|_| ())?;
|
||||
if url.query_pairs().any(|(key, _)| key == "error") {
|
||||
return Err(());
|
||||
}
|
||||
let code = url
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| (key == "code").then(|| value.into_owned()))
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or(())?;
|
||||
let state = url
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or(())?;
|
||||
Ok((code, state))
|
||||
}
|
||||
|
||||
async fn acquire_pkce_session(
|
||||
state: &AppState,
|
||||
oauth_state: &str,
|
||||
request_id: &str,
|
||||
) -> Result<Option<PkceSession>, Response> {
|
||||
let sessions = state.oauth_sessions();
|
||||
let mut sessions = sessions.lock().await;
|
||||
if let Some(session) = sessions.try_acquire(oauth_state) {
|
||||
return Ok(Some(session));
|
||||
}
|
||||
if sessions.is_completed_or_exchanging(oauth_state) {
|
||||
return Ok(None);
|
||||
}
|
||||
Err(AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(
|
||||
40001,
|
||||
"Invalid or expired OAuth session",
|
||||
(),
|
||||
request_id.to_string(),
|
||||
),
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
|
||||
async fn exchange_pkce_session(
|
||||
state: &AppState,
|
||||
repo: &crate::accounts::repository::AccountRepository,
|
||||
oauth_client: &dyn OAuthClient,
|
||||
oauth_state: &str,
|
||||
code: &str,
|
||||
session: PkceSession,
|
||||
request_id: &str,
|
||||
) -> Result<String, Response> {
|
||||
let tokens = match oauth_client
|
||||
.exchange_code(code, &session.code_verifier, &session.redirect_uri)
|
||||
.await
|
||||
{
|
||||
Ok(tokens) => tokens,
|
||||
Err(error) => {
|
||||
state.oauth_sessions().lock().await.release(oauth_state);
|
||||
return Err(oauth_error_response(
|
||||
error,
|
||||
request_id.to_string(),
|
||||
"Token exchange failed",
|
||||
));
|
||||
}
|
||||
};
|
||||
match store_validated_account_import(
|
||||
state,
|
||||
repo,
|
||||
Some(tokens.access_token),
|
||||
tokens.refresh_token,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
state.oauth_sessions().lock().await.complete(oauth_state);
|
||||
Ok(session.return_host)
|
||||
}
|
||||
Err(error) => {
|
||||
state.oauth_sessions().lock().await.complete(oauth_state);
|
||||
Err(validated_account_import_error_response(
|
||||
error,
|
||||
request_id.to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn redirect_to_host(return_host: &str, request_id: String) -> Response {
|
||||
let location = format!("http://{return_host}/");
|
||||
let Ok(location) = HeaderValue::from_str(&location) else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Invalid OAuth return host", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
let mut response = StatusCode::SEE_OTHER.into_response();
|
||||
response.headers_mut().insert(LOCATION, location);
|
||||
response
|
||||
}
|
||||
|
||||
fn manual_account_claims(
|
||||
token: &str,
|
||||
now: DateTime<Utc>,
|
||||
|
||||
+14
-1
@@ -8,7 +8,9 @@ use crate::accounts::{
|
||||
repository::{AccountRepository, AccountRepositoryResult, AccountUsageRepository},
|
||||
};
|
||||
use crate::auth::{
|
||||
api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository, oauth::OAuthClient,
|
||||
api_key::ApiKeyHasher,
|
||||
api_key_repository::ClientApiKeyRepository,
|
||||
oauth::{OAuthClient, PkceSessionStore},
|
||||
refresh::TokenRefresher,
|
||||
};
|
||||
use crate::config::AppConfig;
|
||||
@@ -30,6 +32,7 @@ pub struct AppServices {
|
||||
pub api_key_hasher: Option<ApiKeyHasher>,
|
||||
pub token_refresher: Option<Arc<dyn TokenRefresher>>,
|
||||
pub oauth_client: Option<Arc<dyn OAuthClient>>,
|
||||
pub oauth_sessions: Arc<Mutex<PkceSessionStore>>,
|
||||
pub account_pool: Arc<Mutex<AccountPool>>,
|
||||
}
|
||||
|
||||
@@ -45,6 +48,7 @@ impl AppState {
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
oauth_client: None,
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -61,6 +65,7 @@ impl AppState {
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
oauth_client: None,
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -81,6 +86,7 @@ impl AppState {
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
oauth_client: None,
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -102,6 +108,7 @@ impl AppState {
|
||||
api_key_hasher: Some(api_key_hasher),
|
||||
token_refresher: None,
|
||||
oauth_client: None,
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -127,6 +134,7 @@ impl AppState {
|
||||
api_key_hasher: Some(api_key_hasher),
|
||||
token_refresher: Some(Arc::new(token_refresher)),
|
||||
oauth_client: None,
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -155,6 +163,7 @@ impl AppState {
|
||||
api_key_hasher: Some(api_key_hasher),
|
||||
token_refresher: Some(token_refresher),
|
||||
oauth_client: Some(oauth_client),
|
||||
oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -214,6 +223,10 @@ impl AppState {
|
||||
self.services.oauth_client.clone()
|
||||
}
|
||||
|
||||
pub fn oauth_sessions(&self) -> Arc<Mutex<PkceSessionStore>> {
|
||||
self.services.oauth_sessions.clone()
|
||||
}
|
||||
|
||||
pub async fn reload_account_pool_from_repository(&self) -> AccountRepositoryResult<usize> {
|
||||
let Some(repo) = self.account_repository() else {
|
||||
return Ok(0);
|
||||
|
||||
@@ -52,7 +52,16 @@ struct FailingTokenRefresher {
|
||||
struct StaticOAuthClient {
|
||||
device_response: Result<DeviceCode, OAuthError>,
|
||||
poll_response: Result<TokenPair, OAuthError>,
|
||||
exchange_response: Result<TokenPair, OAuthError>,
|
||||
poll_calls: Arc<Mutex<Vec<String>>>,
|
||||
exchange_calls: Arc<Mutex<Vec<ExchangeCall>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct ExchangeCall {
|
||||
code: String,
|
||||
code_verifier: String,
|
||||
redirect_uri: String,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -79,6 +88,20 @@ impl TokenRefresher for StaticOAuthClient {
|
||||
|
||||
#[async_trait]
|
||||
impl OAuthClient for StaticOAuthClient {
|
||||
async fn exchange_code(
|
||||
&self,
|
||||
code: &str,
|
||||
code_verifier: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<TokenPair, OAuthError> {
|
||||
self.exchange_calls.lock().await.push(ExchangeCall {
|
||||
code: code.to_string(),
|
||||
code_verifier: code_verifier.to_string(),
|
||||
redirect_uri: redirect_uri.to_string(),
|
||||
});
|
||||
self.exchange_response.clone()
|
||||
}
|
||||
|
||||
async fn request_device_code(&self) -> Result<DeviceCode, OAuthError> {
|
||||
self.device_response.clone()
|
||||
}
|
||||
@@ -371,7 +394,9 @@ async fn admin_auth_device_login_should_return_openai_device_code() {
|
||||
interval: 5,
|
||||
}),
|
||||
poll_response: Err(OAuthError::AuthorizationPending),
|
||||
exchange_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
exchange_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let (app, _state, _pool, _dir) =
|
||||
admin_accounts_test_app_with_oauth_client("admin-auth-device-login.sqlite", 38, client)
|
||||
@@ -408,7 +433,9 @@ async fn admin_auth_device_poll_should_return_pending_without_importing_account(
|
||||
let client = StaticOAuthClient {
|
||||
device_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_response: Err(OAuthError::SlowDown),
|
||||
exchange_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_calls: poll_calls.clone(),
|
||||
exchange_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app_with_oauth_client("admin-auth-device-pending.sqlite", 39, client)
|
||||
@@ -456,7 +483,9 @@ async fn admin_auth_device_poll_should_import_tokens_without_returning_secrets()
|
||||
access_token: token.clone(),
|
||||
refresh_token: Some("device-refresh-secret".to_string()),
|
||||
}),
|
||||
exchange_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_calls: poll_calls.clone(),
|
||||
exchange_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let (app, state, pool, _dir) =
|
||||
admin_accounts_test_app_with_oauth_client("admin-auth-device-success.sqlite", 40, client)
|
||||
@@ -513,6 +542,251 @@ async fn admin_auth_device_poll_should_import_tokens_without_returning_secrets()
|
||||
assert_eq!(acquired.email.as_deref(), Some("device@example.com"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_login_start_should_require_admin_session_cookie() {
|
||||
let (app, _state, _pool, _dir) =
|
||||
admin_accounts_test_app("admin-auth-login-start-auth.sqlite", 41).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/login-start")
|
||||
.header("x-request-id", "req_login_start_no_session")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["code"], 40101);
|
||||
assert_eq!(body["requestId"], "req_login_start_no_session");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_login_start_should_return_pkce_auth_url_and_state() {
|
||||
let (app, _state, _pool, _dir) =
|
||||
admin_accounts_test_app("admin-auth-login-start.sqlite", 42).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/login-start")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("host", "127.0.0.1:8080")
|
||||
.header("x-request-id", "req_login_start")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
let state = body["data"]["state"].as_str().unwrap();
|
||||
let auth_url = body["data"]["authUrl"].as_str().unwrap();
|
||||
assert_eq!(state.len(), 32);
|
||||
assert!(auth_url.starts_with("https://auth.openai.com/oauth/authorize?"));
|
||||
assert!(auth_url.contains("response_type=code"));
|
||||
assert!(auth_url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
|
||||
assert!(auth_url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback"));
|
||||
assert!(auth_url.contains("scope=openid%20profile%20email%20offline_access"));
|
||||
assert!(auth_url.contains("code_challenge_method=S256"));
|
||||
assert!(auth_url.contains("originator=codex_cli_rs"));
|
||||
assert!(auth_url.contains(&format!("state={state}")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_code_relay_should_exchange_code_and_import_account() {
|
||||
let token = test_jwt(
|
||||
Some("pkce-account"),
|
||||
Some("pkce-user"),
|
||||
Some("pkce@example.com"),
|
||||
Some("plus"),
|
||||
3600,
|
||||
);
|
||||
let exchange_calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let client = StaticOAuthClient {
|
||||
device_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_response: Err(OAuthError::AuthorizationPending),
|
||||
exchange_response: Ok(TokenPair {
|
||||
access_token: token.clone(),
|
||||
refresh_token: Some("pkce-refresh-secret".to_string()),
|
||||
}),
|
||||
poll_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
exchange_calls: exchange_calls.clone(),
|
||||
};
|
||||
let (app, state, pool, _dir) =
|
||||
admin_accounts_test_app_with_oauth_client("admin-auth-code-relay.sqlite", 43, client).await;
|
||||
let start = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/login-start")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("host", "127.0.0.1:8080")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let start_body = response_json(start).await;
|
||||
let state_value = start_body["data"]["state"].as_str().unwrap();
|
||||
let callback_url =
|
||||
format!("http://localhost:1455/auth/callback?code=oauth-code&state={state_value}");
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/code-relay")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_code_relay")
|
||||
.body(Body::from(
|
||||
json!({ "callbackUrl": callback_url }).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["success"], true);
|
||||
assert!(body["data"].get("accessToken").is_none());
|
||||
assert!(body["data"].get("refreshToken").is_none());
|
||||
let calls = exchange_calls.lock().await;
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].code, "oauth-code");
|
||||
assert!(!calls[0].code_verifier.is_empty());
|
||||
assert_eq!(calls[0].redirect_uri, "http://localhost:1455/auth/callback");
|
||||
drop(calls);
|
||||
let stored = AccountRepository::new(pool.clone(), SecretBox::new([43; 32]))
|
||||
.find_by_chatgpt_identity("pkce-account", Some("pkce-user"))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(stored.email.as_deref(), Some("pkce@example.com"));
|
||||
assert_eq!(
|
||||
stored.refresh_token.unwrap().expose_secret(),
|
||||
"pkce-refresh-secret"
|
||||
);
|
||||
let acquired = state
|
||||
.account_pool()
|
||||
.lock()
|
||||
.await
|
||||
.acquire("gpt-5.5")
|
||||
.unwrap();
|
||||
assert_eq!(acquired.account_id.as_deref(), Some("pkce-account"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_code_relay_should_reject_invalid_callback_url() {
|
||||
let client = StaticOAuthClient {
|
||||
device_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_response: Err(OAuthError::AuthorizationPending),
|
||||
exchange_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
exchange_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let (app, _state, _pool, _dir) = admin_accounts_test_app_with_oauth_client(
|
||||
"admin-auth-code-relay-invalid.sqlite",
|
||||
44,
|
||||
client,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/code-relay")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_code_relay_invalid")
|
||||
.body(Body::from(
|
||||
json!({ "callbackUrl": "not a url" }).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["code"], 40001);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_auth_callback_should_exchange_code_and_redirect_to_return_host() {
|
||||
let token = test_jwt(
|
||||
Some("callback-account"),
|
||||
Some("callback-user"),
|
||||
Some("callback@example.com"),
|
||||
Some("plus"),
|
||||
3600,
|
||||
);
|
||||
let exchange_calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let client = StaticOAuthClient {
|
||||
device_response: Err(OAuthError::Rejected("not used".to_string())),
|
||||
poll_response: Err(OAuthError::AuthorizationPending),
|
||||
exchange_response: Ok(TokenPair {
|
||||
access_token: token,
|
||||
refresh_token: Some("callback-refresh-secret".to_string()),
|
||||
}),
|
||||
poll_calls: Arc::new(Mutex::new(Vec::new())),
|
||||
exchange_calls: exchange_calls.clone(),
|
||||
};
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app_with_oauth_client("admin-auth-callback.sqlite", 45, client).await;
|
||||
let start = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/auth/login-start")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("host", "codex.local:1455")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let start_body = response_json(start).await;
|
||||
let state_value = start_body["data"]["state"].as_str().unwrap();
|
||||
let callback_path = format!("/auth/callback?code=callback-code&state={state_value}");
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri(callback_path)
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SEE_OTHER);
|
||||
assert_eq!(
|
||||
response.headers().get("location").unwrap(),
|
||||
"http://codex.local:1455/"
|
||||
);
|
||||
assert_eq!(exchange_calls.lock().await[0].code, "callback-code");
|
||||
let stored = AccountRepository::new(pool.clone(), SecretBox::new([45; 32]))
|
||||
.find_by_chatgpt_identity("callback-account", Some("callback-user"))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(stored.email.as_deref(), Some("callback@example.com"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_import_should_store_tokens_encrypted_and_list_sanitized_accounts() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -68,6 +68,39 @@ async fn openai_oauth_client_should_request_device_code_with_form_body() {
|
||||
assert!(body.contains("scope=openid+profile+email+offline_access"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_oauth_client_should_exchange_authorization_code_with_pkce_verifier() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "pkce-access",
|
||||
"refresh_token": "pkce-refresh"
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = test_oauth_client(&server);
|
||||
|
||||
let tokens = client
|
||||
.exchange_code(
|
||||
"oauth-code",
|
||||
"pkce-verifier",
|
||||
"http://localhost:1455/auth/callback",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(tokens.access_token, "pkce-access");
|
||||
assert_eq!(tokens.refresh_token.as_deref(), Some("pkce-refresh"));
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let body = String::from_utf8(requests[0].body.clone()).unwrap();
|
||||
assert!(body.contains("grant_type=authorization_code"));
|
||||
assert!(body.contains("client_id=codex-client"));
|
||||
assert!(body.contains("code=oauth-code"));
|
||||
assert!(body.contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback"));
|
||||
assert!(body.contains("code_verifier=pkce-verifier"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_oauth_client_should_map_device_poll_pending_errors() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Reference in New Issue
Block a user