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

package api

import (
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// TestTokenStore_IssueConsume_HappyPath：發放後可消費一次，第二次消費失敗（one-time）。
func TestTokenStore_IssueConsume_HappyPath(t *testing.T) {
	s := NewTokenStore()

	token, expiresAt, err := s.Issue("dev-1")
	if err != nil {
		t.Fatalf("Issue error: %v", err)
	}
	if token == "" {
		t.Fatal("token 不應為空")
	}
	if !expiresAt.After(time.Now()) {
		t.Errorf("expiresAt %v 應在未來", expiresAt)
	}

	// 第一次消費成功
	if err := s.Consume(token, "dev-1"); err != nil {
		t.Fatalf("第一次 Consume 應成功，got %v", err)
	}
	// 第二次消費必失敗（one-time）
	if err := s.Consume(token, "dev-1"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("第二次 Consume 應 ErrTokenInvalid，got %v", err)
	}
}

// TestTokenStore_Consume_DeviceMismatch：deviceId 不符 → ErrTokenInvalid，且 token 已被消費。
func TestTokenStore_Consume_DeviceMismatch(t *testing.T) {
	s := NewTokenStore()
	token, _, _ := s.Issue("dev-1")

	// deviceId 不符 → 失敗
	if err := s.Consume(token, "dev-2"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("deviceId 不符應 ErrTokenInvalid，got %v", err)
	}
	// 即使不符，token 也應已被移除（命中即刪，防以正確 deviceId 重試）
	if err := s.Consume(token, "dev-1"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("不符後 token 應已被消費，got %v", err)
	}
}

// TestTokenStore_Consume_MissingAndEmpty：不存在 / 空字串 token → ErrTokenInvalid。
func TestTokenStore_Consume_MissingAndEmpty(t *testing.T) {
	s := NewTokenStore()
	if err := s.Consume("nonexistent", "dev-1"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("不存在 token 應 ErrTokenInvalid，got %v", err)
	}
	if err := s.Consume("", "dev-1"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("空 token 應 ErrTokenInvalid，got %v", err)
	}
}

// TestTokenStore_Expiry：過期 token 消費失敗。用可注入的 now 模擬時間流逝。
func TestTokenStore_Expiry(t *testing.T) {
	s := NewTokenStore()
	base := time.Now()
	current := base
	s.now = func() time.Time { return current }

	token, _, _ := s.Issue("dev-1")

	// 前進超過 TTL
	current = base.Add(tokenTTL + time.Second)

	if err := s.Consume(token, "dev-1"); !errors.Is(err, ErrTokenInvalid) {
		t.Errorf("過期 token 應 ErrTokenInvalid，got %v", err)
	}
}

// TestTokenStore_Limit：未使用 token 達 32 上限 → ErrTokenLimit；消費一個後可再發。
func TestTokenStore_Limit(t *testing.T) {
	s := NewTokenStore()

	tokens := make([]string, 0, maxUnusedTokens)
	for i := 0; i < maxUnusedTokens; i++ {
		tok, _, err := s.Issue("dev-1")
		if err != nil {
			t.Fatalf("第 %d 個 Issue 不應失敗，got %v", i, err)
		}
		tokens = append(tokens, tok)
	}

	// 第 33 個應被拒
	if _, _, err := s.Issue("dev-1"); !errors.Is(err, ErrTokenLimit) {
		t.Errorf("達上限應 ErrTokenLimit，got %v", err)
	}
	if !s.IsLimitErr(ErrTokenLimit) {
		t.Error("IsLimitErr(ErrTokenLimit) 應為 true")
	}

	// 消費一個後釋放名額，可再發
	if err := s.Consume(tokens[0], "dev-1"); err != nil {
		t.Fatalf("Consume 應成功，got %v", err)
	}
	if _, _, err := s.Issue("dev-1"); err != nil {
		t.Errorf("釋放名額後 Issue 應成功，got %v", err)
	}
}

// TestTokenStore_Limit_ExpiredFreesSlot：過期 token 在 Issue 時被惰性清理，釋放上限名額。
func TestTokenStore_Limit_ExpiredFreesSlot(t *testing.T) {
	s := NewTokenStore()
	base := time.Now()
	current := base
	s.now = func() time.Time { return current }

	for i := 0; i < maxUnusedTokens; i++ {
		if _, _, err := s.Issue("dev-1"); err != nil {
			t.Fatalf("第 %d 個 Issue 失敗: %v", i, err)
		}
	}
	// 全部過期
	current = base.Add(tokenTTL + time.Second)

	// 再 Issue 應觸發惰性清理、成功
	if _, _, err := s.Issue("dev-1"); err != nil {
		t.Errorf("過期清理後 Issue 應成功，got %v", err)
	}
}

// TestTokenStore_ConcurrentConsume_SingleFlight 是 security m2 的關鍵測試：
// 多個 goroutine 同時消費同一 token，必須「恰好一個成功」（防 one-time 失效 / 雙重消費）。
func TestTokenStore_ConcurrentConsume_SingleFlight(t *testing.T) {
	const goroutines = 200

	// 跑多輪，提高抓到 race 的機率。
	for round := 0; round < 50; round++ {
		s := NewTokenStore()
		token, _, _ := s.Issue("dev-1")

		var successCount int32
		var wg sync.WaitGroup
		start := make(chan struct{})
		wg.Add(goroutines)
		for i := 0; i < goroutines; i++ {
			go func() {
				defer wg.Done()
				<-start // 同時起跑，最大化競爭
				if err := s.Consume(token, "dev-1"); err == nil {
					atomic.AddInt32(&successCount, 1)
				}
			}()
		}
		close(start)
		wg.Wait()

		if successCount != 1 {
			t.Fatalf("round %d: 併發消費同一 token 成功數 = %d，必須恰好 1（single-flight 失效）",
				round, successCount)
		}
	}
}

// TestTokenStore_ConcurrentIssue_LimitHeld 是 security m2 的另一半：
// 併發 Issue 時，未使用 token 數不得突破 32 上限。
func TestTokenStore_ConcurrentIssue_LimitHeld(t *testing.T) {
	const goroutines = 200

	for round := 0; round < 30; round++ {
		s := NewTokenStore()

		var wg sync.WaitGroup
		start := make(chan struct{})
		wg.Add(goroutines)
		for i := 0; i < goroutines; i++ {
			go func() {
				defer wg.Done()
				<-start
				_, _, _ = s.Issue("dev-1")
			}()
		}
		close(start)
		wg.Wait()

		if got := s.len(); got > maxUnusedTokens {
			t.Fatalf("round %d: 併發 Issue 後 token 數 = %d，不得超過上限 %d（上限檢查非 atomic）",
				round, got, maxUnusedTokens)
		}
	}
}

// TestTokenStore_TokenUniqueness：連續發放的 token 值不重複。
func TestTokenStore_TokenUniqueness(t *testing.T) {
	s := NewTokenStore()
	seen := make(map[string]bool)
	for i := 0; i < maxUnusedTokens; i++ {
		tok, _, err := s.Issue("dev-1")
		if err != nil {
			t.Fatalf("Issue 失敗: %v", err)
		}
		if seen[tok] {
			t.Fatalf("token 重複: %q", tok)
		}
		seen[tok] = true
	}
}

// TestTokenStore_Cleanup_StopWorks：StartCleanup 回傳的 stop 可正常關閉 goroutine。
func TestTokenStore_Cleanup_StopWorks(t *testing.T) {
	s := NewTokenStore()
	stop := s.StartCleanup()
	// 立即停止不應 panic / deadlock
	stop()
}