iop/apps/edge/internal/authprojection/cache_test.go

package authprojection

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	iop "iop/proto/gen/iop"
)

func projectionFixture(generation uint64, now time.Time, ttl time.Duration, rawToken string) *iop.PrincipalProjection {
	projection := &iop.PrincipalProjection{
		Generation:        generation,
		IssuedAtUnixNano:  now.UnixNano(),
		ExpiresAtUnixNano: now.Add(ttl).UnixNano(),
	}
	if rawToken != "" {
		digest := sha256.Sum256([]byte(rawToken))
		projection.Tokens = []*iop.ProjectedPrincipalToken{{
			TokenDigestSha256: hex.EncodeToString(digest[:]),
			PrincipalRef:      "principal-1",
			PrincipalAlias:    "principal-one",
			TokenRef:          "token-1",
			TokenRevision:     generation,
		}}
		projection.Routes = []*iop.ProjectedPrincipalRoute{{
			RouteId: "route-1", RouteAlias: "model-1", PrincipalRef: "principal-1",
			CredentialSlotRef: "slot-1", ProfileId: "openai", UpstreamModel: "upstream-model",
			ResourceSelector: "default", RouteRevision: generation, CredentialRevision: generation,
		}}
	}
	return projection
}

func TestCacheAppliesOnlyHigherFreshGeneration(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	if state := cache.State(); state != StateUnmanaged {
		t.Fatalf("initial state: got %s", state)
	}
	if err := cache.Apply(projectionFixture(4, now, time.Minute, "token-four")); err != nil {
		t.Fatal(err)
	}
	if generation, ok := cache.Generation(); !ok || generation != 4 {
		t.Fatalf("generation: got %d ok=%v", generation, ok)
	}
	if err := cache.Apply(projectionFixture(4, now, time.Minute, "other-token")); !errors.Is(err, ErrGenerationConflict) {
		t.Fatalf("changed equal generation: got %v", err)
	}
	if err := cache.Apply(projectionFixture(3, now, time.Minute, "other-token")); !errors.Is(err, ErrGenerationNotHigher) {
		t.Fatalf("lower generation: got %v", err)
	}
	if err := cache.Apply(projectionFixture(5, now.Add(-2*time.Minute), time.Minute, "expired")); !errors.Is(err, ErrProjectionExpired) {
		t.Fatalf("expired higher generation: got %v", err)
	}
	if err := cache.Apply(projectionFixture(5, now.Add(time.Second), time.Minute, "future")); !errors.Is(err, ErrInvalidProjection) {
		t.Fatalf("not-yet-issued higher generation: got %v", err)
	}

	digest := sha256.Sum256([]byte("token-four"))
	principal, state, ok := cache.LookupDigest(digest)
	if !ok || state != StateFresh || principal.PrincipalRef != "principal-1" || principal.TokenRef != "token-1" {
		t.Fatalf("lookup: principal=%+v state=%s ok=%v", principal, state, ok)
	}
	if generation, _ := cache.Generation(); generation != 4 {
		t.Fatalf("rejected candidates changed generation to %d", generation)
	}
}

func TestCacheRenewsOnlyIdenticalGeneration(t *testing.T) {
	start := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	var clock atomic.Int64
	clock.Store(start.UnixNano())
	cache := NewCache(DefaultLimits(), func() time.Time { return time.Unix(0, clock.Load()) })
	if err := cache.Apply(projectionFixture(7, start, time.Minute, "renewed-token")); err != nil {
		t.Fatal(err)
	}
	original := cache.current.Load()

	if err := cache.Apply(projectionFixture(7, start, time.Minute, "renewed-token")); err != nil {
		t.Fatalf("idempotent duplicate: %v", err)
	}
	if got := cache.current.Load(); got != original {
		t.Fatal("identical duplicate replaced the installed snapshot")
	}
	if err := cache.Apply(projectionFixture(7, start.Add(-10*time.Second), 40*time.Second, "renewed-token")); err != nil {
		t.Fatalf("older freshness window: %v", err)
	}
	if got := cache.current.Load(); got != original {
		t.Fatal("older freshness window replaced the installed snapshot")
	}

	clock.Store(start.Add(10 * time.Second).UnixNano())
	if err := cache.Apply(projectionFixture(7, start.Add(10*time.Second), time.Minute, "changed-token")); !errors.Is(err, ErrGenerationConflict) {
		t.Fatalf("changed equal generation: got %v", err)
	}

	clock.Store(start.Add(20 * time.Second).UnixNano())
	if err := cache.Apply(projectionFixture(7, start.Add(20*time.Second), time.Minute, "renewed-token")); err != nil {
		t.Fatalf("renewed equal generation: %v", err)
	}
	renewed := cache.current.Load()
	if renewed == original {
		t.Fatal("later freshness window did not replace the installed snapshot")
	}
	if !renewed.issuedAt.After(original.issuedAt) || !renewed.expiresAt.After(original.expiresAt) {
		t.Fatalf("freshness window did not advance: original=%s..%s renewed=%s..%s", original.issuedAt, original.expiresAt, renewed.issuedAt, renewed.expiresAt)
	}
}

func TestCacheRecoversExpiredProjectionWithIdenticalGenerationRenewal(t *testing.T) {
	start := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	var clock atomic.Int64
	clock.Store(start.UnixNano())
	cache := NewCache(DefaultLimits(), func() time.Time { return time.Unix(0, clock.Load()) })
	if err := cache.Apply(projectionFixture(9, start, time.Minute, "recovery-token")); err != nil {
		t.Fatal(err)
	}
	recoveredAt := start.Add(2 * time.Minute)
	clock.Store(recoveredAt.UnixNano())
	if state := cache.State(); state != StateExpired {
		t.Fatalf("state before renewal: got %s", state)
	}
	if err := cache.Apply(projectionFixture(9, recoveredAt, time.Minute, "recovery-token")); err != nil {
		t.Fatalf("same-generation expiry recovery: %v", err)
	}
	if state := cache.State(); state != StateFresh {
		t.Fatalf("state after renewal: got %s", state)
	}
}

func TestCacheExpiryKeepsManagedFailClosed(t *testing.T) {
	start := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	var clock atomic.Int64
	clock.Store(start.UnixNano())
	cache := NewCache(DefaultLimits(), func() time.Time { return time.Unix(0, clock.Load()) })
	if err := cache.Apply(projectionFixture(1, start, time.Minute, "expiring-token")); err != nil {
		t.Fatal(err)
	}
	clock.Store(start.Add(time.Minute).UnixNano())
	if state := cache.State(); state != StateExpired {
		t.Fatalf("expired state: got %s", state)
	}
	digest := sha256.Sum256([]byte("expiring-token"))
	if _, state, ok := cache.LookupDigest(digest); ok || state != StateExpired {
		t.Fatalf("expired lookup: state=%s ok=%v", state, ok)
	}

	recoveredAt := start.Add(2 * time.Minute)
	clock.Store(recoveredAt.UnixNano())
	if err := cache.Apply(projectionFixture(2, recoveredAt, time.Minute, "replacement-token")); err != nil {
		t.Fatal(err)
	}
	if state := cache.State(); state != StateFresh {
		t.Fatalf("recovered state: got %s", state)
	}
}

func TestCacheRevocationSnapshotRemovesVerifier(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(projectionFixture(1, now, time.Minute, "revoked-token")); err != nil {
		t.Fatal(err)
	}
	if err := cache.Apply(projectionFixture(2, now, time.Minute, "")); err != nil {
		t.Fatal(err)
	}
	digest := sha256.Sum256([]byte("revoked-token"))
	if _, state, ok := cache.LookupDigest(digest); ok || state != StateFresh {
		t.Fatalf("revoked lookup: state=%s ok=%v", state, ok)
	}
	if state := cache.State(); state != StateFresh {
		t.Fatalf("revocation snapshot must remain managed and fresh: %s", state)
	}
}

func TestCacheRejectsOversizedSnapshot(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	projection := projectionFixture(1, now, time.Minute, "token-one")
	second := projectionFixture(1, now, time.Minute, "token-two").GetTokens()[0]
	second.TokenRef = "token-2"
	second.PrincipalRef = "principal-2"
	projection.Tokens = append(projection.Tokens, second)

	cache := NewCache(Limits{MaxTokens: 1, MaxRoutes: 4, MaxBytes: 4096}, func() time.Time { return now })
	if err := cache.Apply(projection); !errors.Is(err, ErrProjectionTooLarge) {
		t.Fatalf("token limit: got %v", err)
	}
	byteCache := NewCache(Limits{MaxTokens: 4, MaxRoutes: 4, MaxBytes: 1}, func() time.Time { return now })
	if err := byteCache.Apply(projectionFixture(1, now, time.Minute, "token-one")); !errors.Is(err, ErrProjectionTooLarge) {
		t.Fatalf("byte limit: got %v", err)
	}
}

func TestCacheDeepCopiesAppliedProjection(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	projection := projectionFixture(1, now, time.Minute, "immutable-token")
	if err := cache.Apply(projection); err != nil {
		t.Fatal(err)
	}
	projection.Tokens[0].PrincipalRef = "mutated"
	projection.Tokens[0].TokenDigestSha256 = strings.Repeat("0", sha256.Size*2)
	projection.Routes[0].RouteAlias = "mutated"
	projection.ExpiresAtUnixNano = now.Add(-time.Minute).UnixNano()

	digest := sha256.Sum256([]byte("immutable-token"))
	principal, state, ok := cache.LookupDigest(digest)
	if !ok || state != StateFresh || principal.PrincipalRef != "principal-1" {
		t.Fatalf("mutated input changed cache: principal=%+v state=%s ok=%v", principal, state, ok)
	}
	routes, state := cache.RoutesForPrincipal("principal-1")
	if state != StateFresh || len(routes) != 1 || routes[0].RouteAlias != "model-1" {
		t.Fatalf("mutated route changed cache: routes=%+v state=%s", routes, state)
	}
}

func TestAuthenticatedViewRetainsOneCopiedGeneration(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(projectionFixture(1, now, time.Hour, "view-token")); err != nil {
		t.Fatal(err)
	}
	digest := sha256.Sum256([]byte("view-token"))
	view, ok := cache.AuthenticatedView(digest)
	if !ok || view.Generation != 1 || len(view.Routes) != 1 {
		t.Fatalf("view: %+v ok=%v", view, ok)
	}
	view.Routes[0].RouteID = "caller-mutation"
	if err := cache.Apply(projectionFixture(2, now, time.Hour, "next-token")); err != nil {
		t.Fatal(err)
	}
	if view.Generation != 1 || view.Routes[0].RouteID != "caller-mutation" {
		t.Fatalf("view changed after swap: %+v", view)
	}
	routes, _ := cache.RoutesForPrincipal("principal-1")
	if len(routes) != 1 || routes[0].RouteID != "route-1" {
		t.Fatalf("view mutation changed cache routes: %+v", routes)
	}
}

func TestCacheConcurrentLookupAndApply(t *testing.T) {
	now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(projectionFixture(0, now, time.Hour, "token-0")); err != nil {
		t.Fatal(err)
	}

	var wg sync.WaitGroup
	for reader := 0; reader < 8; reader++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 500; i++ {
				digest := sha256.Sum256([]byte("token-0"))
				_, state, _ := cache.LookupDigest(digest)
				if state == StateUnmanaged {
					t.Errorf("managed cache returned unmanaged state")
					return
				}
				_, _ = cache.RoutesForPrincipal("principal-1")
			}
		}()
	}
	for writer := 0; writer < 4; writer++ {
		writer := writer
		wg.Add(1)
		go func() {
			defer wg.Done()
			for generation := uint64(writer + 1); generation <= 200; generation += 4 {
				err := cache.Apply(projectionFixture(generation, now, time.Hour, "token-0"))
				if err != nil && !errors.Is(err, ErrGenerationNotHigher) {
					t.Errorf("apply generation %d: %v", generation, err)
					return
				}
			}
		}()
	}
	wg.Wait()
	if generation, ok := cache.Generation(); !ok || generation != 200 {
		t.Fatalf("final generation: got %d ok=%v", generation, ok)
	}
}

func TestCacheConcurrentSameGenerationRenewalAndReaders(t *testing.T) {
	start := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC)
	now := start.Add(5 * time.Minute)
	cache := NewCache(DefaultLimits(), func() time.Time { return now })
	if err := cache.Apply(projectionFixture(11, start, 10*time.Minute, "stable-token")); err != nil {
		t.Fatal(err)
	}

	var wg sync.WaitGroup
	for reader := 0; reader < 8; reader++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			digest := sha256.Sum256([]byte("stable-token"))
			for range 500 {
				if _, state, ok := cache.LookupDigest(digest); !ok || state != StateFresh {
					t.Errorf("renewal reader observed state=%s ok=%v", state, ok)
					return
				}
			}
		}()
	}
	for writer := 0; writer < 4; writer++ {
		writer := writer
		wg.Add(1)
		go func() {
			defer wg.Done()
			for offset := writer + 1; offset <= 100; offset += 4 {
				issuedAt := start.Add(time.Duration(offset) * time.Second)
				if err := cache.Apply(projectionFixture(11, issuedAt, 10*time.Minute, "stable-token")); err != nil {
					t.Errorf("same-generation renewal at %s: %v", issuedAt, err)
					return
				}
			}
		}()
	}
	wg.Wait()
	installed := cache.current.Load()
	wantIssuedAt := start.Add(100 * time.Second)
	if !installed.issuedAt.Equal(wantIssuedAt) {
		t.Fatalf("final renewal issued_at: got %s want %s", installed.issuedAt, wantIssuedAt)
	}
}