diff --git a/internal/modules/tcp.go b/internal/modules/tcp.go index 2fe558c..d0e31b5 100644 --- a/internal/modules/tcp.go +++ b/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]...) diff --git a/internal/modules/tcp_test.go b/internal/modules/tcp_test.go index efa1674..a1d355e 100644 --- a/internal/modules/tcp_test.go +++ b/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) }