268 lines
9.2 KiB
Go
268 lines
9.2 KiB
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.
|
|
func renderUnit(exePath, configPath, comfyDir string) string {
|
|
bind := ""
|
|
if comfyDir != "" {
|
|
bind = "BindPaths=" + comfyDir + "\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
|
|
}
|