clavitor/clavis/clavis-vault/lib/tls.go

package lib

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/caddyserver/certmagic"
	"github.com/libdns/cloudflare"
	"github.com/mholt/acmez/v3/acme"
)

// TLSConfig holds configuration for TLS.
type TLSConfig struct {
	Domain  string // e.g. "use1.clavitor.ai"
	CFToken string // Cloudflare API token for DNS-01 challenge
	DataDir string // directory to store certs
	Email   string // ACME account email
}

// LoadTLSConfig reads TLS configuration from environment variables.
func LoadTLSConfig() TLSConfig {
	return TLSConfig{
		Domain:  os.Getenv("TLS_DOMAIN"),
		CFToken: os.Getenv("CF_API_TOKEN"),
		DataDir: os.Getenv("TLS_CERT_DIR"),
		Email:   os.Getenv("TLS_EMAIL"),
	}
}

// ServeTLSWithContext starts an HTTPS server and gracefully shuts down on context cancellation.
// This is the preferred method - it handles SIGTERM properly unlike the blocking ListenAndServeTLS.
func ServeTLSWithContext(ctx context.Context, addr string, handler http.Handler, cfg TLSConfig) error {
	if cfg.Domain != "" && cfg.CFToken != "" {
		return serveLetsEncryptWithContext(ctx, addr, handler, cfg)
	}
	return serveSelfSignedWithContext(ctx, addr, handler, cfg)
}

// serveLetsEncryptWithContext uses certmagic for automatic Let's Encrypt certificates.
func serveLetsEncryptWithContext(ctx context.Context, addr string, handler http.Handler, cfg TLSConfig) error {
	if cfg.DataDir == "" {
		cfg.DataDir = "/opt/clavitor/certs"
	}
	if cfg.Email == "" {
		cfg.Email = "ops@clavitor.ai"
	}

	certmagic.DefaultACME.Agreed = true
	certmagic.DefaultACME.Email = cfg.Email
	certmagic.DefaultACME.DNS01Solver = &certmagic.DNS01Solver{
		DNSManager: certmagic.DNSManager{
			DNSProvider: &cloudflare.Provider{
				APIToken: cfg.CFToken,
			},
		},
	}

	magic := certmagic.NewDefault()
	magic.Storage = &certmagic.FileStorage{Path: cfg.DataDir}

	manageCtx := context.Background()
	if err := magic.ManageSync(manageCtx, []string{cfg.Domain}); err != nil {
		return fmt.Errorf("certmagic manage %s: %w", cfg.Domain, err)
	}

	tlsConfig := magic.TLSConfig()
	tlsConfig.NextProtos = []string{"h2", "http/1.1"}

	server := &http.Server{
		Addr:      addr,
		Handler:   handler,
		TLSConfig: tlsConfig,
	}

	// Start server in goroutine
	go func() {
		log.Printf("Clavitor listening on https://0.0.0.0%s (%s)", addr, cfg.Domain)
		if err := server.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
			log.Fatalf("server: %v", err)
		}
	}()

	// Wait for shutdown signal
	<-ctx.Done()

	// Graceful shutdown with 5-second timeout
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return server.Shutdown(shutdownCtx)
}

// serveSelfSignedWithContext generates a self-signed certificate and serves HTTPS.
func serveSelfSignedWithContext(ctx context.Context, addr string, handler http.Handler, cfg TLSConfig) error {
	certDir := cfg.DataDir
	if certDir == "" {
		certDir = "."
	}
	certPath := filepath.Join(certDir, "clavitor.crt")
	keyPath := filepath.Join(certDir, "clavitor.key")

	if err := ensureSelfSignedCert(certPath, keyPath); err != nil {
		return fmt.Errorf("self-signed cert: %w", err)
	}

	server := &http.Server{
		Addr:    addr,
		Handler: handler,
	}

	// Start server in goroutine
	go func() {
		log.Printf("Clavitor listening on https://0.0.0.0%s (self-signed)", addr)
		if err := server.ListenAndServeTLS(certPath, keyPath); err != nil && err != http.ErrServerClosed {
			log.Fatalf("server: %v", err)
		}
	}()

	// Wait for shutdown signal
	<-ctx.Done()

	// Graceful shutdown with 5-second timeout
	// This tells the tarpit to exit early via the shutdown context
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return server.Shutdown(shutdownCtx)
}

// ListenAndServeTLS starts an HTTPS server (legacy blocking method).
// Deprecated: Use ServeTLSWithContext for graceful shutdown support.
func ListenAndServeTLS(addr string, handler http.Handler, cfg TLSConfig) error {
	if cfg.Domain != "" && cfg.CFToken != "" {
		return listenLetsEncrypt(addr, handler, cfg)
	}
	return listenSelfSigned(addr, handler, cfg)
}

// listenLetsEncrypt uses certmagic for automatic Let's Encrypt certificates.
func listenLetsEncrypt(addr string, handler http.Handler, cfg TLSConfig) error {
	if cfg.DataDir == "" {
		cfg.DataDir = "/opt/clavitor/certs"
	}
	if cfg.Email == "" {
		cfg.Email = "ops@clavitor.ai"
	}

	certmagic.DefaultACME.Agreed = true
	certmagic.DefaultACME.Email = cfg.Email
	certmagic.DefaultACME.DNS01Solver = &certmagic.DNS01Solver{
		DNSManager: certmagic.DNSManager{
			DNSProvider: &cloudflare.Provider{
				APIToken: cfg.CFToken,
			},
		},
	}

	magic := certmagic.NewDefault()
	magic.Storage = &certmagic.FileStorage{Path: cfg.DataDir}

	ctx := context.Background()
	if err := magic.ManageSync(ctx, []string{cfg.Domain}); err != nil {
		return fmt.Errorf("certmagic manage %s: %w", cfg.Domain, err)
	}

	tlsConfig := magic.TLSConfig()
	tlsConfig.NextProtos = []string{"h2", "http/1.1"}

	server := &http.Server{
		Addr:      addr,
		Handler:   handler,
		TLSConfig: tlsConfig,
	}

	log.Printf("Clavitor listening on https://0.0.0.0%s (%s)", addr, cfg.Domain)
	return server.ListenAndServeTLS("", "")
}

// listenSelfSigned generates a self-signed certificate and serves HTTPS.
// Cert is stored in DataDir (or working directory) so it persists across restarts.
// The browser will show a certificate warning — the connection is still encrypted.
func listenSelfSigned(addr string, handler http.Handler, cfg TLSConfig) error {
	certDir := cfg.DataDir
	if certDir == "" {
		certDir = "."
	}
	certPath := filepath.Join(certDir, "clavitor.crt")
	keyPath := filepath.Join(certDir, "clavitor.key")

	if err := ensureSelfSignedCert(certPath, keyPath); err != nil {
		return fmt.Errorf("self-signed cert: %w", err)
	}

	server := &http.Server{
		Addr:    addr,
		Handler: handler,
	}

	log.Printf("Clavitor listening on https://0.0.0.0%s (self-signed)", addr)
	return server.ListenAndServeTLS(certPath, keyPath)
}

// ensureSelfSignedCert creates a self-signed TLS certificate if one doesn't exist.
// Valid for 10 years, covers localhost, 127.0.0.1, and all local network IPs.
func ensureSelfSignedCert(certPath, keyPath string) error {
	if _, err := os.Stat(certPath); err == nil {
		if _, err := os.Stat(keyPath); err == nil {
			return nil // both exist
		}
	}

	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return err
	}

	serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))

	tmpl := x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{Organization: []string{"Clavitor"}, CommonName: "clavitor"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(10 * 365 * 24 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{"localhost"},
		IPAddresses:  []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback},
	}

	// Add all local interface IPs so LAN access works
	if addrs, err := net.InterfaceAddrs(); err == nil {
		for _, a := range addrs {
			if ipnet, ok := a.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
				tmpl.IPAddresses = append(tmpl.IPAddresses, ipnet.IP)
			}
		}
	}

	certDER, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &key.PublicKey, key)
	if err != nil {
		return err
	}

	// Write cert and key as PEM
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return err
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})

	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return err
	}
	if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
		return err
	}

	log.Printf("Generated self-signed TLS certificate: %s", certPath)
	return nil
}

// ChallengeError wraps ACME errors with actionable context.
func ChallengeError(err error) string {
	if acmeErr, ok := err.(acme.Problem); ok {
		return fmt.Sprintf("ACME error: %s (type: %s)", acmeErr.Detail, acmeErr.Type)
	}
	return err.Error()
}

// SelfSignedTLSConfig returns a tls.Config with a fresh self-signed certificate.
// Used by tests that need HTTPS without filesystem cert files.
func SelfSignedTLSConfig() *tls.Config {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	tmpl := x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "test"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(24 * time.Hour),
		DNSNames:     []string{"localhost"},
		IPAddresses:  []net.IP{net.IPv4(127, 0, 0, 1)},
	}
	certDER, _ := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &key.PublicKey, key)
	cert := tls.Certificate{
		Certificate: [][]byte{certDER},
		PrivateKey:  key,
	}
	return &tls.Config{Certificates: []tls.Certificate{cert}}
}