internal/auth/request.go

package auth

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"ds2api/internal/account"
	"ds2api/internal/config"
)

type ctxKey string

const authCtxKey ctxKey = "auth_context"

var (
	ErrUnauthorized = errors.New("unauthorized: missing auth token")
	ErrNoAccount    = errors.New("no accounts configured or all accounts are busy")

	// ErrSessionExpired 表示 DeepSeek 明确拒绝了这个会话 token（401/403 或鉴权类业务码）。
	// 只有这种情况才值得清空本地 token 并回退到密码登录。
	ErrSessionExpired = errors.New("deepseek session rejected")

	// errSessionProbeUnavailable 表示当前没有可用的无密码探针：未注入探针函数，
	// 或本地还没有 token 可以拿去做探测。
	errSessionProbeUnavailable = errors.New("session probe unavailable")
)

type RequestAuth struct {
	UseConfigToken bool
	DeepSeekToken  string
	CallerID       string
	AccountID      string
	TargetAccount  string
	Account        config.Account
	TriedAccounts  map[string]bool
	resolver       *Resolver
}

type LoginFunc func(ctx context.Context, acc config.Account) (string, error)

// RefreshSessionFunc 是不需要密码的会话回填通道：带现有 token 请求
// GET /api/v0/users/current，返回服务端认定的当前 token。
// 明确失效时返回包装了 ErrSessionExpired 的错误，其余错误表示无法判定。
type RefreshSessionFunc func(ctx context.Context, acc config.Account) (string, error)

type Resolver struct {
	Store *config.Store
	Pool  *account.Pool
	Login LoginFunc
	// RefreshSession 为可选的轻量续期通道；未注入时所有刷新都会回退到密码登录
	// （与历史行为一致）。
	RefreshSession RefreshSessionFunc

	mu               sync.Mutex
	tokenRefreshedAt map[string]time.Time
}

func NewResolver(store *config.Store, pool *account.Pool, login LoginFunc) *Resolver {
	return &Resolver{
		Store:            store,
		Pool:             pool,
		Login:            login,
		tokenRefreshedAt: map[string]time.Time{},
	}
}

func (r *Resolver) Determine(req *http.Request) (*RequestAuth, error) {
	callerKey := extractCallerToken(req)
	if callerKey == "" {
		return nil, ErrUnauthorized
	}
	callerID := callerTokenID(callerKey)
	ctx := req.Context()
	if !r.Store.HasAPIKey(callerKey) {
		return &RequestAuth{
			UseConfigToken: false,
			DeepSeekToken:  callerKey,
			CallerID:       callerID,
			resolver:       r,
			TriedAccounts:  map[string]bool{},
		}, nil
	}
	target := strings.TrimSpace(req.Header.Get("X-Ds2-Target-Account"))
	a, err := r.acquireManagedRequestAuth(ctx, callerID, target)
	if err != nil {
		return nil, err
	}
	return a, nil
}

func (r *Resolver) acquireManagedRequestAuth(ctx context.Context, callerID, target string) (*RequestAuth, error) {
	tried := map[string]bool{}
	var lastEnsureErr error
	for {
		if target == "" && len(tried) >= len(r.Store.Accounts()) {
			if lastEnsureErr != nil {
				return nil, lastEnsureErr
			}
			return nil, ErrNoAccount
		}
		acc, ok := r.Pool.AcquireWait(ctx, target, tried)
		if !ok {
			if lastEnsureErr != nil {
				return nil, lastEnsureErr
			}
			return nil, ErrNoAccount
		}

		a := &RequestAuth{
			UseConfigToken: true,
			CallerID:       callerID,
			AccountID:      acc.Identifier(),
			TargetAccount:  target,
			Account:        acc,
			TriedAccounts:  tried,
			resolver:       r,
		}

		if err := r.ensureManagedToken(ctx, a); err != nil {
			lastEnsureErr = err
			tried[a.AccountID] = true
			r.Pool.Release(a.AccountID)
			if target != "" {
				return nil, err
			}
			continue
		}
		return a, nil
	}
}

// DetermineCaller resolves caller identity without acquiring any pooled account.
// Use this for local-cache lookup routes that only need tenant isolation.
func (r *Resolver) DetermineCaller(req *http.Request) (*RequestAuth, error) {
	callerKey := extractCallerToken(req)
	if callerKey == "" {
		return nil, ErrUnauthorized
	}
	callerID := callerTokenID(callerKey)
	a := &RequestAuth{
		UseConfigToken: false,
		CallerID:       callerID,
		resolver:       r,
		TriedAccounts:  map[string]bool{},
	}
	if r == nil || r.Store == nil || !r.Store.HasAPIKey(callerKey) {
		a.DeepSeekToken = callerKey
	}
	return a, nil
}

func WithAuth(ctx context.Context, a *RequestAuth) context.Context {
	return context.WithValue(ctx, authCtxKey, a)
}

func FromContext(ctx context.Context) (*RequestAuth, bool) {
	v := ctx.Value(authCtxKey)
	a, ok := v.(*RequestAuth)
	return a, ok
}

func (r *Resolver) loginAndPersist(ctx context.Context, a *RequestAuth) error {
	token, err := r.Login(ctx, a.Account)
	if err != nil {
		// Detect USER_IS_BANNED: auto-ban the account and promote standby accounts.
		if strings.Contains(strings.ToUpper(err.Error()), "USER_IS_BANNED") {
			r.handleBannedAccount(a.AccountID)
		}
		return err
	}
	a.Account.Token = token
	a.DeepSeekToken = token
	r.markTokenRefreshedNow(a.AccountID)
	return r.Store.UpdateAccountToken(a.AccountID, token)
}

// handleBannedAccount marks an account as banned and promotes standby accounts
// based on the segment-based rule: if active normal accounts fall below half of
// total normal accounts (per segment), promote standby accounts.
func (r *Resolver) handleBannedAccount(accountID string) {
	if r.Store == nil || r.Pool == nil {
		return
	}
	acc, err := r.Store.SetAccountBanned(accountID, true)
	if err != nil {
		config.Logger.Error("[banned] failed to mark account as banned", "account", accountID, "error", err)
		return
	}
	// Only auto-promote if the banned account was a normal (non-standby) account.
	if acc.Role == "standby" {
		return
	}
	config.Logger.Warn("[banned] account banned, checking standby promotion", "account", accountID)

	activeNormal := r.Store.ActiveNormalAccounts()
	totalNormal := r.Store.TotalNormalAccounts()
	// Segment rule: promote when active normal accounts < half of total normal accounts.
	// e.g. 2 normal -> both must be banned; 4 normal -> promote when only 1 active.
	threshold := (totalNormal + 1) / 2 // ceiling of totalNormal/2
	if activeNormal < threshold {
		releaseCount := r.Store.RuntimeBackupReleaseCount()
		promoted := r.Store.PromoteStandbyAccounts(releaseCount)
		if len(promoted) > 0 {
			config.Logger.Info("[banned] promoted standby accounts", "count", len(promoted), "accounts", promoted)
			r.Pool.Reset()
		}
	}
}

// RefreshToken 由鉴权类失败触发。DeepSeek 的 token 是按设备绑定的长期会话凭证，
// 不存在“旧 token 换新 token”的端点，所以这里先用不带密码的 users/current 探针复核：
// 探针确认会话仍有效就直接沿用（顺带收下服务端轮换的新值），只有服务端明确拒绝会话
// 或探针本身不可用时才回退到密码登录。这样既能挡住把非鉴权错误误判成登录失效的情况，
// 又保留了历史行为里的恢复能力。
func (r *Resolver) RefreshToken(ctx context.Context, a *RequestAuth) bool {
	if !a.UseConfigToken || a.AccountID == "" {
		return false
	}
	if err := r.refreshSessionReactive(ctx, a); err != nil {
		config.Logger.Error("[refresh_token] failed", "account", a.AccountID, "error", err)
		return false
	}
	return true
}

// refreshSessionReactive 处理“请求已经拿到鉴权类失败”的场景：探针能证明会话有效就
// 直接返回；否则按历史行为兜底到密码登录（探针不可判定时不清空 token，登录失败还能
// 留着下次再试）。
func (r *Resolver) refreshSessionReactive(ctx context.Context, a *RequestAuth) error {
	err := r.probeSessionToken(ctx, a)
	if err == nil {
		return nil
	}
	if errors.Is(err, ErrSessionExpired) {
		// 会话确认失效：清掉落盘值，避免后续请求继续复用死 token。
		_ = r.Store.UpdateAccountToken(a.AccountID, "")
		a.Account.Token = ""
		a.DeepSeekToken = ""
	} else {
		config.Logger.Warn("[refresh_token] session probe inconclusive, falling back to login",
			"account", a.AccountID, "error", err)
	}
	return r.loginAndPersist(ctx, a)
}

// probeSessionToken 用无密码通道复核当前 token，成功时把服务端值回写内存与配置。
func (r *Resolver) probeSessionToken(ctx context.Context, a *RequestAuth) error {
	if r == nil || r.RefreshSession == nil {
		return errSessionProbeUnavailable
	}
	if a == nil || a.AccountID == "" {
		return errSessionProbeUnavailable
	}
	current := strings.TrimSpace(a.Account.Token)
	if current == "" {
		return errSessionProbeUnavailable
	}
	token, err := r.RefreshSession(ctx, a.Account)
	if err != nil {
		return err
	}
	token = strings.TrimSpace(token)
	if token == "" {
		return errors.New("session probe returned empty token")
	}
	a.Account.Token = token
	a.DeepSeekToken = token
	r.markTokenRefreshedNow(a.AccountID)
	// 以配置存储中的值为准判断是否需要写盘：池里的 Account 可能是较早的快照。
	persisted := current
	if fresh, ok := r.Store.FindAccount(a.AccountID); ok {
		persisted = strings.TrimSpace(fresh.Token)
	}
	if token == persisted {
		// 服务端回显同一 token（抓包与线上实测观察到的主要形态）：无需写盘。
		return nil
	}
	config.Logger.Info("[refresh_token] persisted session token from probe", "account", a.AccountID)
	if err := r.Store.UpdateAccountToken(a.AccountID, token); err != nil {
		config.Logger.Warn("[refresh_token] persist rotated token failed", "account", a.AccountID, "error", err)
	}
	return nil
}

func (r *Resolver) MarkTokenInvalid(a *RequestAuth) {
	if !a.UseConfigToken || a.AccountID == "" {
		return
	}
	a.Account.Token = ""
	a.DeepSeekToken = ""
	r.clearTokenRefreshMark(a.AccountID)
	_ = r.Store.UpdateAccountToken(a.AccountID, "")
}

func (r *Resolver) SwitchAccount(ctx context.Context, a *RequestAuth) bool {
	if !a.UseConfigToken {
		return false
	}
	if strings.TrimSpace(a.TargetAccount) != "" {
		return false
	}
	if a.TriedAccounts == nil {
		a.TriedAccounts = map[string]bool{}
	}
	if a.AccountID != "" {
		a.TriedAccounts[a.AccountID] = true
		r.Pool.Release(a.AccountID)
	}
	for {
		acc, ok := r.Pool.Acquire("", a.TriedAccounts)
		if !ok {
			return false
		}
		a.Account = acc
		a.AccountID = acc.Identifier()
		if err := r.ensureManagedToken(ctx, a); err != nil {
			a.TriedAccounts[a.AccountID] = true
			r.Pool.Release(a.AccountID)
			continue
		}
		return true
	}
}

func (a *RequestAuth) SwitchAccount(ctx context.Context) bool {
	if a == nil || a.resolver == nil {
		return false
	}
	return a.resolver.SwitchAccount(ctx, a)
}

func (r *Resolver) Release(a *RequestAuth) {
	if a == nil || !a.UseConfigToken || a.AccountID == "" {
		return
	}
	r.Pool.Release(a.AccountID)
}

func extractCallerToken(req *http.Request) string {
	authHeader := strings.TrimSpace(req.Header.Get("Authorization"))
	if strings.HasPrefix(strings.ToLower(authHeader), "bearer ") {
		token := strings.TrimSpace(authHeader[7:])
		if token != "" {
			return token
		}
	}
	if key := strings.TrimSpace(req.Header.Get("x-api-key")); key != "" {
		return key
	}
	// Gemini/Google clients commonly send API key via x-goog-api-key.
	if key := strings.TrimSpace(req.Header.Get("x-goog-api-key")); key != "" {
		return key
	}
	// Gemini AI Studio compatibility: allow query key fallback only when no
	// header-based credential is present.
	if key := strings.TrimSpace(req.URL.Query().Get("key")); key != "" {
		return key
	}
	return strings.TrimSpace(req.URL.Query().Get("api_key"))
}

func callerTokenID(token string) string {
	token = strings.TrimSpace(token)
	if token == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(token))
	return "caller:" + hex.EncodeToString(sum[:8])
}

func (r *Resolver) ensureManagedToken(ctx context.Context, a *RequestAuth) error {
	if strings.TrimSpace(a.Account.Token) == "" {
		return r.loginAndPersist(ctx, a)
	}
	if r.shouldForceRefresh(a.AccountID) {
		if err := r.refreshSessionProactively(ctx, a); err != nil {
			return err
		}
	}
	a.DeepSeekToken = a.Account.Token
	return nil
}

// refreshSessionProactively 是定时任务的“体检”：到期时先做一次免费探针，
// 只有会话被明确拒绝（或根本没有探针可用）才回退到密码登录；
// 探针结果不可判定时保留现有 token，不因为一次网络抖动就提交密码。
func (r *Resolver) refreshSessionProactively(ctx context.Context, a *RequestAuth) error {
	err := r.probeSessionToken(ctx, a)
	if err == nil {
		return nil
	}
	if !errors.Is(err, ErrSessionExpired) && !errors.Is(err, errSessionProbeUnavailable) {
		config.Logger.Warn("[token_refresh] session probe inconclusive, keeping current token",
			"account", a.AccountID, "error", err)
		return nil
	}
	config.Logger.Info("[token_refresh] session probe requires login", "account", a.AccountID, "reason", err)
	_ = r.Store.UpdateAccountToken(a.AccountID, "")
	a.Account.Token = ""
	a.DeepSeekToken = ""
	return r.loginAndPersist(ctx, a)
}

// shouldForceRefresh 判断是否到了定期体检的时间点。到期并不等于要重新登录：
// 调用方会用无密码探针复核会话，只有探针明确失败才走密码登录。
func (r *Resolver) shouldForceRefresh(accountID string) bool {
	if r == nil || r.Store == nil {
		return false
	}
	if strings.TrimSpace(accountID) == "" {
		return false
	}
	intervalHours := r.Store.RuntimeTokenRefreshIntervalHours()
	if intervalHours <= 0 {
		return false
	}
	now := time.Now()
	r.mu.Lock()
	defer r.mu.Unlock()
	last, ok := r.tokenRefreshedAt[accountID]
	if !ok || last.IsZero() {
		r.tokenRefreshedAt[accountID] = now
		return false
	}
	// Add jitter: ±10% random variation to avoid all accounts probing at the same time.
	// For a 6h interval, this gives a range of ~5h24m to ~6h36m.
	baseInterval := time.Duration(intervalHours) * time.Hour
	jitterRange := baseInterval / 10 // ±10% → total range is 80%-120% of base
	jitter := time.Duration(rand.Int63n(int64(2*jitterRange+1))) - jitterRange
	effectiveInterval := baseInterval + jitter
	if effectiveInterval < time.Hour {
		effectiveInterval = time.Hour
	}
	return now.Sub(last) >= effectiveInterval
}

func (r *Resolver) markTokenRefreshedNow(accountID string) {
	if strings.TrimSpace(accountID) == "" {
		return
	}
	// Add a random offset (0-30min) to spread out future refreshes across accounts.
	// This prevents all accounts from refreshing at the same wall-clock time.
	offset := time.Duration(rand.Int63n(int64(30 * time.Minute)))
	r.mu.Lock()
	defer r.mu.Unlock()
	r.tokenRefreshedAt[accountID] = time.Now().Add(-offset)
}

func (r *Resolver) clearTokenRefreshMark(accountID string) {
	if strings.TrimSpace(accountID) == "" {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.tokenRefreshedAt, accountID)
}