pmg/proxy/certmanager/generator.go

package certmanager

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/safedep/dry/log"
	"golang.org/x/sync/singleflight"
)

const (
	maxSystemCABundleBytes int64         = 10 * 1024 * 1024
	certExpiryThreshold    time.Duration = 1 * time.Hour
	goosDarwin                           = "darwin"
	goosLinux                            = "linux"
	goosWindows                          = "windows"
)

// CACommonName is the subject CN of the PMG CA. It is the single source of truth
// for the CA's identity: GenerateCA stamps it onto the certificate and the
// truststore package matches on it to find/remove the cert in OS trust stores.
const CACommonName = "SafeDep PMG Proxy CA"

// certManager implements the CertificateManager interface
type certManager struct {
	ca     *Certificate
	cache  CertificateCache
	config CertManagerConfig
	group  singleflight.Group
}

// NewCertificateManagerWithCA creates a new certificate manager with an existing CA certificate
func NewCertificateManagerWithCA(ca *Certificate, config CertManagerConfig) (CertificateManager, error) {
	if ca == nil {
		return nil, fmt.Errorf("CA certificate cannot be nil")
	}

	config.SetDefaults()
	if err := config.Validate(); err != nil {
		return nil, fmt.Errorf("invalid config: %w", err)
	}

	if ca.X509Cert == nil || ca.PrivKey == nil {
		parsedCA, err := parseCertificate(ca)
		if err != nil {
			return nil, fmt.Errorf("failed to parse CA certificate: %w", err)
		}

		ca = parsedCA
	}

	if ca.IsExpired(certExpiryThreshold) {
		return nil, fmt.Errorf("CA certificate is expired")
	}

	return &certManager{
		ca:     ca,
		cache:  NewInMemoryCache(),
		config: config,
	}, nil
}

// GetCA returns the Certificate Authority certificate
func (cm *certManager) GetCA() (*Certificate, error) {
	return cm.ca, nil
}

// GenerateCertForHost creates a certificate for the given hostname.
// Uses caching and singleflight to ensure only one goroutine generates
// a certificate for a given hostname at a time, preventing CPU starvation
// when many concurrent CONNECT requests arrive for the same host.
func (cm *certManager) GenerateCertForHost(hostname string) (*Certificate, error) {
	if cached, found := cm.cache.Get(hostname); found {
		if !cached.IsExpired(certExpiryThreshold) {
			return cached, nil
		}
	}

	result, err, _ := cm.group.Do(hostname, func() (interface{}, error) {
		// Re-check cache: another goroutine in a previous singleflight
		// group may have populated it while we were waiting.
		if cached, found := cm.cache.Get(hostname); found {
			if !cached.IsExpired(certExpiryThreshold) {
				return cached, nil
			}
		}

		cert, err := cm.generateHostCert(hostname)
		if err != nil {
			return nil, fmt.Errorf("failed to generate certificate for %s: %w", hostname, err)
		}

		cm.cache.Set(hostname, cert)
		return cert, nil
	})
	if err != nil {
		return nil, err
	}

	return result.(*Certificate), nil
}

// GetTLSConfig returns a tls.Config for the given hostname
func (cm *certManager) GetTLSConfig(hostname string) (*tls.Config, error) {
	cert, err := cm.GenerateCertForHost(hostname)
	if err != nil {
		return nil, err
	}

	tlsCert, err := tls.X509KeyPair(cert.Certificate, cert.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create tls.Certificate: %w", err)
	}

	return &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		MinVersion:   tls.VersionTLS12,
	}, nil
}

func (cm *certManager) generateHostCert(hostname string) (*Certificate, error) {
	privKey, err := rsa.GenerateKey(rand.Reader, cm.config.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}

	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(cm.config.HostCertValidityDays) * 24 * time.Hour)

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName: hostname,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		DNSNames:              []string{hostname},
	}

	caPrivKey, ok := cm.ca.PrivKey.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("CA private key is not RSA")
	}

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

	x509Cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse generated certificate: %w", err)
	}

	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	return &Certificate{
		Certificate: certPEM,
		PrivateKey:  privKeyPEM,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

// GenerateCA generates a new self-signed CA certificate using the given configuration
func GenerateCA(config CertManagerConfig) (*Certificate, error) {
	config.SetDefaults()
	if err := config.Validate(); err != nil {
		return nil, fmt.Errorf("invalid config: %w", err)
	}

	privKey, err := rsa.GenerateKey(rand.Reader, config.KeySize)
	if err != nil {
		return nil, fmt.Errorf("failed to generate CA private key: %w", err)
	}

	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("failed to generate serial number: %w", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Duration(config.CAValidityDays) * 24 * time.Hour)

	template := &x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			CommonName:   CACommonName,
			Organization: []string{"SafeDep PMG"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}

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

	x509Cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse generated CA certificate: %w", err)
	}

	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})

	return &Certificate{
		Certificate: certPEM,
		PrivateKey:  privKeyPEM,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

func parseCertificate(cert *Certificate) (*Certificate, error) {
	certBlock, _ := pem.Decode(cert.Certificate)
	if certBlock == nil {
		return nil, fmt.Errorf("failed to decode PEM certificate")
	}

	x509Cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse X509 certificate: %w", err)
	}

	keyBlock, _ := pem.Decode(cert.PrivateKey)
	if keyBlock == nil {
		return nil, fmt.Errorf("failed to decode PEM private key")
	}

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

	return &Certificate{
		Certificate: cert.Certificate,
		PrivateKey:  cert.PrivateKey,
		X509Cert:    x509Cert,
		PrivKey:     privKey,
	}, nil
}

// ParseTLSCertificate converts a Certificate to a tls.Certificate
// This is useful for integrating with libraries that expect tls.Certificate
func ParseTLSCertificate(cert *Certificate) (tls.Certificate, error) {
	tlsCert, err := tls.X509KeyPair(cert.Certificate, cert.PrivateKey)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("failed to create X509 key pair: %w", err)
	}

	// Populate Leaf field if we have the parsed X.509 certificate
	// This is important for libraries (like goproxy) that need to inspect the certificate
	if cert.X509Cert != nil {
		tlsCert.Leaf = cert.X509Cert
	}

	return tlsCert, nil
}

// MergeWithSystemCA appends the system CA bundle to certPEM so env-var trust
// injection (SSL_CERT_FILE, NODE_EXTRA_CA_CERTS, ...) covers both the PMG CA and
// the real world. It is best-effort: any problem locating, sizing, or reading the
// system bundle results in certPEM being returned unchanged, so proxy startup is
// never blocked by an optional enhancement.
func MergeWithSystemCA(certPEM []byte) []byte {
	systemBundlePath := firstReadablePath(systemCABundleCandidates()...)
	if systemBundlePath == "" {
		log.Warnf("Skipping system CA bundle merge: No system CA bundle file found")
		return certPEM
	}

	info, err := os.Stat(systemBundlePath)
	if err != nil {
		log.Errorf("Skipping system CA bundle merge: failed to stat %s: %v", systemBundlePath, err)
		return certPEM
	}

	if info.Size() > maxSystemCABundleBytes {
		log.Errorf("Skipping system CA bundle merge: %s is too large (%d bytes > %d bytes)",
			systemBundlePath, info.Size(), maxSystemCABundleBytes)
		return certPEM
	}

	systemBundle, err := os.ReadFile(systemBundlePath)
	if err != nil {
		log.Errorf("Skipping system CA bundle merge: failed to read %s: %v", systemBundlePath, err)
		return certPEM
	}

	const extra = int64(2)
	totalCap := int64(len(certPEM)) + int64(len(systemBundle)) + extra
	if totalCap > maxSystemCABundleBytes {
		log.Errorf("Skipping system CA bundle merge: merged CA would be too large (%d bytes > %d bytes)",
			totalCap, maxSystemCABundleBytes)
		return certPEM
	}

	merged := make([]byte, 0, int(totalCap))
	merged = append(merged, certPEM...)
	if len(merged) > 0 && merged[len(merged)-1] != '\n' {
		merged = append(merged, '\n')
	}
	merged = append(merged, systemBundle...)
	if len(merged) > 0 && merged[len(merged)-1] != '\n' {
		merged = append(merged, '\n')
	}

	return merged
}

// GenerateCAWithSystemCA generates a self-signed CA and merges the system CA
// bundle into the certificate bytes (for env-var trust injection). Signing
// always uses the pure CA via X509Cert/PrivKey.
func GenerateCAWithSystemCA(config CertManagerConfig) (*Certificate, error) {
	caCert, err := GenerateCA(config)
	if err != nil {
		return nil, fmt.Errorf("failed to generate CA: %w", err)
	}

	return &Certificate{
		Certificate: MergeWithSystemCA(caCert.Certificate),
		PrivateKey:  caCert.PrivateKey,
		X509Cert:    caCert.X509Cert,
		PrivKey:     caCert.PrivKey,
	}, nil
}

func firstReadablePath(paths ...string) string {
	for _, path := range paths {
		if path == "" {
			continue
		}

		info, err := os.Stat(path)
		if err != nil || info.IsDir() {
			continue
		}

		f, err := os.Open(path)
		if err != nil {
			continue
		}
		_ = f.Close()

		return path
	}

	return ""
}

func systemCABundleCandidates() []string {
	return systemCABundleCandidatesForOS(runtime.GOOS)
}

func systemCABundleCandidatesForOS(goos string) []string {
	var candidates []string
	appendIfSet := func(key string) {
		if value := os.Getenv(key); value != "" {
			candidates = append(candidates, value)
		}
	}

	appendIfSet("SSL_CERT_FILE")
	appendIfSet("CURL_CA_BUNDLE")

	switch goos {
	case goosDarwin:
		candidates = append(candidates,
			"/opt/homebrew/etc/openssl@3/cert.pem",
			"/usr/local/etc/openssl@3/cert.pem",
			"/etc/ssl/cert.pem",
		)
	case goosLinux:
		candidates = append(candidates,
			"/etc/ssl/certs/ca-certificates.crt",
			"/etc/pki/tls/certs/ca-bundle.crt",
			"/etc/pki/ca-trust/extracted/pem/tls-ca-bundle.pem",
			"/etc/ssl/ca-bundle.pem",
			"/etc/ssl/cert.pem",
		)
	case goosWindows:
		programFiles := os.Getenv("ProgramFiles")
		programFilesX86 := os.Getenv("ProgramFiles(x86)")
		systemRoot := os.Getenv("SystemRoot")

		if programFiles != "" {
			candidates = append(candidates,
				filepath.Join(programFiles, "Git", "mingw64", "ssl", "certs", "ca-bundle.crt"),
				filepath.Join(programFiles, "Git", "usr", "ssl", "certs", "ca-bundle.crt"),
			)
		}
		if programFilesX86 != "" {
			candidates = append(candidates,
				filepath.Join(programFilesX86, "Git", "mingw32", "ssl", "certs", "ca-bundle.crt"),
			)
		}
		if systemRoot != "" {
			candidates = append(candidates,
				filepath.Join(systemRoot, "System32", "curl-ca-bundle.crt"),
			)
		}
	}

	return candidates
}