fix(tcp): stop a cancelled context from being swallowed by the deadline re-arm - #366
Conversation
github-actions
Bot
added
size/xl
pr summary2 files changed (+207 -3)
|
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #366 +/- ##
=======================================
Coverage ? 64.89%
=======================================
Files ? 86
Lines ? 7645
Branches ? 0
=======================================
Hits ? 4961
Misses ? 2301
Partials ? 383 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
a663e79 to
ed6f59a
Compare
vmfunc
left a comment
vmfunc
left a comment
There was a problem hiding this comment.
the cancel-during-write re-arm race is the subtle part here and you handled it right: ctx.Err() before SetWriteDeadline and again at the top of readTCP, with TestExecuteTCPModuleContextCancelDuringWrite actually reproducing it instead of asserting it in the abstract. nice.
one non-blocking thing worth a follow-up: readTCP drains to the read deadline, so any peer that banners and then holds the connection open (ssh, redis after INFO, the mysql handshake, memcached after version... basically all of these) pays the full timeout on every grab. strictly better than the old hang-forever, but at the 10s default that's 10s per module on the happy path. the tests don't surface it because they all close the server side (the real listener defers c.Close, the pipe tests server.Close), so you get EOF instead of the hold-open you hit in the wild. an early return once a banner lands and the read goes idle, or a shorter banner-read window separate from the dial timeout, would keep it blazing.
minor: the ^-anchored banner regexes miss servers that emit a preamble before the ident line (ssh allows it per 4253), since Go ^ is start-of-text. fine to leave.
rest is solid. decodeTCPData bounds check out including the short/malformed \x cases, and validateTCP rejecting status/favicon at load matches what the executor will actually run. modules are low-FP. approving.
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.
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.
ed6f59a to
4229a6d
Compare
TBX3D
changed the title
|
rebased. this branch was carrying the whole tcp executor from before #308 retitled to match what is left: "fix(tcp): stop a cancelled context from being on your follow-up note, agreed that readTCP still drains to the read deadline so |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
This PR hardens TCP module execution against context-cancellation races by ensuring cancellation can’t be “swallowed” by later deadline re-arming, and adds regression tests that reproduce the problematic timing windows.
Changes:
- Added
ctx.Err()guards around TCP writes to prevent overwriting a cancellation-triggeredSetDeadline(now)and to returncontext.Canceledinstead of a raw timeout. - Updated
readTCPto accept acontext.Contextand to avoid re-arming read deadlines after cancellation. - Added multiple targeted tests that reproduce cancellation during write, before write, and while blocked in write.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 8 comments.
| File | Description |
|---|---|
| internal/modules/tcp.go | Adds ctx-aware deadline/IO handling to prevent cancellation races and improve error reporting on cancel. |
| internal/modules/tcp_test.go | Adds regression tests for cancellation timing races and updates readTCP call signature. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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 | ||
| } |
| // 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. |
| func (c *fakeSlowConn) Read([]byte) (int, error) { | ||
| c.mu.Lock() | ||
| dl := c.readDeadline | ||
| c.mu.Unlock() |
| if wait := time.Until(dl); wait > 0 { | ||
| time.Sleep(wait) | ||
| } |
| orig := newTCPConn | ||
| newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return conn, nil } | ||
| t.Cleanup(func() { newTCPConn = orig }) |
| orig := newTCPConn | ||
| newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil } | ||
| t.Cleanup(func() { newTCPConn = orig }) |
| orig := newTCPConn | ||
| newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil } | ||
| t.Cleanup(func() { newTCPConn = orig }) |
| orig := newTCPConn | ||
| newTCPConn = func(context.Context, string, time.Duration) (net.Conn, error) { return client, nil } | ||
| t.Cleanup(func() { newTCPConn = orig }) |
github-actions
Bot
added
size/l
ExecuteTCPModule cancels a run by tripping the connection with a single SetDeadline(now) from its watchdog goroutine. arming a later deadline after that trip silently pushes it back out to the full timeout, so a cancelled scan kept blocking: readTCP called SetReadDeadline unconditionally, and the probe write armed SetWriteDeadline the same way.
check ctx before arming either deadline and again on each read iteration, and prefer ctx.Err() over the raw write error so a write the watchdog trips is reported as a cancellation rather than an i/o timeout.
rebased onto main, which already carries the executor from #308, so this is now just the fix and its tests. TestExecuteTCPModuleContextCancelDuringWrite drives the write-in-flight ordering rather than asserting it in the abstract.