Download internal/auth/request.go from a3216/ds2api: direct link, hf CLI and curl.
- Browser
- Download file 14.7 kB
-
https://huggingface.co/spaces/a3216/ds2api/resolve/ce5a019598af8a3a8fa70d3880098df3e850ecd7/internal/auth/request.go
- Command line
-
hf download hf://spaces/a3216/ds2api@ce5a019598af8a3a8fa70d3880098df3e850ecd7/internal/auth/request.go
-
curl -L -o request.go https://huggingface.co/spaces/a3216/ds2api/resolve/ce5a019598af8a3a8fa70d3880098df3e850ecd7/internal/auth/request.go
14.7 kB
| 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) | |
| } | |