aiops/backend/pkg/agentauth/credential.go

package agentauth

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"
	"sync"
	"time"

	"ezops/model"

	"golang.org/x/crypto/bcrypt"
	"gorm.io/gorm"
)

// ─────────────────────────────────────────────────────────────
// Token 校验缓存
//
// bcrypt 被设计成故意慢（本项目实测单次 Cost=10 约 72ms CPU）。Agent 每 10 秒
// 轮询 3 个接口，每个接口都要校验一次 Token，30 台主机即 9 req/s×72ms ≈ 0.65 核，
// 纯粹烧在重复校验同一条不可变凭证上。
//
// 由于 Token 在轮换前不可变（Issue 才会改写 token_hash），校验结果可以安全缓存：
//   - 缓存键：token 的 SHA-256（不落明文，避免内存里出现可直接使用的凭证）
//   - 有效期：默认 5 分钟，过期后回落一次 bcrypt 并续期
//   - 轮换：Issue 时全量清空，保证「旧 Token 立即失效」语义不变
//
// last_used_at 的写入同样按前缀节流（默认 60s 一次），避免每请求一次 UPDATE。
// // Logan.Li
type cachedCredential struct {
	credential *model.AgentCredential
	expiresAt  time.Time
	// lastWrite 记录最近一次把 last_used_at 落库的时间，用于写入节流。 // Logan.Li
	lastWrite time.Time
}

const (
	defaultCacheTTL      = 5 * time.Minute
	defaultWriteThrottle = 60 * time.Second
	// cacheMaxEntries 上限，防止异常 Token 洪水把内存撑爆。 // Logan.Li
	cacheMaxEntries = 4096
)

var (
	credCacheMu sync.Mutex
	credCache   = make(map[string]*cachedCredential)
	// cacheTTL / writeThrottle 允许测试覆盖。 // Logan.Li
	cacheTTL      = defaultCacheTTL
	writeThrottle = defaultWriteThrottle
)

// cacheKey 只缓存 Token 的哈希，内存中不保留明文 Token。 // Logan.Li
func cacheKey(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

// lookupCache 返回命中的凭证副本；命中时同时按节流更新 last_used_at。 // Logan.Li
func lookupCache(db *gorm.DB, token string) (*model.AgentCredential, bool) {
	key := cacheKey(token)
	credCacheMu.Lock()
	defer credCacheMu.Unlock()

	entry, ok := credCache[key]
	if !ok || time.Now().After(entry.expiresAt) {
		if ok {
			delete(credCache, key)
		}
		return nil, false
	}
	if entry.credential.ExpiresAt != nil && entry.credential.ExpiresAt.Before(time.Now()) {
		delete(credCache, key)
		return nil, false
	}

	// 节流写入 last_used_at：仅当距上次落库超过阈值才真正 UPDATE。 // Logan.Li
	now := time.Now()
	if now.Sub(entry.lastWrite) >= writeThrottle {
		entry.lastWrite = now
		prefix := entry.credential.TokenPrefix
		go func() {
			_ = db.Model(&model.AgentCredential{}).
				Where("token_prefix = ?", prefix).
				Update("last_used_at", now).Error
		}()
	}

	cp := *entry.credential
	cp.LastUsedAt = &now
	return &cp, true
}

// storeCache 写入缓存，必要时按容量上限淘汰过期项。 // Logan.Li
func storeCache(token string, credential *model.AgentCredential) {
	key := cacheKey(token)
	credCacheMu.Lock()
	defer credCacheMu.Unlock()

	if len(credCache) >= cacheMaxEntries {
		now := time.Now()
		for k, v := range credCache {
			if now.After(v.expiresAt) {
				delete(credCache, k)
			}
		}
		// 仍超上限则整体清空，宁可多算几次 bcrypt 也不无限占内存。 // Logan.Li
		if len(credCache) >= cacheMaxEntries {
			credCache = make(map[string]*cachedCredential)
		}
	}

	cp := *credential
	credCache[key] = &cachedCredential{
		credential: &cp,
		expiresAt:  time.Now().Add(cacheTTL),
		lastWrite:  time.Now(),
	}
}

// invalidateCache 清空全部缓存。Token 轮换/撤销后调用，保证旧 Token 立即失效。 // Logan.Li
func invalidateCache() {
	credCacheMu.Lock()
	credCache = make(map[string]*cachedCredential)
	credCacheMu.Unlock()
}

// Invalidate 提供给凭证变更入口（轮换/撤销/删除）显式失效缓存。 // Logan.Li
func Invalidate() { invalidateCache() }

// Issue 为指定主机轮换 Agent Token，旧 Token 会立即失效。 // Logan.Li
func Issue(db *gorm.DB, hostname, issuedBy string) (string, error) {
	hostname = strings.TrimSpace(hostname)
	if hostname == "" {
		return "", fmt.Errorf("hostname is required")
	}
	prefix, err := randomHex(8)
	if err != nil {
		return "", err
	}
	// bcrypt 只处理前 72 字节，Token 总长度控制在 61 字节以内。 // Logan.Li
	secret, err := randomHex(20)
	if err != nil {
		return "", err
	}
	token := "agt_" + prefix + "." + secret
	hash, err := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	now := time.Now()
	credential := model.AgentCredential{Hostname: hostname}
	err = db.Where("hostname = ?", hostname).Assign(map[string]interface{}{
		"token_prefix": "agt_" + prefix,
		"token_hash":   string(hash),
		"status":       "active",
		"issued_by":    issuedBy,
		"last_used_at": nil,
		"expires_at":   nil,
		"rotated_at":   now,
	}).FirstOrCreate(&credential).Error
	if err != nil {
		return "", err
	}
	// 轮换后必须让缓存失效，否则被缓存的旧 Token 还能继续通过校验。 // Logan.Li
	invalidateCache()
	return token, nil
}

// Authenticate 校验 Token 并返回绑定的主机名。 // Logan.Li
//
// 命中缓存时跳过 bcrypt 与凭证查询，只做一次可选节流的 last_used_at 更新；
// 未命中时才回落完整校验路径。 // Logan.Li
func Authenticate(db *gorm.DB, token string) (*model.AgentCredential, error) {
	if credential, ok := lookupCache(db, token); ok {
		return credential, nil
	}

	parts := strings.SplitN(token, ".", 2)
	if len(parts) != 2 || !strings.HasPrefix(parts[0], "agt_") || parts[1] == "" {
		return nil, fmt.Errorf("invalid agent token")
	}
	var credential model.AgentCredential
	if err := db.Where("token_prefix = ? AND status = ?", parts[0], "active").Take(&credential).Error; err != nil {
		return nil, fmt.Errorf("invalid agent token")
	}
	if credential.ExpiresAt != nil && credential.ExpiresAt.Before(time.Now()) {
		return nil, fmt.Errorf("agent token expired")
	}
	if bcrypt.CompareHashAndPassword([]byte(credential.TokenHash), []byte(token)) != nil {
		return nil, fmt.Errorf("invalid agent token")
	}
	now := time.Now()
	_ = db.Model(&credential).Update("last_used_at", now).Error
	credential.LastUsedAt = &now

	// 校验通过后写入缓存，后续请求直接命中。 // Logan.Li
	storeCache(token, &credential)
	return &credential, nil
}

func randomHex(bytes int) (string, error) {
	buf := make([]byte, bytes)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}