fix(tcp): stop a cancelled context from being swallowed by the deadline re-arm (#366)

* fix(tcp): stop the read-deadline re-arm from swallowing a cancelled ctx

the cancel watchdog trips conn.SetDeadline(now) once. if that trip
lands while the probe write is still in flight, readTCP's later
SetReadDeadline(now+timeout) call silently overwrote it, so the read
blocked for the full timeout instead of returning promptly. thread
ctx into readTCP and check it before arming the deadline and on every
read iteration so a cancellation already in flight aborts immediately.

* test(tcp): make the cancel-during-write test exercise the actual race

the old test cancelled ctx before calling ExecuteTCPModule, so the
pre-write ctx.Err() guard returned immediately and the test never
reached readTCP at all; reverting the fix and rerunning it still
passed in ~0ms. cancel from a goroutine mid-write instead, so the
watchdog trips the deadline while readTCP is the next call on the
path, and assert the call returns promptly rather than blocking for
the full timeout. also add a case proving a legitimately slow but
uncancelled connection still completes and matches normally.

Author: TBX3D

internal/modules/tcp.go

@@ -101,14 +101,27 @@ func ExecuteTCPModule(ctx context.Context, target string, def *YAMLModule, opts
 	}()
 
 	if cfg.Data != "" {
+		// same re-arm hazard readTCP guards against: if the cancel watchdog's
+		// SetDeadline(now) trip lands before this SetWriteDeadline call, the
+		// call below would silently push the deadline back out to the full
+		// timeout. Check ctx first so an already-cancelled run never arms it,
+		// and prefer ctx.Err() over the raw write error so a write the
+		// watchdog does trip is reported as a cancellation, not an i/o
+		// timeout.
+		if err := ctx.Err(); err != nil {
+			return result, err
+		}
 		payload := decodeTCPData(cfg.Data)
 		_ = conn.SetWriteDeadline(time.Now().Add(timeout))
 		if _, err := conn.Write([]byte(payload)); err != nil {
+			if ctxErr := ctx.Err(); ctxErr != nil {
+				return result, ctxErr
+			}
 			return nil, fmt.Errorf("tcp write %q: %w", addr, err)
 		}
 	}
 
-	data := readTCP(conn, timeout)
+	data := readTCP(ctx, conn, timeout)
 	if err := ctx.Err(); err != nil {
 		return result, err
 	}
@@ -130,11 +143,24 @@ func ExecuteTCPModule(ctx context.Context, target string, def *YAMLModule, opts
 // cap bounds memory to roughly the limit plus one buffer. A timeout or EOF ends
 // the read normally: a silent or half-open service yields the bytes seen so far
 // rather than an error, leaving the verdict to the matchers.
-func readTCP(conn net.Conn, timeout time.Duration) string {
+//
+// ctx is checked before arming the read deadline and again on every loop
+// iteration. The cancel watchdog in ExecuteTCPModule trips the deadline with a
+// single SetDeadline(now) call, which arming a later deadline here would
+// otherwise silently overwrite (e.g. if the cancel lands while the probe
+// Write is still in flight): a cancelled ctx must abort the read immediately
+// rather than re-arm and block for the full timeout.
+func readTCP(ctx context.Context, conn net.Conn, timeout time.Duration) string {
+	if ctx.Err() != nil {
+		return ""
+	}
 	_ = conn.SetReadDeadline(time.Now().Add(timeout))
 	var out []byte
 	buf := make([]byte, 4096)
 	for len(out) < tcpReadLimit {
+		if ctx.Err() != nil {
+			break
+		}
 		n, err := conn.Read(buf)
 		if n > 0 {
 			out = append(out, buf[:n]...)

internal/modules/tcp_test.go

@@ -19,6 +19,7 @@ import (
 	"net"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 	"time"
 )
@@ -303,6 +304,183 @@ func TestExecuteTCPModuleContextCancel(t *testing.T) {
 	}
 }
 
+// fakeSlowConn is a net.Conn stand-in whose Write yields for a fixed duration
+// before returning success regardless of any deadline, and whose Read blocks
+// until whatever deadline was last armed via SetReadDeadline/SetDeadline, then
+// returns a timeout error. It reproduces a real socket closely enough to prove
+// the cancellation race: a cancel landing while the slow Write is in flight
+// trips the watchdog's SetDeadline(now) before readTCP arms its own (later)
+// read deadline.
+type fakeSlowConn struct {
+	net.Conn
+	mu           sync.Mutex
+	readDeadline time.Time
+	writeYield   time.Duration
+}
+
+func (c *fakeSlowConn) Write(b []byte) (int, error) {
+	time.Sleep(c.writeYield)
+	return len(b), nil
+}
+
+func (c *fakeSlowConn) Read([]byte) (int, error) {
+	c.mu.Lock()
+	dl := c.readDeadline
+	c.mu.Unlock()
+	if dl.IsZero() {
+		dl = time.Now().Add(time.Hour)
+	}
+	if wait := time.Until(dl); wait > 0 {
+		time.Sleep(wait)
+	}
+	return 0, os.ErrDeadlineExceeded
+}
+
+func (c *fakeSlowConn) SetDeadline(t time.Time) error {
+	c.mu.Lock()
+	c.readDeadline = t
+	c.mu.Unlock()
+	return nil
+}
+
+func (c *fakeSlowConn) SetReadDeadline(t time.Time) error { return c.SetDeadline(t) }
+func (c *fakeSlowConn) SetWriteDeadline(time.Time) error  { return nil }
+func (c *fakeSlowConn) Close() error                      { return nil }
+
+// TestExecuteTCPModuleContextCancelDuringWrite reproduces the deadline re-arm
+// race: the ctx is cancelled from another goroutine while the probe Write is
+// still yielding, so the watchdog's conn.SetDeadline(now) trip lands mid-write
+// rather than before ExecuteTCPModule is even entered. The Write then returns
+// its normal success (a fake real socket does not itself enforce the deadline
+// on a write already in flight), so execution reaches readTCP with the ctx
+// already done but no error yet surfaced. Before the fix, readTCP's own
+// SetReadDeadline(now+timeout) call overwrote the watchdog's trip and the read
+// blocked for the full timeout instead of returning promptly.
+func TestExecuteTCPModuleContextCancelDuringWrite(t *testing.T) {
+	conn := &fakeSlowConn{writeYield: 200 * time.Millisecond}
+	orig := newTCPConn
+	newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return conn, nil }
+	t.Cleanup(func() { newTCPConn = orig })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		cancel()
+	}()
+
+	def := tcpDef(&TCPConfig{Port: 22, Data: "PING\r\n", Matchers: []Matcher{tcpWord("x")}})
+	start := time.Now()
+	res, err := ExecuteTCPModule(ctx, "example.com", def, Options{Timeout: 2 * time.Second})
+	elapsed := time.Since(start)
+
+	if elapsed > time.Second {
+		t.Errorf("returned after %v, want prompt (a cancel landing mid-write must not be swallowed by the read-deadline re-arm)", elapsed)
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("err = %v, want context.Canceled", err)
+	}
+	if len(res.Findings) != 0 {
+		t.Errorf("got %d findings on cancel, want 0", len(res.Findings))
+	}
+}
+
+// TestExecuteTCPModuleContextCancelBeforeWrite proves an already-cancelled ctx
+// is caught before the probe write arms its deadline, rather than falling
+// through to SetWriteDeadline and Write regardless. Without the guard this
+// would only fail if the write itself then blocked past the deadline; here it
+// is asserted directly so the guard cannot regress silently.
+func TestExecuteTCPModuleContextCancelBeforeWrite(t *testing.T) {
+	client, server := net.Pipe()
+	t.Cleanup(func() { server.Close() })
+	orig := newTCPConn
+	newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil }
+	t.Cleanup(func() { newTCPConn = orig })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	def := tcpDef(&TCPConfig{Port: 22, Data: "PING\r\n", Matchers: []Matcher{tcpWord("x")}})
+	res, err := ExecuteTCPModule(ctx, "example.com", def, Options{Timeout: 2 * time.Second})
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("err = %v, want context.Canceled", err)
+	}
+	if len(res.Findings) != 0 {
+		t.Errorf("got %d findings on cancel, want 0", len(res.Findings))
+	}
+}
+
+// TestExecuteTCPModuleContextCancelBlocksInWrite reproduces the write-side
+// twin of the read-deadline re-arm race: the probe write blocks on a real
+// synchronous conn (net.Pipe with no reader), and the cancel watchdog's
+// SetDeadline(now) is what has to unblock it. Before the ctx.Err() guard, the
+// watchdog trip could land between goroutine start and the SetWriteDeadline
+// call and be silently overwritten by it, since the watchdog only fires once
+// and never gets a second chance to re-trip; the write would then block for
+// the full timeout instead of returning promptly, and a write that the
+// watchdog did manage to trip surfaced as a raw i/o timeout rather than the
+// cancellation it actually was.
+func TestExecuteTCPModuleContextCancelBlocksInWrite(t *testing.T) {
+	client, server := net.Pipe()
+	t.Cleanup(func() { server.Close() }) // no reader: the probe write blocks until the deadline trips
+
+	orig := newTCPConn
+	newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil }
+	t.Cleanup(func() { newTCPConn = orig })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(20 * time.Millisecond)
+		cancel()
+	}()
+
+	def := tcpDef(&TCPConfig{Port: 22, Data: "PING\r\n", Matchers: []Matcher{tcpWord("x")}})
+	start := time.Now()
+	res, err := ExecuteTCPModule(ctx, "example.com", def, Options{Timeout: 5 * time.Second})
+	if elapsed := time.Since(start); elapsed > 2*time.Second {
+		t.Errorf("returned after %v, want the watchdog to trip the blocked write promptly", elapsed)
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("err = %v, want context.Canceled", err)
+	}
+	if len(res.Findings) != 0 {
+		t.Errorf("got %d findings on cancel, want 0", len(res.Findings))
+	}
+}
+
+// TestExecuteTCPModuleSlowConnCompletesWithoutCancel proves the ctx.Err()
+// guards added for the read-deadline re-arm fix do not clip a legitimately
+// slow but healthy connection: with no cancellation, a banner that trickles
+// in well under the timeout still completes and matches normally.
+func TestExecuteTCPModuleSlowConnCompletesWithoutCancel(t *testing.T) {
+	client, server := net.Pipe()
+	orig := newTCPConn
+	newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil }
+	t.Cleanup(func() { newTCPConn = orig })
+
+	go func() {
+		buf := make([]byte, 4096)
+		_, _ = server.Read(buf)
+		time.Sleep(200 * time.Millisecond)
+		_, _ = server.Write([]byte("+OK ready\r\n"))
+		server.Close()
+	}()
+
+	def := tcpDef(&TCPConfig{Port: 22, Data: "PING\r\n", Matchers: []Matcher{tcpWord("+OK")}})
+	start := time.Now()
+	res, err := ExecuteTCPModule(context.Background(), "example.com", def, Options{Timeout: 2 * time.Second})
+	elapsed := time.Since(start)
+
+	if err != nil {
+		t.Fatalf("ExecuteTCPModule: %v", err)
+	}
+	if elapsed < 200*time.Millisecond {
+		t.Errorf("returned after %v, want it to wait out the slow banner (~200ms)", elapsed)
+	}
+	if len(res.Findings) != 1 {
+		t.Fatalf("got %d findings, want 1 (a healthy slow connection must not be treated as cancelled)", len(res.Findings))
+	}
+}
+
 func TestExecuteTCPModuleDialError(t *testing.T) {
 	orig := newTCPConn
 	newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) {
@@ -430,7 +608,7 @@ func TestReadTCPCapsAtLimit(t *testing.T) {
 	}()
 	defer client.Close()
 
-	got := readTCP(client, time.Second)
+	got := readTCP(context.Background(), client, time.Second)
 	if len(got) < tcpReadLimit {
 		t.Fatalf("read %d bytes, want at least the %d cap", len(got), tcpReadLimit)
 	}