File size: 14,494 Bytes
6d60378
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
// Package auth 解析 WorkBuddy auth 文件(嵌套形/扁平形双形态),
// 提供 refresh 后的原子写回。
package auth

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/linguo2625469/workbuddy2api-panel/internal/logfmt"
)

// Auth 是归一化后的账号凭证(来源可以是插件 OAuth 嵌套形或手写扁平形)。
type Auth struct {
	// mu 串行化 RefreshToken 写与 SaveAtomic 读,防止并发写回半更新 token。
	mu sync.Mutex

	AccessToken  string
	RefreshToken string
	ExpiresAt    int64 // Unix 秒
	Domain       string
	// realm 账号域("cn" / "global"),落盘于 auth.realm(嵌套形)或顶层 realm(扁平形)。
	// 空 = 缺省:Realm() 按 domain 后缀回落,最终恒非空。
	//
	// 命名注记:Go 不允许字段与方法同名,持久化字段用未导出 realm,计算访问器用
	// 导出的 Realm()(跨包调用全部走方法)。Parse/SaveAtomic/login 在包内读写字段。
	realm        string
	UID          string
	EnterpriseID string
	Nickname     string
	FilePath     string // 来源文件;refresh 后原子写回此处

	// DeviceToken 设备风控 Token(X-Device-Token 头),来源 auth 文件的 device_token 键。
	// 缺省为空 = 不注入该头(容器内无桌面端 Turing SDK 的常见部署)。
	// 手写扁平形 auth 文件可直接写 "device_token": "...";插件 OAuth 嵌套形
	// 顶层 device_token 也会被解析(与桌面端共用状态文件的部署方式)。
	DeviceToken string
}

// Lock 供同进程内其他包(upstream.RefreshToken)在改写 Auth 字段期间加锁。
func (a *Auth) Lock() { a.mu.Lock() }

// Unlock 释放 a.Lock 获取的锁。
func (a *Auth) Unlock() { a.mu.Unlock() }

// AccessTokenValue 加锁读取 AccessToken(出站请求头一律经此取值,勿直读字段)。
//
// 为什么必须加锁:RefreshToken 在 a.mu 内改写 AccessToken/RefreshToken/Domain/ExpiresAt
// (client.go「第 2 段(锁内):校验快照一致后写回」),而所有出站请求头构造
// (ChatHeaders / BillingHeaders / fetchEnterpriseModels / fetchV3Models /
// global_models)与调度器的 token 检查都在锁外直读这些字段。生产上两侧真会并发:
// Scheduler.RunKeepaliveNow 定时对**每个**非禁用账号刷新(与是否有在途请求无关),
// 而 handler 正基于**同一个** *auth.Auth 指针构造请求头(Pool.AuthByUID/List 返回的
// 就是池内同一个对象)。无同步直读构成数据竞争(go test -race 实证)。
func (a *Auth) AccessTokenValue() string {
	if a == nil {
		return ""
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.AccessToken
}

// DomainValue 加锁读取 Domain(同 AccessTokenValue:RefreshToken 在锁内改写它)。
func (a *Auth) DomainValue() string {
	if a == nil {
		return ""
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.Domain
}

// RefreshTokenValue 加锁读取 RefreshToken(同 AccessTokenValue:RefreshToken 在锁内
// 改写它)。调度器的「有无凭证」前置守卫(checkin/keepalive/travel 的
// `a.RefreshToken == ""`)必须经此取值,勿直读字段。
func (a *Auth) RefreshTokenValue() string {
	if a == nil {
		return ""
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.RefreshToken
}

// globalEnabled 全局开关:global realm 是否路由(D5 双保险)。
// 默认开启(与 config global.enabled 缺省 true 一致):Realm() 正常按显式 realm/
// domain 判定 global/cn。显式 SetGlobalEnabled(false)(config "enabled": false)关闭
// → 逃生门:纯 CN 部署,即便 auth 文件写了 realm=global 或 domain 为 .workbuddy.ai
// 也恒判 cn——「关了才锁死」的单一闸口集中收敛在 Realm()/IsGlobal() 里。
var globalEnabled atomic.Bool

func init() { globalEnabled.Store(true) }

// SetGlobalEnabled 注入 global realm 路由开关(false = 锁死纯 CN,逃生门)。
func SetGlobalEnabled(enabled bool) { globalEnabled.Store(enabled) }

// GlobalEnabled 报告 global realm 路由开关当前状态(测试/运维观测)。
func GlobalEnabled() bool { return globalEnabled.Load() }

// Realm 返回账号的归一化域:显式 Realm=="global" 或 domain 后缀 .workbuddy.ai → "global",
// 否则 "cn"。显式 global 优先于 domain 回落(D1)。
// 全局开关 SetGlobalEnabled(false) 时恒 "cn"(逃生门:纯 CN 锁定,不影响默认行为)。
// 空 realm + 空 domain → "cn"(老 CN 凭证零回归)。
func (a *Auth) Realm() string {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.realmLocked()
}

// realmLocked Realm 的无锁内部实现:仅限**已持 a.mu** 的调用方使用(sync.Mutex 不可重入,
// 锁内再调 Realm() 会自锁)。realm 由 BackfillRealm 改写、Domain 由 RefreshToken 在锁内
// 改写,故读取必须与写方同锁(理由见 AccessTokenValue 注释)。
func (a *Auth) realmLocked() string {
	if !globalEnabled.Load() {
		return "cn"
	}
	if strings.TrimSpace(a.realm) == "global" || isGlobalDomain(a.Domain) {
		return "global"
	}
	return "cn"
}

// ResolveRealm 归一化 realm(cn/global):显式非空优先,否则按原始 domain 推断
// (isGlobalDomain)。不受逃生门影响(逃生门是路由锁,不应影响标识判定);
// domain 也为空 → "cn"(老 CN 凭证零回归)。
func ResolveRealm(explicit, domain string) string {
	if r := strings.TrimSpace(explicit); r != "" {
		return r
	}
	if isGlobalDomain(domain) {
		return "global"
	}
	return "cn"
}

// BackfillRealm 为缺省 realm 标识的账号持久化补标识:a.realm 为空时按「原始 domain 推断」
// 写回(cn/global),返回 (是否有变更, 归一化后的 realm)。已有标识不动(幂等)。
//
// 注意用 isGlobalDomain(a.Domain) 直接推断,而非 Realm()——Realm() 在逃生门
// (SetGlobalEnabled(false))下恒降级 cn,把 global 账号写死成 cn 会永久污染凭证
// (逃生门是纯 CN 部署的临时锁,不应改写落盘数据)。domain 也为空时写 "cn"(老 CN 凭证)。
func (a *Auth) BackfillRealm() (bool, string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if strings.TrimSpace(a.realm) != "" {
		return false, a.realm
	}
	r := ResolveRealm("", a.Domain)
	a.realm = r
	return true, r
}

// RealmStored 直读持久化的 realm 标识(可能为空 = 未 backfill 的旧文件,Realm() 会 fallback)。
func (a *Auth) RealmStored() string {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.realm
}

// BackfillRealmFor 显式写入 realm 标识(包外登录路径使用:panel login 已知用户选了
// global,直接落盘 realm=global,不依赖 domain 后缀推断)。realm 需为 cn/global,
// 非法值报错(防写脏)。返回是否发生变更。
func BackfillRealmFor(a *Auth, realm string) (bool, error) {
	if a == nil {
		return false, fmt.Errorf("nil auth")
	}
	switch strings.TrimSpace(realm) {
	case "cn", "global":
	default:
		return false, fmt.Errorf("realm must be cn/global, got %q", realm)
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.realm == realm {
		return false, nil
	}
	a.realm = realm
	return true, nil
}

// IsGlobal 报告账号是否属于 global realm(= Realm() == "global")。
func (a *Auth) IsGlobal() bool { return a.Realm() == "global" }

// isGlobalDomain 判定 domain 是否指向 www.workbuddy.ai 家族。
// 同时接受裸域 workbuddy.ai 与任意子域(HasSuffix("www.workbuddy.ai") 或裸域本身)。
func isGlobalDomain(d string) bool {
	d = strings.ToLower(strings.TrimSpace(d))
	return d == "workbuddy.ai" || strings.HasSuffix(d, ".workbuddy.ai")
}

// NeedsRefresh 报告 token 是否将在 within 内过期(或已过期/无 expiry)。
func (a *Auth) NeedsRefresh(within time.Duration) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.ExpiresAt <= 0 {
		return true
	}
	return time.Now().Add(within).Unix() >= a.ExpiresAt
}

// Parse 兼容两种磁盘形态:
//
//	嵌套形 {"auth":{...},"account":{...}}  (插件 OAuth 输出)
//	扁平形 {"accessToken":...,"uid":...}   (手写/旧版)
func Parse(raw []byte) (*Auth, error) {
	if len(raw) == 0 {
		return nil, fmt.Errorf("empty auth storage")
	}
	var probe map[string]json.RawMessage
	if err := json.Unmarshal(raw, &probe); err != nil {
		return nil, fmt.Errorf("storage_parse_error: %w", err)
	}
	var a Auth
	if _, nested := probe["auth"]; nested {
		var n struct {
			Auth struct {
				AccessToken  string `json:"accessToken"`
				RefreshToken string `json:"refreshToken"`
				ExpiresAt    int64  `json:"expiresAt"`
				Domain       string `json:"domain"`
				Realm        string `json:"realm"`
			} `json:"auth"`
			Account struct {
				UID          string `json:"uid"`
				EnterpriseID string `json:"enterpriseId"`
				Nickname     string `json:"nickname"`
			} `json:"account"`
			// DeviceToken 顶层 device_token(嵌套形与扁平形共用;手写时无需嵌进 auth 对象)。
			DeviceToken string `json:"device_token"`
		}
		if err := json.Unmarshal(raw, &n); err != nil {
			return nil, fmt.Errorf("storage_parse_error: %w", err)
		}
		a = Auth{
			AccessToken:  n.Auth.AccessToken,
			RefreshToken: n.Auth.RefreshToken,
			ExpiresAt:    n.Auth.ExpiresAt,
			Domain:       n.Auth.Domain,
			realm:        n.Auth.Realm,
			UID:          n.Account.UID,
			EnterpriseID: n.Account.EnterpriseID,
			Nickname:     n.Account.Nickname,
			DeviceToken:  n.DeviceToken,
		}
	} else {
		var f struct {
			AccessToken  string `json:"accessToken"`
			RefreshToken string `json:"refreshToken"`
			ExpiresAt    int64  `json:"expiresAt"`
			Domain       string `json:"domain"`
			Realm        string `json:"realm"`
			UID          string `json:"uid"`
			EnterpriseID string `json:"enterpriseId"`
			Nickname     string `json:"nickname"`
			DeviceToken  string `json:"device_token"`
		}
		if err := json.Unmarshal(raw, &f); err != nil {
			return nil, fmt.Errorf("storage_parse_error: %w", err)
		}
		a = Auth{
			AccessToken:  f.AccessToken,
			RefreshToken: f.RefreshToken,
			ExpiresAt:    f.ExpiresAt,
			Domain:       f.Domain,
			realm:        f.Realm,
			UID:          f.UID,
			EnterpriseID: f.EnterpriseID,
			Nickname:     f.Nickname,
			DeviceToken:  f.DeviceToken,
		}
	}
	if strings.TrimSpace(a.AccessToken) == "" {
		return nil, fmt.Errorf("parse_error: missing accessToken")
	}
	return &a, nil
}

// SaveAtomic 以嵌套形原子写回 FilePath(tmp + rename),保持嵌套形(插件可读)格式。
// 全程持 a.mu:防止与 RefreshToken 修改 token 字段并发,杜绝写回半更新。
// 防御:accessToken 为空时拒绝写回,避免误用空凭证覆盖有效文件。
func (a *Auth) SaveAtomic() error {
	a.mu.Lock()
	defer a.mu.Unlock()
	if strings.TrimSpace(a.AccessToken) == "" {
		return fmt.Errorf("save refused: empty accessToken (uid=%s)", a.UID)
	}
	if a.FilePath == "" {
		return fmt.Errorf("no FilePath set")
	}
	doc := map[string]any{
		"auth": map[string]any{
			"accessToken":  a.AccessToken,
			"refreshToken": a.RefreshToken,
			"expiresAt":    a.ExpiresAt,
			"domain":       a.Domain,
			"realm":        a.realm,
		},
		"account": map[string]any{
			"uid":          a.UID,
			"enterpriseId": a.EnterpriseID,
			"nickname":     a.Nickname,
		},
	}
	// DeviceToken 非空才写回顶层 device_token:避免在无该字段的旧文件里引入空键
	// (保持与插件 OAuth 输出形状一致,插件读取忽略未知键)。
	if a.DeviceToken != "" {
		doc["device_token"] = a.DeviceToken
	}
	raw, err := json.MarshalIndent(doc, "", "  ")
	if err != nil {
		return err
	}
	tmp := a.FilePath + ".tmp"
	if err := os.WriteFile(tmp, raw, 0o600); err != nil {
		// Docker bind-mount 权限问题的典型现场:容器内 app 用户(uid 10001)
		// 对宿主机挂载目录无写权限。给出可操作指引而不是裸 syscall 错误。
		msg := fmt.Sprintf("写入 %s 失败: %v", tmp, err)
		if errors.Is(err, fs.ErrPermission) {
			msg += "\n(Docker 部署:容器内用户对宿主机挂载目录无写权限。解法任选:" +
				"1) 以本机 uid 运行容器:PUID=$(id -u) PGID=$(id -g) docker compose up -d;" +
				"2) sudo chown -R 10001:10001 ./auths ./data ./config.json;" +
				"3) compose 设 user: \"0:0\" 以 root 运行)"
		}
		return errors.New(msg)
	}
	return os.Rename(tmp, a.FilePath)
}

// LoadDir 扫描并解析 dir 下 workbuddy*.json;解析失败的文件静默跳过(启动日志由调用方统计)。
// 顺带做 realm 标识存量迁移:对空 realm 的 auth 自动 backfill(原始 domain 推断)并 SaveAtomic
// 落盘,一次性把旧文件补上 realm 键。单个文件写失败不阻断启动(log WARN 继续),
// 避免历史 auth 目录个别文件不可写时整个服务起不来。
func LoadDir(dir string) ([]*Auth, error) {
	files, err := filepath.Glob(filepath.Join(dir, "workbuddy*.json"))
	if err != nil {
		return nil, err
	}
	// seenUID 重复 UID 检测:同 UID 出现在多个文件时(双 realm 同名 UID 概率近零)
	// 打 WARN 告警含两文件路径,由「后载入者胜出」保持现状行为(不改变加载结果)。
	seenUID := make(map[string]string, len(files))
	var out []*Auth
	for _, f := range files {
		raw, err := os.ReadFile(f)
		if err != nil {
			continue
		}
		a, err := Parse(raw)
		if err != nil {
			continue
		}
		a.FilePath = f
		if prev, ok := seenUID[a.UID]; ok {
			log.Printf("WARN: uid %s duplicated across %s and %s — 后者覆盖(不同 realm 同名 UID?)",
				logfmt.Label(a.UID, a.Nickname), prev, f)
		}
		seenUID[a.UID] = f
		if a.RealmStored() == "" {
			if changed, r := a.BackfillRealm(); changed {
				if err := a.SaveAtomic(); err != nil {
					log.Printf("WARN: auth %s realm backfill save: %v", logfmt.Label(a.UID, a.Nickname), err)
				} else if r == "global" {
					log.Printf("auth %s 存量迁移: 补 realm=global(domain=%s)", logfmt.Label(a.UID, a.Nickname), a.Domain)
				}
			}
		}
		out = append(out, a)
	}
	return out, nil
}