gomail/internal/tlsutil/acme_manager.go

package tlsutil

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"gomail/internal/acme"
	"gomail/internal/crypto"
	"gomail/internal/db"
)

// renewalMargin is how far before expiry a certificate is renewed.
const renewalMargin = 30 * 24 * time.Hour

// ACMEManager obtains and caches ACME certificates per domain, encrypted at
// rest (same HKDF-per-record scheme as everything else), and serves them
// via a SNI-aware tls.Config.GetCertificate callback so a single listener
// can present the right certificate for whichever domain a client connects
// to. A background loop renews any certificate within renewalMargin of
// expiry.
type ACMEManager struct {
	database     *db.DB
	mk           *crypto.MasterKey
	directoryURL string
	contactEmail string
	responder    *acme.ChallengeResponder

	mu    sync.RWMutex
	cache map[string]*tls.Certificate
}

func NewACMEManager(database *db.DB, mk *crypto.MasterKey, directoryURL, contactEmail string, responder *acme.ChallengeResponder) *ACMEManager {
	return &ACMEManager{
		database: database, mk: mk, directoryURL: directoryURL, contactEmail: contactEmail,
		responder: responder, cache: make(map[string]*tls.Certificate),
	}
}

// TLSConfig returns a tls.Config whose GetCertificate looks up the right
// cert per SNI, obtaining one on first use if none is cached yet.
func (m *ACMEManager) TLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion: tls.VersionTLS12,
		GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			return m.CertificateFor(hello.ServerName)
		},
	}
}

// CertificateFor returns a cached certificate for domain, obtaining one via
// ACME (and caching it, in memory and encrypted in the DB) if not already
// cached or if the cached one is expired/near expiry.
func (m *ACMEManager) CertificateFor(domain string) (*tls.Certificate, error) {
	m.mu.RLock()
	cached, ok := m.cache[domain]
	m.mu.RUnlock()
	if ok {
		return cached, nil
	}

	if stored, err := m.loadFromDB(domain); err == nil {
		m.mu.Lock()
		m.cache[domain] = stored
		m.mu.Unlock()
		return stored, nil
	}

	cert, err := m.obtainAndStore(domain)
	if err != nil {
		return nil, err
	}
	return cert, nil
}

func (m *ACMEManager) loadFromDB(domain string) (*tls.Certificate, error) {
	row, err := m.database.GetTLSCert(domain)
	if err != nil {
		return nil, err
	}
	if row.CertPEMEnc == nil || row.KeyPEMEnc == nil {
		return nil, fmt.Errorf("no cert material stored for %s", domain)
	}
	if row.ExpiresAt != nil && time.Now().UTC().After(row.ExpiresAt.Add(-renewalMargin)) {
		return nil, fmt.Errorf("stored cert for %s is expired or near expiry", domain)
	}

	certPEM, err := crypto.Decrypt(m.mk, row.ID, "tls-cert", row.CertPEMEnc)
	if err != nil {
		return nil, fmt.Errorf("decrypting cert: %w", err)
	}
	keyPEM, err := crypto.Decrypt(m.mk, row.ID, "tls-key", row.KeyPEMEnc)
	if err != nil {
		return nil, fmt.Errorf("decrypting key: %w", err)
	}

	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("parsing stored cert/key: %w", err)
	}
	return &cert, nil
}

func (m *ACMEManager) obtainAndStore(domain string) (*tls.Certificate, error) {
	accountKey, err := m.loadOrCreateAccountKey(domain)
	if err != nil {
		return nil, fmt.Errorf("account key: %w", err)
	}

	slog.Info("obtaining ACME certificate", "domain", domain, "directory", m.directoryURL)
	certPEM, keyPEM, err := acme.Obtain(m.directoryURL, m.contactEmail, []string{domain}, accountKey, m.responder)
	if err != nil {
		return nil, fmt.Errorf("ACME obtain for %s: %w", domain, err)
	}

	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("parsing obtained cert/key: %w", err)
	}

	var expiresAt *time.Time
	if len(cert.Certificate) > 0 {
		if leaf, err := x509.ParseCertificate(cert.Certificate[0]); err == nil {
			expiresAt = &leaf.NotAfter
		}
	}

	existing, _ := m.database.GetTLSCert(domain)
	recordID := domain
	if existing != nil {
		recordID = existing.ID
	}
	encCert, err := crypto.Encrypt(m.mk, recordID, "tls-cert", certPEM)
	if err != nil {
		return nil, fmt.Errorf("encrypting cert: %w", err)
	}
	encKey, err := crypto.Encrypt(m.mk, recordID, "tls-key", keyPEM)
	if err != nil {
		return nil, fmt.Errorf("encrypting key: %w", err)
	}

	if err := m.database.UpsertTLSCert(&db.TLSCert{
		ID: recordID, Domain: domain, CertPEMEnc: encCert, KeyPEMEnc: encKey, ExpiresAt: expiresAt,
	}); err != nil {
		return nil, fmt.Errorf("storing cert: %w", err)
	}

	m.mu.Lock()
	m.cache[domain] = &cert
	m.mu.Unlock()

	slog.Info("ACME certificate obtained and stored", "domain", domain, "expires_at", expiresAt)
	return &cert, nil
}

func (m *ACMEManager) loadOrCreateAccountKey(domain string) (*acme.AccountKey, error) {
	row, err := m.database.GetTLSCert(domain)
	if err == nil && row.ACMEAccountKeyEnc != nil {
		plain, decErr := crypto.Decrypt(m.mk, row.ID, "acme-account-key", row.ACMEAccountKeyEnc)
		if decErr == nil {
			if key, parseErr := acme.ParseAccountKeyPEM(plain); parseErr == nil {
				return key, nil
			}
		}
	}

	key, err := acme.GenerateAccountKey()
	if err != nil {
		return nil, err
	}
	keyPEM, err := key.MarshalPEM()
	if err != nil {
		return nil, err
	}

	recordID := domain
	if row != nil {
		recordID = row.ID
	}
	encKey, err := crypto.Encrypt(m.mk, recordID, "acme-account-key", keyPEM)
	if err != nil {
		return nil, err
	}
	if err := m.database.SetACMEAccountKey(domain, encKey); err != nil {
		return nil, err
	}
	return key, nil
}

// StartRenewalLoop runs a background check (default: daily) and renews any
// domain whose cached/stored certificate is within renewalMargin of expiry.
// domains is the full set this instance is responsible for — typically all
// active hosted domains plus the server's own hostname.
func (m *ACMEManager) StartRenewalLoop(ctx context.Context, domains []string, checkInterval time.Duration) {
	ticker := time.NewTicker(checkInterval)
	defer ticker.Stop()

	checkAndRenew := func() {
		for _, domain := range domains {
			row, err := m.database.GetTLSCert(domain)
			needsRenewal := err != nil || row.ExpiresAt == nil || time.Now().UTC().After(row.ExpiresAt.Add(-renewalMargin))
			if !needsRenewal {
				continue
			}
			slog.Info("renewing ACME certificate", "domain", domain)
			m.mu.Lock()
			delete(m.cache, domain) // force re-obtain, not a stale in-memory hit
			m.mu.Unlock()
			if _, err := m.obtainAndStore(domain); err != nil {
				slog.Error("ACME renewal failed", "domain", domain, "err", err)
			}
		}
	}

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			checkAndRenew()
		}
	}
}