URBM/internal/platform/workloads.go

package platform

import (
	"context"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"git.casaderoll.de/michael/urbm/internal/model"
)

type WorkloadManager struct {
	RuntimeDir  string
	FindmntPath string
	LsblkPath   string
	Stat        func(string) (os.FileInfo, error)
}

type Workload struct {
	ID     string `json:"id"`
	Name   string `json:"name"`
	Status string `json:"status"`
}

type WorkloadDiscovery struct {
	Available bool       `json:"available"`
	Message   string     `json:"message,omitempty"`
	Items     []Workload `json:"items"`
}

type Prepared struct {
	Sources     []string
	Stopped     []model.Source
	MetadataDir string
	Kind        model.JobType
}

func (w *WorkloadManager) Discover(ctx context.Context, kind model.JobType) (WorkloadDiscovery, error) {
	switch kind {
	case model.JobDocker:
		if _, err := os.Stat("/var/run/docker.sock"); err != nil {
			return WorkloadDiscovery{Available: false, Message: "Docker is not enabled or is not running on this Unraid server.", Items: []Workload{}}, nil
		}
		output, err := exec.CommandContext(ctx, "docker", "ps", "-a", "--format", "{{json .}}").Output()
		if err != nil {
			return WorkloadDiscovery{}, fmt.Errorf("environment: list Docker containers: %w", err)
		}
		var result []Workload
		for _, line := range strings.Split(strings.TrimSpace(string(output)), "\n") {
			if strings.TrimSpace(line) == "" {
				continue
			}
			var item struct {
				ID     string `json:"ID"`
				Names  string `json:"Names"`
				Status string `json:"Status"`
			}
			if err := json.Unmarshal([]byte(line), &item); err != nil {
				return WorkloadDiscovery{}, fmt.Errorf("environment: decode Docker container list: %w", err)
			}
			result = append(result, Workload{ID: item.Names, Name: item.Names, Status: item.Status})
		}
		message := ""
		if len(result) == 0 {
			message = "Docker is active, but no containers are configured."
		}
		return WorkloadDiscovery{Available: true, Message: message, Items: result}, nil
	case model.JobVM:
		if !unraidServiceEnabled("/boot/config/domain.cfg") {
			return WorkloadDiscovery{Available: false, Message: "The Unraid VM Manager is not enabled. Enable VMs under Settings > VM Manager to select virtual machines.", Items: []Workload{}}, nil
		}
		output, err := exec.CommandContext(ctx, "virsh", "list", "--all", "--name").Output()
		if err != nil {
			return WorkloadDiscovery{}, fmt.Errorf("environment: VM Manager is enabled but virtual machines could not be listed: %w", err)
		}
		var result []Workload
		for _, name := range strings.Split(strings.TrimSpace(string(output)), "\n") {
			name = strings.TrimSpace(name)
			if name == "" {
				continue
			}
			state, _ := exec.CommandContext(ctx, "virsh", "domstate", name).Output()
			result = append(result, Workload{ID: name, Name: name, Status: strings.TrimSpace(string(state))})
		}
		message := ""
		if len(result) == 0 {
			message = "The VM Manager is active, but no virtual machines are configured."
		}
		return WorkloadDiscovery{Available: true, Message: message, Items: result}, nil
	default:
		return WorkloadDiscovery{}, fmt.Errorf("validation: unsupported workload type %q", kind)
	}
}

func unraidServiceEnabled(path string) bool {
	data, err := os.ReadFile(path)
	if err != nil {
		return false
	}
	for _, line := range strings.Split(string(data), "\n") {
		parts := strings.SplitN(strings.TrimSpace(line), "=", 2)
		if len(parts) == 2 && parts[0] == "SERVICE" {
			return strings.Trim(strings.TrimSpace(parts[1]), "\"'") == "enable"
		}
	}
	return false
}

func (w *WorkloadManager) Prepare(ctx context.Context, job model.Job) (Prepared, error) {
	prepared := Prepared{Kind: job.Type}
	for _, source := range job.Sources {
		if source.Path != "" {
			prepared.Sources = append(prepared.Sources, source.Path)
		}
	}
	if job.Type != model.JobDocker && job.Type != model.JobVM {
		return prepared, nil
	}
	metadataDir, err := os.MkdirTemp(w.RuntimeDir, "metadata-"+job.ID+"-")
	if err != nil {
		return prepared, err
	}
	prepared.MetadataDir = metadataDir
	prepared.Sources = append(prepared.Sources, metadataDir)
	if job.Type == model.JobDocker {
		if _, err := os.Stat("/boot/config/plugins/dockerMan/templates-user"); err == nil {
			prepared.Sources = append(prepared.Sources, "/boot/config/plugins/dockerMan/templates-user")
		}
	}
	for _, source := range job.Sources {
		if source.WorkloadID == "" {
			continue
		}
		discovered, err := w.captureMetadata(ctx, job.Type, source.WorkloadID, metadataDir)
		if err != nil {
			w.Cleanup(context.Background(), prepared)
			return Prepared{}, err
		}
		prepared.Sources = append(prepared.Sources, discovered...)
		if job.Consistency.Mode == "stop" {
			if err := w.stop(ctx, job.Type, source.WorkloadID, time.Duration(job.ShutdownSecs)*time.Second); err != nil {
				w.Cleanup(context.Background(), prepared)
				return Prepared{}, err
			}
			prepared.Stopped = append(prepared.Stopped, source)
		}
	}
	prepared.Sources = uniqueStrings(prepared.Sources)
	return prepared, nil
}

func (w *WorkloadManager) FlashDevice(ctx context.Context) (string, error) {
	findmnt := w.FindmntPath
	if findmnt == "" {
		findmnt = "findmnt"
	}
	output, err := exec.CommandContext(ctx, findmnt, "--noheadings", "--output", "SOURCE", "--target", "/boot").Output()
	if err != nil {
		return "", fmt.Errorf("source: determine USB flash partition: %w", err)
	}
	partition := strings.TrimSpace(string(output))
	if index := strings.IndexByte(partition, '['); index >= 0 {
		partition = partition[:index]
	}
	partition = filepath.Clean(partition)
	if !strings.HasPrefix(partition, "/dev/") || strings.ContainsAny(partition, "\r\n\x00") {
		return "", fmt.Errorf("source: /boot is not backed by a device: %q", partition)
	}

	lsblk := w.LsblkPath
	if lsblk == "" {
		lsblk = "lsblk"
	}
	output, err = exec.CommandContext(ctx, lsblk, "--noheadings", "--output", "PKNAME", "--", partition).Output()
	if err != nil {
		return "", fmt.Errorf("source: determine USB flash device for %s: %w", partition, err)
	}
	parent := strings.TrimSpace(string(output))
	device := partition
	if parent != "" {
		if strings.ContainsAny(parent, "/\r\n\x00") {
			return "", fmt.Errorf("source: invalid USB flash parent device %q", parent)
		}
		device = filepath.Join("/dev", parent)
	}
	stat := w.Stat
	if stat == nil {
		stat = os.Stat
	}
	info, err := stat(device)
	if err != nil {
		return "", fmt.Errorf("source: access USB flash device %s: %w", device, err)
	}
	if info.Mode()&os.ModeDevice == 0 {
		return "", fmt.Errorf("source: USB flash source %s is not a block device", device)
	}
	return device, nil
}

func (w *WorkloadManager) Cleanup(ctx context.Context, prepared Prepared) error {
	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, 5*time.Second)
		defer cancel()
	}
	var first error
	for i := len(prepared.Stopped) - 1; i >= 0; i-- {
		source := prepared.Stopped[i]
		if err := w.start(ctx, prepared.Kind, source.WorkloadID); err != nil && first == nil {
			first = err
		}
	}
	if prepared.MetadataDir != "" {
		_ = os.RemoveAll(prepared.MetadataDir)
	}
	return first
}

func (w *WorkloadManager) captureMetadata(ctx context.Context, kind model.JobType, id, dir string) ([]string, error) {
	var cmd *exec.Cmd
	var name string
	if kind == model.JobDocker {
		cmd, name = exec.CommandContext(ctx, "docker", "inspect", id), "docker-"+safeName(id)+".json"
	} else {
		cmd, name = exec.CommandContext(ctx, "virsh", "dumpxml", id), "vm-"+safeName(id)+".xml"
	}
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("source: capture metadata for %s: %w", id, err)
	}
	if err := os.WriteFile(filepath.Join(dir, name), output, 0600); err != nil {
		return nil, err
	}
	if kind == model.JobDocker {
		var inspected []struct {
			Mounts []struct {
				Type   string `json:"Type"`
				Source string `json:"Source"`
			} `json:"Mounts"`
		}
		if err := json.Unmarshal(output, &inspected); err != nil {
			return nil, fmt.Errorf("source: decode docker metadata: %w", err)
		}
		var paths []string
		if len(inspected) > 0 {
			for _, mount := range inspected[0].Mounts {
				if mount.Type == "bind" && filepath.IsAbs(mount.Source) {
					paths = append(paths, mount.Source)
				}
			}
		}
		return paths, nil
	}
	var domain struct {
		OS struct {
			NVRAM string `xml:"nvram"`
		} `xml:"os"`
		Devices struct {
			Disks []struct {
				Device string `xml:"device,attr"`
				Source struct {
					File string `xml:"file,attr"`
					Dev  string `xml:"dev,attr"`
				} `xml:"source"`
			} `xml:"disk"`
		} `xml:"devices"`
	}
	if err := xml.Unmarshal(output, &domain); err != nil {
		return nil, fmt.Errorf("source: decode vm metadata: %w", err)
	}
	paths := []string{}
	if filepath.IsAbs(strings.TrimSpace(domain.OS.NVRAM)) {
		paths = append(paths, strings.TrimSpace(domain.OS.NVRAM))
	}
	for _, disk := range domain.Devices.Disks {
		if disk.Device != "disk" && disk.Device != "lun" {
			continue
		}
		path := disk.Source.File
		if path == "" {
			path = disk.Source.Dev
		}
		if filepath.IsAbs(path) {
			paths = append(paths, path)
		}
	}
	return paths, nil
}

func (w *WorkloadManager) stop(ctx context.Context, kind model.JobType, id string, timeout time.Duration) error {
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	if kind == model.JobDocker {
		seconds := fmt.Sprint(int(timeout.Seconds()))
		if err := exec.CommandContext(ctx, "docker", "stop", "--time", seconds, id).Run(); err != nil {
			return fmt.Errorf("source: stop docker %s: %w", id, err)
		}
		return nil
	}
	if err := exec.CommandContext(ctx, "virsh", "shutdown", id).Run(); err != nil {
		return fmt.Errorf("source: shutdown vm %s: %w", id, err)
	}
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		output, _ := exec.CommandContext(ctx, "virsh", "domstate", id).Output()
		if strings.Contains(strings.ToLower(string(output)), "shut off") {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
	return fmt.Errorf("source: vm %s did not stop before timeout", id)
}

func (w *WorkloadManager) start(ctx context.Context, kind model.JobType, id string) error {
	if kind == model.JobDocker {
		return exec.CommandContext(ctx, "docker", "start", id).Run()
	}
	return exec.CommandContext(ctx, "virsh", "start", id).Run()
}

func safeName(value string) string {
	value = strings.Map(func(r rune) rune {
		if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '-' || r == '_' {
			return r
		}
		return '_'
	}, value)
	if value == "" {
		return "workload"
	}
	return value
}

func uniqueStrings(values []string) []string {
	seen := make(map[string]struct{}, len(values))
	result := make([]string, 0, len(values))
	for _, value := range values {
		value = filepath.Clean(value)
		if _, exists := seen[value]; exists {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result
}