onesvm-browser-server/server/internal/store/keys.go

// 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
}