gpu-turnstile/internal/service/service_linux.go

//go:build linux

// Package service integrates gpu-turnstile with systemd on Linux: running
// under a unit with readiness notification and watchdog, plus
// install/remove helpers that manage a hardened system unit.
package service

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"syscall"
)

// Name matches the Windows service name; the systemd unit is Name + ".service".
const Name = "gpu-turnstile"

// unitPath is where Install writes the unit file.
const unitPath = "/etc/systemd/system/" + Name + ".service"

// stateDir holds the installed binary (and staged updates); the unit's
// StateDirectory= directive makes systemd own it and grant the dynamic
// user write access. etcConfig is the config file the unit loads.
const (
	stateDir  = "/var/lib/" + Name
	etcConfig = "/etc/" + Name + ".env"
)

// IsService reports whether the process was started by systemd.
func IsService() bool { return os.Getenv("INVOCATION_ID") != "" }

// Elevated is always true on Linux: there is no UAC equivalent; privilege
// errors surface from the failing operation with a "run as root" hint.
func Elevated() bool { return true }

// RelaunchElevated is a Windows-only concept (UAC).
func RelaunchElevated([]string) (int, error) {
	return 0, fmt.Errorf("elevated relaunch is only supported on Windows")
}

// Run executes run with SIGINT/SIGTERM cancellation (which is how systemctl
// stop signals the process) and tells systemd when the shutdown begins.
func Run(run func(ctx context.Context) error) error {
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	defer NotifyStopping()
	return run(ctx)
}

// renderUnit builds the hardened systemd unit: Type=notify so systemctl
// start blocks until the listeners are bound, a 30s watchdog, and
// restart-on-failure with a 5s delay — which is also what brings up a
// staged update after the updater exits with a non-zero code.
//
// Sandboxing mirrors the Windows virtual account: DynamicUser=yes gives
// the service a transient per-service UID with no login and no home, the
// filesystem is read-only except StateDirectory (the install dir, so
// self-updates can rewrite the binary), and the usual no-privilege-escalation
// directives apply. The proxy needs nothing but outbound TCP/UDP and the
// notify socket, so it loses nothing. A managed ComfyUI (comfyDir) gets a
// BindPaths hole through ProtectHome/ProtectSystem: it reads its venv and
// writes output/temp/user data under COMFY_DIR. A venv whose base
// interpreter (pyvenv.cfg home) lives outside COMFY_DIR gets an additional
// read-only bind.
func renderUnit(exePath, configPath, comfyDir string) string {
	bind := ""
	if comfyDir != "" {
		bind = "BindPaths=" + comfyDir + "\n"
		if home := comfyVenvHome(comfyDir); home != "" {
			bind += "BindReadOnlyPaths=" + home + "\n"
		}
	}
	return fmt.Sprintf(`[Unit]
Description=gpu-turnstile GPU arbitration proxy for Ollama and ComfyUI
After=network-online.target
Wants=network-online.target

[Service]
Type=notify
WatchdogSec=30s
ExecStart=%q -config %q
Restart=on-failure
RestartSec=5s

DynamicUser=yes
StateDirectory=%s
%sProtectSystem=strict
ProtectHome=yes
PrivateTmp=yes
NoNewPrivileges=yes
ProtectKernelTunables=yes
ProtectKernelModules=yes
ProtectKernelLogs=yes
ProtectControlGroups=yes
ProtectClock=yes
RestrictNamespaces=yes
RestrictSUIDSGID=yes
RestrictRealtime=yes
LockPersonality=yes
MemoryDenyWriteExecute=yes
CapabilityBoundingSet=
AmbientCapabilities=
RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6
SystemCallFilter=@system-service
SystemCallErrorNumber=EPERM

[Install]
WantedBy=multi-user.target
`, exePath, configPath, Name, bind)
}

// copyFile copies src to dst, creating dst with the given mode.
func copyFile(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
	if err != nil {
		return err
	}
	defer out.Close()
	if _, err := io.Copy(out, in); err != nil {
		return err
	}
	return out.Close()
}

// Install copies the current executable into /var/lib/gpu-turnstile, makes
// sure /etc/gpu-turnstile.env exists (copied from the given config file if
// provided), writes the hardened unit, then enables and starts it. With
// copyBin=false the current executable location and config path are
// registered as-is instead. When the config sets COMFY_DIR, the unit gets a
// BindPaths= for it so the sandboxed service can reach the managed ComfyUI
// even under /home. Needs root.
//
// Re-running install converges an existing unit instead of failing: it is
// stopped first if active, the installed binary is replaced only when the
// content differs, the unit file is rewritten (followed by daemon-reload)
// only when it changed, and the service is started again only if it was
// active before.
//
// The binary lives in the StateDirectory rather than /usr/local/sbin on
// purpose: replacing a running binary needs write access to its
// *directory*, and granting the sandboxed service user write access to a
// shared system directory would let a compromised service overwrite other
// binaries. /var/lib/gpu-turnstile is exclusively ours.
func Install(configPath string, copyBin bool, version string) error {
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	if abs, absErr := filepath.Abs(exe); absErr == nil {
		exe = abs
	}
	unit := Name + ".service"
	_, statErr := os.Stat(unitPath)
	fresh := os.IsNotExist(statErr)
	wasRunning := exec.Command("systemctl", "is-active", "--quiet", unit).Run() == nil
	if wasRunning {
		if out, err := exec.Command("systemctl", "stop", unit).CombinedOutput(); err != nil {
			return fmt.Errorf("systemctl stop (run as root): %w (%s)", err, out)
		}
	}

	cfg := etcConfig
	if copyBin {
		if err := os.MkdirAll(stateDir, 0o755); err != nil {
			return fmt.Errorf("create %s (run as root): %w", stateDir, err)
		}
		installedExe := filepath.Join(stateDir, Name)
		if exe != installedExe {
			if same, _ := sameFileContent(exe, installedExe); !same {
				if err := copyFile(exe, installedExe, 0o755); err != nil {
					return fmt.Errorf("install binary to %s: %w", installedExe, err)
				}
			}
		}
		exe = installedExe
		if _, err := os.Stat(etcConfig); os.IsNotExist(err) && configPath != "" {
			copyFile(configPath, etcConfig, 0o644) //nolint:errcheck // best effort
		}
		// Missing configs get a fully commented sample (LOG_FILE stays
		// commented — stderr goes to the journal on Linux);
		// installer-written ones from older versions get new settings
		// appended; anything without CFG_VER is invalid and gets replaced
		// (backup kept as .bak).
		if err := syncEnvFile(etcConfig, version, ""); err != nil {
			return err
		}
	} else if configPath != "" {
		if abs, absErr := filepath.Abs(configPath); absErr == nil {
			cfg = abs
		}
	}
	rendered := renderUnit(exe, cfg, configuredValue(cfg, "COMFY_DIR"))
	if old, _ := os.ReadFile(unitPath); string(old) != rendered {
		if err := os.WriteFile(unitPath, []byte(rendered), 0o644); err != nil {
			return fmt.Errorf("write %s (run as root): %w", unitPath, err)
		}
		if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil {
			return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out)
		}
	}
	if fresh {
		if out, err := exec.Command("systemctl", "enable", "--now", unit).CombinedOutput(); err != nil {
			return fmt.Errorf("systemctl enable --now: %w (%s)", err, out)
		}
		return nil
	}
	if exec.Command("systemctl", "is-enabled", "--quiet", unit).Run() != nil {
		if out, err := exec.Command("systemctl", "enable", unit).CombinedOutput(); err != nil {
			return fmt.Errorf("systemctl enable: %w (%s)", err, out)
		}
	}
	if wasRunning {
		if out, err := exec.Command("systemctl", "start", unit).CombinedOutput(); err != nil {
			return fmt.Errorf("systemctl start: %w (%s)", err, out)
		}
	}
	return nil
}

// sameFileContent reports whether two files hold identical bytes. A missing
// destination is simply "different".
func sameFileContent(a, b string) (bool, error) {
	ba, err := os.ReadFile(a)
	if err != nil {
		return false, err
	}
	bb, err := os.ReadFile(b)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return bytes.Equal(ba, bb), nil
}

// Remove stops and disables the service and deletes the unit file and the
// installed binary. The config file in /etc is left in place (user data).
func Remove() error {
	exec.Command("systemctl", "disable", "--now", Name+".service").Run() // ignore: may not exist
	if err := os.Remove(unitPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("remove %s: %w", unitPath, err)
	}
	os.RemoveAll(stateDir) // installed binary + staged updates; ignore error
	if out, err := exec.Command("systemctl", "daemon-reload").CombinedOutput(); err != nil {
		return fmt.Errorf("systemctl daemon-reload: %w (%s)", err, out)
	}
	return nil
}

// RestartIfRunning restarts the systemd unit when it is active (used after
// a forced update staged a new binary). Reports whether a restart
// happened. An inactive or missing unit is not an error. Needs root.
func RestartIfRunning() (bool, error) {
	if err := exec.Command("systemctl", "is-active", "--quiet", Name+".service").Run(); err != nil {
		return false, nil // inactive or not installed
	}
	if out, err := exec.Command("systemctl", "restart", Name+".service").CombinedOutput(); err != nil {
		return false, fmt.Errorf("systemctl restart (run as root): %w (%s)", err, out)
	}
	return true, nil
}