1 files changed, 356 insertions, 0 deletions
diff --git a/internal/clients/handlers/basehandler_test.go b/internal/clients/handlers/basehandler_test.go
new file mode 100644
index 0000000..e070e48
--- /dev/null
+++ b/internal/clients/handlers/basehandler_test.go
@@ -0,0 +1,356 @@
+package handlers
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
+)
+
+func TestParseAuthKeyMessage(t *testing.T) {
+	tests := []struct {
+		name     string
+		message  string
+		wantAuth bool
+		wantOK   bool
+		wantInfo string
+	}{
+		{
+			name:     "server formatted success",
+			message:  fmt.Sprintf("SERVER%s%s%sAUTHKEY OK\n", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+			wantAuth: true,
+			wantOK:   true,
+		},
+		{
+			name:     "server formatted error",
+			message:  fmt.Sprintf("SERVER%s%s%sAUTHKEY ERR feature disabled\n", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+			wantAuth: true,
+			wantOK:   false,
+			wantInfo: "feature disabled",
+		},
+		{
+			name:     "plain response success",
+			message:  "AUTHKEY OK",
+			wantAuth: true,
+			wantOK:   true,
+		},
+		{
+			name:     "not an authkey message",
+			message:  fmt.Sprintf("SERVER%s%s%ssome other message", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+			wantAuth: false,
+			wantOK:   false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			gotAuth, gotOK, gotInfo := parseAuthKeyMessage(tc.message)
+			if gotAuth != tc.wantAuth {
+				t.Fatalf("Unexpected auth marker: got %v want %v", gotAuth, tc.wantAuth)
+			}
+			if gotOK != tc.wantOK {
+				t.Fatalf("Unexpected ok marker: got %v want %v", gotOK, tc.wantOK)
+			}
+			if gotInfo != tc.wantInfo {
+				t.Fatalf("Unexpected info: got %q want %q", gotInfo, tc.wantInfo)
+			}
+		})
+	}
+}
+
+func TestHandleCapabilitiesMessage(t *testing.T) {
+	handler := baseHandler{
+		done:           internal.NewDone(),
+		capabilities:   make(map[string]struct{}),
+		capabilitiesCh: make(chan struct{}),
+		sessionAcks:    make(chan SessionAck, 1),
+	}
+
+	handler.handleHiddenMessage(".syn capabilities query-update-v1 feature-two")
+
+	if !handler.HasCapability(protocol.CapabilityQueryUpdateV1) {
+		t.Fatalf("expected handler to track %q", protocol.CapabilityQueryUpdateV1)
+	}
+	if !handler.HasCapability("feature-two") {
+		t.Fatalf("expected handler to track feature-two")
+	}
+	if handler.WaitForCapabilities(10*time.Millisecond) != true {
+		t.Fatalf("expected capabilities wait to succeed")
+	}
+
+	capabilities := handler.Capabilities()
+	if len(capabilities) != 2 {
+		t.Fatalf("unexpected capabilities: %#v", capabilities)
+	}
+}
+
+func TestWaitForCapabilitiesTimeout(t *testing.T) {
+	handler := baseHandler{
+		done:           internal.NewDone(),
+		capabilities:   make(map[string]struct{}),
+		capabilitiesCh: make(chan struct{}),
+		sessionAcks:    make(chan SessionAck, 1),
+	}
+
+	if handler.WaitForCapabilities(5 * time.Millisecond) {
+		t.Fatalf("expected capabilities wait to time out")
+	}
+}
+
+func TestFormatServerErrorMessage(t *testing.T) {
+	got := formatServerErrorMessage("srv1", "journal file targets require server capability journal-v1")
+	// No trailing newline: the message is emitted via the diagnostic (Log) sink,
+	// which appends the newline itself (see ReportServerError / RawLog).
+	want := "SERVER|srv1|ERROR|journal file targets require server capability journal-v1"
+	if got != want {
+		t.Fatalf("formatServerErrorMessage() = %q, want %q", got, want)
+	}
+}
+
+func TestParseSessionAckMessage(t *testing.T) {
+	tests := []struct {
+		name    string
+		message string
+		want    SessionAck
+		wantOK  bool
+	}{
+		{
+			name:    "start ok",
+			message: ".syn session start ok 7",
+			want: SessionAck{
+				Action:     "start",
+				Generation: 7,
+			},
+			wantOK: true,
+		},
+		{
+			name:    "update ok",
+			message: ".syn session update ok 8",
+			want: SessionAck{
+				Action:     "update",
+				Generation: 8,
+			},
+			wantOK: true,
+		},
+		{
+			name:    "error",
+			message: ".syn session err query sessions not supported yet",
+			want: SessionAck{
+				Action: "error",
+				Error:  "query sessions not supported yet",
+			},
+			wantOK: true,
+		},
+		{
+			name:    "invalid",
+			message: ".syn session start ok nope",
+			wantOK:  false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got, ok := parseSessionAckMessage(tc.message)
+			if ok != tc.wantOK {
+				t.Fatalf("unexpected ok flag: got %v want %v", ok, tc.wantOK)
+			}
+			if !tc.wantOK {
+				return
+			}
+			if got != tc.want {
+				t.Fatalf("unexpected ack: got %#v want %#v", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestHandleSessionAckMessage(t *testing.T) {
+	handler := baseHandler{
+		done:        internal.NewDone(),
+		sessionAcks: make(chan SessionAck, 1),
+	}
+
+	handler.handleHiddenMessage(".syn session update ok 4")
+
+	ack, ok := handler.WaitForSessionAck(10 * time.Millisecond)
+	if !ok {
+		t.Fatalf("expected session ack")
+	}
+	if ack.Action != "update" || ack.Generation != 4 {
+		t.Fatalf("unexpected session ack: %#v", ack)
+	}
+}
+
+func TestHandleCloseConnectionAcknowledgesBeforeShutdown(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	handler := baseHandler{
+		done:     internal.NewDone(),
+		server:   "server-under-test",
+		commands: make(chan string, 1),
+	}
+
+	handler.handleHiddenMessage(".syn close connection")
+
+	select {
+	case command := <-handler.commands:
+		if command == "" {
+			t.Fatal("expected close acknowledgement command")
+		}
+	case <-time.After(10 * time.Millisecond):
+		t.Fatal("expected close acknowledgement command to be queued")
+	}
+
+	select {
+	case <-handler.Done():
+	default:
+		t.Fatal("expected handler to be shut down after close acknowledgement")
+	}
+}
+
+// TestReadDrainsAckBeforeEOF verifies that Read() always delivers the
+// '.ack close connection' command before returning io.EOF when Done() fires
+// concurrently. Without the priority-select fix, Go's non-deterministic select
+// would randomly pick the Done() case and drop the queued ack ~50% of the time.
+func TestReadDrainsAckBeforeEOF(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	// Run many iterations to catch the race reliably even with -race.
+	const iterations = 500
+	for i := 0; i < iterations; i++ {
+		// Use an unbuffered commands channel so the ack write and Done() close
+		// are truly concurrent, exposing the original non-deterministic select.
+		handler := baseHandler{
+			done:     internal.NewDone(),
+			server:   "server-under-test",
+			commands: make(chan string, 1),
+		}
+
+		// handleHiddenMessage calls SendMessage (queues ack) then Shutdown (closes Done).
+		// After this call both handler.commands and handler.Done() are ready.
+		handler.handleHiddenMessage(".syn close connection")
+
+		buf := make([]byte, 4096)
+		n, err := handler.Read(buf)
+		if err != nil {
+			// Done() won the select — the ack was dropped. This is the bug.
+			t.Fatalf("iteration %d: Read returned io.EOF before draining the ack", i)
+		}
+		if n == 0 {
+			t.Fatalf("iteration %d: Read returned 0 bytes without an error", i)
+		}
+
+		// Second Read must now return EOF since Done() is closed and commands is empty.
+		_, err = handler.Read(buf)
+		if err == nil {
+			t.Fatalf("iteration %d: second Read should return io.EOF", i)
+		}
+	}
+}
+
+// newCommandReadTestHandler returns a handler for exercising Read directly
+// and registers a cleanup that unblocks any Read still waiting on the
+// commands channel, so a regression cannot leak a blocked goroutine.
+func newCommandReadTestHandler(t *testing.T) *baseHandler {
+	t.Helper()
+
+	handler := &baseHandler{
+		done:     internal.NewDone(),
+		server:   "server-under-test",
+		commands: make(chan string, 2),
+	}
+	t.Cleanup(handler.done.Shutdown)
+	return handler
+}
+
+// readCommandsWithin drains wantLen bytes from the handler on a separate
+// goroutine and fails the test when the reads do not finish within the
+// timeout. The deadline matters: under a regression that drops the remainder
+// of a partially delivered command, the next Read would block forever on the
+// empty commands channel; the timeout turns that hang into a clean assertion
+// (and shuts the handler down so the reader goroutine exits via Done).
+func readCommandsWithin(t *testing.T, handler *baseHandler, bufSize, wantLen int) []byte {
+	t.Helper()
+
+	type result struct {
+		data []byte
+		err  error
+	}
+	resultCh := make(chan result, 1)
+	go func() {
+		var got []byte
+		p := make([]byte, bufSize)
+		for len(got) < wantLen {
+			n, err := handler.Read(p)
+			if err != nil {
+				resultCh <- result{got, err}
+				return
+			}
+			got = append(got, p[:n]...)
+		}
+		resultCh <- result{got, nil}
+	}()
+
+	select {
+	case res := <-resultCh:
+		if res.err != nil {
+			t.Fatalf("Read() error after %d of %d bytes: %v",
+				len(res.data), wantLen, res.err)
+		}
+		return res.data
+	case <-time.After(5 * time.Second):
+		handler.done.Shutdown()
+		t.Fatalf("Read stalled: timed out waiting for %d bytes", wantLen)
+		return nil
+	}
+}
+
+// TestBaseHandlerReadLargeCommandAcrossMultipleReads verifies that a command
+// frame larger than the caller's buffer (e.g. a huge base64-encoded regex or
+// MapReduce query) is delivered completely across multiple Read calls instead
+// of being truncated at the buffer boundary.
+func TestBaseHandlerReadLargeCommandAcrossMultipleReads(t *testing.T) {
+	handler := newCommandReadTestHandler(t)
+
+	command := "grep " + strings.Repeat("x", 1000) + ";"
+	handler.commands <- command
+
+	got := readCommandsWithin(t, handler, 32, len(command))
+	if string(got) != command {
+		t.Fatalf("large command corrupted across reads:\ngot  %q\nwant %q",
+			got, command)
+	}
+}
+
+// TestBaseHandlerReadDrainsRemainderBeforeNextCommand verifies that the
+// remainder of a partially delivered command is fully drained before the next
+// queued command starts, so frames are never interleaved, and that an
+// exact-fit read leaves no stale remainder behind.
+func TestBaseHandlerReadDrainsRemainderBeforeNextCommand(t *testing.T) {
+	handler := newCommandReadTestHandler(t)
+
+	first := "first command payload;"
+	second := "second;"
+	handler.commands <- first
+	handler.commands <- second
+
+	// len(first) = 22, so first fills exactly two 11-byte reads (the second
+	// of which is the remainder), then second must start fresh.
+	got := readCommandsWithin(t, handler, 11, len(first)+len(second))
+	if string(got) != first+second {
+		t.Fatalf("commands interleaved or corrupted:\ngot  %q\nwant %q",
+			got, first+second)
+	}
+}