at-container-registry/pkg/auth/oauth/keys.go

package oauth

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
)

// KeyType represents the elliptic curve algorithm for key generation
type KeyType int

const (
	// KeyTypeP256 uses NIST P-256 (ES256) - standard for OAuth/OIDC
	KeyTypeP256 KeyType = iota
	// KeyTypeK256 uses secp256k1 (ES256K) - used for ATProto PDS signing
	KeyTypeK256
)

// GenerateOrLoadKey generates a new key pair or loads an existing one
// Supports both P-256 (OAuth) and K-256 (ATProto PDS) key types
func GenerateOrLoadKey(keyPath string, keyType KeyType) (atcrypto.PrivateKey, error) {
	// Ensure directory exists
	dir := filepath.Dir(keyPath)
	if err := os.MkdirAll(dir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create key directory: %w", err)
	}

	// Check if key already exists
	if _, err := os.Stat(keyPath); err == nil {
		// Key exists, load it
		return loadKey(keyPath, keyType)
	}

	// Key doesn't exist, generate new one
	return generateKey(keyPath, keyType)
}

// generateKey creates a new key pair of the specified type
func generateKey(keyPath string, keyType KeyType) (atcrypto.PrivateKey, error) {
	var privateKey atcrypto.PrivateKey
	var keyName string
	var err error

	switch keyType {
	case KeyTypeP256:
		privateKey, err = atcrypto.GeneratePrivateKeyP256()
		keyName = "P-256"
	case KeyTypeK256:
		privateKey, err = atcrypto.GeneratePrivateKeyK256()
		keyName = "K-256"
	default:
		return nil, fmt.Errorf("unsupported key type: %d", keyType)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to generate %s key: %w", keyName, err)
	}

	// Serialize key to bytes
	exportableKey, ok := privateKey.(atcrypto.PrivateKeyExportable)
	if !ok {
		return nil, fmt.Errorf("key does not support export")
	}
	keyBytes := exportableKey.Bytes()

	// Write to file with restrictive permissions
	if err := os.WriteFile(keyPath, keyBytes, 0600); err != nil {
		return nil, fmt.Errorf("failed to write key file: %w", err)
	}

	slog.Info("Generated new signing key", "type", keyName, "path", keyPath)
	return privateKey, nil
}

// loadKey loads an existing private key from disk
func loadKey(keyPath string, keyType KeyType) (atcrypto.PrivateKey, error) {
	// Read key bytes
	keyBytes, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read key file: %w", err)
	}

	var privateKey atcrypto.PrivateKey
	var keyName string

	switch keyType {
	case KeyTypeP256:
		privateKey, err = atcrypto.ParsePrivateBytesP256(keyBytes)
		keyName = "P-256"
	case KeyTypeK256:
		// Check for old PEM format (migration path)
		if IsPEMFormat(keyBytes) {
			slog.Warn("Detected old P-256 PEM key, replacing with K-256")
			return generateKey(keyPath, keyType)
		}
		privateKey, err = atcrypto.ParsePrivateBytesK256(keyBytes)
		keyName = "K-256"
	default:
		return nil, fmt.Errorf("unsupported key type: %d", keyType)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to parse %s private key: %w", keyName, err)
	}

	slog.Info("Loaded existing signing key", "type", keyName, "path", keyPath)
	return privateKey, nil
}

// IsPEMFormat checks if bytes are in PEM format (for migration detection)
// Exported for testing and migration utilities
func IsPEMFormat(data []byte) bool {
	return len(data) > 10 && string(data[:5]) == "-----"
}

// GenerateOrLoadClientKey generates a new P256 key pair or loads an existing one
// This is a convenience wrapper for OAuth client keys
func GenerateOrLoadClientKey(keyPath string) (*atcrypto.PrivateKeyP256, error) {
	key, err := GenerateOrLoadKey(keyPath, KeyTypeP256)
	if err != nil {
		return nil, err
	}

	p256Key, ok := key.(*atcrypto.PrivateKeyP256)
	if !ok {
		return nil, fmt.Errorf("expected P-256 key, got different type")
	}

	return p256Key, nil
}

// GenerateOrLoadPDSKey generates a new K256 key pair or loads an existing one
// This is a convenience wrapper for ATProto PDS signing keys
func GenerateOrLoadPDSKey(keyPath string) (*atcrypto.PrivateKeyK256, error) {
	key, err := GenerateOrLoadKey(keyPath, KeyTypeK256)
	if err != nil {
		return nil, err
	}

	k256Key, ok := key.(*atcrypto.PrivateKeyK256)
	if !ok {
		return nil, fmt.Errorf("expected K-256 key, got different type")
	}

	return k256Key, nil
}

// GenerateKeyID generates a stable key ID from a P256 public key
// Uses the first 8 characters of the hex-encoded SHA256 hash of the public key bytes
func GenerateKeyID(privateKey *atcrypto.PrivateKeyP256) (string, error) {
	// Get public key
	pubKey, err := privateKey.PublicKey()
	if err != nil {
		return "", fmt.Errorf("failed to get public key: %w", err)
	}

	// Get public key bytes
	pubKeyBytes := pubKey.Bytes()

	// Hash public key bytes
	hash := sha256.Sum256(pubKeyBytes)

	// Return first 8 characters of hex-encoded hash
	return hex.EncodeToString(hash[:])[:8], nil
}

// PrivateKeyToMultibase converts a P256 private key to multibase format
// Required by indigo's SetClientSecret() API
func PrivateKeyToMultibase(key *atcrypto.PrivateKeyP256) string {
	return key.Multibase()
}

// MultibaseToPrivateKey parses a multibase-encoded P256 private key
func MultibaseToPrivateKey(encoded string) (*atcrypto.PrivateKeyP256, error) {
	// ParsePrivateMultibase returns PrivateKeyExportable interface
	key, err := atcrypto.ParsePrivateMultibase(encoded)
	if err != nil {
		return nil, fmt.Errorf("failed to parse multibase key: %w", err)
	}

	// Type assert to P256 key
	p256Key, ok := key.(*atcrypto.PrivateKeyP256)
	if !ok {
		return nil, fmt.Errorf("expected P-256 key, got different key type")
	}

	return p256Key, nil
}