at-container-registry/pkg/auth/cache_test.go

package auth

import (
	"encoding/base64"
	"fmt"
	"testing"
	"time"
)

func TestGetServiceToken_NotCached(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:test123"
	holdDID := "did:web:hold.example.com"

	token, expiresAt := GetServiceToken(did, holdDID)
	if token != "" {
		t.Errorf("Expected empty token for uncached entry, got %q", token)
	}
	if !expiresAt.IsZero() {
		t.Error("Expected zero time for uncached entry")
	}
}

func TestSetServiceToken_ManualExpiry(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:test123"
	holdDID := "did:web:hold.example.com"
	token := "invalid_jwt_token" // Will fall back to 50s default

	// This should succeed with default 50s TTL since JWT parsing will fail
	err := SetServiceToken(did, holdDID, token)
	if err != nil {
		t.Fatalf("SetServiceToken() error = %v", err)
	}

	// Verify token was cached
	cachedToken, expiresAt := GetServiceToken(did, holdDID)
	if cachedToken != token {
		t.Errorf("Expected token %q, got %q", token, cachedToken)
	}
	if expiresAt.IsZero() {
		t.Error("Expected non-zero expiry time")
	}

	// Expiry should be approximately 50s from now (with 10s margin subtracted in some cases)
	expectedExpiry := time.Now().Add(50 * time.Second)
	diff := expiresAt.Sub(expectedExpiry)
	if diff < -5*time.Second || diff > 5*time.Second {
		t.Errorf("Expiry time off by %v (expected ~50s from now)", diff)
	}
}

func TestGetServiceToken_Expired(t *testing.T) {
	// Manually insert an expired token by reaching into the cache. Set()
	// would apply the safety margin which is the opposite of what we
	// want here — we want a guaranteed-stale entry.
	did := "did:plc:test123"
	holdDID := "did:web:hold.example.com"
	cacheKey := did + ":" + holdDID

	defaultCache.mu.Lock()
	defaultCache.tokens[cacheKey] = &serviceTokenEntry{
		token:     "expired_token",
		expiresAt: time.Now().Add(-1 * time.Hour),
	}
	defaultCache.mu.Unlock()

	// Try to get - should return empty since expired
	token, expiresAt := GetServiceToken(did, holdDID)
	if token != "" {
		t.Errorf("Expected empty token for expired entry, got %q", token)
	}
	if !expiresAt.IsZero() {
		t.Error("Expected zero time for expired entry")
	}

	// Verify token was removed from cache
	defaultCache.mu.RLock()
	_, exists := defaultCache.tokens[cacheKey]
	defaultCache.mu.RUnlock()

	if exists {
		t.Error("Expected expired token to be removed from cache")
	}
}

func TestInvalidateServiceToken(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:test123"
	holdDID := "did:web:hold.example.com"
	token := "test_token"

	err := SetServiceToken(did, holdDID, token)
	if err != nil {
		t.Fatalf("SetServiceToken() error = %v", err)
	}

	cachedToken, _ := GetServiceToken(did, holdDID)
	if cachedToken != token {
		t.Fatal("Token should be cached")
	}

	InvalidateServiceToken(did, holdDID)

	cachedToken, _ = GetServiceToken(did, holdDID)
	if cachedToken != "" {
		t.Error("Expected token to be invalidated")
	}
}

func TestCleanExpiredTokens(t *testing.T) {
	defaultCache.Clear()

	defaultCache.mu.Lock()
	defaultCache.tokens["expired:hold1"] = &serviceTokenEntry{
		token:     "expired1",
		expiresAt: time.Now().Add(-1 * time.Hour),
	}
	defaultCache.tokens["valid:hold2"] = &serviceTokenEntry{
		token:     "valid1",
		expiresAt: time.Now().Add(1 * time.Hour),
	}
	defaultCache.mu.Unlock()

	CleanExpiredTokens()

	defaultCache.mu.RLock()
	_, expiredExists := defaultCache.tokens["expired:hold1"]
	_, validExists := defaultCache.tokens["valid:hold2"]
	defaultCache.mu.RUnlock()

	if expiredExists {
		t.Error("Expected expired token to be removed")
	}
	if !validExists {
		t.Error("Expected valid token to remain")
	}
}

func TestGetCacheStats(t *testing.T) {
	defaultCache.Clear()

	defaultCache.mu.Lock()
	defaultCache.tokens["did1:hold1"] = &serviceTokenEntry{
		token:     "token1",
		expiresAt: time.Now().Add(1 * time.Hour),
	}
	defaultCache.tokens["did2:hold2"] = &serviceTokenEntry{
		token:     "token2",
		expiresAt: time.Now().Add(1 * time.Hour),
	}
	defaultCache.mu.Unlock()

	stats := GetCacheStats()
	if stats == nil {
		t.Fatal("Expected non-nil stats")
	}

	totalEntries, ok := stats["total_entries"].(int)
	if !ok {
		t.Fatalf("Expected total_entries in stats map, got: %v", stats)
	}

	if totalEntries != 2 {
		t.Errorf("Expected 2 entries, got %d", totalEntries)
	}

	validTokens, ok := stats["valid_tokens"].(int)
	if !ok {
		t.Fatal("Expected valid_tokens in stats map")
	}

	if validTokens != 2 {
		t.Errorf("Expected 2 valid tokens, got %d", validTokens)
	}
}

func TestCache_StructAPI_IsolatedInstances(t *testing.T) {
	// A freshly constructed Cache must not share state with defaultCache —
	// otherwise tests that expect isolation would silently fail.
	c1 := NewCache()
	c2 := NewCache()

	did := "did:plc:alice"
	holdDID := "did:web:hold.test"

	c1.mu.Lock()
	c1.tokens[did+":"+holdDID] = &serviceTokenEntry{
		token:     "tok-c1",
		expiresAt: time.Now().Add(1 * time.Hour),
	}
	c1.mu.Unlock()

	if tok, _ := c1.Get(did, holdDID); tok != "tok-c1" {
		t.Errorf("c1.Get() = %q, want tok-c1", tok)
	}
	if tok, _ := c2.Get(did, holdDID); tok != "" {
		t.Errorf("c2.Get() = %q, want empty (instances must not share state)", tok)
	}
	if tok, _ := GetServiceToken(did, holdDID); tok == "tok-c1" {
		t.Error("defaultCache should not see writes to c1")
	}

	c1.Clear()
	if tok, _ := c1.Get(did, holdDID); tok != "" {
		t.Errorf("c1.Get() after Clear() = %q, want empty", tok)
	}
}

func TestCache_PackageFunctionsDelegateToDefault(t *testing.T) {
	// The package-level wrappers must route to defaultCache.
	defaultCache.Clear()

	did := "did:plc:bob"
	holdDID := "did:web:hold.test"

	if err := SetServiceToken(did, holdDID, "wrapper-tok"); err != nil {
		t.Fatalf("SetServiceToken: %v", err)
	}

	tok, exp := DefaultCache().Get(did, holdDID)
	if tok != "wrapper-tok" {
		t.Errorf("DefaultCache().Get() = %q, want wrapper-tok", tok)
	}
	if exp.IsZero() {
		t.Error("DefaultCache().Get() expiry is zero")
	}

	InvalidateServiceToken(did, holdDID)
	if tok, _ := DefaultCache().Get(did, holdDID); tok != "" {
		t.Errorf("after InvalidateServiceToken, DefaultCache().Get() = %q, want empty", tok)
	}
}

// testServiceToken builds an unsigned JWT whose exp claim is expiresAt. Only
// the payload is meaningful: the cache reads exp without verifying anything.
func testServiceToken(expiresAt time.Time) string {
	payload := fmt.Sprintf(`{"exp":%d}`, expiresAt.Unix())
	return "header." + base64.RawURLEncoding.EncodeToString([]byte(payload)) + ".signature"
}

func TestSetServiceToken_AppliesSafetyMargin(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:margin"
	holdDID := "did:web:hold.example.com"

	realExp := time.Now().Add(5 * time.Minute)
	if err := SetServiceToken(did, holdDID, testServiceToken(realExp)); err != nil {
		t.Fatalf("SetServiceToken() error = %v", err)
	}

	_, expiresAt := GetServiceToken(did, holdDID)
	if expiresAt.IsZero() {
		t.Fatal("expected the token to be cached")
	}

	want := realExp.Add(-ServiceTokenSafetyMargin)
	if diff := expiresAt.Sub(want); diff < -2*time.Second || diff > 2*time.Second {
		t.Errorf("cached expiry off by %v (want exp minus %v)", diff, ServiceTokenSafetyMargin)
	}

	// The point of the margin: the cache must stop serving the token at least
	// distribution's 60s JWT leeway before the hold would reject it.
	if got := realExp.Sub(expiresAt); got < 60*time.Second {
		t.Errorf("cache serves the token until %v before its real exp, want >= 60s", got)
	}
}

func TestSetServiceToken_ShortGrantKeepsPositiveTTL(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:shortgrant"
	holdDID := "did:web:hold.example.com"

	// A PDS that grants far less than the safety margin. Subtracting the margin
	// outright would cache an already-expired entry and make every request
	// re-mint, so the cache keeps half the remaining life instead.
	realExp := time.Now().Add(20 * time.Second)
	if err := SetServiceToken(did, holdDID, testServiceToken(realExp)); err != nil {
		t.Fatalf("SetServiceToken() error = %v", err)
	}

	token, expiresAt := GetServiceToken(did, holdDID)
	if token == "" {
		t.Fatal("short-lived token should still be cached, not dropped on sight")
	}
	if !expiresAt.After(time.Now()) {
		t.Fatalf("cached expiry %v is not in the future", expiresAt)
	}
	if !expiresAt.Before(realExp) {
		t.Errorf("cached expiry %v should be before the token's real exp %v", expiresAt, realExp)
	}
}

func TestSetServiceToken_UnparsableExpUsesFallbackTTL(t *testing.T) {
	defaultCache.Clear()

	did := "did:plc:unparsable"
	holdDID := "did:web:hold.example.com"

	if err := SetServiceToken(did, holdDID, "not-a-jwt"); err != nil {
		t.Fatalf("SetServiceToken() error = %v", err)
	}

	_, expiresAt := GetServiceToken(did, holdDID)
	want := time.Now().Add(unparsableTokenTTL)
	if diff := expiresAt.Sub(want); diff < -5*time.Second || diff > 5*time.Second {
		t.Errorf("expiry off by %v (want ~%v from now)", diff, unparsableTokenTTL)
	}
}

func TestServiceTokenExpiry(t *testing.T) {
	want := time.Now().Add(3 * time.Minute).Truncate(time.Second)

	got, err := ServiceTokenExpiry(testServiceToken(want))
	if err != nil {
		t.Fatalf("ServiceTokenExpiry() error = %v", err)
	}
	if !got.Equal(want) {
		t.Errorf("ServiceTokenExpiry() = %v, want %v", got, want)
	}

	if _, err := ServiceTokenExpiry("not-a-jwt"); err == nil {
		t.Error("expected an error for a token that is not a JWT")
	}
}