Cyrene_v2/internal/auth/apikey.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"log/slog"
	"sync"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// maskEllipsis 是脱敏标识中间的省略符；maskEdgeLen 是两端保留的十六进制位数。
const (
	maskEllipsis = "••••"
	maskEdgeLen  = 4
)

// APIKey 管理一张有界、常量的 API key 表。按裁决 15.4.3（主人原话：「key 始终给
// 超级管理员权限」），表里**不再绑定角色**：任何已登记且未吊销的 key 鉴权后一律
// owner 级 —— key 用于管理员调试，甚至生产环境救火，因此不做权限分级。
//
// 表里存的是 key 的 **SHA-256 哈希**（十六进制），不是明文：明文只在 Generate 的
// 返回值里出现一次，既不驻留内存也不落库、不落日志。这样「内存副本」与「api_keys
// 表」的凭据口径完全一致（迁移 0016_api_keys），重启后把库里未吊销的哈希载回内存
// 即可复原，无需也不可能还原明文。
//
// store 为 nil 时是本包原有的纯内存行为（测试与未接库的构建路径），行为不变。
type APIKey struct {
	mu    sync.RWMutex
	cfg   config.AuthConfig
	keys  map[string]struct{}
	store KeyStore
	log   *slog.Logger
}

// NewAPIKey 返回一个空的 key 管理器。cfg 暂只为构造签名的一致性而接收（暂无
// 影响 key 行为的字段）；落盘载体由 AttachKeyStore 单独注入，因此未接库的调用方
// （既有测试、纯内存用法）拿到的就是原样的内存表。
func NewAPIKey(cfg config.AuthConfig) *APIKey {
	return &APIKey{
		cfg:  cfg,
		keys: make(map[string]struct{}),
		log:  logger.WithSource(slog.Default(), source),
	}
}

// AttachKeyStore 把可选的落盘载体接到 key 管理器（迁移 0016_api_keys；K1）：
//
//  1. 记下 store —— 此后 Generate/Revoke/RevokeAll 先落库、成功后才改内存；
//  2. 立即 LoadUnrevoked 把库里**未吊销**的哈希载入内存，于是重启后已签发的
//     key 仍能通过 Verify（这就是本修复要的行为）。
//
// store 为 nil 是合法的空操作（保持纯内存，即既有行为）。Load 失败时返回错误，
// 但**store 依然被接上**：写路径随后会**响亮地**失败（Generate 返回错误），而不是
// 悄悄退回「只在内存里成功」—— 那正是 K1 要消灭的静默失效（签发端点报 201、重启
// 后全部 401）。此处与 capability/registry.AttachPolicyStore「Load 失败就不接线」
// 刻意不同：那边默认值本身就是最严（fail-closed），这边不接线的默认是**不落盘**。
func (k *APIKey) AttachKeyStore(ctx context.Context, store KeyStore) error {
	if k == nil || store == nil {
		return nil
	}
	k.mu.Lock()
	k.store = store
	k.mu.Unlock()

	hashes, err := store.LoadUnrevoked(ctx)
	if err != nil {
		return fmt.Errorf("auth: load api key store: %w", err)
	}
	k.mu.Lock()
	defer k.mu.Unlock()
	for _, hash := range hashes {
		if hash != "" {
			k.keys[hash] = struct{}{}
		}
	}
	return nil
}

// Generate 铸造一个新 key 并登记进表，返回**仅此一次**可见的明文（绝不记录到
// 日志、绝不落库：库里只写它的哈希）。已接上 store 时**先落库后改内存**（照
// capability/registry.SetPolicy 的既有纪律）：落库失败 ⇒ 内存不变 + 返回错误，
// 绝不出现「面板说签发成功、重启却不认」的假成功。
//
// role 是可选形参：传入时只校验它是不是已知角色名（未知即报 ErrInvalidRole），
// **不参与存储、不影响授权** —— 按裁决 15.4.3，key 表已无角色维度，任何有效 key
// 都是超管级（见 Lookup）。
func (k *APIKey) Generate(role ...Role) (string, error) {
	if len(role) > 0 && !role[0].Valid() {
		return "", fmt.Errorf("%w: %q", ErrInvalidRole, role[0])
	}
	raw := make([]byte, keyLength)
	if _, err := rand.Read(raw); err != nil {
		return "", fmt.Errorf("auth: generate api key: %w", err)
	}
	key := base64.RawURLEncoding.EncodeToString(raw)
	hash := hashAPIKey(key)

	if err := k.persistInsert(hash); err != nil {
		return "", err
	}
	k.mu.Lock()
	k.keys[hash] = struct{}{}
	k.mu.Unlock()
	return key, nil
}

// Verify 报告 key 是否是一把已登记且未吊销的 API key。
func (k *APIKey) Verify(key string) bool {
	_, ok := k.Lookup(key)
	return ok
}

// Lookup 报告 key 是否已登记，并返回其授权角色。按裁决 15.4.3，**任何**已登记的
// key 都是超管（owner）级。key 不在表中时 ok 为 false，调用方必须据此拒绝。
//
// 判定就是一次按哈希的 map 查表：未命中**不回查库**（启动时已载入全部有效哈希），
// 因此鉴权路径既不碰磁盘、也不做明文逐字节比较，天然 fail-closed（未知/已吊销一律
// 不在表里）。被吊销的 key 因此绝不会回落到 Bearer（见 WithAPIKey）。
func (k *APIKey) Lookup(key string) (Role, bool) {
	k.mu.RLock()
	defer k.mu.RUnlock()
	_, ok := k.keys[hashAPIKey(key)]
	return RoleOwner, ok
}

// Revoke 吊销一把 key：**先落库标记 revoked_at，成功后才移出内存**。落库失败时
// 内存保持不变（这把 key 在本进程内仍然有效）并记一条告警日志 —— 绝不出现「本进程
// 说已吊销、重启后又活回来」的反向不一致。既有签名（无 ctx、无 error，网关的 404
// 与计数出口依赖它）保持不变，因此这里用 context.Background() 并把失败降级为日志。
func (k *APIKey) Revoke(key string) {
	hash := hashAPIKey(key)
	if err := k.persistRevoke(hash); err != nil {
		return
	}
	k.mu.Lock()
	delete(k.keys, hash)
	k.mu.Unlock()
}

// RevokeAll 移除当前已登记的每一个 API key 并返回移除数量。它是管理端的「全部
// 吊销」：调用之后表为空，此前签发的 key 一律无法认证。已接上 store 时先把库中
// 所有未吊销行标记为已吊销（跨重启同样失效），成功后才清空内存；落库失败则内存
// 不动、返回 0（「一个都没废掉」是诚实的回答，而不是清空内存后重启又复活）。
func (k *APIKey) RevokeAll() int {
	if err := k.persistRevokeAll(); err != nil {
		return 0
	}
	k.mu.Lock()
	defer k.mu.Unlock()
	n := len(k.keys)
	k.keys = make(map[string]struct{})
	return n
}

// RevokeByMask 移除脱敏标识（与 List 返回的 id 同一口径）等于 mask 的每一个 key，
// 有至少一把被移除时返回 true。掩码是有损的，因此按「为每条记录重算掩码」的方式
// 匹配；表很小，线性扫描足够。空掩码是空操作。
//
// 顺序仍是「先落库、后改内存」：只有落库成功的那些才移出内存，失败的留在表里并
// 记告警（否则重启后它们会带着未标记的行复活）。SQL 在管理器锁之外执行，避免在
// 持锁时做 IO。
func (k *APIKey) RevokeByMask(mask string) bool {
	if mask == "" {
		return false
	}
	matches := k.hashesMatchingMask(mask)
	if len(matches) == 0 {
		return false
	}
	persisted := make([]string, 0, len(matches))
	for _, hash := range matches {
		if err := k.persistRevoke(hash); err != nil {
			continue
		}
		persisted = append(persisted, hash)
	}
	if len(persisted) == 0 {
		return false
	}
	k.mu.Lock()
	defer k.mu.Unlock()
	for _, hash := range persisted {
		delete(k.keys, hash)
	}
	return true
}

// List 返回面板可见的每一条 key 摘要。绝不返回密钥材料：每条只给脱敏标识（哈希
// 的首尾各 4 位十六进制），面板据此辨认与吊销。内存表不记创建时间，因此每条
// CreatedAt 均为 nil（库里有 created_at 列，但目前不回读、不在本结构上暴露）。
// 每条的角色恒为 owner（裁决 15.4.3：key 不做角色分级）。
func (k *APIKey) List() []UserSummary {
	k.mu.RLock()
	defer k.mu.RUnlock()
	out := make([]UserSummary, 0, len(k.keys))
	for hash := range k.keys {
		out = append(out, UserSummary{ID: maskHash(hash), Username: "apikey", Role: RoleOwner})
	}
	return out
}

// hashesMatchingMask 返回当前表里脱敏标识等于 mask 的全部哈希（读锁内完成）。
func (k *APIKey) hashesMatchingMask(mask string) []string {
	k.mu.RLock()
	defer k.mu.RUnlock()
	out := make([]string, 0, 1)
	for hash := range k.keys {
		if maskHash(hash) == mask {
			out = append(out, hash)
		}
	}
	return out
}

// persistInsert 把新 key 的哈希写进已接上的 store（未接上是空操作）。它刻意在
// 管理器锁之外执行 SQL，调用方只在它成功后才改内存。
func (k *APIKey) persistInsert(hash string) error {
	store := k.storeSnapshot()
	if store == nil {
		return nil
	}
	if err := store.Insert(context.Background(), hash); err != nil {
		return fmt.Errorf("auth: persist api key: %w", err)
	}
	return nil
}

// persistRevoke 在库里把哈希对应的行标记为已吊销（未接 store 是空操作）。失败时
// 记一条告警并返回错误，由调用方决定不动内存。
func (k *APIKey) persistRevoke(hash string) error {
	store := k.storeSnapshot()
	if store == nil {
		return nil
	}
	if err := store.Revoke(context.Background(), hash); err != nil {
		k.log.Warn("api key revoke not persisted, key stays valid until restart", "err", err)
		return fmt.Errorf("auth: revoke api key: %w", err)
	}
	return nil
}

// persistRevokeAll 在库里标记全部未吊销行为已吊销（未接 store 是空操作）。
func (k *APIKey) persistRevokeAll() error {
	store := k.storeSnapshot()
	if store == nil {
		return nil
	}
	if _, err := store.RevokeAll(context.Background()); err != nil {
		k.log.Warn("api key revoke-all not persisted, keys stay valid until restart", "err", err)
		return fmt.Errorf("auth: revoke all api keys: %w", err)
	}
	return nil
}

// storeSnapshot 在锁内读一次 store 指针，避免与 AttachKeyStore 的写入竞态。
func (k *APIKey) storeSnapshot() KeyStore {
	k.mu.RLock()
	defer k.mu.RUnlock()
	return k.store
}

// MaskKey 返回一把 key 的脱敏标识：明文 → SHA-256 → 首尾各 4 位十六进制。它是
// maskHash 的导出形态，供调用方（如管理端签发 handler 返回的 id、审计链的 Sub）
// 算出与 List 页面一致的标识。
//
// 与修复前的旧口径（直接取明文首尾 4 位）相比，这个 id **不再泄露明文的任何一位**，
// 而且只由哈希决定 —— 同一把 key 在任何一次重启后都得到同一个 id，面板上记下的 id
// 跨重启仍然可吊销。
func MaskKey(key string) string { return maskHash(hashAPIKey(key)) }

// maskHash 把 64 位十六进制哈希渲染成脱敏标识（首尾各 maskEdgeLen 位）。它是
// MaskKey / List / RevokeByMask 三者共用的 id 口径，改动它等于改动面板与审计链
// 看到的标识，因此集中在这里。
func maskHash(hash string) string {
	if len(hash) <= maskEdgeLen*2 {
		return maskEllipsis
	}
	return hash[:maskEdgeLen] + maskEllipsis + hash[len(hash)-maskEdgeLen:]
}

// hashAPIKey 返回 key 的十六进制 SHA-256 —— 内存表与 api_keys 表都只认这一形式。
// key 是 32 字节随机秘密（256 位熵），没有低熵输入可供爆破，因此无需加盐；查表由
// map 与唯一索引完成，不做明文逐字节比较。
func hashAPIKey(key string) string {
	sum := sha256.Sum256([]byte(key))
	return hex.EncodeToString(sum[:])
}