From ae595526daabe4757a4a82eda8ab85587dc8f6a0 Mon Sep 17 00:00:00 2001
From: pjay <ipriyanshujain@gmail.com>
Date: Sat, 18 Apr 2026 14:16:15 +0700
Subject: [PATCH] fix(agent): race in readWithDeadline clobbers conn deadline
 (#7)

* refactor(docs): inline pandoc build into Makefile, drop scripts dir

* fix(agent): wait for deadline watcher before returning

readWithDeadline's watcher goroutine could clobber the conn's read
deadline with time.Unix(1, 0) after the main function reset it to zero.
When the Accept ctx was canceled shortly after Accept returned, the
watcher raced with close(done) in select and sometimes picked ctx.Done()
even though we were already done reading, leaving the conn unusable for
the next read (instant i/o timeout on step 1 snapshot).

Synchronize on the watcher's exit before resetting the deadline so it
can never override the reset.

* test(agent): cover readWithDeadline race on Accept ctx cancel

Drives Accept with a short-timeout ctx, cancels it right after Accept
returns, then does a Snapshot. Reliably fails without the readWithDeadline
synchronization fix (watcher goroutine overwrites the deadline to past).
---
 internal/agent/server.go      |  7 ++++--
 internal/agent/server_test.go | 47 +++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/internal/agent/server.go b/internal/agent/server.go
index 64c9d97..af3d0ad 100644
--- a/internal/agent/server.go
+++ b/internal/agent/server.go
@@ -92,11 +92,11 @@ func (c *Conn) Close() error {
 func readWithDeadline(ctx context.Context, conn net.Conn) (Message, error) {
 	if deadline, ok := ctx.Deadline(); ok {
 		_ = conn.SetReadDeadline(deadline)
-		defer conn.SetReadDeadline(time.Time{})
 	}
 	done := make(chan struct{})
-	defer close(done)
+	exited := make(chan struct{})
 	go func() {
+		defer close(exited)
 		select {
 		case <-ctx.Done():
 			_ = conn.SetReadDeadline(time.Unix(1, 0))
@@ -104,6 +104,9 @@ func readWithDeadline(ctx context.Context, conn net.Conn) (Message, error) {
 		}
 	}()
 	message, err := ReadMessage(conn)
+	close(done)
+	<-exited
+	_ = conn.SetReadDeadline(time.Time{})
 	if err != nil && ctx.Err() != nil {
 		return Message{}, ctx.Err()
 	}
diff --git a/internal/agent/server_test.go b/internal/agent/server_test.go
index 7f825f6..1e63c24 100644
--- a/internal/agent/server_test.go
+++ b/internal/agent/server_test.go
@@ -249,6 +249,53 @@ func TestConn_CloseSendsGoodbye(t *testing.T) {
 	}
 }
 
+// TestConn_SnapshotAfterAcceptContextCancel guards against a race in
+// readWithDeadline where the watcher goroutine from Accept could clobber the
+// conn's read deadline with a past time after Accept returned, causing the
+// next read on the same conn (Snapshot) to time out instantly.
+func TestConn_SnapshotAfterAcceptContextCancel(t *testing.T) {
+	for iteration := range 50 {
+		server := newLoopbackServer(t)
+
+		clientDone := make(chan struct{})
+		go func() {
+			defer close(clientDone)
+			client, err := net.Dial("tcp", server.Addr().String())
+			if err != nil {
+				return
+			}
+			defer client.Close()
+			if err := WriteMessage(client, Hello("0.0.1", "android", "com.x")); err != nil {
+				return
+			}
+			msg, err := ReadMessage(client)
+			if err != nil {
+				return
+			}
+			_ = WriteMessage(client, State(msg.ID, map[string]json.RawMessage{"ok": json.RawMessage(`true`)}))
+		}()
+
+		acceptCtx, acceptCancel := context.WithTimeout(context.Background(), time.Second)
+		conn, err := server.Accept(acceptCtx)
+		acceptCancel()
+		if err != nil {
+			t.Fatalf("iteration %d: Accept: %v", iteration, err)
+		}
+
+		snapCtx, snapCancel := context.WithTimeout(context.Background(), 2*time.Second)
+		state, err := conn.Snapshot(snapCtx)
+		snapCancel()
+		if err != nil {
+			t.Fatalf("iteration %d: Snapshot: %v", iteration, err)
+		}
+		if string(state.Snapshots["ok"]) != `true` {
+			t.Errorf("iteration %d: unexpected snapshots: %v", iteration, state.Snapshots)
+		}
+		conn.Close()
+		<-clientDone
+	}
+}
+
 func TestConn_SnapshotTimesOutIfSDKSilent(t *testing.T) {
 	server := newLoopbackServer(t)