mailgosend/internal/tls/acme.go

package tls

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/go-acme/lego/v4/certcrypto"
	"github.com/go-acme/lego/v4/certificate"
	"github.com/go-acme/lego/v4/challenge"
	"github.com/go-acme/lego/v4/challenge/http01"
	"github.com/go-acme/lego/v4/lego"
	"github.com/go-acme/lego/v4/providers/dns/cloudflare"
	"github.com/go-acme/lego/v4/providers/dns/digitalocean"
	"github.com/go-acme/lego/v4/providers/dns/hetzner"
	"github.com/go-acme/lego/v4/providers/dns/route53"
	"github.com/go-acme/lego/v4/registration"
)

// ACME manages Let's Encrypt certificate issuance and renewal.
type ACME struct {
	cfg      ACMEConfig
	client   *lego.Client
	account  *acmeAccount
	cacheDir string
}

// ACMEConfig holds all ACME-related settings from app config.
type ACMEConfig struct {
	Email       string
	CacheDir    string
	Staging     bool
	Domains     []string // e.g. ["example.com", "*.example.com"]
	Mode        string   // "dns01" | "http01"
	DNSProvider string   // "cloudflare" | "route53" | "digitalocean" | "hetzner" | ...
}

// acmeAccount implements lego's registration.User interface.
type acmeAccount struct {
	Email        string                 `json:"email"`
	Registration *registration.Resource `json:"registration"`
	key          crypto.PrivateKey
	keyPEM       []byte
}

func (a *acmeAccount) GetEmail() string                        { return a.Email }
func (a *acmeAccount) GetRegistration() *registration.Resource { return a.Registration }
func (a *acmeAccount) GetPrivateKey() crypto.PrivateKey        { return a.key }

// NewACME initialises the ACME client. Does not yet obtain a cert.
// Call ObtainOrRenew() to get/refresh the certificate.
func NewACME(cfg ACMEConfig) (*ACME, error) {
	if cfg.Email == "" {
		return nil, fmt.Errorf("ACME_EMAIL required for Let's Encrypt")
	}
	if len(cfg.Domains) == 0 {
		return nil, fmt.Errorf("ACME_DOMAINS must not be empty")
	}
	if err := os.MkdirAll(cfg.CacheDir, 0700); err != nil {
		return nil, fmt.Errorf("acme cache dir: %w", err)
	}

	a := &ACME{cfg: cfg, cacheDir: cfg.CacheDir}

	acc, err := a.loadOrCreateAccount()
	if err != nil {
		return nil, fmt.Errorf("acme account: %w", err)
	}
	a.account = acc

	legoConfig := lego.NewConfig(a.account)
	if cfg.Staging {
		legoConfig.CADirURL = lego.LEDirectoryStaging
	} else {
		legoConfig.CADirURL = lego.LEDirectoryProduction
	}
	legoConfig.Certificate.KeyType = certcrypto.RSA2048

	client, err := lego.NewClient(legoConfig)
	if err != nil {
		return nil, fmt.Errorf("lego client: %w", err)
	}
	a.client = client

	// Configure challenge provider.
	if err := a.setProvider(); err != nil {
		return nil, err
	}

	// Register account if not already registered.
	if acc.Registration == nil {
		reg, err := client.Registration.Register(registration.RegisterOptions{
			TermsOfServiceAgreed: true,
		})
		if err != nil {
			return nil, fmt.Errorf("acme register: %w", err)
		}
		a.account.Registration = reg
		if err := a.saveAccount(); err != nil {
			return nil, fmt.Errorf("save acme account: %w", err)
		}
	}

	return a, nil
}

// setProvider configures the ACME challenge based on cfg.Mode and cfg.DNSProvider.
func (a *ACME) setProvider() error {
	switch a.cfg.Mode {
	case "http01":
		// Lego manages an ephemeral HTTP server on port 80.
		return a.client.Challenge.SetHTTP01Provider(http01.NewProviderServer("", "80"))

	case "dns01":
		provider, err := a.buildDNSProvider()
		if err != nil {
			return fmt.Errorf("dns01 provider %q: %w", a.cfg.DNSProvider, err)
		}
		return a.client.Challenge.SetDNS01Provider(provider)

	default:
		return fmt.Errorf("unknown ACME mode %q (want dns01 or http01)", a.cfg.Mode)
	}
}

// buildDNSProvider returns the lego DNS provider for cfg.DNSProvider.
// Credentials are read from env vars set by config.exportProviderEnv().
func (a *ACME) buildDNSProvider() (challenge.Provider, error) {
	switch a.cfg.DNSProvider {
	case "cloudflare":
		return cloudflare.NewDNSProvider()

	case "route53":
		return route53.NewDNSProvider()

	case "digitalocean":
		return digitalocean.NewDNSProvider()

	case "hetzner":
		return hetzner.NewDNSProvider()

	default:
		// Generic: lego supports 90+ providers; if the user sets the correct
		// env vars and the provider name matches a lego provider, it WILL work
		// through the lego plugin system. For unlisted providers, document that
		// the user must set env vars matching the lego provider docs.
		return nil, fmt.Errorf(
			"provider %q not built-in; set lego env vars and use a supported provider name.\n"+
				"Built-in: cloudflare, route53, digitalocean, hetzner.\n"+
				"Full list: https://go-acme.github.io/lego/dns/",
			a.cfg.DNSProvider,
		)
	}
}

// ObtainOrRenew returns a valid *tls.Certificate, obtaining or renewing as needed.
// Uses cached cert on disk if it has >30 days remaining.
func (a *ACME) ObtainOrRenew() (*tls.Certificate, error) {
	cached, err := a.loadCachedCert()
	if err == nil && cached != nil {
		return cached, nil
	}

	return a.obtain()
}

// obtain requests a new certificate from Let's Encrypt.
func (a *ACME) obtain() (*tls.Certificate, error) {
	req := certificate.ObtainRequest{
		Domains: a.cfg.Domains,
		Bundle:  true, // include full chain in cert PEM
	}

	res, err := a.client.Certificate.Obtain(req)
	if err != nil {
		return nil, fmt.Errorf("acme obtain %v: %w", a.cfg.Domains, err)
	}

	if err := a.saveCert(res); err != nil {
		// Non-fatal: cert is in memory, just can't persist.
		fmt.Printf("[acme] WARNING: could not cache cert: %v\n", err)
	}

	return parseCert(res.Certificate, res.PrivateKey)
}

// RenewalLoop blocks forever, renewing the cert 30 days before expiry.
// manager.UpdateCert() is called on each renewal to hot-swap the TLS config.
// Call as a goroutine. stopCh receives when it should quit.
func (a *ACME) RenewalLoop(manager *Manager, stopCh <-chan struct{}) {
	ticker := time.NewTicker(12 * time.Hour)
	defer ticker.Stop()

	for {
		select {
		case <-stopCh:
			return
		case <-ticker.C:
			cached, err := a.loadCachedCert()
			if err != nil || cached == nil {
				// No cert or corrupt cache — re-obtain.
				cert, err := a.obtain()
				if err != nil {
					fmt.Printf("[acme] renewal obtain error: %v\n", err)
					continue
				}
				manager.UpdateCert(cert)
				fmt.Printf("[acme] cert renewed for %v\n", a.cfg.Domains)
			}
			// loadCachedCert returns nil if cert expires in < 30 days → triggers re-obtain above.
		}
	}
}

// ---- Cert persistence ----

func (a *ACME) certPath() string  { return filepath.Join(a.cacheDir, "cert.pem") }
func (a *ACME) keyPath() string   { return filepath.Join(a.cacheDir, "key.pem") }
func (a *ACME) accPath() string   { return filepath.Join(a.cacheDir, "account.json") }
func (a *ACME) accKeyPath() string { return filepath.Join(a.cacheDir, "account.key") }

func (a *ACME) saveCert(res *certificate.Resource) error {
	if err := os.WriteFile(a.certPath(), res.Certificate, 0600); err != nil {
		return err
	}
	return os.WriteFile(a.keyPath(), res.PrivateKey, 0600)
}

// loadCachedCert returns the cached cert if it has >30 days remaining, nil otherwise.
func (a *ACME) loadCachedCert() (*tls.Certificate, error) {
	certPEM, err := os.ReadFile(a.certPath())
	if err != nil {
		return nil, err
	}
	keyPEM, err := os.ReadFile(a.keyPath())
	if err != nil {
		return nil, err
	}

	cert, err := parseCert(certPEM, keyPEM)
	if err != nil {
		return nil, err
	}

	// Check expiry — renew if <30 days left.
	if cert.Leaf != nil && time.Until(cert.Leaf.NotAfter) < 30*24*time.Hour {
		return nil, nil // signal: needs renewal
	}
	// Parse leaf if not already parsed.
	if cert.Leaf == nil {
		leaf, err := x509.ParseCertificate(cert.Certificate[0])
		if err == nil && time.Until(leaf.NotAfter) < 30*24*time.Hour {
			return nil, nil
		}
	}

	return cert, nil
}

func parseCert(certPEM, keyPEM []byte) (*tls.Certificate, error) {
	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("parse cert: %w", err)
	}
	// Pre-parse leaf for expiry checks.
	if len(cert.Certificate) > 0 {
		leaf, err := x509.ParseCertificate(cert.Certificate[0])
		if err == nil {
			cert.Leaf = leaf
		}
	}
	return &cert, nil
}

// ---- Account persistence ----

func (a *ACME) loadOrCreateAccount() (*acmeAccount, error) {
	// Try loading existing account.
	accData, errAcc := os.ReadFile(a.accPath())
	keyData, errKey := os.ReadFile(a.accKeyPath())

	if errAcc == nil && errKey == nil {
		acc := &acmeAccount{}
		if err := json.Unmarshal(accData, acc); err != nil {
			return nil, fmt.Errorf("parse account: %w", err)
		}
		key, err := parseECKey(keyData)
		if err != nil {
			return nil, fmt.Errorf("parse account key: %w", err)
		}
		acc.key = key
		acc.keyPEM = keyData
		return acc, nil
	}

	// Generate new account key.
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("gen account key: %w", err)
	}
	keyPEM, err := encodeECKey(key)
	if err != nil {
		return nil, err
	}

	acc := &acmeAccount{
		Email: a.cfg.Email,
		key:   key,
		keyPEM: keyPEM,
	}
	return acc, nil
}

func (a *ACME) saveAccount() error {
	data, err := json.MarshalIndent(a.account, "", "  ")
	if err != nil {
		return err
	}
	if err := os.WriteFile(a.accPath(), data, 0600); err != nil {
		return err
	}
	return os.WriteFile(a.accKeyPath(), a.account.keyPEM, 0600)
}

// ---- EC key helpers ----

func encodeECKey(key *ecdsa.PrivateKey) ([]byte, error) {
	der, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("marshal ec key: %w", err)
	}
	return pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: der}), nil
}

func parseECKey(pemData []byte) (*ecdsa.PrivateKey, error) {
	block, _ := pem.Decode(pemData)
	if block == nil {
		return nil, fmt.Errorf("no PEM block in account key")
	}
	return x509.ParseECPrivateKey(block.Bytes)
}