From ed64fd428a4795bfe74a6be25d25c452a5804003 Mon Sep 17 00:00:00 2001
From: Rain Seven <128443127+RAiNY7Study@users.noreply.github.com>
Date: Sun, 16 Aug 2026 21:25:50 +0800
Subject: [PATCH] fix: restart the active systemd unit after update (#36)

Co-authored-by: Meng Meng <227010654+MengMengCode@users.noreply.github.com>
---
 internal/update/systemd_test.go | 42 ++++++++++++++++++++++++++
 internal/update/update.go       | 52 +++++++++++++++++++++++++++++++--
 2 files changed, 92 insertions(+), 2 deletions(-)
 create mode 100644 internal/update/systemd_test.go

diff --git a/internal/update/systemd_test.go b/internal/update/systemd_test.go
new file mode 100644
index 0000000..4bc3f2b
--- /dev/null
+++ b/internal/update/systemd_test.go
@@ -0,0 +1,42 @@
+package update
+
+import (
+	"io"
+	"log/slog"
+	"testing"
+)
+
+func TestDetectSystemdUnitUsesExplicitOverride(t *testing.T) {
+	t.Setenv("VOCAT_SYSTEMD_UNIT", "vocat-test.service")
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	if got := detectSystemdUnit(logger); got != "vocat-test.service" {
+		t.Fatalf("detectSystemdUnit() = %q, want vocat-test.service", got)
+	}
+}
+
+func TestSystemdUnitFromCgroup(t *testing.T) {
+	tests := []struct {
+		name string
+		data string
+		want string
+	}{
+		{name: "cgroup v2", data: "0::/system.slice/vocat-test.service\n", want: "vocat-test.service"},
+		{name: "legacy", data: "1:name=systemd:/system.slice/vocat@dji.service\n", want: "vocat@dji.service"},
+		{name: "no service", data: "0::/user.slice/user-1000.slice/session-1.scope\n", want: ""},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if got := systemdUnitFromCgroup(test.data); got != test.want {
+				t.Fatalf("systemdUnitFromCgroup() = %q, want %q", got, test.want)
+			}
+		})
+	}
+}
+
+func TestValidSystemdUnitRejectsArgumentsAndPaths(t *testing.T) {
+	for _, value := range []string{"vocat", "../vocat.service", "vocat.service --now", "vocat.service/other"} {
+		if validSystemdUnit.MatchString(value) {
+			t.Fatalf("validSystemdUnit unexpectedly accepted %q", value)
+		}
+	}
+}
diff --git a/internal/update/update.go b/internal/update/update.go
index 745dedb..53b7782 100644
--- a/internal/update/update.go
+++ b/internal/update/update.go
@@ -20,7 +20,9 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -320,17 +322,63 @@ func RestartService(logger *slog.Logger) error {
 	if _, err := exec.LookPath("systemctl"); err != nil {
 		return fmt.Errorf("neither /etc/init.d/vocat nor systemctl is available")
 	}
+	unit := detectSystemdUnit(logger)
 	// Queue the restart and let systemctl exit before systemd stops this unit.
 	// A blocking restart command becomes part of vocat.service's own cgroup and
 	// waits for that same cgroup to terminate, creating a stop-timeout cycle.
-	cmd := exec.Command("systemctl", "restart", "--no-block", "vocat")
+	cmd := exec.Command("systemctl", "restart", "--no-block", unit)
 	if out, err := cmd.CombinedOutput(); err != nil {
 		logger.Warn("systemctl restart failed", "error", err, "output", string(out))
-		return fmt.Errorf("systemctl restart vocat: %w", err)
+		return fmt.Errorf("systemctl restart %s: %w", unit, err)
 	}
 	return nil
 }
 
+var validSystemdUnit = regexp.MustCompile(`^[A-Za-z0-9_.@:-]+\.service$`)
+
+func detectSystemdUnit(logger *slog.Logger) string {
+	if configured := strings.TrimSpace(os.Getenv("VOCAT_SYSTEMD_UNIT")); validSystemdUnit.MatchString(configured) {
+		return configured
+	}
+	if data, err := os.ReadFile("/proc/self/cgroup"); err == nil {
+		if unit := systemdUnitFromCgroup(string(data)); unit != "" {
+			return unit
+		}
+	}
+	// Some cgroup namespaces hide the unit name. Query loaded services and
+	// identify the unit whose MainPID is this process before falling back.
+	list := exec.Command("systemctl", "list-units", "--type=service", "--all", "--no-legend", "--plain")
+	if output, err := list.Output(); err == nil {
+		pid := strconv.Itoa(os.Getpid())
+		for _, line := range strings.Split(string(output), "\n") {
+			fields := strings.Fields(line)
+			if len(fields) == 0 || !validSystemdUnit.MatchString(fields[0]) {
+				continue
+			}
+			show := exec.Command("systemctl", "show", fields[0], "--property=MainPID", "--value")
+			if value, showErr := show.Output(); showErr == nil && strings.TrimSpace(string(value)) == pid {
+				return fields[0]
+			}
+		}
+	}
+	if logger != nil {
+		logger.Warn("could not identify the current systemd unit; using vocat.service", "hint", "set VOCAT_SYSTEMD_UNIT for a custom unit")
+	}
+	return "vocat.service"
+}
+
+func systemdUnitFromCgroup(data string) string {
+	for _, line := range strings.Split(data, "\n") {
+		for _, part := range strings.Split(line, "/") {
+			part = strings.TrimSpace(part)
+			if validSystemdUnit.MatchString(part) {
+				return part
+			}
+		}
+	}
+	return ""
+}
+
 // resolveDefaultTarget returns the conventional install path when present,
 // falling back to the running executable. This lets `vocat update` "just work"
 // on the standard systemd host without flags.