uncloud/internal/machine/caddyconfig/caddyfile.go

package caddyconfig

import (
	"bytes"
	"cmp"
	"context"
	"errors"
	"fmt"
	"log/slog"
	"maps"
	"net"
	"slices"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/psviderski/uncloud/internal/machine/store"
	"github.com/psviderski/uncloud/pkg/api"
)

const (
	caddyfileHeaderFmt = `# Caddyfile autogenerated by Uncloud on machine '%s' (DO NOT EDIT): %s
# Automatically updated on service or health status changes.
# Docs: https://uncloud.run/docs/concepts/ingress/overview
`
	caddyfileTemplate = `# Health check endpoint to verify Caddy reachability on this machine.
http:// {
	handle {{.VerifyPath}} {
		respond "{{.VerifyResponse}}" 200
	}
	log
}

(common_proxy) {
	# Retry failed requests up to lb_retries times against other available upstreams.
	lb_retries 3
	# Upstreams are marked unhealthy for fail_duration after a failed request (passive health checking).
	fail_duration 30s
}
{{- if or .HTTPHostUpstreams .HTTPSHostUpstreams }}

# Sites generated from service ports.{{end}}
{{- range $hostname, $upstreams := .HTTPHostUpstreams}}

http://{{$hostname}} {
	reverse_proxy {{join $upstreams " "}} {
		import common_proxy
	}
	log
}{{end}}
{{- range $hostname, $upstreams := .HTTPSHostUpstreams}}

https://{{$hostname}} {
	reverse_proxy {{join $upstreams " "}} {
		import common_proxy
	}
	log
}{{end}}
`
	bootstrapMarker       = "# Uncloud bootstrap for Caddy container "
	legacyBootstrapMarker = "# NOTE: User-defined configs for services were skipped because Caddy is not running"
)

// CaddyfileGenerator generates a Caddyfile configuration for the Caddy reverse proxy.
// It combines generated routes from published ports with user-defined Caddyfile snippets from service specs (x-caddy).
type CaddyfileGenerator struct {
	// machineID is the unique identifier of the machine where the controller is running.
	machineID string
	// machineName is the human-friendly name of the machine.
	machineName string
	validator   CaddyfileValidator
	log         *slog.Logger
}

// CaddyfileValidator checks a candidate Caddyfile without loading it. Validate returns InvalidCaddyfileError only
// when the candidate is known to be invalid. Other errors mean the check could not be completed. A successful check
// does not guarantee that the configuration will load.
type CaddyfileValidator interface {
	Validate(ctx context.Context, caddyfile string) error
}

// InvalidCaddyfileError means validation completed and the candidate Caddyfile was rejected.
type InvalidCaddyfileError struct{ Message string }

func (e *InvalidCaddyfileError) Error() string { return e.Message }

func NewCaddyfileGenerator(
	machineID, machineName string, validator CaddyfileValidator, log *slog.Logger,
) *CaddyfileGenerator {
	if log == nil {
		log = slog.Default()
	}
	return &CaddyfileGenerator{
		machineID:   machineID,
		machineName: machineName,
		validator:   validator,
		log:         log,
	}
}

// Generate creates a Caddyfile for the local Caddy container from the supplied healthy application containers.
// Application service ports provide the generated sites. The Caddy container's custom Caddy config (x-caddy), if
// defined, provides the global block even when that container is unhealthy. When application custom Caddy configs
// are included, the newest container for each service provides its config.
//
// The resulting Caddyfile has this structure:
//
//	[caddy custom Caddy config (global)]
//	[generated sites from application service ports]
//	[service-a custom Caddy config]
//	...
//	[service-z custom Caddy config]
//
// Bootstrap mode keeps the global custom Caddy config and generated sites, but omits application custom Caddy configs.
// It skips validation because Caddy may not be running yet. In this case, Caddy validates the config when it starts.
// Global template errors still stop generation. In full mode, invalid globals stop generation, while invalid
// application custom Caddy configs are logged and skipped.
func (g *CaddyfileGenerator) Generate(
	ctx context.Context,
	caddyCtr api.ServiceContainer,
	records []store.ContainerRecord,
	bootstrap bool,
) (string, error) {
	records = slices.Clone(records)
	// Sort records by local machine first, then by service name and creation time. Placing containers on the local
	// machine first lets user-defined Caddy configs pair this ordering with the "first" lb_policy to always send
	// traffic to the same-host replica (skipping the cross-machine hop) and only fall back to remote upstreams when
	// the local one is unhealthy.
	// The service name, creation time, and container ID tiebreakers keep the file stable across regenerations.
	slices.SortStableFunc(records, func(a, b store.ContainerRecord) int {
		return cmp.Or(
			g.localMachineRank(a.MachineID)-g.localMachineRank(b.MachineID),
			strings.Compare(a.Container.ServiceName(), b.Container.ServiceName()),
			a.Container.CreatedTime().Compare(b.Container.CreatedTime()),
			strings.Compare(a.Container.ID, b.Container.ID),
		)
	})

	containers := make([]api.ServiceContainer, len(records))
	for i, cr := range records {
		containers[i] = cr.Container
	}

	caddyfile, err := g.generateBaseFromPorts(containers)
	if err != nil {
		return "", fmt.Errorf("generate base Caddyfile from service ports: %w", err)
	}

	caddyfileHeader := fmt.Sprintf(caddyfileHeaderFmt, g.machineName, time.Now().UTC().Format(time.RFC3339))
	upstreams := serviceUpstreams(containers)
	// Track validation errors for reporting.
	var configErrors []string

	if caddyCtr.ServiceSpec.CaddyConfig() != "" {
		// Render the custom global Caddy config as a Go template with the upstreams.
		tmplCtx := templateContext{
			Name:      caddyCtr.ServiceName(),
			Upstreams: upstreams,
		}
		renderedConfig, err := renderCaddyfile(tmplCtx, caddyCtr.ServiceSpec.CaddyConfig())
		if err != nil {
			return "", fmt.Errorf(
				"render template directives in user-defined global Caddy config from Caddy container %s: %w",
				caddyCtr.ShortID(), err)
		}
		caddyfileCandidate := fmt.Sprintf("# User-defined global config from service '%s'.\n%s\n\n%s",
			caddyCtr.ServiceName(), renderedConfig, caddyfile)
		if !bootstrap && g.validator != nil {
			if err = g.validator.Validate(ctx, caddyfileCandidate); err != nil {
				return "", fmt.Errorf("validate user-defined global Caddy config from Caddy container %s: %w",
					caddyCtr.ShortID(), err)
			}
		}
		caddyfile = caddyfileCandidate
	}

	if bootstrap {
		return caddyfileHeader + bootstrapMarker + caddyCtr.ID + "\n\n" + caddyfile, nil
	}

	// There could be multiple service containers for the same service with different custom Caddy configs, for example,
	// if the service has been partially updated. The most recent container for each service defines the current custom
	// Caddy config for that service.
	latestServiceContainers := make(map[string]api.ServiceContainer, len(containers))
	for _, ctr := range containers {
		if latest, ok := latestServiceContainers[ctr.ServiceName()]; ok {
			if ctr.CreatedTime().Compare(latest.CreatedTime()) > 0 {
				latestServiceContainers[ctr.ServiceName()] = ctr
			}
		} else {
			latestServiceContainers[ctr.ServiceName()] = ctr
		}
	}
	sortedServiceNames := slices.Sorted(maps.Keys(latestServiceContainers))

	// Append a custom Caddy config for each service to the Caddyfile and validate it. If the config for a service
	// is invalid, skip it but continue processing other services to ensure the Caddyfile remains valid.
	for _, serviceName := range sortedServiceNames {
		// Skip the caddy container as we already processed it as the global config.
		if serviceName == CaddyServiceName {
			continue
		}

		ctr := latestServiceContainers[serviceName]
		if ctr.ServiceSpec.CaddyConfig() == "" {
			continue
		}

		// Render the template actions in the service's Caddy config.
		tmplCtx := templateContext{
			Name:      serviceName,
			Upstreams: upstreams,
		}
		renderedConfig, err := renderCaddyfile(tmplCtx, ctr.ServiceSpec.CaddyConfig())
		if err != nil {
			g.log.Error("Failed to render template directives in user-defined Caddy config for service, skipping it.",
				"service", serviceName, "err", err)
			configErrors = append(configErrors,
				fmt.Sprintf("service '%s': failed to render template: %v", serviceName, err))
			continue
		}

		caddyfileCandidate := fmt.Sprintf("%s\n# User-defined config for service '%s'.\n%s\n",
			caddyfile, serviceName, renderedConfig)
		if g.validator != nil {
			if err = g.validator.Validate(ctx, caddyfileCandidate); err != nil {
				if _, ok := errors.AsType[*InvalidCaddyfileError](err); !ok {
					return "", fmt.Errorf("validate Caddy config for service '%s': %w", serviceName, err)
				}
				g.log.Error("User-defined Caddy config for service is invalid, skipping it.",
					"service", serviceName, "err", err)
				configErrors = append(configErrors,
					fmt.Sprintf("service '%s': validation failed: %v", serviceName, err))
				continue
			}
		}
		caddyfile = caddyfileCandidate
	}

	// Keep skipped-snippet errors in the saved file so an operator can see why routes are missing after a restart.
	if len(configErrors) > 0 {
		var errorsComment strings.Builder
		errorsComment.WriteString("# Skipped invalid user-defined configs:\n")
		for _, e := range configErrors {
			errorsComment.WriteString(fmt.Sprintf("# - %s\n", e))
		}

		caddyfile += "\n" + errorsComment.String()
	}

	return caddyfileHeader + "\n" + caddyfile, nil
}

// localMachineRank is a sorting function for containers that puts running on the local machine first.
func (g *CaddyfileGenerator) localMachineRank(machineID string) int {
	if g.machineID == machineID {
		return 0
	}
	return 1
}

func (g *CaddyfileGenerator) generateBaseFromPorts(containers []api.ServiceContainer) (string, error) {
	httpHostUpstreams, httpsHostUpstreams := httpUpstreamsFromPorts(containers)

	funcs := template.FuncMap{"join": strings.Join}
	tmpl, err := template.New("Caddyfile").Funcs(funcs).Parse(caddyfileTemplate)
	if err != nil {
		return "", fmt.Errorf("parse Caddyfile template: %w", err)
	}

	data := struct {
		VerifyPath         string
		VerifyResponse     string
		HTTPHostUpstreams  map[string][]string
		HTTPSHostUpstreams map[string][]string
	}{
		VerifyPath:         VerifyPath,
		VerifyResponse:     g.machineID,
		HTTPHostUpstreams:  httpHostUpstreams,
		HTTPSHostUpstreams: httpsHostUpstreams,
	}

	var buf bytes.Buffer
	if err = tmpl.Execute(&buf, data); err != nil {
		return "", fmt.Errorf("execute Caddyfile template: %w", err)
	}

	return buf.String(), nil
}

// httpUpstreamsFromPorts extracts upstreams for HTTP and HTTPS protocols from the published ports of the provided
// service containers. It's expected that all containers are healthy.
func httpUpstreamsFromPorts(containers []api.ServiceContainer) (map[string][]string, map[string][]string) {
	// Maps hostnames to lists of upstreams (container IP:port pairs).
	httpHostUpstreams := make(map[string][]string)
	httpsHostUpstreams := make(map[string][]string)
	for _, ctr := range containers {
		ip := ctr.UncloudNetworkIP()
		if !ip.IsValid() {
			// Container is not connected to the uncloud Docker network (could be host network).
			continue
		}
		log := slog.With("container", ctr.ID)

		ports, err := ctr.ServicePorts()
		if err != nil {
			log.Error("Failed to parse service ports for container.", "err", err)
			continue
		}

		for _, port := range ports {
			if port.Mode != api.PortModeIngress {
				continue
			}

			switch port.Protocol {
			case api.ProtocolHTTP:
				upstream := net.JoinHostPort(ip.String(), strconv.Itoa(int(port.ContainerPort)))
				httpHostUpstreams[port.Hostname] = append(httpHostUpstreams[port.Hostname], upstream)
			case api.ProtocolHTTPS:
				upstream := net.JoinHostPort(ip.String(), strconv.Itoa(int(port.ContainerPort)))
				httpsHostUpstreams[port.Hostname] = append(httpsHostUpstreams[port.Hostname], upstream)
			default:
				// TODO: implement L4 ingress routing for TCP and UDP.
				log.Error("Unsupported protocol for ingress port.", "port", port)
				continue
			}
		}
	}

	return httpHostUpstreams, httpsHostUpstreams
}

// serviceUpstreams creates a map of service names to their container IPs.
// Only includes containers connected to the uncloud Docker network.
func serviceUpstreams(containers []api.ServiceContainer) map[string][]string {
	upstreams := make(map[string][]string)
	for _, ctr := range containers {
		ip := ctr.UncloudNetworkIP()
		if !ip.IsValid() {
			// Container is not connected to the uncloud Docker network (could be host network).
			continue
		}

		serviceName := ctr.ServiceName()
		upstreams[serviceName] = append(upstreams[serviceName], ip.String())
	}

	return upstreams
}

// renderCaddyfile renders a Caddyfile template with the upstreams function and data.
func renderCaddyfile(tmplCtx templateContext, caddyfile string) (string, error) {
	funcs := template.FuncMap{
		"upstreams": upstreamsTemplateFn(tmplCtx),
	}

	tmpl, err := template.New("Caddyfile").Funcs(funcs).Parse(caddyfile)
	if err != nil {
		return "", fmt.Errorf("parse config as Go template: %w", err)
	}

	var buf bytes.Buffer
	if err = tmpl.Execute(&buf, tmplCtx); err != nil {
		return "", fmt.Errorf("execute template: %w", err)
	}

	return buf.String(), nil
}