李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
468 lines
16 KiB
Go
468 lines
16 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"log/slog"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// OAuthRefreshExecutor 各平台实现的 OAuth 刷新执行器
|
||
// TokenRefresher 接口的超集：增加了 CacheKey 方法用于分布式锁
|
||
type OAuthRefreshExecutor interface {
|
||
TokenRefresher
|
||
|
||
// CacheKey 返回用于分布式锁的缓存键（与 TokenProvider 使用的一致）
|
||
CacheKey(account *Account) string
|
||
}
|
||
|
||
// GrokOAuthRefreshSuccessRepository is the persistence boundary for a
|
||
// provider-issued Grok credential rotation. Implementations must compare the
|
||
// complete credential document and proxy used by the upstream attempt, and
|
||
// atomically publish scheduler invalidation with a successful update.
|
||
type GrokOAuthRefreshSuccessRepository interface {
|
||
UpdateGrokOAuthCredentialsIfUnchanged(
|
||
ctx context.Context,
|
||
id int64,
|
||
expectedCredentials map[string]any,
|
||
expectedProxyID *int64,
|
||
credentials map[string]any,
|
||
) (bool, error)
|
||
}
|
||
|
||
const (
|
||
defaultRefreshLockTTL = 60 * time.Second
|
||
defaultRefreshLockReleaseTimeout = 2 * time.Second
|
||
defaultRefreshPostPersistCleanupTimeout = 2 * time.Second
|
||
)
|
||
|
||
var (
|
||
errOAuthRefreshAccountRereadFailed = errors.New("oauth refresh account reread failed")
|
||
errOAuthRefreshAccountStateChanged = errors.New("oauth refresh account state changed")
|
||
errOAuthRefreshCredentialPersist = errors.New("oauth refresh credential persistence failed")
|
||
)
|
||
|
||
type oauthRefreshRequestPathKey struct{}
|
||
|
||
func withOAuthRefreshRequestPath(ctx context.Context) context.Context {
|
||
return context.WithValue(ctx, oauthRefreshRequestPathKey{}, true)
|
||
}
|
||
|
||
func isOAuthRefreshRequestPath(ctx context.Context) bool {
|
||
requestPath, _ := ctx.Value(oauthRefreshRequestPathKey{}).(bool)
|
||
return requestPath
|
||
}
|
||
|
||
type contextMutex struct {
|
||
token chan struct{}
|
||
}
|
||
|
||
// Keep the request-path credential mutation lock API introduced by #4212
|
||
// while sharing the context-aware mutex implementation used by pool refresh.
|
||
type oauthRefreshLocalLock = contextMutex
|
||
|
||
func newOAuthRefreshLocalLock() *oauthRefreshLocalLock {
|
||
return newContextMutex()
|
||
}
|
||
|
||
type oauthRefreshStateUnavailableError struct {
|
||
err error
|
||
}
|
||
|
||
func (e *oauthRefreshStateUnavailableError) Error() string {
|
||
return "OAuth refresh account state is unavailable"
|
||
}
|
||
|
||
func (e *oauthRefreshStateUnavailableError) Unwrap() error {
|
||
if e == nil {
|
||
return nil
|
||
}
|
||
return e.err
|
||
}
|
||
|
||
func newContextMutex() *contextMutex {
|
||
return &contextMutex{token: make(chan struct{}, 1)}
|
||
}
|
||
|
||
func (m *contextMutex) Lock(ctx context.Context) error {
|
||
select {
|
||
case m.token <- struct{}{}:
|
||
return nil
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
}
|
||
}
|
||
|
||
func (m *contextMutex) Unlock() {
|
||
<-m.token
|
||
}
|
||
|
||
// OAuthRefreshResult 统一刷新结果
|
||
type OAuthRefreshResult struct {
|
||
Refreshed bool // 实际执行了刷新
|
||
NewCredentials map[string]any // 刷新后的 credentials（nil 表示未刷新）
|
||
Account *Account // 成功时为最新 account；刷新错误时为实际尝试的凭据快照
|
||
LockHeld bool // 锁被其他 worker 持有（未执行刷新）
|
||
}
|
||
|
||
func snapshotOAuthRefreshAccount(account *Account) *Account {
|
||
if account == nil {
|
||
return nil
|
||
}
|
||
snapshot := *account
|
||
snapshot.Credentials = shallowCopyMap(account.Credentials)
|
||
if account.ProxyID != nil {
|
||
proxyID := *account.ProxyID
|
||
snapshot.ProxyID = &proxyID
|
||
}
|
||
return &snapshot
|
||
}
|
||
|
||
// OAuthRefreshAPI 统一的 OAuth Token 刷新入口
|
||
// 封装分布式锁、进程内互斥锁、DB 重读、已刷新检查、竞争恢复等通用逻辑
|
||
type OAuthRefreshAPI struct {
|
||
accountRepo AccountRepository
|
||
tokenCache GeminiTokenCache // 可选，nil = 无分布式锁
|
||
lockTTL time.Duration
|
||
localLocks sync.Map // key: cacheKey string -> value: *contextMutex
|
||
}
|
||
|
||
// NewOAuthRefreshAPI 创建统一刷新 API
|
||
// 可选传入 lockTTL 覆盖默认的 60s 分布式锁 TTL
|
||
func NewOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiTokenCache, lockTTL ...time.Duration) *OAuthRefreshAPI {
|
||
ttl := defaultRefreshLockTTL
|
||
if len(lockTTL) > 0 && lockTTL[0] > 0 {
|
||
ttl = lockTTL[0]
|
||
}
|
||
return &OAuthRefreshAPI{
|
||
accountRepo: accountRepo,
|
||
tokenCache: tokenCache,
|
||
lockTTL: ttl,
|
||
}
|
||
}
|
||
|
||
// getLocalLock 返回指定 cacheKey 的进程内互斥锁
|
||
func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *contextMutex {
|
||
actual, _ := api.localLocks.LoadOrStore(cacheKey, newContextMutex())
|
||
mu, ok := actual.(*contextMutex)
|
||
if !ok {
|
||
mu = newContextMutex()
|
||
api.localLocks.Store(cacheKey, mu)
|
||
}
|
||
return mu
|
||
}
|
||
|
||
// RefreshIfNeeded 在分布式锁保护下按需刷新 OAuth token
|
||
//
|
||
// 流程:
|
||
// 1. 获取分布式锁
|
||
// 2. 从 DB 重读最新 account（防止使用过时的 refresh_token）
|
||
// 3. 二次检查是否仍需刷新
|
||
// 4. 调用 executor.Refresh() 执行平台特定刷新逻辑
|
||
// 5. 设置 _token_version + 更新 DB
|
||
// 6. 释放锁
|
||
func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||
ctx context.Context,
|
||
account *Account,
|
||
executor OAuthRefreshExecutor,
|
||
refreshWindow time.Duration,
|
||
) (*OAuthRefreshResult, error) {
|
||
if api == nil || api.accountRepo == nil {
|
||
return nil, errors.New("oauth refresh account repository is not configured")
|
||
}
|
||
if account == nil {
|
||
return nil, errors.New("oauth refresh account is nil")
|
||
}
|
||
if executor == nil {
|
||
return nil, errors.New("oauth refresh executor is nil")
|
||
}
|
||
requestPath := isOAuthRefreshRequestPath(ctx)
|
||
cacheKey := executor.CacheKey(account)
|
||
|
||
// 0. 获取进程内互斥锁（防止同一进程内的并发刷新竞争）
|
||
localMu := api.getLocalLock(cacheKey)
|
||
if err := localMu.Lock(ctx); err != nil {
|
||
return nil, fmt.Errorf("oauth refresh local lock: %w", err)
|
||
}
|
||
defer localMu.Unlock()
|
||
|
||
// 1. 获取分布式锁
|
||
if api.tokenCache != nil {
|
||
acquired, lockErr := api.tokenCache.AcquireRefreshLock(ctx, cacheKey, api.lockTTL)
|
||
if lockErr != nil {
|
||
// Redis 错误，降级为无锁刷新（进程内互斥锁仍生效）
|
||
slog.Warn("oauth_refresh_lock_failed_degraded",
|
||
"account_id", account.ID,
|
||
"cache_key", cacheKey,
|
||
"error", lockErr,
|
||
)
|
||
} else if !acquired {
|
||
// 锁被其他 worker 持有
|
||
return &OAuthRefreshResult{LockHeld: true}, nil
|
||
} else {
|
||
defer api.releaseRefreshLock(ctx, cacheKey)
|
||
}
|
||
}
|
||
|
||
// 2. 从 DB 重读最新 account（锁保护下，确保使用最新的 refresh_token）
|
||
freshAccount, err := api.accountRepo.GetByID(ctx, account.ID)
|
||
if err != nil {
|
||
if requestPath {
|
||
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
|
||
}
|
||
return nil, &oauthRefreshStateUnavailableError{err: err}
|
||
}
|
||
if freshAccount == nil {
|
||
if requestPath {
|
||
return nil, fmt.Errorf("%w: account not found", errOAuthRefreshAccountStateChanged)
|
||
}
|
||
return nil, &oauthRefreshStateUnavailableError{err: fmt.Errorf("account not found")}
|
||
}
|
||
if freshAccount.ID != account.ID {
|
||
return nil, fmt.Errorf("%w: account identity mismatch", errOAuthRefreshAccountRereadFailed)
|
||
}
|
||
if !freshAccount.IsActive() {
|
||
if requestPath {
|
||
return nil, fmt.Errorf("%w: account is not active", errOAuthRefreshAccountStateChanged)
|
||
}
|
||
return &OAuthRefreshResult{Account: freshAccount}, nil
|
||
}
|
||
if requestPath && freshAccount.Platform == PlatformGrok {
|
||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil {
|
||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount)
|
||
}
|
||
}
|
||
if !executor.CanRefresh(freshAccount) {
|
||
if requestPath && freshAccount.IsGrokOAuth() && strings.TrimSpace(freshAccount.GetGrokRefreshToken()) == "" {
|
||
return nil, withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, freshAccount)
|
||
}
|
||
if requestPath {
|
||
return nil, fmt.Errorf("%w: account is no longer refreshable", errOAuthRefreshAccountStateChanged)
|
||
}
|
||
return &OAuthRefreshResult{Account: freshAccount}, nil
|
||
}
|
||
|
||
// 3. 二次检查是否仍需刷新（另一条路径可能已刷新）
|
||
if !executor.NeedsRefresh(freshAccount, refreshWindow) {
|
||
return &OAuthRefreshResult{
|
||
Account: freshAccount,
|
||
}, nil
|
||
}
|
||
|
||
// 4. 执行平台特定刷新逻辑
|
||
attemptedAccount := snapshotOAuthRefreshAccount(freshAccount)
|
||
newCredentials, refreshErr := executor.Refresh(ctx, freshAccount)
|
||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||
// A provider implementation may ignore cancellation and return late
|
||
// credentials. Never persist them after the attempt/cycle boundary.
|
||
return nil, ctxErr
|
||
}
|
||
if refreshErr != nil {
|
||
// 竞争恢复：invalid_grant 可能是另一个 worker 已消费了旧 refresh_token
|
||
// 重新读取 DB，如果 refresh_token 已更新则说明是竞争，返回成功
|
||
if isInvalidGrantError(refreshErr) {
|
||
if recoveredAccount, recovered := api.tryRecoverFromRefreshRace(ctx, freshAccount); recovered {
|
||
if requestPath && recoveredAccount.Platform == PlatformGrok {
|
||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(recoveredAccount); eligibilityErr != nil {
|
||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, recoveredAccount)
|
||
}
|
||
}
|
||
slog.Info("oauth_refresh_race_recovered",
|
||
"account_id", freshAccount.ID,
|
||
"platform", freshAccount.Platform,
|
||
)
|
||
return &OAuthRefreshResult{
|
||
Account: recoveredAccount,
|
||
}, nil
|
||
}
|
||
}
|
||
// Preserve the exact account snapshot used by the failed upstream call.
|
||
// Callers can then conditionally mutate only that credential version and
|
||
// avoid quarantining a concurrently reauthorized account.
|
||
result := &OAuthRefreshResult{Account: attemptedAccount}
|
||
if requestPath && attemptedAccount.Platform == PlatformGrok {
|
||
return result, withGrokCredentialFailureSnapshot(refreshErr, attemptedAccount)
|
||
}
|
||
return result, refreshErr
|
||
}
|
||
|
||
// 5. 设置版本号 + 更新 DB
|
||
if newCredentials != nil {
|
||
newCredentials["_token_version"] = time.Now().UnixMilli()
|
||
if freshAccount.IsGrokOAuth() {
|
||
conditionalRepo, ok := api.accountRepo.(GrokOAuthRefreshSuccessRepository)
|
||
if !ok {
|
||
return nil, &providerConfigurationRefreshError{
|
||
err: fmt.Errorf("grok OAuth refresh success CAS repository is not configured"),
|
||
}
|
||
}
|
||
applied, updateErr := conditionalRepo.UpdateGrokOAuthCredentialsIfUnchanged(
|
||
ctx,
|
||
freshAccount.ID,
|
||
attemptedAccount.Credentials,
|
||
attemptedAccount.ProxyID,
|
||
newCredentials,
|
||
)
|
||
if updateErr != nil {
|
||
slog.Error("oauth_refresh_update_failed",
|
||
"account_id", freshAccount.ID,
|
||
"platform", freshAccount.Platform,
|
||
"error", updateErr,
|
||
)
|
||
// The provider may have rotated and consumed the refresh token.
|
||
// Retrying after an ambiguous local persistence result can turn a
|
||
// healthy account into invalid_grant, so contain this provider cycle.
|
||
return nil, &providerCycleContainmentRefreshError{
|
||
err: fmt.Errorf("OAuth refresh succeeded but credential persistence failed: %w", updateErr),
|
||
}
|
||
}
|
||
if !applied {
|
||
currentAccount, readErr := api.accountRepo.GetByID(ctx, freshAccount.ID)
|
||
if readErr != nil || currentAccount == nil {
|
||
if readErr == nil {
|
||
readErr = fmt.Errorf("account not found after Grok OAuth success CAS miss")
|
||
}
|
||
return nil, &providerCycleContainmentRefreshError{
|
||
err: fmt.Errorf("grok OAuth success CAS lost and current state is unavailable: %w", readErr),
|
||
}
|
||
}
|
||
slog.Info("oauth_refresh_success_cas_skipped_stale_credentials",
|
||
"account_id", freshAccount.ID,
|
||
"platform", freshAccount.Platform,
|
||
)
|
||
return &OAuthRefreshResult{Account: currentAccount}, nil
|
||
}
|
||
durableAccount, readErr := api.loadGrokDurableAccountAfterPersist(ctx, cacheKey, freshAccount.ID)
|
||
if readErr != nil || durableAccount == nil {
|
||
if readErr == nil {
|
||
readErr = fmt.Errorf("account not found after Grok OAuth success CAS")
|
||
}
|
||
return nil, &providerCycleContainmentRefreshError{
|
||
err: fmt.Errorf("grok OAuth success persisted but durable account state is unavailable: %w", readErr),
|
||
}
|
||
}
|
||
// The CAS changes credentials only. A concurrent admin or scheduler
|
||
// mutation may have changed status, schedulability, or cooldown fields
|
||
// while the provider call was in flight. Return the durable row so
|
||
// post-refresh cache publication cannot restore that stale snapshot.
|
||
freshAccount = durableAccount
|
||
} else if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil {
|
||
slog.Error("oauth_refresh_update_failed",
|
||
"account_id", freshAccount.ID,
|
||
"error", updateErr,
|
||
)
|
||
return nil, fmt.Errorf("%w: %v", errOAuthRefreshCredentialPersist, updateErr)
|
||
}
|
||
}
|
||
|
||
if requestPath && freshAccount.Platform == PlatformGrok {
|
||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil {
|
||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount)
|
||
}
|
||
}
|
||
|
||
return &OAuthRefreshResult{
|
||
Refreshed: true,
|
||
NewCredentials: newCredentials,
|
||
Account: freshAccount,
|
||
}, nil
|
||
}
|
||
|
||
func (api *OAuthRefreshAPI) releaseRefreshLock(parent context.Context, cacheKey string) {
|
||
cleanupParent := context.Background()
|
||
if parent != nil {
|
||
cleanupParent = context.WithoutCancel(parent)
|
||
}
|
||
ctx, cancel := context.WithTimeout(cleanupParent, defaultRefreshLockReleaseTimeout)
|
||
defer cancel()
|
||
if err := api.tokenCache.ReleaseRefreshLock(ctx, cacheKey); err != nil {
|
||
slog.Warn("oauth_refresh_lock_release_failed", "cache_key", cacheKey, "error", err)
|
||
}
|
||
}
|
||
|
||
func (api *OAuthRefreshAPI) loadGrokDurableAccountAfterPersist(parent context.Context, cacheKey string, accountID int64) (*Account, error) {
|
||
cleanupParent := context.Background()
|
||
if parent != nil {
|
||
cleanupParent = context.WithoutCancel(parent)
|
||
}
|
||
ctx, cancel := context.WithTimeout(cleanupParent, defaultRefreshPostPersistCleanupTimeout)
|
||
defer cancel()
|
||
|
||
// A successful rotation can revoke the access token still cached from the
|
||
// pre-rotation credential document. Trigger deletion at the commit boundary,
|
||
// even if the attempt/parent context was canceled immediately after CAS.
|
||
if api.tokenCache != nil {
|
||
if err := api.tokenCache.DeleteAccessToken(ctx, cacheKey); err != nil {
|
||
slog.Warn("oauth_refresh_post_persist_cache_delete_failed",
|
||
"account_id", accountID,
|
||
"cache_key", cacheKey,
|
||
"error", err,
|
||
)
|
||
}
|
||
}
|
||
|
||
return api.accountRepo.GetByID(ctx, accountID)
|
||
}
|
||
|
||
// isInvalidGrantError 检查错误是否为 invalid_grant
|
||
func isInvalidGrantError(err error) bool {
|
||
return err != nil && strings.Contains(strings.ToLower(err.Error()), "invalid_grant")
|
||
}
|
||
|
||
// tryRecoverFromRefreshRace 在 invalid_grant 错误后尝试竞争恢复
|
||
// 重新读取 DB，如果 refresh_token 已改变（说明另一个 worker 成功刷新），则返回更新后的 account
|
||
func (api *OAuthRefreshAPI) tryRecoverFromRefreshRace(ctx context.Context, usedAccount *Account) (*Account, bool) {
|
||
if api.accountRepo == nil {
|
||
return nil, false
|
||
}
|
||
reReadAccount, err := api.accountRepo.GetByID(ctx, usedAccount.ID)
|
||
if err != nil || reReadAccount == nil {
|
||
return nil, false
|
||
}
|
||
usedRT := usedAccount.GetCredential("refresh_token")
|
||
currentRT := reReadAccount.GetCredential("refresh_token")
|
||
if usedRT == "" || currentRT == "" {
|
||
return nil, false
|
||
}
|
||
// refresh_token 不同 → 另一个 worker 已成功刷新
|
||
if usedRT != currentRT {
|
||
return reReadAccount, true
|
||
}
|
||
return nil, false
|
||
}
|
||
|
||
// MergeCredentials 将旧 credentials 中不存在于新 map 的字段保留到新 map 中
|
||
func MergeCredentials(oldCreds, newCreds map[string]any) map[string]any {
|
||
if newCreds == nil {
|
||
newCreds = make(map[string]any)
|
||
}
|
||
for k, v := range oldCreds {
|
||
if _, exists := newCreds[k]; !exists {
|
||
newCreds[k] = v
|
||
}
|
||
}
|
||
return newCreds
|
||
}
|
||
|
||
// BuildClaudeAccountCredentials 为 Claude 平台构建 OAuth credentials map
|
||
// 消除 Claude 平台没有 BuildAccountCredentials 方法的问题
|
||
func BuildClaudeAccountCredentials(tokenInfo *TokenInfo) map[string]any {
|
||
creds := map[string]any{
|
||
"access_token": tokenInfo.AccessToken,
|
||
"token_type": tokenInfo.TokenType,
|
||
"expires_in": strconv.FormatInt(tokenInfo.ExpiresIn, 10),
|
||
"expires_at": strconv.FormatInt(tokenInfo.ExpiresAt, 10),
|
||
}
|
||
if tokenInfo.RefreshToken != "" {
|
||
creds["refresh_token"] = tokenInfo.RefreshToken
|
||
}
|
||
if tokenInfo.Scope != "" {
|
||
creds["scope"] = tokenInfo.Scope
|
||
}
|
||
return creds
|
||
}
|