1 files changed, 252 insertions, 0 deletions
diff --git a/internal/server/handlers/readcommand_sendmessage_test.go b/internal/server/handlers/readcommand_sendmessage_test.go
new file mode 100644
index 0000000..a10be90
--- /dev/null
+++ b/internal/server/handlers/readcommand_sendmessage_test.go
@@ -0,0 +1,252 @@
+package handlers
+
+// Tests for readCommand.sendServerMessage.
+//
+// Leak context: the session's serverMessages channel (capacity 10 in the real
+// ServerHandler) is drained only by baseHandler.Read. Once the client
+// disconnects and the session shuts down, Read stops draining, so a bare
+// blocking send would pin its goroutine forever. The warn paths hit this in
+// practice: readGlob warns once per retry and readFileIfPermissions warns once
+// per permission-denied file, and a single glob may fan out to up to
+// MaxGlobTargets goroutines. sendServerMessage must therefore abandon the send
+// when the per-command context is cancelled (the context is cancelled on
+// command completion and handler shutdown via baseHandler.newCommandContext).
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/omode"
+)
+
+// newSendMessageTestCommand returns a readCommand whose server messages
+// channel has the given capacity, so tests can control exactly when the
+// channel is full.
+func newSendMessageTestCommand(channelCapacity int) (*readCommand, *globCapTestServer) {
+	srv := newGlobCapTestServer(10)
+	srv.serverMessage = make(chan string, channelCapacity)
+	return newReadCommand(srv, omode.CatClient), srv
+}
+
+// fillServerMessagesChannel fills the channel to capacity so any further
+// send would block.
+func fillServerMessagesChannel(t *testing.T, ch chan string) {
+	t.Helper()
+	for i := 0; i < cap(ch); i++ {
+		select {
+		case ch <- fmt.Sprintf("filler %d", i):
+		default:
+			t.Fatalf("channel unexpectedly full while pre-filling (i=%d, cap=%d)", i, cap(ch))
+		}
+	}
+}
+
+// TestSendServerMessageReturnsOnCancelledContext verifies the core leak fix:
+// with a full channel and an already-cancelled context, sendServerMessage must
+// return instead of blocking forever — and the abandoned message must never
+// surface on the channel afterwards.
+func TestSendServerMessageReturnsOnCancelledContext(t *testing.T) {
+	resetServerLogger(t)
+
+	cmd, srv := newSendMessageTestCommand(2)
+	fillServerMessagesChannel(t, srv.serverMessage)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	returned := make(chan struct{})
+	go func() {
+		cmd.sendServerMessage(ctx, "must not block")
+		close(returned)
+	}()
+
+	select {
+	case <-returned:
+		// Fixed behavior: the send was abandoned on ctx.Done().
+	case <-time.After(2 * time.Second):
+		t.Fatal("sendServerMessage blocked on a full channel despite cancelled context (goroutine leak)")
+	}
+
+	// An abandoned send must be dropped for good: draining the channel now
+	// must yield only the filler messages, never the abandoned one.
+	drained := 0
+	for {
+		select {
+		case raw := <-srv.serverMessage:
+			if !strings.HasPrefix(raw, "filler ") {
+				t.Fatalf("abandoned message leaked onto the channel: %q", raw)
+			}
+			drained++
+		default:
+			if want := cap(srv.serverMessage); drained != want {
+				t.Fatalf("expected %d filler messages on the channel, drained %d", want, drained)
+			}
+			return
+		}
+	}
+}
+
+// TestSendServerMessageUnblocksExactlyOnCancel verifies the send stays pending
+// while the context is live (no message is dropped prematurely) and is
+// abandoned as soon as the context is cancelled.
+func TestSendServerMessageUnblocksExactlyOnCancel(t *testing.T) {
+	resetServerLogger(t)
+
+	cmd, srv := newSendMessageTestCommand(1)
+	fillServerMessagesChannel(t, srv.serverMessage)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	started := make(chan struct{})
+	returned := make(chan struct{})
+	go func() {
+		close(started)
+		cmd.sendServerMessage(ctx, "pending message")
+		close(returned)
+	}()
+
+	// Wait until the sender goroutine has been scheduled before starting the
+	// "still pending" window, so the check below cannot pass vacuously just
+	// because the goroutine never ran.
+	<-started
+
+	// While ctx is live and the channel is full, the send must still be
+	// pending — sendServerMessage must not silently drop the message.
+	select {
+	case <-returned:
+		t.Fatal("sendServerMessage returned although channel is full and context is live")
+	case <-time.After(100 * time.Millisecond):
+		// Still pending, as expected.
+	}
+
+	cancel()
+
+	select {
+	case <-returned:
+		// Abandoned on cancel, as expected.
+	case <-time.After(2 * time.Second):
+		t.Fatal("sendServerMessage did not return after context cancellation")
+	}
+}
+
+// TestSendServerMessageDeliversWhenChannelHasRoom verifies the happy path:
+// with room in the channel and a live context, the message is delivered with
+// the command's generation encoded and a trailing newline appended.
+func TestSendServerMessageDeliversWhenChannelHasRoom(t *testing.T) {
+	resetServerLogger(t)
+
+	cmd, srv := newSendMessageTestCommand(2)
+	cmd.generation = 7
+
+	cmd.sendServerMessage(context.Background(), "hello")
+
+	select {
+	case raw := <-srv.serverMessage:
+		generation, message := decodeGeneratedMessage(raw)
+		if generation != 7 {
+			t.Errorf("expected generation 7, got %d", generation)
+		}
+		if message != "hello\n" {
+			t.Errorf("expected message %q, got %q", "hello\n", message)
+		}
+	default:
+		t.Fatal("expected a message on the server messages channel, got none")
+	}
+}
+
+// TestSendServerMessageDeliversToDrainedChannel verifies that a send blocked
+// on a full channel completes normally (message delivered, not abandoned)
+// once a consumer drains the channel while the context stays live. This
+// mirrors the healthy-session case where baseHandler.Read keeps draining.
+func TestSendServerMessageDeliversToDrainedChannel(t *testing.T) {
+	resetServerLogger(t)
+
+	cmd, srv := newSendMessageTestCommand(1)
+	fillServerMessagesChannel(t, srv.serverMessage)
+
+	returned := make(chan struct{})
+	go func() {
+		cmd.sendServerMessage(context.Background(), "queued")
+		close(returned)
+	}()
+
+	// Drain the filler message; the pending send must now complete.
+	<-srv.serverMessage
+
+	select {
+	case <-returned:
+	case <-time.After(2 * time.Second):
+		t.Fatal("sendServerMessage did not complete after the channel was drained")
+	}
+
+	select {
+	case raw := <-srv.serverMessage:
+		if _, message := decodeGeneratedMessage(raw); message != "queued\n" {
+			t.Errorf("expected message %q, got %q", "queued\n", message)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("expected the queued message to be delivered")
+	}
+}
+
+// TestSendServerMessageReleasedByHandlerShutdown exercises the production
+// wiring end-to-end: the per-command context comes from
+// baseHandler.newCommandContext, whose watcher goroutine cancels it when the
+// handler's done channel is shut down (the client-disconnect path). A
+// warn-sender stuck on a full serverMessages channel must be released by
+// done.Shutdown() alone — this is exactly the goroutine leak the fix removes.
+func TestSendServerMessageReleasedByHandlerShutdown(t *testing.T) {
+	resetServerLogger(t)
+
+	handler := &baseHandler{
+		done:           internal.NewDone(),
+		serverMessages: make(chan string, 1),
+	}
+
+	// Wire the readCommand's server messages channel to the real handler's
+	// channel so the stuck send targets the same channel baseHandler.Read
+	// would drain in production.
+	srv := newGlobCapTestServer(10)
+	srv.serverMessage = handler.serverMessages
+	cmd := newReadCommand(srv, omode.CatClient)
+
+	fillServerMessagesChannel(t, handler.serverMessages)
+
+	ctx, cancel := handler.newCommandContext(context.Background())
+	defer cancel()
+
+	started := make(chan struct{})
+	returned := make(chan struct{})
+	go func() {
+		close(started)
+		cmd.sendServerMessage(ctx, "stuck warn message")
+		close(returned)
+	}()
+
+	// Ensure the sender is scheduled and pending before shutting down, so the
+	// release below is attributable to done.Shutdown() rather than the send
+	// never having started.
+	<-started
+	select {
+	case <-returned:
+		t.Fatal("sendServerMessage returned although channel is full and handler is not shut down")
+	case <-time.After(100 * time.Millisecond):
+		// Still pending, as expected.
+	}
+
+	// Simulate the session teardown after a client disconnect: nothing drains
+	// serverMessages anymore, only the done channel fires.
+	handler.done.Shutdown()
+
+	select {
+	case <-returned:
+		// Released via newCommandContext's done watcher cancelling ctx.
+	case <-time.After(2 * time.Second):
+		t.Fatal("sendServerMessage not released by handler shutdown (goroutine leak)")
+	}
+}