mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-02 04:20:50 +08:00
287 lines
9.1 KiB
Go
287 lines
9.1 KiB
Go
// Package codex provides authentication and token management for OpenAI's Codex API.
|
|
// It handles the OAuth2 flow, including generating authorization URLs, exchanging
|
|
// authorization codes for tokens, and refreshing expired tokens. The package also
|
|
// defines data structures for storing and managing Codex authentication credentials.
|
|
package codex
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
const (
|
|
openaiAuthURL = "https://auth.openai.com/oauth/authorize"
|
|
openaiTokenURL = "https://auth.openai.com/oauth/token"
|
|
openaiClientID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
|
redirectURI = "http://localhost:1455/auth/callback"
|
|
)
|
|
|
|
// CodexAuth handles the OpenAI OAuth2 authentication flow.
|
|
// It manages the HTTP client and provides methods for generating authorization URLs,
|
|
// exchanging authorization codes for tokens, and refreshing access tokens.
|
|
type CodexAuth struct {
|
|
httpClient *http.Client
|
|
}
|
|
|
|
// NewCodexAuth creates a new CodexAuth service instance.
|
|
// It initializes an HTTP client with proxy settings from the provided configuration.
|
|
func NewCodexAuth(cfg *config.Config) *CodexAuth {
|
|
return &CodexAuth{
|
|
httpClient: util.SetProxy(cfg, &http.Client{}),
|
|
}
|
|
}
|
|
|
|
// GenerateAuthURL creates the OAuth authorization URL with PKCE (Proof Key for Code Exchange).
|
|
// It constructs the URL with the necessary parameters, including the client ID,
|
|
// response type, redirect URI, scopes, and PKCE challenge.
|
|
func (o *CodexAuth) GenerateAuthURL(state string, pkceCodes *PKCECodes) (string, error) {
|
|
if pkceCodes == nil {
|
|
return "", fmt.Errorf("PKCE codes are required")
|
|
}
|
|
|
|
params := url.Values{
|
|
"client_id": {openaiClientID},
|
|
"response_type": {"code"},
|
|
"redirect_uri": {redirectURI},
|
|
"scope": {"openid email profile offline_access"},
|
|
"state": {state},
|
|
"code_challenge": {pkceCodes.CodeChallenge},
|
|
"code_challenge_method": {"S256"},
|
|
"prompt": {"login"},
|
|
"id_token_add_organizations": {"true"},
|
|
"codex_cli_simplified_flow": {"true"},
|
|
}
|
|
|
|
authURL := fmt.Sprintf("%s?%s", openaiAuthURL, params.Encode())
|
|
return authURL, nil
|
|
}
|
|
|
|
// ExchangeCodeForTokens exchanges an authorization code for access and refresh tokens.
|
|
// It performs an HTTP POST request to the OpenAI token endpoint with the provided
|
|
// authorization code and PKCE verifier.
|
|
func (o *CodexAuth) ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *PKCECodes) (*CodexAuthBundle, error) {
|
|
if pkceCodes == nil {
|
|
return nil, fmt.Errorf("PKCE codes are required for token exchange")
|
|
}
|
|
|
|
// Prepare token exchange request
|
|
data := url.Values{
|
|
"grant_type": {"authorization_code"},
|
|
"client_id": {openaiClientID},
|
|
"code": {code},
|
|
"redirect_uri": {redirectURI},
|
|
"code_verifier": {pkceCodes.CodeVerifier},
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", openaiTokenURL, strings.NewReader(data.Encode()))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create token request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := o.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("token exchange request failed: %w", err)
|
|
}
|
|
defer func() {
|
|
_ = resp.Body.Close()
|
|
}()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to read token response: %w", err)
|
|
}
|
|
// log.Debugf("Token response: %s", string(body))
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
// Parse token response
|
|
var tokenResp struct {
|
|
AccessToken string `json:"access_token"`
|
|
RefreshToken string `json:"refresh_token"`
|
|
IDToken string `json:"id_token"`
|
|
TokenType string `json:"token_type"`
|
|
ExpiresIn int `json:"expires_in"`
|
|
}
|
|
|
|
if err = json.Unmarshal(body, &tokenResp); err != nil {
|
|
return nil, fmt.Errorf("failed to parse token response: %w", err)
|
|
}
|
|
|
|
// Extract account ID from ID token
|
|
claims, err := ParseJWTToken(tokenResp.IDToken)
|
|
if err != nil {
|
|
log.Warnf("Failed to parse ID token: %v", err)
|
|
}
|
|
|
|
accountID := ""
|
|
email := ""
|
|
if claims != nil {
|
|
accountID = claims.GetAccountID()
|
|
email = claims.GetUserEmail()
|
|
}
|
|
|
|
// Create token data
|
|
tokenData := CodexTokenData{
|
|
IDToken: tokenResp.IDToken,
|
|
AccessToken: tokenResp.AccessToken,
|
|
RefreshToken: tokenResp.RefreshToken,
|
|
AccountID: accountID,
|
|
Email: email,
|
|
Expire: time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339),
|
|
}
|
|
|
|
// Create auth bundle
|
|
bundle := &CodexAuthBundle{
|
|
TokenData: tokenData,
|
|
LastRefresh: time.Now().Format(time.RFC3339),
|
|
}
|
|
|
|
return bundle, nil
|
|
}
|
|
|
|
// RefreshTokens refreshes an access token using a refresh token.
|
|
// This method is called when an access token has expired. It makes a request to the
|
|
// token endpoint to obtain a new set of tokens.
|
|
func (o *CodexAuth) RefreshTokens(ctx context.Context, refreshToken string) (*CodexTokenData, error) {
|
|
if refreshToken == "" {
|
|
return nil, fmt.Errorf("refresh token is required")
|
|
}
|
|
|
|
data := url.Values{
|
|
"client_id": {openaiClientID},
|
|
"grant_type": {"refresh_token"},
|
|
"refresh_token": {refreshToken},
|
|
"scope": {"openid profile email"},
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", openaiTokenURL, strings.NewReader(data.Encode()))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create refresh request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := o.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("token refresh request failed: %w", err)
|
|
}
|
|
defer func() {
|
|
_ = resp.Body.Close()
|
|
}()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to read refresh response: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("token refresh failed with status %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var tokenResp struct {
|
|
AccessToken string `json:"access_token"`
|
|
RefreshToken string `json:"refresh_token"`
|
|
IDToken string `json:"id_token"`
|
|
TokenType string `json:"token_type"`
|
|
ExpiresIn int `json:"expires_in"`
|
|
}
|
|
|
|
if err = json.Unmarshal(body, &tokenResp); err != nil {
|
|
return nil, fmt.Errorf("failed to parse refresh response: %w", err)
|
|
}
|
|
|
|
// Extract account ID from ID token
|
|
claims, err := ParseJWTToken(tokenResp.IDToken)
|
|
if err != nil {
|
|
log.Warnf("Failed to parse refreshed ID token: %v", err)
|
|
}
|
|
|
|
accountID := ""
|
|
email := ""
|
|
if claims != nil {
|
|
accountID = claims.GetAccountID()
|
|
email = claims.Email
|
|
}
|
|
|
|
return &CodexTokenData{
|
|
IDToken: tokenResp.IDToken,
|
|
AccessToken: tokenResp.AccessToken,
|
|
RefreshToken: tokenResp.RefreshToken,
|
|
AccountID: accountID,
|
|
Email: email,
|
|
Expire: time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339),
|
|
}, nil
|
|
}
|
|
|
|
// CreateTokenStorage creates a new CodexTokenStorage from a CodexAuthBundle.
|
|
// It populates the storage struct with token data, user information, and timestamps.
|
|
func (o *CodexAuth) CreateTokenStorage(bundle *CodexAuthBundle) *CodexTokenStorage {
|
|
storage := &CodexTokenStorage{
|
|
IDToken: bundle.TokenData.IDToken,
|
|
AccessToken: bundle.TokenData.AccessToken,
|
|
RefreshToken: bundle.TokenData.RefreshToken,
|
|
AccountID: bundle.TokenData.AccountID,
|
|
LastRefresh: bundle.LastRefresh,
|
|
Email: bundle.TokenData.Email,
|
|
Expire: bundle.TokenData.Expire,
|
|
}
|
|
|
|
return storage
|
|
}
|
|
|
|
// RefreshTokensWithRetry refreshes tokens with a built-in retry mechanism.
|
|
// It attempts to refresh the tokens up to a specified maximum number of retries,
|
|
// with an exponential backoff strategy to handle transient network errors.
|
|
func (o *CodexAuth) RefreshTokensWithRetry(ctx context.Context, refreshToken string, maxRetries int) (*CodexTokenData, error) {
|
|
var lastErr error
|
|
|
|
for attempt := 0; attempt < maxRetries; attempt++ {
|
|
if attempt > 0 {
|
|
// Wait before retry
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-time.After(time.Duration(attempt) * time.Second):
|
|
}
|
|
}
|
|
|
|
tokenData, err := o.RefreshTokens(ctx, refreshToken)
|
|
if err == nil {
|
|
return tokenData, nil
|
|
}
|
|
|
|
lastErr = err
|
|
log.Warnf("Token refresh attempt %d failed: %v", attempt+1, err)
|
|
}
|
|
|
|
return nil, fmt.Errorf("token refresh failed after %d attempts: %w", maxRetries, lastErr)
|
|
}
|
|
|
|
// UpdateTokenStorage updates an existing CodexTokenStorage with new token data.
|
|
// This is typically called after a successful token refresh to persist the new credentials.
|
|
func (o *CodexAuth) UpdateTokenStorage(storage *CodexTokenStorage, tokenData *CodexTokenData) {
|
|
storage.IDToken = tokenData.IDToken
|
|
storage.AccessToken = tokenData.AccessToken
|
|
storage.RefreshToken = tokenData.RefreshToken
|
|
storage.AccountID = tokenData.AccountID
|
|
storage.LastRefresh = time.Now().Format(time.RFC3339)
|
|
storage.Email = tokenData.Email
|
|
storage.Expire = tokenData.Expire
|
|
}
|