llmgw/internal/store/claims_test.go

package store

import (
	"context"
	"errors"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

func claimFixture(t *testing.T) (*DB, string, string) {
	t.Helper()
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "c.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	tree(t, s)

	secret, k, err := NewKey("key1", "n", "proj1")
	if err != nil {
		t.Fatal(err)
	}
	k.CreatedAt = time.Now()
	if err := s.CreateKey(ctx, k); err != nil {
		t.Fatal(err)
	}

	token, err := NewClaimToken()
	if err != nil {
		t.Fatal(err)
	}
	if err := s.CreateClaim(ctx, k.ID, token, secret, ClaimTTL); err != nil {
		t.Fatal(err)
	}
	return s, token, secret
}

func TestAClaimDeliversTheKeyExactlyOnce(t *testing.T) {
	s, token, secret := claimFixture(t)
	ctx := context.Background()

	if _, err := s.PeekClaim(ctx, token); err != nil {
		t.Fatalf("a fresh link is not valid: %v", err)
	}
	got, err := s.ConsumeClaim(ctx, token)
	if err != nil {
		t.Fatal(err)
	}
	if got != secret {
		t.Fatal("the delivered key is not the one that was issued")
	}

	if _, err := s.ConsumeClaim(ctx, token); !errors.Is(err, ErrClaimed) {
		t.Errorf("the link worked twice: %v", err)
	}
	if _, err := s.PeekClaim(ctx, token); !errors.Is(err, ErrClaimed) {
		t.Errorf("a used link still looks valid: %v", err)
	}
}

func TestTheDatabaseNeverHoldsAUsableKey(t *testing.T) {
	// The property that makes a pending claim safe to store at all: what is in
	// the table cannot be opened by anything that can read the table.
	s, token, secret := claimFixture(t)

	var sealed []byte
	if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
		t.Fatal(err)
	}
	if len(sealed) == 0 {
		t.Fatal("nothing was stored")
	}
	if string(sealed) == secret {
		t.Fatal("the key is sitting in the database in the clear")
	}
	// And the token itself is not there either, so nothing in the database can
	// derive the key that opens it.
	var count int
	if err := s.db.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM key_claims WHERE token_hash = ?`, []byte(token)).
		Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 0 {
		t.Fatal("the claim token was stored, not just its hash")
	}

	if _, err := open("wrong token entirely", sealed); !errors.Is(err, ErrSealed) {
		t.Errorf("a wrong token opened the blob: %v", err)
	}
}

func TestAClaimedLinkLeavesNothingBehind(t *testing.T) {
	s, token, _ := claimFixture(t)
	if _, err := s.ConsumeClaim(context.Background(), token); err != nil {
		t.Fatal(err)
	}
	var sealed []byte
	if err := s.db.QueryRowContext(context.Background(), `SELECT sealed FROM key_claims`).Scan(&sealed); err != nil {
		t.Fatal(err)
	}
	if len(sealed) != 0 {
		t.Errorf("a used claim still holds %d bytes", len(sealed))
	}
}

func TestOnlyOneOfTwoSimultaneousClaimsWins(t *testing.T) {
	// Two people opening the same link at the same moment must produce one key,
	// not two copies of it.
	s, token, _ := claimFixture(t)

	var wg sync.WaitGroup
	var mu sync.Mutex
	won := 0
	for range 16 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if _, err := s.ConsumeClaim(context.Background(), token); err == nil {
				mu.Lock()
				won++
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
	if won != 1 {
		t.Errorf("%d callers each received the key", won)
	}
}

func TestAnExpiredLinkIsRefused(t *testing.T) {
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "c.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	tree(t, s)

	secret, k, _ := NewKey("key1", "n", "proj1")
	k.CreatedAt = time.Now()
	must(t, s.CreateKey(ctx, k))

	token, _ := NewClaimToken()
	must(t, s.CreateClaim(ctx, k.ID, token, secret, ClaimTTL))
	// Age it directly. CreateClaim treats a non-positive TTL as "use the
	// default", which is the right thing for a caller and no way to build the
	// row this test needs.
	if _, err := s.db.ExecContext(context.Background(), `UPDATE key_claims SET expires_at = ?`,
		time.Now().Add(-time.Hour).UnixMilli()); err != nil {
		t.Fatal(err)
	}

	if _, err := s.PeekClaim(ctx, token); !errors.Is(err, ErrClaimed) {
		t.Errorf("an expired link looks valid: %v", err)
	}
	if _, err := s.ConsumeClaim(ctx, token); !errors.Is(err, ErrClaimed) {
		t.Errorf("an expired link delivered a key: %v", err)
	}
}

func TestAnUnknownLinkLooksLikeAUsedOne(t *testing.T) {
	// Distinguishing them would confirm that a token was once valid.
	s, token, _ := claimFixture(t)
	ctx := context.Background()
	must(t, func() error { _, err := s.ConsumeClaim(ctx, token); return err }())

	unknown, _ := NewClaimToken()
	_, usedErr := s.PeekClaim(ctx, token)
	_, unknownErr := s.PeekClaim(ctx, unknown)
	if usedErr.Error() != unknownErr.Error() {
		t.Errorf("used says %q, unknown says %q", usedErr, unknownErr)
	}
}

func TestPurgeRemovesWhatCannotBeUsed(t *testing.T) {
	s, token, _ := claimFixture(t)
	ctx := context.Background()
	must(t, func() error { _, err := s.ConsumeClaim(ctx, token); return err }())

	n, err := s.PurgeClaims(ctx, time.Now())
	if err != nil {
		t.Fatal(err)
	}
	if n != 1 {
		t.Errorf("purged %d rows", n)
	}
}

func TestSealRoundTrips(t *testing.T) {
	token, _ := NewClaimToken()
	sealed, err := seal(token, []byte("sk-llmgw-secret"))
	if err != nil {
		t.Fatal(err)
	}
	got, err := open(token, sealed)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != "sk-llmgw-secret" {
		t.Errorf("round trip gave %q", got)
	}
	// Tampering must fail rather than produce something.
	sealed[len(sealed)-1] ^= 0xff
	if _, err := open(token, sealed); !errors.Is(err, ErrSealed) {
		t.Errorf("a tampered blob opened: %v", err)
	}
}