gpu-turnstile/internal/service/service_windows.go

//go:build windows

// Package service integrates gpu-turnstile with the Windows Service
// Control Manager: running as a service with graceful stop, plus
// install/remove helpers. Installed services always run as the virtual
// account NT SERVICE\gpu-turnstile — a per-service low-privilege identity
// managed by the SCM, with no password and no admin rights.
package service

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/mgr"
)

// Name is the Windows service name.
const Name = "gpu-turnstile"

// virtualAccount is the per-service identity the service runs as. The SCM
// manages it: no password, automatic "log on as a service" right, gone
// when the service is removed.
const virtualAccount = `NT SERVICE\` + Name

// installDirs returns the canonical install (Program Files) and data
// (ProgramData) directories.
func installDirs() (install, data string) {
	pf := os.Getenv("ProgramFiles")
	if pf == "" {
		pf = `C:\Program Files`
	}
	pd := os.Getenv("ProgramData")
	if pd == "" {
		pd = `C:\ProgramData`
	}
	return filepath.Join(pf, Name), filepath.Join(pd, Name)
}

// IsService reports whether the process is running as a Windows service.
func IsService() bool {
	isSvc, err := svc.IsWindowsService()
	return err == nil && isSvc
}

// Run executes run as a Windows service. SCM Stop and Shutdown cancel the
// context passed to run, triggering the same graceful shutdown as SIGTERM
// in interactive mode.
func Run(run func(ctx context.Context) error) error {
	return svc.Run(Name, &handler{run: run})
}

type handler struct {
	run func(ctx context.Context) error
}

func (h *handler) Execute(_ []string, requests <-chan svc.ChangeRequest, status chan<- svc.Status) (bool, uint32) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	status <- svc.Status{State: svc.StartPending}
	errCh := make(chan error, 1)
	go func() { errCh <- h.run(ctx) }()
	status <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown}

	for {
		select {
		case err := <-errCh:
			status <- svc.Status{State: svc.Stopped}
			if err != nil {
				return true, 1
			}
			return false, 0
		case c := <-requests:
			switch c.Cmd {
			case svc.Interrogate:
				status <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				status <- svc.Status{State: svc.StopPending}
				cancel()
			}
		}
	}
}

// Install registers gpu-turnstile as an auto-start Windows service running
// as the NT SERVICE\gpu-turnstile virtual account, whose binPath loads the
// given config file. With copyBin it first creates the canonical layout —
// the binary is copied into %ProgramFiles%\gpu-turnstile and the config
// next to it (an existing config there is kept), %ProgramData%\gpu-turnstile
// is created for logs — and registers that copy; with copyBin=false the
// current executable location is registered as-is. Recovery actions restart
// the service after 5s on failure — this is also what brings up a staged
// update after the updater exits with a non-zero code. After registering,
// the virtual account is granted modify access to the install and data
// directories (self-updates rewrite the exe), read access to the
// config file if it lives elsewhere, and — when the config sets COMFY_DIR —
// recursive modify access to the managed ComfyUI's install tree, which may
// live inside a user profile the account otherwise cannot enter. The grants
// must come after CreateService: the virtual account's SID only exists once
// the service is registered.
//
// Re-running install on an existing service converges instead of failing:
// the service is stopped first if running (so the binary can be replaced),
// the installed binary is refreshed only when the content differs, the
// registration is updated only where it drifted, and the service is
// started again only if it was running before.
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
	}
	if configPath != "" {
		if abs, absErr := filepath.Abs(configPath); absErr == nil {
			configPath = abs
		}
	}

	m, err := mgr.Connect()
	if err != nil {
		return fmt.Errorf("connect to service manager (run as administrator): %w", err)
	}
	defer m.Disconnect()

	// An existing service is converged, not an error. Stop it first so the
	// binary copy can be replaced, and remember whether to start it again.
	var s *mgr.Service
	wasRunning := false
	if existing, openErr := m.OpenService(Name); openErr == nil {
		s = existing
		defer s.Close()
		if st, qErr := s.Query(); qErr == nil &&
			(st.State == svc.Running || st.State == svc.StartPending) {
			wasRunning = true
			fmt.Println("stopping the running gpu-turnstile service")
			if err := stopAndWait(s); err != nil {
				return err
			}
		}
	}

	installDir, dataDir := installDirs()
	if copyBin {
		targetCfg := filepath.Join(installDir, "gpu-turnstile.env")
		if !strings.EqualFold(filepath.Dir(exe), installDir) {
			if err := os.MkdirAll(installDir, 0o755); err != nil {
				return fmt.Errorf("create %s: %w", installDir, err)
			}
			installedExe := filepath.Join(installDir, "gpu-turnstile.exe")
			if same, _ := sameFileContent(exe, installedExe); !same {
				fmt.Printf("installing %s\n", installedExe)
				if err := copyFile(exe, installedExe); err != nil {
					return fmt.Errorf("copy binary to %s: %w", installedExe, err)
				}
			}
			exe = installedExe
			if configPath != "" && !strings.EqualFold(configPath, targetCfg) {
				if _, statErr := os.Stat(targetCfg); os.IsNotExist(statErr) {
					copyFile(configPath, targetCfg) //nolint:errcheck // best effort
				}
			}
		}
		// A service has no console: without LOG_FILE the output vanishes, so
		// the sample pre-wires it to ProgramData. Missing configs get a
		// fresh sample; installer-written ones from older versions get new
		// settings appended; anything without CFG_VER is invalid and gets
		// replaced (backup kept as .bak).
		logPath := filepath.Join(dataDir, "gpu-turnstile.log")
		if err := syncEnvFile(targetCfg, version, logPath); err != nil {
			return err
		}
		configPath = targetCfg
	}
	binPath := fmt.Sprintf(`"%s" -config "%s"`, exe, configPath)

	if s == nil {
		s, err = m.CreateService(Name, binPath, mgr.Config{
			StartType:        mgr.StartAutomatic,
			DisplayName:      "gpu-turnstile",
			Description:      "GPU arbitration proxy for Ollama and ComfyUI",
			ServiceStartName: virtualAccount,
		})
		if err != nil {
			return fmt.Errorf("create service: %w", err)
		}
		defer s.Close()

		if err := ensureRecovery(s); err != nil {
			s.Delete() // roll back so a retry starts clean
			return err
		}
		if err := grantAll(exe, configPath); err != nil {
			s.Delete() // roll back so a retry starts clean
			return err
		}
		// Best effort: start now instead of waiting for the next boot. A
		// missing config (no consumer URLs) fails the start; the service stays
		// registered and can be started once the config exists.
		fmt.Println("starting the gpu-turnstile service")
		s.Start()
		return nil
	}

	// Existing service: update the registration only where it drifted.
	cur, err := s.Config()
	if err != nil {
		return fmt.Errorf("query service config: %w", err)
	}
	const displayName = "gpu-turnstile"
	const description = "GPU arbitration proxy for Ollama and ComfyUI"
	if cur.BinaryPathName != binPath ||
		cur.StartType != mgr.StartAutomatic ||
		cur.ServiceStartName != virtualAccount ||
		cur.DisplayName != displayName ||
		cur.Description != description {
		upd := cur
		upd.BinaryPathName = binPath
		upd.StartType = mgr.StartAutomatic
		upd.ServiceStartName = virtualAccount
		upd.DisplayName = displayName
		upd.Description = description
		if err := s.UpdateConfig(upd); err != nil {
			return fmt.Errorf("update service config: %w", err)
		}
	}
	if err := ensureRecovery(s); err != nil {
		return err
	}
	// Idempotent: re-assert the virtual account's ACLs (granting an
	// existing ACE is a no-op).
	if err := grantAll(exe, configPath); err != nil {
		return err
	}
	if wasRunning {
		fmt.Println("starting the gpu-turnstile service")
		if err := s.Start(); err != nil {
			return fmt.Errorf("start service: %w", err)
		}
	}
	return nil
}

// ensureRecovery makes sure the service restarts 5s after a failure — the
// mechanism that also brings up a staged update. The current actions are
// queried first so an up-to-date service is left untouched.
func ensureRecovery(s *mgr.Service) error {
	want := mgr.RecoveryAction{Type: mgr.ServiceRestart, Delay: 5 * time.Second}
	actions, err := s.RecoveryActions()
	onFailure, flagErr := s.RecoveryActionsOnNonCrashFailures()
	if err == nil && flagErr == nil && onFailure && len(actions) == 3 {
		ok := true
		for _, a := range actions {
			if a.Type != want.Type || a.Delay != want.Delay {
				ok = false
			}
		}
		if ok {
			return nil
		}
	}
	if err := s.SetRecoveryActions([]mgr.RecoveryAction{want, want, want}, 24*60*60); err != nil {
		return fmt.Errorf("set recovery actions: %w", err)
	}
	if err := s.SetRecoveryActionsOnNonCrashFailures(true); err != nil {
		return fmt.Errorf("set failure actions flag: %w", err)
	}
	return nil
}

// stopAndWait stops the service and waits up to 30s for the stopped state.
func stopAndWait(s *mgr.Service) error {
	if _, err := s.Control(svc.Stop); err != nil {
		return fmt.Errorf("stop service: %w", err)
	}
	deadline := time.Now().Add(30 * time.Second)
	for {
		st, err := s.Query()
		if err != nil {
			return fmt.Errorf("query service: %w", err)
		}
		if st.State == svc.Stopped {
			return nil
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("service did not stop within 30s")
		}
		time.Sleep(300 * time.Millisecond)
	}
}

// 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
}

// copyFile copies src to dst (0755 on the new file).
func copyFile(src, dst string) 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, 0o755)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

// grantAll gives the virtual account every ACL the service needs: modify
// on the install and ProgramData directories and the LOG_FILE directory
// (if configured elsewhere), read on a config file outside the install
// directory.
func grantAll(exe, configPath string) error {
	_, dataDir := installDirs()
	if err := os.MkdirAll(dataDir, 0o755); err != nil {
		return fmt.Errorf("create %s: %w", dataDir, err)
	}
	if err := grantAccess(dataDir, "(OI)(CI)(M)"); err != nil {
		return err
	}
	exeDir := filepath.Dir(exe)
	if err := grantAccess(exeDir, "(OI)(CI)(M)"); err != nil {
		return err
	}
	if configPath != "" && !strings.HasPrefix(strings.ToLower(configPath), strings.ToLower(exeDir)+`\`) {
		if err := grantAccess(configPath, "(R)"); err != nil {
			return err
		}
	}
	if logFile := configuredValue(configPath, "LOG_FILE"); logFile != "" {
		dir := filepath.Dir(logFile)
		if err := os.MkdirAll(dir, 0o755); err == nil {
			if err := grantAccess(dir, "(OI)(CI)(M)"); err != nil {
				return err
			}
		}
	}
	// A managed ComfyUI whose install lives somewhere the virtual account
	// may not go (a user profile) needs an explicit grant — recursively,
	// since ComfyUI also writes output/temp/user data next to its code. A
	// missing directory is skipped: the startup warning covers it.
	if comfyDir := configuredValue(configPath, "COMFY_DIR"); comfyDir != "" {
		if _, err := os.Stat(comfyDir); err == nil {
			if err := grantAccessTree(comfyDir, "(OI)(CI)(M)"); err != nil {
				return err
			}
			// uv venvs (Comfy-Desktop) redirect to a base interpreter that
			// can live outside COMFY_DIR; read+execute suffices for it.
			if home := comfyVenvHome(comfyDir); home != "" {
				if err := grantAccessTree(home, "(OI)(CI)(RX)"); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// Remove stops (if running) and unregisters the service. The virtual
// account ceases to exist with it; the ACL grants on the install and log
// directories are left in place (harmless without the account).
func Remove() error {
	m, err := mgr.Connect()
	if err != nil {
		return fmt.Errorf("connect to service manager (run as administrator): %w", err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(Name)
	if err != nil {
		return fmt.Errorf("open service: %w", err)
	}
	defer s.Close()
	s.Control(svc.Stop) // ignore error: may already be stopped
	if err := s.Delete(); err != nil {
		return fmt.Errorf("delete service: %w", err)
	}
	return nil
}

var procShellExecuteExW = windows.NewLazySystemDLL("shell32.dll").NewProc("ShellExecuteExW")

const seeMaskNoCloseProcess = 0x40

// shellExecuteInfo mirrors SHELLEXECUTEINFOW (64-bit layout).
type shellExecuteInfo struct {
	cbSize       uint32
	fMask        uint32
	hwnd         uintptr
	lpVerb       *uint16
	lpFile       *uint16
	lpParameters *uint16
	lpDirectory  *uint16
	nShow        int32
	_            int32
	hInstApp     uintptr
	lpIDList     unsafe.Pointer
	lpClass      *uint16
	hkeyClass    uintptr
	dwHotKey     uint32
	_            uint32
	hIcon        uintptr
	hProcess     windows.Handle
}

// Elevated reports whether the current process token is UAC-elevated.
func Elevated() bool {
	var token windows.Token
	if err := windows.OpenProcessToken(windows.CurrentProcess(), windows.TOKEN_QUERY, &token); err != nil {
		return false
	}
	defer token.Close()
	return token.IsElevated()
}

// RelaunchElevated re-runs the current executable elevated via the UAC
// "runas" verb with the given arguments, waits for the child, and returns
// its exit code. The child gets a fresh console window for its output.
func RelaunchElevated(args []string) (int, error) {
	exe, err := os.Executable()
	if err != nil {
		return 0, err
	}
	quoted := make([]string, len(args))
	for i, a := range args {
		quoted[i] = syscall.EscapeArg(a)
	}
	cwd, _ := os.Getwd()
	verb, _ := windows.UTF16PtrFromString("runas")
	exeP, _ := windows.UTF16PtrFromString(exe)
	params, _ := windows.UTF16PtrFromString(strings.Join(quoted, " "))
	dir, _ := windows.UTF16PtrFromString(cwd)
	info := shellExecuteInfo{
		fMask:        seeMaskNoCloseProcess,
		lpVerb:       verb,
		lpFile:       exeP,
		lpParameters: params,
		lpDirectory:  dir,
		nShow:        windows.SW_NORMAL,
	}
	info.cbSize = uint32(unsafe.Sizeof(info))
	r, _, callErr := procShellExecuteExW.Call(uintptr(unsafe.Pointer(&info)))
	if r == 0 {
		if errors.Is(callErr, syscall.Errno(1223)) { // ERROR_CANCELLED
			return 0, ErrUserCancelled
		}
		return 0, fmt.Errorf("ShellExecuteEx: %w", callErr)
	}
	defer windows.CloseHandle(windows.Handle(info.hProcess))
	windows.WaitForSingleObject(windows.Handle(info.hProcess), windows.INFINITE)
	var code uint32
	if err := windows.GetExitCodeProcess(windows.Handle(info.hProcess), &code); err != nil {
		return 0, err
	}
	return int(code), nil
}

// grantAccess gives the virtual account the icacls permission set (e.g.
// "(OI)(CI)(M)") on path.
func grantAccess(path, perms string) error {
	return runIcacls(path, perms, false)
}

// grantAccessTree is grantAccess with /T: the ACE is applied to the
// existing tree, not just inherited by children created later. Needed when
// the tree already exists, e.g. a ComfyUI install in a user profile. On a
// large tree (a venv has tens of thousands of files) this takes minutes,
// so it says what it is doing instead of looking hung.
func grantAccessTree(path, perms string) error {
	return runIcacls(path, perms, true)
}

func runIcacls(path, perms string, recursive bool) error {
	args := []string{path, "/grant", virtualAccount + ":" + perms}
	if recursive {
		fmt.Printf("granting %s %s access to %s (large trees can take minutes)\n", virtualAccount, perms, path)
		args = append(args, "/T")
	}
	start := time.Now()
	out, err := exec.Command("icacls", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("grant %s access to %s: %w (%s)", virtualAccount, path, err, strings.TrimSpace(string(out)))
	}
	if recursive {
		fmt.Printf("access granted in %s\n", time.Since(start).Round(time.Second))
	}
	return nil
}

// RestartIfRunning restarts the service when it is installed and running
// (used after a forced update staged a new binary). Reports whether a
// restart happened. A service that is not installed or not running is not
// an error. Needs elevation.
func RestartIfRunning() (bool, error) {
	m, err := mgr.Connect()
	if err != nil {
		return false, fmt.Errorf("connect to service manager (run as administrator): %w", err)
	}
	defer m.Disconnect()
	s, err := m.OpenService(Name)
	if err != nil {
		return false, nil // not installed
	}
	defer s.Close()
	st, err := s.Query()
	if err != nil {
		return false, fmt.Errorf("query service: %w", err)
	}
	if st.State != svc.Running && st.State != svc.StartPending {
		return false, nil
	}
	if err := stopAndWait(s); err != nil {
		return false, err
	}
	if err := s.Start(); err != nil {
		return false, fmt.Errorf("start service: %w", err)
	}
	return true, nil
}