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

package token

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"log/slog"
	"math/big"
	"os"
	"path/filepath"
	"strings"
	"time"

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

// Issuer handles JWT token creation and signing
type Issuer struct {
	privateKey  *rsa.PrivateKey
	publicKey   *rsa.PublicKey
	certificate []byte // DER-encoded certificate
	issuer      string
	service     string
	expiration  time.Duration
}

// NewIssuer creates a new JWT issuer
func NewIssuer(privateKeyPath, issuer, service string, expiration time.Duration) (*Issuer, error) {
	privateKey, err := loadOrGenerateKey(privateKeyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to load private key: %w", err)
	}

	// Load the certificate for x5c header
	certPath := strings.TrimSuffix(privateKeyPath, ".pem") + ".crt"
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read certificate: %w", err)
	}

	// Parse PEM to get DER-encoded certificate
	block, _ := pem.Decode(certPEM)
	if block == nil || block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("failed to decode certificate PEM")
	}

	return &Issuer{
		privateKey:  privateKey,
		publicKey:   &privateKey.PublicKey,
		certificate: block.Bytes, // DER-encoded certificate
		issuer:      issuer,
		service:     service,
		expiration:  expiration,
	}, nil
}

// NewIssuerFromKey creates a JWT issuer from pre-loaded key material.
// certDER is the DER-encoded X.509 certificate for the x5c JWT header.
func NewIssuerFromKey(privateKey *rsa.PrivateKey, certDER []byte, issuer, service string, expiration time.Duration) *Issuer {
	return &Issuer{
		privateKey:  privateKey,
		publicKey:   &privateKey.PublicKey,
		certificate: certDER,
		issuer:      issuer,
		service:     service,
		expiration:  expiration,
	}
}

// Issue creates and signs a new JWT token
func (i *Issuer) Issue(subject string, access []auth.AccessEntry, authMethod string) (string, error) {
	claims := NewClaims(subject, i.issuer, i.service, i.expiration, access, authMethod)

	slog.Debug("Creating JWT token",
		"issuer", i.issuer,
		"service", i.service,
		"subject", subject,
		"access", access)

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)

	// Add x5c header - embeds the certificate chain in the JWT
	// This is base64-encoded DER certificate(s)
	certChain := []string{
		base64.StdEncoding.EncodeToString(i.certificate),
	}
	token.Header["x5c"] = certChain

	signedToken, err := token.SignedString(i.privateKey)
	if err != nil {
		return "", fmt.Errorf("failed to sign token: %w", err)
	}

	slog.Debug("Successfully signed token with x5c header")

	return signedToken, nil
}

// PublicKey returns the public key for token verification
func (i *Issuer) PublicKey() *rsa.PublicKey {
	return i.publicKey
}

// Expiration returns the token expiration duration
func (i *Issuer) Expiration() time.Duration {
	return i.expiration
}

// loadOrGenerateKey loads an existing RSA private key or generates a new one
func loadOrGenerateKey(path string) (*rsa.PrivateKey, error) {
	// Try to load existing key
	if _, err := os.Stat(path); err == nil {
		keyData, err := os.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("failed to read key file: %w", err)
		}

		block, _ := pem.Decode(keyData)
		if block == nil || block.Type != "RSA PRIVATE KEY" {
			return nil, fmt.Errorf("failed to decode PEM block containing private key")
		}

		privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse private key: %w", err)
		}

		// Ensure certificate exists
		certPath := strings.TrimSuffix(path, ".pem") + ".crt"
		if _, err := os.Stat(certPath); os.IsNotExist(err) {
			// Certificate doesn't exist, generate it
			if err := generateCertificate(privateKey, certPath); err != nil {
				return nil, fmt.Errorf("failed to generate certificate: %w", err)
			}
		}

		return privateKey, nil
	}

	// Generate new key
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}

	// Ensure directory exists
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create key directory: %w", err)
	}

	// Save key to file
	keyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: keyBytes,
	})

	if err := os.WriteFile(path, keyPEM, 0600); err != nil {
		return nil, fmt.Errorf("failed to write private key: %w", err)
	}

	// Also generate a self-signed certificate for the public key
	certPath := strings.TrimSuffix(path, ".pem") + ".crt"
	if err := generateCertificate(privateKey, certPath); err != nil {
		return nil, fmt.Errorf("failed to generate certificate: %w", err)
	}

	return privateKey, nil
}

// generateCertificate creates a self-signed certificate for JWT validation
func generateCertificate(privateKey *rsa.PrivateKey, certPath string) error {
	// Create certificate template
	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"ATCR"},
			CommonName:   "ATCR Token Signing Certificate",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(10 * 365 * 24 * time.Hour), // 10 years
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	// Create self-signed certificate
	certBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return fmt.Errorf("failed to create certificate: %w", err)
	}

	// Encode certificate to PEM
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})

	// Write certificate to file
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return fmt.Errorf("failed to write certificate: %w", err)
	}

	return nil
}