diff --git a/internal/upstream/client.go b/internal/upstream/client.go index ed4137ba..79549feb 100644 --- a/internal/upstream/client.go +++ b/internal/upstream/client.go @@ -12,6 +12,7 @@ import ( "log" "net/http" "net/url" + "os" "regexp" "sort" "strconv" @@ -681,6 +682,13 @@ type Client struct { // false 时即便用户 auth 写了 realm=global 也**不**路由到 global base—— // chatBase/billingBase 返回 CN base,路径也走 CN(双保险,与 auth.Realm() 的开关闸呼应)。 GlobalEnabled bool + + // refreshMu/refreshInFlight 同账号 refresh 单飞(按 UID 分桶):同一时刻只允许 + // 一个 refresh 在途,其余调用等待并复用首个结果。上游对 refreshToken 反重放, + // 并发第二发会带着已被消费的 RT 撞 12153 并整链作废(issue #215:账号从此恒 + // 401,只能人工重登)。单飞在这里是保号手段,不只是省流量。 + refreshMu sync.Mutex + refreshInFlight map[string]*refreshFlight } // New 生产默认值。Transport 由 newTransport() 集中构造(连接层加固:真正禁 h2 / @@ -901,13 +909,92 @@ const refreshTokenExpiresInMax = 10 * 365 * 24 * time.Hour // RefreshToken 刷新 access token;成功时更新 a 的字段(缺省值保留旧值), // 调用方负责 SaveAtomic。 // -// 并发安全模型(两段式,缩小持锁窗口): -// - 锁内仅做「读 refreshToken 快照」与「校验未变后写回新 token」两小段内存操作; -// - 网络 I/O(doJSON)在**锁外**执行,带 30s ctx 超时——避免上游 hang 时长时间 -// 独占 a.mu,阻塞同账号的 SaveAtomic / 其他刷新(issue:持锁 120s I/O)。 -// - 写回前重新校验快照一致性:若锁外期间另一 goroutine 已完成刷新(refreshToken -// 已变),本次结果直接采用(新 token 已生效),不再重复写回。 +// 并发安全模型(单飞 + 两段式): +// - **单飞**:同一账号(UID)同一时刻只允许一个 refresh 在途,其余调用等待并 +// 原样复用首个结果。上游对 refreshToken 有反重放:RT 被消费后再用会撞 12153 +// (refresh token is invalid)并把整条 RT 链作废,账号从此恒 401、只能人工 +// 重新登录(issue #215:scheduler keepalive / handler 临近过期刷新 / signin +// 会并发触发同一账号)。所以单飞在这里是保号手段,而不只是省流量。 +// - 首个调用内部仍是两段式:锁内读快照 → 锁外 I/O(30s)→ 锁内校验快照一致后 +// 写回;写回前若发现另一刷新已完成(两 token 同时变化)则直接采用。 +// - 12153 兜底:若磁盘 auth 文件已被其他进程(server + signin 双进程部署)刷新 +// 过,则采纳磁盘 token 重试一次——进程内单飞挡不住跨进程竞态。 func (c *Client) RefreshToken(a *auth.Auth) error { + key := a.UID + if key == "" { + // 无 UID 的临时 auth:退化为按对象指针去重,仍保证同一对象串行。 + key = fmt.Sprintf("%p", a) + } + c.refreshMu.Lock() + if c.refreshInFlight == nil { + c.refreshInFlight = make(map[string]*refreshFlight) + } + if f, ok := c.refreshInFlight[key]; ok { + c.refreshMu.Unlock() + <-f.done + return f.err + } + f := &refreshFlight{done: make(chan struct{})} + c.refreshInFlight[key] = f + c.refreshMu.Unlock() + + err := c.refreshTokenOnce(a) + if err != nil && isSessionDead(err) && adoptNewerTokensFromDisk(a) { + // 跨进程竞态:磁盘 RT 已被别的进程 rotate 过,用新 RT 再试一次。 + err = c.refreshTokenOnce(a) + } + + // 先写结果再关 channel:等待者 <-f.done 之后读到的一定是最终 err。 + f.err = err + close(f.done) + c.refreshMu.Lock() + delete(c.refreshInFlight, key) + c.refreshMu.Unlock() + return err +} + +// refreshFlight 是同一账号一次 refresh 的单飞票据:done 关闭即结果就绪。 +type refreshFlight struct { + done chan struct{} + err error +} + +// isSessionDead 判定错误是否为上游 12153(会话 / RT 失效)。 +func isSessionDead(err error) bool { + var ue *Error + return errors.As(err, &ue) && ue.Kind == ErrSessionDead +} + +// adoptNewerTokensFromDisk 在 refresh 失败后读回 auth 文件:若磁盘 RT 与内存不同 +// (另一进程已刷新并写回),采纳磁盘 token 并返回 true,供调用方重试一次。 +// 读文件失败、解析失败、RT 一致均返回 false(不改内存状态),保持失败语义。 +func adoptNewerTokensFromDisk(a *auth.Auth) bool { + if a.FilePath == "" { + return false + } + raw, err := os.ReadFile(a.FilePath) + if err != nil { + return false + } + disk, err := auth.Parse(raw) + if err != nil || strings.TrimSpace(disk.RefreshToken) == "" { + return false + } + a.Lock() + defer a.Unlock() + if disk.RefreshToken == a.RefreshToken { + return false + } + a.AccessToken = disk.AccessToken + a.RefreshToken = disk.RefreshToken + if disk.ExpiresAt > 0 { + a.ExpiresAt = disk.ExpiresAt + } + return true +} + +// refreshTokenOnce 是单次刷新本体(调用方 RefreshToken 已保证同账号串行)。 +func (c *Client) refreshTokenOnce(a *auth.Auth) error { // 第 1 段(锁内):读快照。 a.Lock() rtSnapshot := a.RefreshToken diff --git a/internal/upstream/refresh_singleflight_test.go b/internal/upstream/refresh_singleflight_test.go new file mode 100644 index 00000000..2af21f56 --- /dev/null +++ b/internal/upstream/refresh_singleflight_test.go @@ -0,0 +1,174 @@ +package upstream + +import ( + "io" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/linguo2625469/workbuddy2api-panel/internal/auth" +) + +// TestRefreshTokenSingleFlightPerAccount 验证同一账号并发 refresh 只打一次上游: +// 上游对 refreshToken 反重放,第二发带已消费的 RT 会撞 12153 并把整条 RT 链作废 +// (issue #215,账号只能人工重登)。单飞失败时本用例会看到 hits>1 且集中报 12153。 +func TestRefreshTokenSingleFlightPerAccount(t *testing.T) { + var hits, inFlight, maxInFlight int32 + c := testClient(func(r *http.Request) (*http.Response, error) { + atomic.AddInt32(&hits, 1) + n := atomic.AddInt32(&inFlight, 1) + for { + old := atomic.LoadInt32(&maxInFlight) + if n <= old || atomic.CompareAndSwapInt32(&maxInFlight, old, n) { + break + } + } + time.Sleep(50 * time.Millisecond) // 拉长窗口:无单飞时 8 个调用必然并发 + atomic.AddInt32(&inFlight, -1) + return jsonResp(200, `{"code":0,"msg":"ok","data":{"accessToken":"at1","refreshToken":"rt1","expiresIn":3600}}`), nil + }) + a := &auth.Auth{UID: "uid-a", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1} + + const callers = 8 + var wg sync.WaitGroup + errs := make([]error, callers) + for i := 0; i < callers; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + errs[i] = c.RefreshToken(a) + }(i) + } + wg.Wait() + + for i, err := range errs { + if err != nil { + t.Fatalf("caller %d: %v", i, err) + } + } + if got := atomic.LoadInt32(&hits); got != 1 { + t.Errorf("upstream refresh hits = %d, want 1 (single-flight)", got) + } + if got := atomic.LoadInt32(&maxInFlight); got != 1 { + t.Errorf("max concurrent refresh = %d, want 1", got) + } + if a.AccessToken != "at1" || a.RefreshToken != "rt1" { + t.Errorf("tokens not updated: at=%q rt=%q", a.AccessToken, a.RefreshToken) + } +} + +// TestRefreshTokenSingleFlightIsPerAccount 验证单飞按账号隔离:不同 UID 的刷新 +// 并行执行,不被同一把锁串行化(否则 keepalive 全量刷新会退化成串行慢刷)。 +func TestRefreshTokenSingleFlightIsPerAccount(t *testing.T) { + var inFlight, maxInFlight int32 + c := testClient(func(r *http.Request) (*http.Response, error) { + n := atomic.AddInt32(&inFlight, 1) + for { + old := atomic.LoadInt32(&maxInFlight) + if n <= old || atomic.CompareAndSwapInt32(&maxInFlight, old, n) { + break + } + } + time.Sleep(50 * time.Millisecond) + atomic.AddInt32(&inFlight, -1) + return jsonResp(200, `{"code":0,"data":{"accessToken":"at1","refreshToken":"rt1","expiresIn":3600}}`), nil + }) + accounts := []*auth.Auth{ + {UID: "uid-a", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1}, + {UID: "uid-b", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1}, + {UID: "uid-c", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1}, + } + var wg sync.WaitGroup + for _, a := range accounts { + wg.Add(1) + go func(a *auth.Auth) { + defer wg.Done() + if err := c.RefreshToken(a); err != nil { + t.Errorf("refresh %s: %v", a.UID, err) + } + }(a) + } + wg.Wait() + if got := atomic.LoadInt32(&maxInFlight); got < 2 { + t.Errorf("max concurrent refresh across accounts = %d, want >= 2 (per-account single-flight)", got) + } +} + +// TestRefreshAdoptsNewerDiskTokenAfterSessionDead 验证跨进程兜底:内存 RT 已被上游 +// 判 12153,但磁盘上的 auth 文件已被另一进程(如 signin 工具)rotate 过时,采纳 +// 磁盘 token 重试一次即可自救,而不是直接把账号判死。 +func TestRefreshAdoptsNewerDiskTokenAfterSessionDead(t *testing.T) { + dir := t.TempDir() + fp := filepath.Join(dir, "uid-a.json") + disk := `{"auth":{"accessToken":"at-disk","refreshToken":"rt-disk","expiresAt":1999999999,"domain":"copilot.tencent.com","realm":"cn"},"account":{"uid":"uid-a"}}` + if err := os.WriteFile(fp, []byte(disk), 0o600); err != nil { + t.Fatal(err) + } + + var mu sync.Mutex + var seen []string + c := testClient(func(r *http.Request) (*http.Response, error) { + rt := r.Header.Get("X-Refresh-Token") + mu.Lock() + seen = append(seen, rt) + mu.Unlock() + if rt != "rt-disk" { + return &http.Response{ + StatusCode: 401, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"code":12153,"msg":"refresh token is invalid"}`)), + }, nil + } + return jsonResp(200, `{"code":0,"data":{"accessToken":"at2","refreshToken":"rt2","expiresIn":3600}}`), nil + }) + a := &auth.Auth{UID: "uid-a", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1, FilePath: fp} + + if err := c.RefreshToken(a); err != nil { + t.Fatalf("refresh should recover from disk token: %v", err) + } + if a.RefreshToken != "rt2" || a.AccessToken != "at2" { + t.Errorf("tokens not updated after retry: at=%q rt=%q", a.AccessToken, a.RefreshToken) + } + mu.Lock() + defer mu.Unlock() + if len(seen) != 2 || seen[0] != "rt0" || seen[1] != "rt-disk" { + t.Errorf("refresh attempts = %v, want [rt0 rt-disk]", seen) + } +} + +// TestRefreshDoesNotRetryWhenDiskTokenUnchanged 验证兜底不改变失败语义:磁盘 RT 与 +// 内存一致(无跨进程刷新)时,12153 原样返回,不会无限重试打上游。 +func TestRefreshDoesNotRetryWhenDiskTokenUnchanged(t *testing.T) { + dir := t.TempDir() + fp := filepath.Join(dir, "uid-a.json") + disk := `{"auth":{"accessToken":"at0","refreshToken":"rt0","expiresAt":1,"domain":"copilot.tencent.com","realm":"cn"},"account":{"uid":"uid-a"}}` + if err := os.WriteFile(fp, []byte(disk), 0o600); err != nil { + t.Fatal(err) + } + var hits int32 + c := testClient(func(r *http.Request) (*http.Response, error) { + atomic.AddInt32(&hits, 1) + return &http.Response{ + StatusCode: 401, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"code":12153,"msg":"refresh token is invalid"}`)), + }, nil + }) + a := &auth.Auth{UID: "uid-a", AccessToken: "at0", RefreshToken: "rt0", ExpiresAt: 1, FilePath: fp} + + err := c.RefreshToken(a) + if err == nil { + t.Fatal("want error") + } + if !isSessionDead(err) { + t.Errorf("want ErrSessionDead, got %v", err) + } + if got := atomic.LoadInt32(&hits); got != 1 { + t.Errorf("refresh attempts = %d, want 1 (no retry when disk token unchanged)", got) + } +}