233 lines
8.0 KiB
Go
233 lines
8.0 KiB
Go
// Copyright (c) 2026 Storm Dragon
|
|
// Authors: Storm Dragon, Angus Kola
|
|
// SPDX-License-Identifier: GPL-3.0-or-later
|
|
|
|
package ocr
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"errors"
|
|
"hash/crc32"
|
|
"image"
|
|
"image/color"
|
|
"io"
|
|
"math"
|
|
"os"
|
|
"os/exec"
|
|
"strconv"
|
|
"strings"
|
|
"syscall"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func capture_request() CaptureRequest {
|
|
return CaptureRequest{Kind: ObjectSource, Bounds: image.Rect(-20, 10, 30, 50), Foreground: 123, Generation: 7}
|
|
}
|
|
|
|
func capture_frame(t *testing.T, header CaptureHeader, pixels image.Image) []byte {
|
|
t.Helper()
|
|
var output bytes.Buffer
|
|
if err := WriteCaptureResponse(&output, header, pixels); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return output.Bytes()
|
|
}
|
|
|
|
func TestCaptureRequestRejectsUnownedOrInvalidSources(t *testing.T) {
|
|
base := capture_request()
|
|
if err := base.Validate(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
mutations := []func(*CaptureRequest){
|
|
func(r *CaptureRequest) { r.Foreground = 0 },
|
|
func(r *CaptureRequest) { r.Foreground = 1 << 32 },
|
|
func(r *CaptureRequest) { r.Generation = 0 },
|
|
func(r *CaptureRequest) { r.Kind = "unknown" },
|
|
func(r *CaptureRequest) { r.Bounds = image.Rectangle{} },
|
|
func(r *CaptureRequest) { r.Bounds.Min.X = math.MinInt32 },
|
|
func(r *CaptureRequest) { r.Bounds = image.Rect(0, 0, MaxPixels+1, 1) },
|
|
func(r *CaptureRequest) { r.Window = 42 },
|
|
func(r *CaptureRequest) { r.Kind = WindowSource },
|
|
}
|
|
for _, mutate := range mutations {
|
|
request := base
|
|
mutate(&request)
|
|
if !errors.Is(request.Validate(), ErrSource) {
|
|
t.Fatalf("invalid request accepted: %+v", request)
|
|
}
|
|
}
|
|
for _, kind := range []SourceKind{WindowSource, DesktopSource, ClipboardSource} {
|
|
request := base
|
|
request.Kind, request.Bounds = kind, image.Rectangle{}
|
|
if err := request.Validate(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
data, _ := json.Marshal(base)
|
|
request, err := ReadCaptureRequest(bytes.NewReader(data))
|
|
if err != nil || request != base {
|
|
t.Fatalf("request roundtrip: %+v, %v", request, err)
|
|
}
|
|
for _, invalid := range [][]byte{nil, []byte("{}"), append(append([]byte(nil), data...), data...), []byte(`{"Private":"generic"}`), bytes.Repeat([]byte{' '}, maxCaptureHeader+1)} {
|
|
if _, err := ReadCaptureRequest(bytes.NewReader(invalid)); err == nil {
|
|
t.Fatal("invalid private request accepted")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCaptureFrameClippingAndGeneration(t *testing.T) {
|
|
request := capture_request()
|
|
header := CaptureHeader{Request: request, SourceWindow: request.Foreground,
|
|
RootBounds: image.Rect(0, 0, 100, 100), SourceBounds: image.Rect(-10, 0, 80, 80), ScreenBounds: image.Rect(0, 10, 30, 50)}
|
|
pixels := image.NewRGBA(image.Rect(0, 0, 30, 40))
|
|
pixels.Set(1, 2, color.RGBA{12, 34, 56, 255})
|
|
frame := capture_frame(t, header, pixels)
|
|
result, err := decode_capture(frame, request)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if color.RGBAModel.Convert(result.Pixels.At(1, 2)) != pixels.At(1, 2) {
|
|
t.Fatal("pixels changed")
|
|
}
|
|
if _, err := decode_capture(append(append([]byte(nil), frame...), 'x'), request); !errors.Is(err, ErrCapture) {
|
|
t.Fatalf("trailing payload accepted: %v", err)
|
|
}
|
|
mapped, ok := result.Transform.ScreenRectangle(image.Rect(1, 2, 3, 4))
|
|
if !ok || mapped != image.Rect(1, 12, 3, 14) {
|
|
t.Fatalf("mapping: %v, %v", mapped, ok)
|
|
}
|
|
request.Generation++
|
|
if _, err := decode_capture(frame, request); !errors.Is(err, ErrCapture) {
|
|
t.Fatalf("old generation accepted: %v", err)
|
|
}
|
|
request.Generation--
|
|
for _, change := range []func(*CaptureHeader){
|
|
func(h *CaptureHeader) { h.ScreenBounds.Min.X = -1 },
|
|
func(h *CaptureHeader) { h.SourceBounds = image.Rect(50, 50, 80, 80) },
|
|
func(h *CaptureHeader) { h.SourceWindow++ },
|
|
func(h *CaptureHeader) { h.ScreenBounds.Max.X++ },
|
|
} {
|
|
bad := header
|
|
change(&bad)
|
|
if _, err := decode_capture(capture_frame(t, bad, pixels), request); err == nil {
|
|
t.Fatalf("invalid mapping accepted: %+v", bad)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCaptureClipboardDoesNotMapToScreen(t *testing.T) {
|
|
request := CaptureRequest{Kind: ClipboardSource, Foreground: 123, Generation: 1}
|
|
header := CaptureHeader{Request: request}
|
|
pixels := image.NewRGBA(image.Rect(0, 0, 10, 10))
|
|
result, err := decode_capture(capture_frame(t, header, pixels), request)
|
|
if err != nil || result.Transform.HasScreenMapping {
|
|
t.Fatalf("clipboard: %+v, %v", result, err)
|
|
}
|
|
header.ScreenBounds = image.Rect(0, 0, 10, 10)
|
|
if _, err := decode_capture(capture_frame(t, header, pixels), request); err == nil {
|
|
t.Fatal("clipboard gained coordinates")
|
|
}
|
|
}
|
|
|
|
func TestCaptureProtocolErrorsAndPixelLimit(t *testing.T) {
|
|
request := capture_request()
|
|
for status, want := range map[string]error{"unavailable": ErrCapture, "source": ErrSource, "changed": ErrSourceChanged, "limit": ErrLimit, "Private diagnostics": ErrCapture} {
|
|
frame := capture_frame(t, CaptureHeader{Request: request, Status: status}, nil)
|
|
_, err := decode_capture(frame, request)
|
|
if !errors.Is(err, want) || strings.Contains(err.Error(), "Private") {
|
|
t.Fatalf("status %q: %v", status, err)
|
|
}
|
|
if _, err := decode_capture(append(frame, 'x'), request); !errors.Is(err, ErrCapture) {
|
|
t.Fatal("failure with payload accepted")
|
|
}
|
|
}
|
|
clipboard := CaptureRequest{Kind: ClipboardSource, Foreground: 123, Generation: 1}
|
|
frame := capture_frame(t, CaptureHeader{Request: clipboard}, image.NewRGBA(image.Rect(0, 0, 1, 1)))
|
|
start := bytes.IndexByte(frame, '\n') + 1
|
|
// A valid CRC and malicious IHDR are rejected before allocating pixel data.
|
|
binary.BigEndian.PutUint32(frame[start+16:start+20], 10000)
|
|
binary.BigEndian.PutUint32(frame[start+20:start+24], 10000)
|
|
binary.BigEndian.PutUint32(frame[start+29:start+33], crc32.ChecksumIEEE(frame[start+12:start+29]))
|
|
if _, err := decode_capture(frame, clipboard); !errors.Is(err, ErrLimit) {
|
|
t.Fatalf("oversize PNG: %v", err)
|
|
}
|
|
for _, invalid := range [][]byte{nil, []byte("{}\nnot a png"), bytes.Repeat([]byte{'x'}, maxCaptureHeader+1)} {
|
|
if _, err := decode_capture(invalid, request); !errors.Is(err, ErrCapture) {
|
|
t.Fatalf("invalid response: %v", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCapturerMissingHelperAndCancellation(t *testing.T) {
|
|
request := capture_request()
|
|
if _, err := NewCapturer("/nonexistent-ocr-capture-helper").Capture(context.Background(), request); !errors.Is(err, ErrCapture) {
|
|
t.Fatalf("missing helper: %v", err)
|
|
}
|
|
executable, err := os.Executable()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
readyRead, readyWrite, err := os.Pipe()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer readyRead.Close()
|
|
defer readyWrite.Close()
|
|
capturer := NewCapturer(executable)
|
|
capturer.command = func(ctx context.Context, args ...string) *exec.Cmd {
|
|
command := exec.CommandContext(ctx, executable, "-test.run=^TestCaptureChildProcess$", "--", args[0])
|
|
command.Env = append(os.Environ(), "CTHULHU_CAPTURE_TEST_CHILD=1")
|
|
command.ExtraFiles = []*os.File{readyWrite}
|
|
return command
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
completed := make(chan error, 1)
|
|
go func() { _, err := capturer.Capture(ctx, request); completed <- err }()
|
|
_ = readyRead.SetReadDeadline(time.Now().Add(5 * time.Second))
|
|
data := make([]byte, 64)
|
|
count, err := readyRead.Read(data)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
pid, err := strconv.Atoi(strings.TrimSpace(string(data[:count])))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
waitContext, waitCancel := context.WithCancel(context.Background())
|
|
waitCancel()
|
|
if _, err := capturer.Capture(waitContext, request); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("waiting cancellation: %v", err)
|
|
}
|
|
cancel()
|
|
select {
|
|
case err := <-completed:
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("running cancellation: %v", err)
|
|
}
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("capture did not join cancelled helper")
|
|
}
|
|
if err := syscall.Kill(pid, 0); !errors.Is(err, syscall.ESRCH) {
|
|
t.Fatalf("helper still alive or unreaped: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCaptureChildProcess(t *testing.T) {
|
|
if os.Getenv("CTHULHU_CAPTURE_TEST_CHILD") != "1" {
|
|
return
|
|
}
|
|
if _, err := io.ReadAll(os.Stdin); err != nil {
|
|
os.Exit(20)
|
|
}
|
|
ready := os.NewFile(3, "ready")
|
|
_, _ = ready.WriteString(strconv.Itoa(os.Getpid()) + "\n")
|
|
_ = ready.Close()
|
|
time.Sleep(time.Hour)
|
|
os.Exit(0)
|
|
}
|