feat(gateway): per-key 用量配额(token + 请求数)与重置周期

问题:密钥控制只能限制模型范围。实测发现三个缺陷,其中前两个让
per-model token_quota 在真实链路上从未生效:

1. 桶键不含 key。scopeTokens 调 WindowTokens(model, source, win),
   桶键是 model / source::model,与调用方无关。实测两把 key 各用
   1000 token,窗口报 2000 —— A key 的额度被 B key 消耗。
2. 无 source pin 的桶永远是空的。真实记录 Source 总被填上,桶键存成
   "deepseek::m1",而无 pin 的查询找 "m1" —— 读到 0,永远 < quota,
   配额形同虚设。实测 WindowTokens("m1","",1h)=0 而 pinned=2000。
3. AUTO scope 走 KeyTokens(key),是全时段累计、永不重置。实测 30 天
   前的 200 token 仍计入 1 小时配额(报 210 而非 10)。配了
   period: hour 也不会每小时归零。

生产 5 把 user key 全是 token_quota: 0,所以前两条一直没暴露。

改动:
- Stats 新增 per-key 小时桶 keyModelHour(key → model → hour)与
  keyHour(key 总量)、keyReqHour(请求数),retention 40 天,与既有
  modelHour 对齐以覆盖最长的 month 窗口;LoadAudit 走 aggregateLocked,
  所以窗口用量跨重启存活。modelHour 保持 key-blind:它服务的是 AUTO
  槽位配额(限制整个网关对某槽位的消耗),语义不同,不应被 per-key
  改造污染。
- 每个请求写两份模型桶:裸 model 与 source::model。无 pin 的 scope
  条目读前者,有 pin 的读后者。
- GWKey 新增 TokenQuota / ReqQuota / Period / Hours:整钥配额,
  跨该 key 所有模型共享一份预算;ReqQuota 覆盖持续请求量(源上的
  RPM 只管突发)。
- 配额耗尽返回 429 + Retry-After(rate_limit_exceeded),而不是 403:
  403 让客户端以为这把 key 永远不能用该模型,直接放弃;429 + 等待
  才能在窗口重置后自动恢复。模型越权仍是 403。
- admin key 永不受配额限制 —— 否则操作者会把自己锁在门外。
- 周期词表在写入时校验,拼错的 period 被拒绝而不是静默当成永不过期
  (那与操作者输入的意图正好相反)。
- PUT /api/keys 的配额字段是指针:省略=保留原值,显式 0=解除限制。
  否则只改模型范围就会悄悄清空预算。

判据 3 个文件 24 例,9 个变异全部被抓:key 隔离、pin 桶缺失、
AUTO 周期、key-blind 退化、429→403、admin 被限、PUT 清空配额、
Validate 失效、pinned 桶缺失。前三个变异最初漏网 —— 判据只测了
Stats 层没测接线,补了走真实 HTTP 的接线层与 API 层判据后抓住。
端到端验证:真实进程 + 加密配置往返,配额字段与 enc:v1 密钥均正常。

(cherry picked from commit 5306251840)
This commit is contained in:
JianFeeeee
2026-09-27 17:23:36 +08:00
parent a7355debed
commit 9c3aabb7f9
9 changed files with 1151 additions and 41 deletions

View File

@ -250,6 +250,11 @@ func (c *Core) FindKey(key string) (config.GWKey, bool) {
// CreateKey builds a new random gateway key and persists it to config.yaml.
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
return c.CreateKeyWithQuota(name, role, models, note, config.KeyQuota{})
}
// CreateKeyWithQuota is CreateKey plus the key-wide token/request caps.
func (c *Core) CreateKeyWithQuota(name, role string, models []config.ModelScope, note string, q config.KeyQuota) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
models = cleanScopes(models)
@ -257,6 +262,9 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
if _, err := rand.Read(key); err != nil {
return config.GWKey{}, err
}
if err := q.Validate(); err != nil {
return config.GWKey{}, err
}
rec := config.GWKey{
Key: "sk-gw-" + hex.EncodeToString(key),
Role: role,
@ -265,9 +273,8 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
Note: note,
CreatedAt: time.Now().Unix(),
}
if rec.Role == "" {
rec.Role = "user"
}
rec.Role = config.NormalizeRole(rec.Role)
rec.ApplyQuota(q)
c.cfg.Keys = append(c.cfg.Keys, rec)
if err := c.saveConfig(); err != nil {
return config.GWKey{}, err
@ -277,8 +284,20 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
// UpdateKey mutates a key's name/role/model scope and persists it.
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
return c.UpdateKeyWithQuota(key, name, role, models, note, nil)
}
// UpdateKeyWithQuota is UpdateKey plus the key-wide caps. quota == nil leaves
// the existing caps untouched, so a caller that only edits the model scope
// does not silently clear a key's budget.
func (c *Core) UpdateKeyWithQuota(key, name, role string, models []config.ModelScope, note string, quota *config.KeyQuota) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
if quota != nil {
if err := quota.Validate(); err != nil {
return config.GWKey{}, err
}
}
for i, k := range c.cfg.Keys {
if k.Key == key {
if name != "" {
@ -293,6 +312,9 @@ func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, not
c.cfg.Keys[i].Models = cleanScopes(models)
}
c.cfg.Keys[i].Note = note
if quota != nil {
c.cfg.Keys[i].ApplyQuota(*quota)
}
if err := c.saveConfig(); err != nil {
return config.GWKey{}, err
}