at-container-registry/pkg/auth/token/issuer_test.go

package token

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"atcr.io/pkg/auth"
	"github.com/golang-jwt/jwt/v5"
)

// Shared test key to avoid generating a new RSA key for each test
// Generating a 2048-bit RSA key takes ~0.15s, so reusing one key saves significant time
var (
	issuerSharedTestKeyPath string
	issuerSharedTestKeyOnce sync.Once
	issuerSharedTestKeyDir  string
)

// getSharedTestKey returns a shared RSA key and its file path for all tests
// The key is generated once and reused across all tests in this package
func getIssuerSharedTestKey(t *testing.T) string {
	issuerSharedTestKeyOnce.Do(func() {
		// Create a persistent temp directory for the shared key
		var err error
		issuerSharedTestKeyDir, err = os.MkdirTemp("", "atcr-issuer-test-keys-*")
		if err != nil {
			t.Fatalf("Failed to create test key directory: %v", err)
		}

		issuerSharedTestKeyPath = filepath.Join(issuerSharedTestKeyDir, "test-key.pem")

		// Generate the key once (this is the expensive operation we want to avoid repeating)
		_, err = NewIssuer(issuerSharedTestKeyPath, "atcr.io", "registry", 15*time.Minute)
		if err != nil {
			t.Fatalf("Failed to generate shared test key: %v", err)
		}
	})

	return issuerSharedTestKeyPath
}

func TestNewIssuer_GeneratesKey(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "private-key.pem")

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	if issuer == nil {
		t.Fatal("Expected non-nil issuer")
	}

	// Verify key file was created
	if _, err := os.Stat(keyPath); os.IsNotExist(err) {
		t.Error("Expected private key file to be created")
	}

	// Verify certificate file was created
	certPath := filepath.Join(tmpDir, "private-key.crt")
	if _, err := os.Stat(certPath); os.IsNotExist(err) {
		t.Error("Expected certificate file to be created")
	}

	// Verify key file permissions (should be 0600)
	info, err := os.Stat(keyPath)
	if err != nil {
		t.Fatalf("Failed to stat key file: %v", err)
	}
	mode := info.Mode()
	if mode.Perm() != 0600 {
		t.Errorf("Expected key file permissions 0600, got %04o", mode.Perm())
	}

	// Verify issuer fields
	if issuer.issuer != "atcr.io" {
		t.Errorf("Expected issuer %q, got %q", "atcr.io", issuer.issuer)
	}

	if issuer.service != "registry" {
		t.Errorf("Expected service %q, got %q", "registry", issuer.service)
	}

	if issuer.expiration != 15*time.Minute {
		t.Errorf("Expected expiration %v, got %v", 15*time.Minute, issuer.expiration)
	}

	if issuer.privateKey == nil {
		t.Error("Expected private key to be set")
	}

	if issuer.publicKey == nil {
		t.Error("Expected public key to be set")
	}

	if issuer.certificate == nil {
		t.Error("Expected certificate to be set")
	}
}

func TestNewIssuer_LoadsExistingKey(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "private-key.pem")

	// First create - generates key
	issuer1, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("First NewIssuer() error = %v", err)
	}

	// Second create - should load existing key
	issuer2, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("Second NewIssuer() error = %v", err)
	}

	// Compare public keys - should be the same
	if issuer1.publicKey.N.Cmp(issuer2.publicKey.N) != 0 {
		t.Error("Expected same public key when loading existing key")
	}
	if issuer1.publicKey.E != issuer2.publicKey.E {
		t.Error("Expected same public key exponent when loading existing key")
	}
}

func TestIssuer_Issue(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	subject := "did:plc:user123"
	access := []auth.AccessEntry{
		{
			Type:    "repository",
			Name:    "alice/myapp",
			Actions: []string{"pull", "push"},
		},
	}

	token, err := issuer.Issue(subject, access, AuthMethodOAuth)
	if err != nil {
		t.Fatalf("Issue() error = %v", err)
	}

	if token == "" {
		t.Fatal("Expected non-empty token")
	}

	// Token should be a JWT (3 parts separated by dots)
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		t.Errorf("Expected JWT with 3 parts, got %d parts", len(parts))
	}
}

func TestIssuer_IssueWithExpiration_HonorsCallerDuration(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	// Issuer baked-in expiration is 15 min; per-call should override.
	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	tokenString, err := issuer.IssueWithExpiration(
		"did:plc:user123",
		[]auth.AccessEntry{{Type: "repository", Name: "alice/myapp", Actions: []string{"pull"}}},
		AuthMethodOAuth,
		2*time.Minute,
	)
	if err != nil {
		t.Fatalf("IssueWithExpiration() error = %v", err)
	}

	parts := strings.Split(tokenString, ".")
	if len(parts) != 3 {
		t.Fatalf("expected 3 JWT parts, got %d", len(parts))
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		t.Fatalf("decode payload: %v", err)
	}
	var claims struct {
		Exp int64 `json:"exp"`
		Iat int64 `json:"iat"`
	}
	if err := json.Unmarshal(payload, &claims); err != nil {
		t.Fatalf("unmarshal claims: %v", err)
	}
	delta := claims.Exp - claims.Iat
	if delta != 120 {
		t.Errorf("expected exp - iat = 120, got %d", delta)
	}
}

func TestIssuer_Issue_EmptyAccess(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	token, err := issuer.Issue("did:plc:user123", nil, AuthMethodOAuth)
	if err != nil {
		t.Fatalf("Issue() error = %v", err)
	}

	if token == "" {
		t.Fatal("Expected non-empty token even with nil access")
	}
}

func TestIssuer_Issue_ValidateToken(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	subject := "did:plc:user123"
	access := []auth.AccessEntry{
		{
			Type:    "repository",
			Name:    "alice/myapp",
			Actions: []string{"pull", "push"},
		},
	}

	tokenString, err := issuer.Issue(subject, access, AuthMethodOAuth)
	if err != nil {
		t.Fatalf("Issue() error = %v", err)
	}

	// Parse and validate the token
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (any, error) {
		return issuer.publicKey, nil
	})
	if err != nil {
		t.Fatalf("Failed to parse token: %v", err)
	}

	if !token.Valid {
		t.Error("Expected token to be valid")
	}

	claims, ok := token.Claims.(*Claims)
	if !ok {
		t.Fatal("Failed to cast claims to *Claims")
	}

	// Verify claims
	if claims.Subject != subject {
		t.Errorf("Expected subject %q, got %q", subject, claims.Subject)
	}

	if claims.Issuer != "atcr.io" {
		t.Errorf("Expected issuer %q, got %q", "atcr.io", claims.Issuer)
	}

	if len(claims.Audience) != 1 || claims.Audience[0] != "registry" {
		t.Errorf("Expected audience [%q], got %v", "registry", claims.Audience)
	}

	if len(claims.Access) != 1 {
		t.Errorf("Expected 1 access entry, got %d", len(claims.Access))
	}

	if len(claims.Access) > 0 {
		if claims.Access[0].Type != "repository" {
			t.Errorf("Expected type %q, got %q", "repository", claims.Access[0].Type)
		}
		if claims.Access[0].Name != "alice/myapp" {
			t.Errorf("Expected name %q, got %q", "alice/myapp", claims.Access[0].Name)
		}
		if len(claims.Access[0].Actions) != 2 {
			t.Errorf("Expected 2 actions, got %d", len(claims.Access[0].Actions))
		}
	}

	// Verify expiration is set and reasonable
	if claims.ExpiresAt == nil {
		t.Fatal("Expected ExpiresAt to be set")
	}

	expiresIn := time.Until(claims.ExpiresAt.Time)
	if expiresIn < 14*time.Minute || expiresIn > 16*time.Minute {
		t.Errorf("Expected expiration around 15 minutes, got %v", expiresIn)
	}
}

func TestIssuer_Issue_X5CHeader(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	tokenString, err := issuer.Issue("did:plc:user123", nil, "oauth")
	if err != nil {
		t.Fatalf("Issue() error = %v", err)
	}

	// Parse token to inspect header
	token, _, err := jwt.NewParser().ParseUnverified(tokenString, &Claims{})
	if err != nil {
		t.Fatalf("Failed to parse token: %v", err)
	}

	// Check x5c header exists
	x5c, ok := token.Header["x5c"]
	if !ok {
		t.Fatal("Expected x5c header in token")
	}

	// x5c should be a slice of base64-encoded certificates
	x5cSlice, ok := x5c.([]any)
	if !ok {
		t.Fatal("Expected x5c to be a slice")
	}

	if len(x5cSlice) != 1 {
		t.Errorf("Expected 1 certificate in x5c chain, got %d", len(x5cSlice))
	}

	// Decode and verify certificate
	certStr, ok := x5cSlice[0].(string)
	if !ok {
		t.Fatal("Expected certificate to be a string")
	}

	certBytes, err := base64.StdEncoding.DecodeString(certStr)
	if err != nil {
		t.Fatalf("Failed to decode certificate: %v", err)
	}

	// Parse certificate
	cert, err := x509.ParseCertificate(certBytes)
	if err != nil {
		t.Fatalf("Failed to parse certificate: %v", err)
	}

	// Verify certificate is self-signed and matches our public key
	if cert.Subject.CommonName != "ATCR Token Signing Certificate" {
		t.Errorf("Expected CN %q, got %q", "ATCR Token Signing Certificate", cert.Subject.CommonName)
	}

	// Verify certificate's public key matches issuer's public key
	certPubKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		t.Fatal("Expected RSA public key in certificate")
	}

	if certPubKey.N.Cmp(issuer.publicKey.N) != 0 {
		t.Error("Certificate public key doesn't match issuer public key")
	}
}

func TestIssuer_PublicKey(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	pubKey := issuer.PublicKey()
	if pubKey == nil {
		t.Fatal("Expected non-nil public key")
	}

	// Verify it's a valid RSA public key
	if pubKey.N == nil {
		t.Error("Expected public key modulus to be set")
	}

	if pubKey.E == 0 {
		t.Error("Expected public key exponent to be set")
	}
}

func TestIssuer_Expiration(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	expiration := 30 * time.Minute
	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", expiration)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	if issuer.Expiration() != expiration {
		t.Errorf("Expected expiration %v, got %v", expiration, issuer.Expiration())
	}
}

func TestIssuer_ConcurrentIssue(t *testing.T) {
	keyPath := getIssuerSharedTestKey(t)

	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	// Issue tokens concurrently
	const numGoroutines = 10
	var wg sync.WaitGroup

	tokens := make([]string, numGoroutines)
	errors := make([]error, numGoroutines)

	for i := range numGoroutines {
		wg.Go(func() {
			subject := "did:plc:user" + string(rune('0'+i))
			token, err := issuer.Issue(subject, nil, AuthMethodOAuth)
			tokens[i] = token
			errors[i] = err
		})
	}

	wg.Wait()

	// Verify all tokens were issued successfully
	for i, err := range errors {
		if err != nil {
			t.Errorf("Goroutine %d: Issue() error = %v", i, err)
		}
	}

	for i, token := range tokens {
		if token == "" {
			t.Errorf("Goroutine %d: Expected non-empty token", i)
		}
	}
}

func TestNewIssuer_InvalidCertificate(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "private-key.pem")

	// First generate key + cert
	_, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("First NewIssuer() error = %v", err)
	}

	// Corrupt the certificate file
	certPath := filepath.Join(tmpDir, "private-key.crt")
	err = os.WriteFile(certPath, []byte("invalid certificate data"), 0644)
	if err != nil {
		t.Fatalf("Failed to corrupt certificate: %v", err)
	}

	// Try to create issuer again - should fail
	_, err = NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err == nil {
		t.Error("Expected error when certificate is invalid")
	}

	if !strings.Contains(err.Error(), "certificate") {
		t.Errorf("Expected error message to mention certificate, got: %v", err)
	}
}

func TestNewIssuer_MissingCertificate(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "private-key.pem")

	// First generate key + cert
	_, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("First NewIssuer() error = %v", err)
	}

	// Delete certificate but keep key
	certPath := filepath.Join(tmpDir, "private-key.crt")
	err = os.Remove(certPath)
	if err != nil {
		t.Fatalf("Failed to remove certificate: %v", err)
	}

	// Try to create issuer - should regenerate certificate
	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() should regenerate certificate, got error: %v", err)
	}

	if issuer == nil {
		t.Fatal("Expected non-nil issuer")
	}

	// Verify certificate was regenerated
	if _, err := os.Stat(certPath); os.IsNotExist(err) {
		t.Error("Expected certificate to be regenerated")
	}
}

func TestLoadOrGenerateKey_InvalidPEM(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "invalid-key.pem")

	// Write invalid PEM data
	err := os.WriteFile(keyPath, []byte("not a valid PEM file"), 0600)
	if err != nil {
		t.Fatalf("Failed to write invalid PEM: %v", err)
	}

	// Try to load - should fail
	_, err = NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err == nil {
		t.Error("Expected error when loading invalid PEM")
	}
}

func TestGenerateCertificate_ValidCertificate(t *testing.T) {
	tmpDir := t.TempDir()
	keyPath := filepath.Join(tmpDir, "private-key.pem")
	certPath := filepath.Join(tmpDir, "private-key.crt")

	// Generate issuer (which generates key and cert)
	issuer, err := NewIssuer(keyPath, "atcr.io", "registry", 15*time.Minute)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}

	// Read and parse the certificate
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		t.Fatalf("Failed to read certificate: %v", err)
	}

	block, _ := pem.Decode(certPEM)
	if block == nil || block.Type != "CERTIFICATE" {
		t.Fatal("Failed to decode certificate PEM")
	}

	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		t.Fatalf("Failed to parse certificate: %v", err)
	}

	// Verify certificate properties
	if cert.Subject.CommonName != "ATCR Token Signing Certificate" {
		t.Errorf("Expected CN %q, got %q", "ATCR Token Signing Certificate", cert.Subject.CommonName)
	}

	if len(cert.Subject.Organization) == 0 || cert.Subject.Organization[0] != "ATCR" {
		t.Error("Expected Organization to be ATCR")
	}

	// Verify key usage
	if cert.KeyUsage&x509.KeyUsageDigitalSignature == 0 {
		t.Error("Expected certificate to have DigitalSignature key usage")
	}

	// Verify validity period (should be 10 years)
	validityPeriod := cert.NotAfter.Sub(cert.NotBefore)
	expectedPeriod := 10 * 365 * 24 * time.Hour
	if validityPeriod < expectedPeriod-24*time.Hour || validityPeriod > expectedPeriod+24*time.Hour {
		t.Errorf("Expected validity period around 10 years, got %v", validityPeriod)
	}

	// Verify certificate's public key matches issuer's public key
	certPubKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		t.Fatal("Expected RSA public key in certificate")
	}

	if certPubKey.N.Cmp(issuer.publicKey.N) != 0 {
		t.Error("Certificate public key doesn't match issuer public key")
	}

	// Verify certificate is self-signed
	if err := cert.CheckSignature(cert.SignatureAlgorithm, cert.RawTBSCertificate, cert.Signature); err != nil {
		t.Errorf("Certificate is not properly self-signed: %v", err)
	}
}

func TestIssuer_DifferentExpirations(t *testing.T) {
	expirations := []time.Duration{
		1 * time.Minute,
		15 * time.Minute,
		1 * time.Hour,
		24 * time.Hour,
	}

	for _, expiration := range expirations {
		t.Run(expiration.String(), func(t *testing.T) {
			keyPath := getIssuerSharedTestKey(t)

			issuer, err := NewIssuer(keyPath, "atcr.io", "registry", expiration)
			if err != nil {
				t.Fatalf("NewIssuer() error = %v", err)
			}

			tokenString, err := issuer.Issue("did:plc:user123", nil, AuthMethodOAuth)
			if err != nil {
				t.Fatalf("Issue() error = %v", err)
			}

			// Parse token and verify expiration
			token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (any, error) {
				return issuer.publicKey, nil
			})
			if err != nil {
				t.Fatalf("Failed to parse token: %v", err)
			}

			claims, ok := token.Claims.(*Claims)
			if !ok {
				t.Fatal("Failed to cast claims")
			}

			expiresIn := time.Until(claims.ExpiresAt.Time)
			// Allow 2 second tolerance for test execution time
			if expiresIn < expiration-2*time.Second || expiresIn > expiration+2*time.Second {
				t.Errorf("Expected expiration around %v, got %v", expiration, expiresIn)
			}
		})
	}
}