at-container-registry/pkg/appview/crypto_keys.go

package appview

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

	"atcr.io/pkg/appview/db"
	"github.com/bluesky-social/indigo/atproto/atcrypto"
)

// loadOAuthKey loads the OAuth P-256 key from the DB, generating one if absent.
func loadOAuthKey(database *sql.DB) (*atcrypto.PrivateKeyP256, error) {
	data, err := db.GetCryptoKey(database, "oauth_p256")
	if err != nil {
		return nil, fmt.Errorf("failed to query crypto_keys: %w", err)
	}
	if data != nil {
		key, err := atcrypto.ParsePrivateBytesP256(data)
		if err != nil {
			return nil, fmt.Errorf("failed to parse OAuth key from database: %w", err)
		}
		slog.Info("Loaded OAuth P-256 key from database")
		return key, nil
	}

	p256Key, err := atcrypto.GeneratePrivateKeyP256()
	if err != nil {
		return nil, fmt.Errorf("failed to generate OAuth P-256 key: %w", err)
	}

	keyBytes := p256Key.Bytes()
	if err := db.PutCryptoKey(database, "oauth_p256", keyBytes); err != nil {
		return nil, fmt.Errorf("failed to store generated OAuth key in database: %w", err)
	}

	// Re-read rather than trusting the key we just generated. If another
	// instance won the race to write first, PutCryptoKey kept theirs, and the
	// stored key is the one the published JWKS will advertise. Using ours would
	// mean signing client assertions with a key nobody can verify.
	stored, err := db.GetCryptoKey(database, "oauth_p256")
	if err != nil {
		return nil, fmt.Errorf("failed to re-read OAuth key after storing: %w", err)
	}
	if stored == nil {
		return nil, fmt.Errorf("OAuth key missing immediately after storing it")
	}
	key, err := atcrypto.ParsePrivateBytesP256(stored)
	if err != nil {
		return nil, fmt.Errorf("failed to parse stored OAuth key: %w", err)
	}
	if bytes.Equal(stored, keyBytes) {
		slog.Info("Generated new OAuth P-256 key and stored in database")
	} else {
		slog.Info("Another instance stored an OAuth P-256 key first, using theirs")
	}

	return key, nil
}

// loadJWTKeyAndCert loads the JWT RSA key from the DB and generates a self-signed
// certificate. The cert is always regenerated and written to certPath on disk
// because the distribution library reads it via os.Open().
func loadJWTKeyAndCert(database *sql.DB, certPath string) (*rsa.PrivateKey, []byte, error) {
	rsaKey, err := loadRSAKey(database)
	if err != nil {
		return nil, nil, err
	}

	certDER, err := generateAndWriteCert(rsaKey, certPath)
	if err != nil {
		return nil, nil, err
	}

	return rsaKey, certDER, nil
}

// loadRSAKey loads the RSA private key from the DB, generating one if absent.
func loadRSAKey(database *sql.DB) (*rsa.PrivateKey, error) {
	data, err := db.GetCryptoKey(database, "jwt_rsa")
	if err != nil {
		return nil, fmt.Errorf("failed to query crypto_keys: %w", err)
	}
	if data != nil {
		key, err := parseRSAKeyPEM(data)
		if err != nil {
			return nil, fmt.Errorf("failed to parse RSA key from database: %w", err)
		}
		slog.Info("Loaded JWT RSA key from database")
		return key, nil
	}

	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("failed to generate RSA key: %w", err)
	}

	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(rsaKey),
	})
	if err := db.PutCryptoKey(database, "jwt_rsa", keyPEM); err != nil {
		return nil, fmt.Errorf("failed to store generated RSA key in database: %w", err)
	}

	// Re-read for the same reason as loadOAuthKey: whoever wrote first owns the
	// key, and the certificate written to disk is derived from whatever we
	// return here. Keeping our own copy after losing the race would mean issuing
	// registry JWTs that do not match the advertised certificate.
	stored, err := db.GetCryptoKey(database, "jwt_rsa")
	if err != nil {
		return nil, fmt.Errorf("failed to re-read JWT RSA key after storing: %w", err)
	}
	if stored == nil {
		return nil, fmt.Errorf("JWT RSA key missing immediately after storing it")
	}
	key, err := parseRSAKeyPEM(stored)
	if err != nil {
		return nil, fmt.Errorf("failed to parse stored JWT RSA key: %w", err)
	}
	if bytes.Equal(stored, keyPEM) {
		slog.Info("Generated new JWT RSA key and stored in database")
	} else {
		slog.Info("Another instance stored a JWT RSA key first, using theirs")
	}

	return key, nil
}

func parseRSAKeyPEM(data []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(data)
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("failed to decode PEM block containing RSA private key")
	}
	return x509.ParsePKCS1PrivateKey(block.Bytes)
}

// generateAndWriteCert creates a self-signed certificate from the RSA key and writes
// it to certPath. Returns the DER-encoded certificate bytes for the JWT x5c header.
func generateAndWriteCert(rsaKey *rsa.PrivateKey, certPath string) ([]byte, error) {
	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),
		KeyUsage:              x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &rsaKey.PublicKey, rsaKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate: %w", err)
	}

	// Write cert to disk for distribution library
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	dir := filepath.Dir(certPath)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create cert directory: %w", err)
	}
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return nil, fmt.Errorf("failed to write certificate: %w", err)
	}

	slog.Info("Generated JWT signing certificate", "path", certPath)
	return certDER, nil
}