visionA/local-agent/server/internal/api/token_store.go

package api

import (
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"sync"
	"time"
)

// ADR-019 §2.4：one-time upload token store。
//
// 設計（已經 security review 議題 2 通過）：
//   - token = crypto/rand 32 bytes → base64url（禁 math/rand）
//   - TTL 120s、one-time（consume 即刪）、綁 deviceId
//   - 記憶體 store（map + 單一 sync.Mutex，不持久化）
//   - 未使用上限 32 個（防記憶體 DoS）
//   - 比對用 crypto/subtle.ConstantTimeCompare（防 timing attack）
//   - consume single-flight：查存在 + 比對 + 刪除三步在同一 Lock 內完成（security m2，防 race）
const (
	tokenTTL           = 120 * time.Second
	maxUnusedTokens    = 32
	tokenCleanupPeriod = 60 * time.Second
	tokenRandBytes     = 32
)

// 明確的錯誤，供 handler 對應到 api-spec §6.5 的錯誤碼。
var (
	// ErrTokenLimit：未使用 token 達 32 上限（→ 429 LOCAL_TOKEN_LIMIT）。
	ErrTokenLimit = errors.New("local token limit reached")
	// ErrTokenInvalid：token 不存在 / 過期 / 已使用 / deviceId 不符（→ 401 LOCAL_TOKEN_INVALID）。
	ErrTokenInvalid = errors.New("local token invalid")
)

// tokenEntry 是一筆未消費的 token 記錄。
type tokenEntry struct {
	deviceID  string
	expiresAt time.Time
}

// TokenStore 是執行緒安全的 one-time token 記憶體 store。
//
// 併發正確性核心：所有讀寫都在單一 mu 內完成。
// Consume 是 single-flight——「查存在 + ConstantTimeCompare + 刪除」在同一 Lock()
// 內原子完成，兩個併發 consume 同一 token 不可能都成功（防 one-time 失效）。
type TokenStore struct {
	mu     sync.Mutex
	tokens map[string]tokenEntry
	now    func() time.Time // 可注入，方便測試過期邏輯
}

// NewTokenStore 建立 store。now 預設為 time.Now。
func NewTokenStore() *TokenStore {
	return &TokenStore{
		tokens: make(map[string]tokenEntry),
		now:    time.Now,
	}
}

// Issue 產生一個新 token 綁定 deviceID，single-flight 持鎖完成
// 「清過期 + 查 len < 32 + 插入」。達上限回 ErrTokenLimit。
//
// token 值以 crypto/rand 產生（32 bytes → base64url RawURL）。
func (s *TokenStore) Issue(deviceID string) (string, time.Time, error) {
	// 先在鎖外產生亂數（crypto/rand 可能較慢，避免長時間持鎖）。
	buf := make([]byte, tokenRandBytes)
	if _, err := rand.Read(buf); err != nil {
		return "", time.Time{}, err
	}
	token := base64.RawURLEncoding.EncodeToString(buf)

	s.mu.Lock()
	defer s.mu.Unlock()

	// 惰性清理過期 token，順便為上限計算釋放名額。
	s.pruneExpiredLocked()

	if len(s.tokens) >= maxUnusedTokens {
		return "", time.Time{}, ErrTokenLimit
	}

	expiresAt := s.now().Add(tokenTTL)
	s.tokens[token] = tokenEntry{deviceID: deviceID, expiresAt: expiresAt}
	return token, expiresAt, nil
}

// Consume 驗證並消費一個 token（one-time）。single-flight 持鎖：
// 查存在 + 比對 deviceID + 未過期 + 刪除，全部在同一 Lock 內完成。
//
// 成功 → 回 nil（token 已從 store 移除，不可再用）。
// 失敗（不存在 / 過期 / deviceId 不符）→ 回 ErrTokenInvalid。
//
// deviceID 比對用 ConstantTimeCompare（雖然 deviceID 非高機密，維持一致的常數時間比對紀律）。
func (s *TokenStore) Consume(token, deviceID string) error {
	if token == "" {
		return ErrTokenInvalid
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	entry, ok := s.tokens[token]
	if !ok {
		return ErrTokenInvalid
	}

	// 不論後續成功與否，one-time 語意要求「命中即刪」——刪除放在最前面，
	// 確保兩個併發 consume 只有第一個拿到 entry、第二個 map 查不到。
	delete(s.tokens, token)

	// 過期檢查（惰性）。
	if !s.now().Before(entry.expiresAt) {
		return ErrTokenInvalid
	}

	// deviceID 綁定檢查（常數時間比對）。
	if subtle.ConstantTimeCompare([]byte(entry.deviceID), []byte(deviceID)) != 1 {
		return ErrTokenInvalid
	}

	return nil
}

// IsLimitErr 回報 err 是否為「達 token 上限」（給 handler 對應 429）。
// 讓 handlers 套件不需 import sentinel error 即可判斷。
func (s *TokenStore) IsLimitErr(err error) bool {
	return errors.Is(err, ErrTokenLimit)
}

// pruneExpiredLocked 移除所有已過期的 token。呼叫端必須已持有 mu。
func (s *TokenStore) pruneExpiredLocked() {
	now := s.now()
	for tok, entry := range s.tokens {
		if !now.Before(entry.expiresAt) {
			delete(s.tokens, tok)
		}
	}
}

// pruneExpired 是背景 goroutine 用的加鎖版本。
func (s *TokenStore) pruneExpired() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneExpiredLocked()
}

// len 回傳目前未使用 token 數（測試用）。
func (s *TokenStore) len() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.tokens)
}

// StartCleanup 啟動背景清理 goroutine，每 tokenCleanupPeriod 掃一次過期 token。
// 惰性清理（Consume/Issue 時）+ 背景清理雙保險。
// 回傳 stop 函式（給測試 / graceful shutdown 用）。
func (s *TokenStore) StartCleanup() (stop func()) {
	ticker := time.NewTicker(tokenCleanupPeriod)
	done := make(chan struct{})
	go func() {
		for {
			select {
			case <-ticker.C:
				s.pruneExpired()
			case <-done:
				ticker.Stop()
				return
			}
		}
	}()
	return func() { close(done) }
}