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

package appview

import (
	"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 with priority: DB → file → generate.
// Keys loaded from file or newly generated are stored in the DB.
func loadOAuthKey(database *sql.DB, keyPath string) (*atcrypto.PrivateKeyP256, error) {
	// Try database first
	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
	}

	// Try file fallback
	if keyPath != "" {
		if fileData, err := os.ReadFile(keyPath); err == nil {
			key, err := atcrypto.ParsePrivateBytesP256(fileData)
			if err != nil {
				return nil, fmt.Errorf("failed to parse OAuth key from file %s: %w", keyPath, err)
			}
			// Migrate to database
			if err := db.PutCryptoKey(database, "oauth_p256", fileData); err != nil {
				return nil, fmt.Errorf("failed to store OAuth key in database: %w", err)
			}
			slog.Info("Migrated OAuth P-256 key from file to database", "path", keyPath)
			return key, nil
		}
	}

	// Generate new key
	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)
	}
	slog.Info("Generated new OAuth P-256 key and stored in database")

	return p256Key, nil
}

// loadJWTKeyAndCert loads the JWT RSA key from DB (with file fallback) 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, keyPath, certPath string) (*rsa.PrivateKey, []byte, error) {
	rsaKey, err := loadRSAKey(database, keyPath)
	if err != nil {
		return nil, nil, err
	}

	// Generate cert and write to disk for distribution library
	certDER, err := generateAndWriteCert(rsaKey, certPath)
	if err != nil {
		return nil, nil, err
	}

	return rsaKey, certDER, nil
}

// loadRSAKey loads the RSA private key with priority: DB → file → generate.
func loadRSAKey(database *sql.DB, keyPath string) (*rsa.PrivateKey, error) {
	// Try database first
	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
	}

	// Try file fallback
	if keyPath != "" {
		if fileData, err := os.ReadFile(keyPath); err == nil {
			key, err := parseRSAKeyPEM(fileData)
			if err != nil {
				return nil, fmt.Errorf("failed to parse RSA key from file %s: %w", keyPath, err)
			}
			// Migrate to database
			if err := db.PutCryptoKey(database, "jwt_rsa", fileData); err != nil {
				return nil, fmt.Errorf("failed to store RSA key in database: %w", err)
			}
			slog.Info("Migrated JWT RSA key from file to database", "path", keyPath)
			return key, nil
		}
	}

	// Generate new key
	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)
	}
	slog.Info("Generated new JWT RSA key and stored in database")

	return rsaKey, 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
}