anchorage/internal/pkg/token/rotation_test.go

package token_test

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	jwtPkg "github.com/golang-jwt/jwt/v5"

	"anchorage/internal/pkg/ids"
	"anchorage/internal/pkg/token"
)

// keyBytes produces a deterministic 32-byte key. Using fmt'd values
// rather than crypto/rand keeps the test reproducible.
func keyBytes(seed byte) []byte {
	b := make([]byte, 48)
	for i := range b {
		b[i] = seed
	}
	return b
}

const (
	issuer = "https://anchorage.test/"
	aud    = "anchorage"
)

// TestRotationOverlap walks the canonical rotation procedure:
//
//  1. Start with signer_A (key_A primary).
//  2. Add key_B as secondary → signer_A+B (A primary, B loaded).
//  3. Flip primary to B → signer_A+B (B primary, A retained).
//  4. Drop key_A → signer_B (only B).
//
// Tokens minted at each stage must keep verifying at least until their
// minting key is dropped, and new mints must use the primary key's
// kid.
func TestRotationOverlap(t *testing.T) {
	keyA := keyBytes(0xA0)
	keyB := keyBytes(0xB0)

	makeSigner := func(t *testing.T, keys []token.SigningKey) *token.Signer {
		t.Helper()
		s, err := token.NewSigner(keys, issuer, aud, nil)
		if err != nil {
			t.Fatalf("NewSigner: %v", err)
		}
		return s
	}
	mintHere := func(t *testing.T, s *token.Signer) string {
		t.Helper()
		raw, _, err := s.Mint(context.Background(), ids.MustNewOrg(), ids.MustNewUser(), "member", nil, time.Hour)
		if err != nil {
			t.Fatalf("Mint: %v", err)
		}
		return raw
	}
	parseOK := func(t *testing.T, s *token.Signer, raw, label string) {
		t.Helper()
		if _, err := s.Parse(context.Background(), raw); err != nil {
			t.Errorf("%s should parse: %v", label, err)
		}
	}
	parseFail := func(t *testing.T, s *token.Signer, raw, label string) {
		t.Helper()
		if _, err := s.Parse(context.Background(), raw); err == nil {
			t.Errorf("%s should NOT parse", label)
		}
	}

	// --- stage 1: only A ---
	stage1 := makeSigner(t, []token.SigningKey{
		{ID: "A", Key: keyA, Primary: true},
	})
	tokenA := mintHere(t, stage1)
	if stage1.PrimaryKeyID() != "A" {
		t.Fatalf("stage1 primary = %q", stage1.PrimaryKeyID())
	}

	// --- stage 2: A primary, B added as secondary ---
	stage2 := makeSigner(t, []token.SigningKey{
		{ID: "A", Key: keyA, Primary: true},
		{ID: "B", Key: keyB},
	})
	parseOK(t, stage2, tokenA, "stage2: old A-minted token")
	tokenA2 := mintHere(t, stage2)
	// Both stage2 tokens are still A-minted (primary unchanged).
	parseOK(t, stage2, tokenA2, "stage2: new mint still on key A")

	// --- stage 3: B primary, A retained ---
	stage3 := makeSigner(t, []token.SigningKey{
		{ID: "A", Key: keyA},
		{ID: "B", Key: keyB, Primary: true},
	})
	parseOK(t, stage3, tokenA, "stage3: A-era token still verifies")
	tokenB := mintHere(t, stage3)
	parseOK(t, stage3, tokenB, "stage3: new B-minted token")
	if stage3.PrimaryKeyID() != "B" {
		t.Fatalf("stage3 primary = %q", stage3.PrimaryKeyID())
	}

	// --- stage 4: drop A ---
	stage4 := makeSigner(t, []token.SigningKey{
		{ID: "B", Key: keyB, Primary: true},
	})
	parseFail(t, stage4, tokenA, "stage4: key A dropped → A-minted tokens rejected")
	parseOK(t, stage4, tokenB, "stage4: B-minted still verifies")
}

// TestRotationRejectsUnknownKid confirms an attacker who knows the
// mint algorithm can't trick anchorage into accepting a kid that isn't
// configured. Silent "try all keys" would defeat the whole feature.
func TestRotationRejectsUnknownKid(t *testing.T) {
	s, _ := token.NewSigner([]token.SigningKey{
		{ID: "real-key", Key: keyBytes(0x11), Primary: true},
	}, issuer, aud, nil)

	// Mint a token with a signer that happens to use the same bytes but
	// a DIFFERENT kid. Verification against the first signer must fail.
	attacker, _ := token.NewSigner([]token.SigningKey{
		{ID: "attacker-key", Key: keyBytes(0x11), Primary: true},
	}, issuer, aud, nil)
	raw, _, err := attacker.Mint(context.Background(), ids.MustNewOrg(), ids.MustNewUser(), "sysadmin", nil, time.Hour)
	if err != nil {
		t.Fatalf("attacker mint: %v", err)
	}
	_, err = s.Parse(context.Background(), raw)
	if err == nil {
		t.Fatal("expected unknown-kid token to be rejected")
	}
	if !strings.Contains(err.Error(), "unknown kid") {
		t.Errorf("error should mention unknown kid, got: %v", err)
	}
}

// TestRotationRejectsMissingKid confirms that a token without a `kid`
// header is rejected outright. anchorage is new software — there is
// no pre-rotation token population to accept via a fallback.
func TestRotationRejectsMissingKid(t *testing.T) {
	s, err := token.NewSigner([]token.SigningKey{
		{ID: "A", Key: keyBytes(0xAA), Primary: true},
	}, issuer, aud, nil)
	if err != nil {
		t.Fatalf("NewSigner: %v", err)
	}
	// Hand-craft a JWT signed with the right bytes but no `kid` header.
	raw := jwtNoKid(t, keyBytes(0xAA))
	_, err = s.Parse(context.Background(), raw)
	if err == nil {
		t.Fatal("expected missing-kid token to be rejected")
	}
	if !strings.Contains(err.Error(), "kid") {
		t.Errorf("error should mention kid, got: %v", err)
	}
}

// jwtNoKid mints a minimally-valid HS256 JWT without a `kid` header,
// used by TestRotationRejectsMissingKid to prove Parse rejects the
// pre-rotation shape.
func jwtNoKid(t *testing.T, key []byte) string {
	t.Helper()
	tok := jwtPkg.NewWithClaims(jwtPkg.SigningMethodHS256, &token.Claims{
		RegisteredClaims: jwtPkg.RegisteredClaims{
			ID:        ids.MustNewToken().String(),
			Issuer:    issuer,
			Audience:  jwtPkg.ClaimStrings{aud},
			IssuedAt:  jwtPkg.NewNumericDate(time.Now()),
			NotBefore: jwtPkg.NewNumericDate(time.Now()),
			ExpiresAt: jwtPkg.NewNumericDate(time.Now().Add(time.Hour)),
		},
	})
	// Deliberately do NOT set tok.Header["kid"].
	signed, err := tok.SignedString(key)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	return signed
}

// TestNewSignerRejectsInvalidConfig covers the validation surface.
func TestNewSignerRejectsInvalidConfig(t *testing.T) {
	tests := []struct {
		name string
		keys []token.SigningKey
	}{
		{"empty", nil},
		{"no-primary", []token.SigningKey{{ID: "A", Key: keyBytes(0xAA)}}},
		{"two-primaries", []token.SigningKey{
			{ID: "A", Key: keyBytes(0xAA), Primary: true},
			{ID: "B", Key: keyBytes(0xBB), Primary: true},
		}},
		{"duplicate-id", []token.SigningKey{
			{ID: "A", Key: keyBytes(0xAA), Primary: true},
			{ID: "A", Key: keyBytes(0xBB)},
		}},
		{"empty-id", []token.SigningKey{
			{ID: "", Key: keyBytes(0xAA), Primary: true},
		}},
		{"short-key", []token.SigningKey{
			{ID: "A", Key: []byte("too-short"), Primary: true},
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if _, err := token.NewSigner(tt.keys, issuer, aud, nil); err == nil {
				t.Error("expected error")
			}
		})
	}
}

// TestLoadSigningKeysEmptyReturnsDevFallback confirms zero-config runs
// land on the built-in dev key.
func TestLoadSigningKeysEmptyReturnsDevFallback(t *testing.T) {
	keys, err := token.LoadSigningKeys(nil)
	if err != nil {
		t.Fatalf("LoadSigningKeys: %v", err)
	}
	if len(keys) != 1 || keys[0].ID != token.DevKeyID || !keys[0].Primary {
		t.Fatalf("empty load should produce single primary DevKeyID entry, got %+v", keys)
	}
}

// TestLoadSigningKeysRotation verifies multi-key loading from config.
func TestLoadSigningKeysRotation(t *testing.T) {
	dir := t.TempDir()
	writeKey := func(name string, seed byte) string {
		p := filepath.Join(dir, name)
		if err := os.WriteFile(p, keyBytes(seed), 0o600); err != nil {
			t.Fatalf("write %s: %v", name, err)
		}
		return p
	}
	pathA := writeKey("jwt.a.key", 0xAA)
	pathB := writeKey("jwt.b.key", 0xBB)

	keys, err := token.LoadSigningKeys([]token.KeyFileSpec{
		{ID: "2025-10", Path: pathA},
		{ID: "2026-04", Path: pathB, Primary: true},
	})
	if err != nil {
		t.Fatalf("LoadSigningKeys: %v", err)
	}
	if len(keys) != 2 {
		t.Fatalf("want 2 keys, got %d", len(keys))
	}
	if keys[1].ID != "2026-04" || !keys[1].Primary {
		t.Errorf("primary entry mis-marked: %+v", keys)
	}
}

// TestLoadSigningKeysRejectsNoPrimary makes sure a malformed config
// can't silently produce a signer that can't mint.
func TestLoadSigningKeysRejectsNoPrimary(t *testing.T) {
	dir := t.TempDir()
	p := filepath.Join(dir, "k")
	_ = os.WriteFile(p, keyBytes(0x55), 0o600)
	if _, err := token.LoadSigningKeys([]token.KeyFileSpec{
		{ID: "A", Path: p}, {ID: "B", Path: p},
	}); err == nil {
		t.Error("expected error when no primary is set")
	}
}