at-container-registry/deploy/upcloud/cloudinit.go

package main

import (
	"bytes"
	_ "embed"
	"fmt"
	"strings"
	"text/template"

	"go.yaml.in/yaml/v3"
)

//go:embed systemd/appview.service.tmpl
var appviewServiceTmpl string

//go:embed systemd/hold.service.tmpl
var holdServiceTmpl string

//go:embed systemd/scanner.service.tmpl
var scannerServiceTmpl string

//go:embed configs/appview.yaml.tmpl
var appviewConfigTmpl string

//go:embed configs/hold.yaml.tmpl
var holdConfigTmpl string

//go:embed configs/scanner.yaml.tmpl
var scannerConfigTmpl string

//go:embed systemd/labeler.service.tmpl
var labelerServiceTmpl string

//go:embed configs/labeler.yaml.tmpl
var labelerConfigTmpl string

//go:embed configs/cloudinit.sh.tmpl
var cloudInitTmpl string

// ConfigValues holds values injected into config YAML templates.
// Only truly dynamic/computed values belong here — deployment-specific
// values like client_name, owner_did, etc. are literal in the templates.
type ConfigValues struct {
	// S3 / Object Storage
	S3Endpoint  string
	S3Region    string
	S3Bucket    string
	S3AccessKey string
	S3SecretKey string

	// Infrastructure (computed from zone + config)
	Zone          string // e.g. "us-chi1"
	HoldDomain    string // e.g. "us-chi1.cove.seamark.dev"
	HoldDid       string // e.g. "did:web:us-chi1.cove.seamark.dev"
	LabelerDomain string // e.g. "labeler.seamark.dev"
	BasePath      string // e.g. "/var/lib/seamark"

	// Scanner (auto-generated shared secret)
	ScannerSecret string // hex-encoded 32-byte secret; empty disables scanning
	// ScannerHoldURL is the hold's WebSocket address as seen from the scanner
	// server, e.g. "ws://10.0.1.3:8080" over the private network.
	ScannerHoldURL string
}

// renderConfig executes a Go template with the given values.
func renderConfig(tmplStr string, vals *ConfigValues) (string, error) {
	t, err := template.New("config").Parse(tmplStr)
	if err != nil {
		return "", fmt.Errorf("parse config template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, vals); err != nil {
		return "", fmt.Errorf("render config template: %w", err)
	}
	return buf.String(), nil
}

// serviceUnitParams holds values for rendering systemd service unit templates.
type serviceUnitParams struct {
	DisplayName string // e.g. "Seamark"
	User        string // e.g. "seamark"
	BinaryPath  string // e.g. "/opt/seamark/bin/seamark-appview"
	ConfigPath  string // e.g. "/etc/seamark/appview.yaml"
	DataDir     string // e.g. "/var/lib/seamark"
	ServiceName string // e.g. "seamark-appview"
}

func renderServiceUnit(tmplStr string, p serviceUnitParams) (string, error) {
	t, err := template.New("service").Parse(tmplStr)
	if err != nil {
		return "", fmt.Errorf("parse service template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render service template: %w", err)
	}
	return buf.String(), nil
}

// scannerServiceUnitParams holds values for rendering the scanner systemd unit.
// Extends the standard fields with the memory limits, which are sized to the
// scanner's own server (see scannerPlanSpec) rather than shared with a hold.
type scannerServiceUnitParams struct {
	DisplayName string // e.g. "Seamark"
	User        string // e.g. "seamark"
	BinaryPath  string // e.g. "/opt/seamark/bin/seamark-scanner"
	ConfigPath  string // e.g. "/etc/seamark/scanner.yaml"
	DataDir     string // e.g. "/var/lib/seamark"
	ServiceName string // e.g. "seamark-scanner"
	GoMemLimit  string // GOMEMLIMIT, e.g. "2560MiB"
	MemoryHigh  string // systemd MemoryHigh, e.g. "3G"
	MemoryMax   string // systemd MemoryMax, e.g. "3500M"
}

// scannerMemoryLimits sizes the unit's memory knobs from the plan's RAM. The
// OS and page cache keep roughly the last 500 MiB; MemoryMax sits just under
// that line, MemoryHigh 500 MiB below it so reclaim starts before the kill,
// and the Go soft limit another 500 MiB below so the collector gets the first
// try. On the 4 GB scanner plan that is 2560MiB / 3G / 3500M.
func scannerMemoryLimits(memoryMB int) (goMemLimit, memoryHigh, memoryMax string) {
	maxMB := memoryMB - 596
	highMB := maxMB - 428
	goMB := highMB - 512
	if goMB < 512 {
		goMB = 512
	}
	return fmt.Sprintf("%dMiB", goMB), fmt.Sprintf("%dM", highMB), fmt.Sprintf("%dM", maxMB)
}

// scannerUnitParams builds the scanner unit parameters for a deployment.
func scannerUnitParams(naming Naming) scannerServiceUnitParams {
	goLimit, high, maxMem := scannerMemoryLimits(scannerPlanSpec.MemoryMB)
	return scannerServiceUnitParams{
		DisplayName: naming.DisplayName(),
		User:        naming.SystemUser(),
		BinaryPath:  naming.InstallDir() + "/bin/" + naming.Scanner(),
		ConfigPath:  naming.ScannerConfigPath(),
		DataDir:     naming.BasePath(),
		ServiceName: naming.Scanner(),
		GoMemLimit:  goLimit,
		MemoryHigh:  high,
		MemoryMax:   maxMem,
	}
}

func renderScannerServiceUnit(p scannerServiceUnitParams) (string, error) {
	t, err := template.New("scanner-service").Parse(scannerServiceTmpl)
	if err != nil {
		return "", fmt.Errorf("parse scanner service template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render scanner service template: %w", err)
	}
	return buf.String(), nil
}

// labelerServiceUnitParams holds values for rendering the labeler systemd unit.
type labelerServiceUnitParams struct {
	DisplayName        string // e.g. "Seamark"
	User               string // e.g. "seamark"
	BinaryPath         string // e.g. "/opt/seamark/bin/seamark-labeler"
	ConfigPath         string // e.g. "/etc/seamark/labeler.yaml"
	DataDir            string // e.g. "/var/lib/seamark"
	ServiceName        string // e.g. "seamark-labeler"
	AppviewServiceName string // e.g. "seamark-appview" (After= dependency)
}

func renderLabelerServiceUnit(p labelerServiceUnitParams) (string, error) {
	t, err := template.New("labeler-service").Parse(labelerServiceTmpl)
	if err != nil {
		return "", fmt.Errorf("parse labeler service template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render labeler service template: %w", err)
	}
	return buf.String(), nil
}

// generateAppviewCloudInit generates the cloud-init user-data script for the appview server.
// When withLabeler is true, a second phase is appended that creates labeler data
// directories and installs a labeler systemd service. Binaries are deployed separately via SCP.
func generateAppviewCloudInit(cfg *InfraConfig, vals *ConfigValues, withLabeler bool) (string, error) {
	naming := cfg.Naming()

	configYAML, err := renderConfig(appviewConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("appview config: %w", err)
	}

	serviceUnit, err := renderServiceUnit(appviewServiceTmpl, serviceUnitParams{
		DisplayName: naming.DisplayName(),
		User:        naming.SystemUser(),
		BinaryPath:  naming.InstallDir() + "/bin/" + naming.Appview(),
		ConfigPath:  naming.AppviewConfigPath(),
		DataDir:     naming.BasePath(),
		ServiceName: naming.Appview(),
	})
	if err != nil {
		return "", fmt.Errorf("appview service unit: %w", err)
	}

	script, err := generateCloudInit(cloudInitParams{
		BinaryName:  naming.Appview(),
		ServiceUnit: serviceUnit,
		ConfigYAML:  configYAML,
		ConfigPath:  naming.AppviewConfigPath(),
		ServiceName: naming.Appview(),
		DataDir:     naming.BasePath(),
		InstallDir:  naming.InstallDir(),
		SystemUser:  naming.SystemUser(),
		ConfigDir:   naming.ConfigDir(),
		LogFile:     naming.LogFile(),
		DisplayName: naming.DisplayName(),
	})
	if err != nil {
		return "", err
	}

	if !withLabeler {
		return script, nil
	}

	// Render labeler config YAML
	labelerConfigYAML, err := renderConfig(labelerConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("labeler config: %w", err)
	}

	// Append labeler setup phase
	labelerUnit, err := renderLabelerServiceUnit(labelerServiceUnitParams{
		DisplayName:        naming.DisplayName(),
		User:               naming.SystemUser(),
		BinaryPath:         naming.InstallDir() + "/bin/" + naming.Labeler(),
		ConfigPath:         naming.LabelerConfigPath(),
		DataDir:            naming.BasePath(),
		ServiceName:        naming.Labeler(),
		AppviewServiceName: naming.Appview(),
	})
	if err != nil {
		return "", fmt.Errorf("labeler service unit: %w", err)
	}

	// Escape single quotes for heredoc embedding
	labelerUnit = strings.ReplaceAll(labelerUnit, "'", "'\\''")
	labelerConfigYAML = strings.ReplaceAll(labelerConfigYAML, "'", "'\\''")

	labelerPhase := fmt.Sprintf(`
# === Labeler Setup ===

# Labeler data dirs
mkdir -p %s
chown -R %s:%s %s

# Labeler config
cat > %s << 'CFGEOF'
%s
CFGEOF

# Labeler systemd service
cat > /etc/systemd/system/%s.service << 'SVCEOF'
%s
SVCEOF
systemctl daemon-reload
systemctl enable %s

echo "=== Labeler setup complete ==="
`,
		naming.LabelerDataDir(),
		naming.SystemUser(), naming.SystemUser(), naming.LabelerDataDir(),
		naming.LabelerConfigPath(),
		labelerConfigYAML,
		naming.Labeler(),
		labelerUnit,
		naming.Labeler(),
	)

	return script + labelerPhase, nil
}

// generateHoldCloudInit generates the cloud-init user-data script for the hold
// server. Binaries are deployed separately via SCP. The scanner has its own
// server (generateScannerCloudInit) and no longer rides along here.
func generateHoldCloudInit(cfg *InfraConfig, vals *ConfigValues) (string, error) {
	naming := cfg.Naming()

	configYAML, err := renderConfig(holdConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("hold config: %w", err)
	}

	serviceUnit, err := renderServiceUnit(holdServiceTmpl, serviceUnitParams{
		DisplayName: naming.DisplayName(),
		User:        naming.SystemUser(),
		BinaryPath:  naming.InstallDir() + "/bin/" + naming.Hold(),
		ConfigPath:  naming.HoldConfigPath(),
		DataDir:     naming.BasePath(),
		ServiceName: naming.Hold(),
	})
	if err != nil {
		return "", fmt.Errorf("hold service unit: %w", err)
	}

	return generateCloudInit(cloudInitParams{
		BinaryName:  naming.Hold(),
		ServiceUnit: serviceUnit,
		ConfigYAML:  configYAML,
		ConfigPath:  naming.HoldConfigPath(),
		ServiceName: naming.Hold(),
		DataDir:     naming.BasePath(),
		InstallDir:  naming.InstallDir(),
		SystemUser:  naming.SystemUser(),
		ConfigDir:   naming.ConfigDir(),
		LogFile:     naming.LogFile(),
		DisplayName: naming.DisplayName(),
	})
}

// generateScannerCloudInit generates the cloud-init user-data script for the
// scanner's own server. It is the standard base script (packages, swap,
// service user, config, unit) plus the vulnerability database and extraction
// directories the scanner needs. The binary is deployed separately via SCP.
// vals.ScannerHoldURL must already point at the hold's private address.
func generateScannerCloudInit(cfg *InfraConfig, vals *ConfigValues) (string, error) {
	if vals.ScannerHoldURL == "" {
		return "", fmt.Errorf("scanner cloud-init: hold URL is empty (hold server must exist first)")
	}
	naming := cfg.Naming()

	configYAML, err := renderConfig(scannerConfigTmpl, vals)
	if err != nil {
		return "", fmt.Errorf("scanner config: %w", err)
	}

	serviceUnit, err := renderScannerServiceUnit(scannerUnitParams(naming))
	if err != nil {
		return "", fmt.Errorf("scanner service unit: %w", err)
	}

	script, err := generateCloudInit(cloudInitParams{
		BinaryName:  naming.Scanner(),
		ServiceUnit: serviceUnit,
		ConfigYAML:  configYAML,
		ConfigPath:  naming.ScannerConfigPath(),
		ServiceName: naming.Scanner(),
		DataDir:     naming.BasePath(),
		InstallDir:  naming.InstallDir(),
		SystemUser:  naming.SystemUser(),
		ConfigDir:   naming.ConfigDir(),
		LogFile:     naming.LogFile(),
		DisplayName: naming.DisplayName(),
	})
	if err != nil {
		return "", err
	}

	return script + scannerDirsScript(naming), nil
}

// scannerDirsScript creates the scanner's data directories. Shared between
// cloud-init (first boot) and update (a server whose directories were removed).
func scannerDirsScript(naming Naming) string {
	return fmt.Sprintf(`
# Scanner data dirs (vulnerability database, layer extraction scratch)
mkdir -p %s/vulndb %s/tmp
chown -R %s:%s %s
`,
		naming.ScannerDataDir(), naming.ScannerDataDir(),
		naming.SystemUser(), naming.SystemUser(), naming.ScannerDataDir(),
	)
}

type cloudInitParams struct {
	BinaryName  string
	ServiceUnit string
	ConfigYAML  string
	ConfigPath  string
	ServiceName string
	DataDir     string
	InstallDir  string
	SystemUser  string
	ConfigDir   string
	LogFile     string
	DisplayName string
}

func generateCloudInit(p cloudInitParams) (string, error) {
	// Escape single quotes in embedded content for heredoc safety
	p.ServiceUnit = strings.ReplaceAll(p.ServiceUnit, "'", "'\\''")
	p.ConfigYAML = strings.ReplaceAll(p.ConfigYAML, "'", "'\\''")

	t, err := template.New("cloudinit").Parse(cloudInitTmpl)
	if err != nil {
		return "", fmt.Errorf("parse cloudinit template: %w", err)
	}
	var buf bytes.Buffer
	if err := t.Execute(&buf, p); err != nil {
		return "", fmt.Errorf("render cloudinit template: %w", err)
	}
	return buf.String(), nil
}

// syncServiceUnit compares a rendered systemd service unit against what's on
// the server. If they differ, it writes the new unit file. If the unit is
// missing entirely, it installs it and runs `systemctl enable` so the service
// starts on boot. Returns true if the unit was created or updated (caller
// should daemon-reload before restart).
func syncServiceUnit(name, ip, serviceName, renderedUnit string) (bool, error) {
	unitPath := "/etc/systemd/system/" + serviceName + ".service"

	remote, err := runSSH(ip, fmt.Sprintf("cat %s 2>/dev/null || echo '__MISSING__'", unitPath), false)
	if err != nil {
		fmt.Printf("  service unit sync: could not reach %s (%v)\n", name, err)
		return false, nil
	}
	remote = strings.TrimSpace(remote)
	rendered := strings.TrimSpace(renderedUnit)

	if remote == "__MISSING__" {
		// First-time install: write file, daemon-reload, and enable so the
		// service comes up on boot. The caller's restart will start it.
		script := fmt.Sprintf("cat > %s << 'SVCEOF'\n%s\nSVCEOF\nsystemctl daemon-reload\nsystemctl enable %s",
			unitPath, rendered, serviceName)
		if _, err := runSSH(ip, script, false); err != nil {
			return false, fmt.Errorf("install service unit: %w", err)
		}
		fmt.Printf("  service unit: %s installed and enabled\n", name)
		return true, nil
	}

	if remote == rendered {
		fmt.Printf("  service unit: %s up to date\n", name)
		return false, nil
	}

	// Write the updated unit file
	script := fmt.Sprintf("cat > %s << 'SVCEOF'\n%s\nSVCEOF", unitPath, rendered)
	if _, err := runSSH(ip, script, false); err != nil {
		return false, fmt.Errorf("write service unit: %w", err)
	}
	fmt.Printf("  service unit: %s updated\n", name)
	return true, nil
}

// syncConfigKeys fetches the existing config from a server and merges in any
// missing keys from the rendered template. Existing values are never overwritten.
//
// installIfMissing is true only during provision. On an update a missing
// config is never a first install: the service has been running with one, so
// the file has been moved or deleted, and writing the template in its place
// would hand the service empty identity fields (owner DID, DID, keys) that it
// refuses to start with. That is exactly what happened to the labeler on
// 2026-09-09, and it only surfaced at the next restart three days later.
func syncConfigKeys(name, ip, configPath, templateYAML string, installIfMissing bool) error {
	remote, err := runSSH(ip, fmt.Sprintf("cat %s 2>/dev/null || echo '__MISSING__'", configPath), false)
	if err != nil {
		if !installIfMissing {
			return fmt.Errorf("read %s config %s: %w", name, configPath, err)
		}
		fmt.Printf("  config sync: could not reach %s (%v)\n", name, err)
		return nil
	}
	remote = strings.TrimSpace(remote)

	if remote == "__MISSING__" {
		if !installIfMissing {
			return fmt.Errorf("%s config %s is missing on the server; refusing to install the template during an update because its identity fields are empty. Restore the file from the newest predeploy-* backup under the data directory and rerun", name, configPath)
		}
		// First-time install: write the rendered template as-is. Subsequent
		// runs use the merge-keys path below to preserve operator edits.
		dir := configPath[:strings.LastIndex(configPath, "/")]
		if _, err := runSSH(ip, fmt.Sprintf("mkdir -p %s", dir), false); err != nil {
			return fmt.Errorf("create config dir: %w", err)
		}
		script := fmt.Sprintf("cat > %s << 'CFGEOF'\n%s\nCFGEOF", configPath, strings.TrimRight(templateYAML, "\n"))
		if _, err := runSSH(ip, script, false); err != nil {
			return fmt.Errorf("write initial config: %w", err)
		}
		fmt.Printf("  config sync: %s installed\n", name)
		return nil
	}

	// Parse both into yaml.Node trees
	var templateDoc yaml.Node
	if err := yaml.Unmarshal([]byte(templateYAML), &templateDoc); err != nil {
		return fmt.Errorf("parse template yaml: %w", err)
	}
	var existingDoc yaml.Node
	if err := yaml.Unmarshal([]byte(remote), &existingDoc); err != nil {
		return fmt.Errorf("parse remote yaml: %w", err)
	}

	// Unwrap document nodes to get the root mapping
	templateRoot := unwrapDocNode(&templateDoc)
	existingRoot := unwrapDocNode(&existingDoc)
	if templateRoot == nil || existingRoot == nil {
		fmt.Printf("  config sync: %s skipped (unexpected YAML structure)\n", name)
		return nil
	}

	added := mergeYAMLNodes(templateRoot, existingRoot)
	if !added {
		fmt.Printf("  config sync: %s up to date\n", name)
		return nil
	}

	// Marshal the modified tree back
	merged, err := yaml.Marshal(&existingDoc)
	if err != nil {
		return fmt.Errorf("marshal merged yaml: %w", err)
	}

	// Write back to server
	script := fmt.Sprintf("cat > %s << 'CFGEOF'\n%sCFGEOF", configPath, string(merged))
	if _, err := runSSH(ip, script, false); err != nil {
		return fmt.Errorf("write merged config: %w", err)
	}
	fmt.Printf("  config sync: %s updated with new keys\n", name)
	return nil
}

// unwrapDocNode returns the root mapping node, unwrapping a DocumentNode wrapper if present.
func unwrapDocNode(n *yaml.Node) *yaml.Node {
	if n.Kind == yaml.DocumentNode && len(n.Content) > 0 {
		return n.Content[0]
	}
	if n.Kind == yaml.MappingNode {
		return n
	}
	return nil
}

// mergeYAMLNodes recursively adds keys from base into existing that are not
// already present. Existing values are never overwritten. Returns true if any
// new keys were added.
func mergeYAMLNodes(base, existing *yaml.Node) bool {
	if base.Kind != yaml.MappingNode || existing.Kind != yaml.MappingNode {
		return false
	}

	added := false
	for i := 0; i+1 < len(base.Content); i += 2 {
		baseKey := base.Content[i]
		baseVal := base.Content[i+1]

		// Look for this key in existing
		found := false
		for j := 0; j+1 < len(existing.Content); j += 2 {
			if existing.Content[j].Value == baseKey.Value {
				found = true
				// If both are mappings, recurse to merge sub-keys
				if baseVal.Kind == yaml.MappingNode && existing.Content[j+1].Kind == yaml.MappingNode {
					if mergeYAMLNodes(baseVal, existing.Content[j+1]) {
						added = true
					}
				}
				break
			}
		}

		if !found {
			// Append the missing key+value pair
			existing.Content = append(existing.Content, baseKey, baseVal)
			added = true
		}
	}

	return added
}