// keys.go:consumers / api_keys / quota_ledger repository。 // 存储形状(design-arch §2.1):明文 key 不落库,只存 prefix + hash(SHA-256+salt)。 package store import ( "database/sql" "errors" "fmt" "strings" "time" "onesvm.com/onesvm/browser-server/internal/config" ) // ErrNotFound 行不存在。 var ErrNotFound = errors.New("store: not found") // Consumer 注册主体。 type Consumer struct { ID int64 Name string Contact string Status string CreatedAt time.Time } // ApiKey 访问密钥(hash 存 SHA-256(salt+plaintext),见 auth 包)。 type ApiKey struct { ID int64 ConsumerID int64 Prefix string Name string Hash string Salt string Scopes []string RPM int DailyQuota int MonthlyQuota int ConcurrentSessions int Status string ExpiresAt *time.Time CreatedAt time.Time } const keyCols = `id, consumer_id, prefix, name, hash, salt, scopes, rpm, daily_quota, monthly_quota, concurrent_sessions, status, expires_at, created_at` func scanKey(row interface{ Scan(...any) error }) (*ApiKey, error) { var k ApiKey var scopes, createdAt string var expires sql.NullString if err := row.Scan(&k.ID, &k.ConsumerID, &k.Prefix, &k.Name, &k.Hash, &k.Salt, &scopes, &k.RPM, &k.DailyQuota, &k.MonthlyQuota, &k.ConcurrentSessions, &k.Status, &expires, &createdAt); err != nil { return nil, err } k.Scopes = strings.Split(scopes, ",") if expires.Valid && expires.String != "" { t, err := time.Parse(time.RFC3339, expires.String) if err != nil { return nil, fmt.Errorf("store: api_keys.expires_at 解析: %w", err) } tz := t.In(config.TZ) k.ExpiresAt = &tz } ct, err := time.Parse(time.RFC3339, createdAt) if err != nil { return nil, fmt.Errorf("store: api_keys.created_at 解析: %w", err) } k.CreatedAt = ct return &k, nil } // CreateConsumer 建注册主体。 func (d *DB) CreateConsumer(name, contact string) (int64, error) { res, err := d.raw.Exec( `INSERT INTO consumers(name, contact, status, created_at) VALUES(?, ?, 'enabled', ?)`, name, contact, config.Now().Format(time.RFC3339)) if err != nil { return 0, fmt.Errorf("store: CreateConsumer: %w", err) } return res.LastInsertId() } // CreateKey 写入密钥记录(hash/salt 由调用方 auth 包生成,本层不碰明文)。 func (d *DB) CreateKey(k *ApiKey) (int64, error) { res, err := d.raw.Exec( `INSERT INTO api_keys(consumer_id, prefix, name, hash, salt, scopes, rpm, daily_quota, monthly_quota, concurrent_sessions, status, expires_at, created_at) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, k.ConsumerID, k.Prefix, k.Name, k.Hash, k.Salt, strings.Join(k.Scopes, ","), k.RPM, k.DailyQuota, k.MonthlyQuota, k.ConcurrentSessions, k.Status, formatOptTime(k.ExpiresAt), config.Now().Format(time.RFC3339)) if err != nil { return 0, fmt.Errorf("store: CreateKey: %w", err) } return res.LastInsertId() } // KeyByHash 按 hash 查密钥(auth 校验路径)。 func (d *DB) KeyByHash(hash string) (*ApiKey, error) { row := d.raw.QueryRow(`SELECT `+keyCols+` FROM api_keys WHERE hash = ?`, hash) k, err := scanKey(row) if errors.Is(err, sql.ErrNoRows) { return nil, ErrNotFound } return k, err } // KeyByID 按 id 查密钥(配额路径)。 func (d *DB) KeyByID(id int64) (*ApiKey, error) { row := d.raw.QueryRow(`SELECT `+keyCols+` FROM api_keys WHERE id = ?`, id) k, err := scanKey(row) if errors.Is(err, sql.ErrNoRows) { return nil, ErrNotFound } return k, err } // SetKeyStatus 吊销/禁用/启用。 func (d *DB) SetKeyStatus(id int64, status string) error { res, err := d.raw.Exec(`UPDATE api_keys SET status = ? WHERE id = ?`, status, id) if err != nil { return fmt.Errorf("store: SetKeyStatus: %w", err) } if n, _ := res.RowsAffected(); n == 0 { return ErrNotFound } return nil } // formatOptTime 可空时间转存储串。 func formatOptTime(t *time.Time) any { if t == nil { return nil } return t.In(config.TZ).Format(time.RFC3339) } // ---------- quota_ledger ---------- // ledgerWindow 计量窗口标识:day:YYYY-MM-DD / month:YYYY-MM。 func ledgerWindow(kind string, now time.Time) string { t := now.In(config.TZ) if kind == "month" { return "month:" + t.Format("2006-01") } return "day:" + t.Format("2006-01-02") } // QuotaReserve 配额预扣(事务内条件 UPDATE WHERE used+reserved < limit,防超卖)。 // kind: "day"|"month"。返回剩余可用额度;额度不足返回 ErrNotFound 语义外的专用错误由调用方判断。 func (d *DB) QuotaReserve(keyID int64, limit int, now time.Time) (int, error) { win := ledgerWindow("day", now) tx, err := d.raw.Begin() if err != nil { return 0, fmt.Errorf("store: QuotaReserve begin: %w", err) } defer tx.Rollback() //nolint:errcheck // 显式 Commit 后 rollback 为幂等空操作 if _, err := tx.Exec( `INSERT OR IGNORE INTO quota_ledger(key_id, window, used, reserved) VALUES(?, ?, 0, 0)`, keyID, win); err != nil { return 0, fmt.Errorf("store: QuotaReserve init: %w", err) } res, err := tx.Exec( `UPDATE quota_ledger SET reserved = reserved + 1 WHERE key_id = ? AND window = ? AND used + reserved < ?`, keyID, win, limit) if err != nil { return 0, fmt.Errorf("store: QuotaReserve update: %w", err) } n, _ := res.RowsAffected() if n == 0 { return 0, ErrNotFound // 额度不足或 key 不存在 } var remaining int if err := tx.QueryRow( `SELECT ? - used - reserved FROM quota_ledger WHERE key_id = ? AND window = ?`, limit, keyID, win).Scan(&remaining); err != nil { return 0, fmt.Errorf("store: QuotaReserve remaining: %w", err) } if err := tx.Commit(); err != nil { return 0, fmt.Errorf("store: QuotaReserve commit: %w", err) } return remaining, nil } // QuotaSettle 结算:预扣转实耗(预留位 -1,used +1)。 func (d *DB) QuotaSettle(keyID int64, now time.Time) error { win := ledgerWindow("day", now) _, err := d.raw.Exec( `UPDATE quota_ledger SET reserved = MAX(reserved - 1, 0), used = used + 1 WHERE key_id = ? AND window = ?`, keyID, win) if err != nil { return fmt.Errorf("store: QuotaSettle: %w", err) } return nil } // QuotaRelease 释放预扣(任务失败回滚额度)。 func (d *DB) QuotaRelease(keyID int64, now time.Time) error { win := ledgerWindow("day", now) _, err := d.raw.Exec( `UPDATE quota_ledger SET reserved = MAX(reserved - 1, 0) WHERE key_id = ? AND window = ?`, keyID, win) if err != nil { return fmt.Errorf("store: QuotaRelease: %w", err) } return nil } // QuotaUsed 查询某 key 当日已耗(观测用)。 func (d *DB) QuotaUsed(keyID int64, now time.Time) (int, error) { win := ledgerWindow("day", now) var used int err := d.raw.QueryRow( `SELECT used FROM quota_ledger WHERE key_id = ? AND window = ?`, keyID, win).Scan(&used) if errors.Is(err, sql.ErrNoRows) { return 0, nil } return used, err }