feat: DTail fork — server/client feature development

Squashed development of the snonux/dtail fork's product code (internal/, cmd/)
since diverging from mimecast/dtail. Major areas:

- Read/output path: the former "turbo" channel-less path is now the single,
  default server-side read/output path for cat/grep/tail and MapReduce; the old
  channel-based path and its config/env toggles were removed.
- MapReduce: single aggregate implementation (server + serverless) fed directly
  by a processor pipeline, with input-exhausted finalization via the shutdown
  coordinator; high-concurrency and data-race fixes.
- Journal source reads (journal:unit.service) via journalctl, Linux-gated behind
  a journal-v1 capability.
- Auth-key fast reconnect: in-memory per-user public-key cache with TTL/max-keys,
  registered over an authenticated session (AUTHKEY), checked before
  authorized_keys.
- Interactive query reload (--interactive-query) with SESSION START/UPDATE
  generation boundaries and capability negotiation.
- Client-side deadlines: --timeout / --shutdownAfter as context deadlines;
  follow shutdown handling.
- Client logging: diagnostics-only daily log by default, opt-in payload tee via
  --log-payload.
- Numerous correctness fixes (buffer-pool double-recycle races, EOF-sentinel
  leaks, glob-expansion cap, TOCTOU in CSV parsing) with accompanying unit tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

1 files changed, 568 insertions, 0 deletions

diff --git a/internal/server/handlers/sessioncommand_test.go b/internal/server/handlers/sessioncommand_test.go
new file mode 100644
index 0000000..0d32e87
--- /dev/null
+++ b/internal/server/handlers/sessioncommand_test.go
@@ -0,0 +1,568 @@
+package handlers
+
+import (
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"slices"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/io/line"
+	"github.com/mimecast/dtail/internal/lcontext"
+	maprserver "github.com/mimecast/dtail/internal/mapr/server"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/protocol"
+	"github.com/mimecast/dtail/internal/session"
+	sshserver "github.com/mimecast/dtail/internal/ssh/server"
+	userserver "github.com/mimecast/dtail/internal/user/server"
+)
+
+func TestNewServerHandlerSendsAdvertisedServerCapabilities(t *testing.T) {
+	resetServerLogger(t)
+
+	originalCapabilities := advertisedServerCapabilities
+	advertisedServerCapabilities = strings.Join([]string{
+		protocol.CapabilityQueryUpdateV1,
+		protocol.CapabilityJournalV1,
+	}, " ")
+	t.Cleanup(func() {
+		advertisedServerCapabilities = originalCapabilities
+	})
+
+	handler := NewServerHandler(
+		&userserver.User{Name: "session-capability-user"},
+		make(chan struct{}, 1),
+		make(chan struct{}, 1),
+		&config.ServerConfig{AuthKeyEnabled: true},
+		sshserver.NewAuthKeyStore(time.Hour, 5),
+	)
+
+	message := readServerMessage(t, handler.serverMessages)
+	if !strings.HasPrefix(message, protocol.HiddenCapabilitiesPrefix) {
+		t.Fatalf("unexpected capability advertisement: %q", message)
+	}
+	if want := protocol.HiddenCapabilitiesPrefix + advertisedServerCapabilities; message != want {
+		t.Fatalf("capability advertisement = %q, want %q", message, want)
+	}
+
+	capabilities := strings.Fields(strings.TrimPrefix(message, protocol.HiddenCapabilitiesPrefix))
+	if !slices.Contains(capabilities, protocol.CapabilityQueryUpdateV1) {
+		t.Fatalf("expected %q capability in %q", protocol.CapabilityQueryUpdateV1, message)
+	}
+	if !slices.Contains(capabilities, protocol.CapabilityJournalV1) {
+		t.Fatalf("expected %q capability in %q", protocol.CapabilityJournalV1, message)
+	}
+}
+
+func TestServerCapabilitiesAdvertisesJournalOnlyOnLinuxWithJournalctl(t *testing.T) {
+	tests := []struct {
+		name                string
+		goos                string
+		journalctlAvailable bool
+		want                []string
+	}{
+		{
+			name:                "linux with journalctl",
+			goos:                "linux",
+			journalctlAvailable: true,
+			want: []string{
+				protocol.CapabilityQueryUpdateV1,
+				protocol.CapabilityJournalV1,
+			},
+		},
+		{
+			name:                "linux without journalctl",
+			goos:                "linux",
+			journalctlAvailable: false,
+			want:                []string{protocol.CapabilityQueryUpdateV1},
+		},
+		{
+			name:                "non linux with journalctl",
+			goos:                "freebsd",
+			journalctlAvailable: true,
+			want:                []string{protocol.CapabilityQueryUpdateV1},
+		},
+		{
+			name:                "non linux without journalctl",
+			goos:                "darwin",
+			journalctlAvailable: false,
+			want:                []string{protocol.CapabilityQueryUpdateV1},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := strings.Fields(serverCapabilities(tt.goos, tt.journalctlAvailable))
+			if !slices.Equal(got, tt.want) {
+				t.Fatalf("capabilities = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestHandleSessionCommandStartStoresSpec(t *testing.T) {
+	handler := newSessionTestHandler("session-start-user")
+	readServerMessage(t, handler.serverMessages)
+
+	spec := session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	payload := mustSessionPayload(t, spec)
+
+	commandFinished := false
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", payload}, func() {
+		commandFinished = true
+	})
+
+	if !commandFinished {
+		t.Fatalf("expected commandFinished callback")
+	}
+	if !handler.sessionState.activeSession() {
+		t.Fatalf("expected session state to become active")
+	}
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckStartOKPrefix+" 1" {
+		t.Fatalf("unexpected session start message: %q", message)
+	}
+}
+
+func TestHandleSessionCommandUpdateCancelsPreviousGenerationImmediately(t *testing.T) {
+	handler, recorder := newSessionDispatchTestHandler("session-update-cancel-user")
+	readServerMessage(t, handler.serverMessages)
+	t.Cleanup(func() {
+		if handler.sessionState.cancel != nil {
+			handler.sessionState.cancel()
+		}
+		recorder.wg.Wait()
+	})
+
+	startPayload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app-a.log"},
+		Regex: "ERROR",
+	})
+	updatePayload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app-b.log"},
+		Regex: "WARN",
+	})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", startPayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckStartOKPrefix+" 1" {
+		t.Fatalf("unexpected session start ack: %q", message)
+	}
+
+	first := recorder.waitForStart(t)
+	if !strings.Contains(first.command, "/var/log/app-a.log") {
+		t.Fatalf("expected first command to target app-a.log, got %q", first.command)
+	}
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "UPDATE", updatePayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckUpdateOKPrefix+" 2" {
+		t.Fatalf("unexpected session update ack: %q", message)
+	}
+
+	waitForContextDone(first.ctx, t)
+
+	second := recorder.waitForStart(t)
+	if !strings.Contains(second.command, "/var/log/app-b.log") {
+		t.Fatalf("expected second command to target app-b.log, got %q", second.command)
+	}
+	select {
+	case <-second.ctx.Done():
+		t.Fatalf("expected replacement generation context to remain active")
+	default:
+	}
+}
+
+func TestHandleSessionCommandUpdateRequiresActiveSession(t *testing.T) {
+	handler := newSessionTestHandler("session-update-user")
+	readServerMessage(t, handler.serverMessages)
+
+	spec := session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	payload := mustSessionPayload(t, spec)
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "UPDATE", payload}, func() {})
+
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckErrorPrefix+"session not started" {
+		t.Fatalf("unexpected session update error: %q", message)
+	}
+}
+
+func TestHandleSessionCommandRejectsInvalidPayload(t *testing.T) {
+	handler := newSessionTestHandler("session-invalid-user")
+	readServerMessage(t, handler.serverMessages)
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", "not-base64"}, func() {})
+
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckErrorPrefix+"invalid session payload" {
+		t.Fatalf("unexpected invalid payload message: %q", message)
+	}
+}
+
+func TestHandleSessionCommandStartDispatchesQueryWorkload(t *testing.T) {
+	handler, recorder := newQuerySessionDispatchTestHandler("session-query-user")
+	readServerMessage(t, handler.serverMessages)
+
+	payload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Query: "from STATS select count(*)",
+		Regex: ".",
+	})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", payload}, func() {})
+
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckStartOKPrefix+" 1" {
+		t.Fatalf("unexpected query-session ack: %q", message)
+	}
+
+	first := recorder.waitForStart(t)
+	if !strings.HasPrefix(first.command, "map:") {
+		t.Fatalf("expected map command first, got %q", first.command)
+	}
+	if !strings.Contains(first.command, "from STATS select count(*)") {
+		t.Fatalf("expected map command to contain query, got %q", first.command)
+	}
+
+	second := recorder.waitForStart(t)
+	if !strings.HasPrefix(second.command, "tail:") {
+		t.Fatalf("expected tail command second, got %q", second.command)
+	}
+	if !strings.Contains(second.command, "/var/log/app.log") {
+		t.Fatalf("expected tail command to contain file, got %q", second.command)
+	}
+}
+
+func TestHandleSessionCommandRejectsInvalidSerializedOptions(t *testing.T) {
+	handler := newSessionTestHandler("session-options-user")
+	readServerMessage(t, handler.serverMessages)
+
+	payload := mustSessionPayload(t, session.Spec{
+		Mode:    omode.TailClient,
+		Files:   []string{"/var/log/app.log"},
+		Options: "badoption",
+		Regex:   "ERROR",
+	})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", payload}, func() {})
+
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckErrorPrefix+"invalid session spec" {
+		t.Fatalf("unexpected invalid options error: %q", message)
+	}
+}
+
+func TestHandleSessionCommandRejectsInvalidQuerySession(t *testing.T) {
+	handler := newSessionTestHandler("session-invalid-query-user")
+	readServerMessage(t, handler.serverMessages)
+
+	payload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Query: "select from",
+		Regex: ".",
+	})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", payload}, func() {})
+
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckErrorPrefix+"invalid session spec" {
+		t.Fatalf("unexpected invalid query-session error: %q", message)
+	}
+}
+
+func TestHandleAckCommandCloseConnectionConcurrentDoesNotPanic(t *testing.T) {
+	handler := newSessionTestHandler("ack-close-user")
+
+	const workers = 16
+	start := make(chan struct{})
+	panicCh := make(chan any, workers)
+	var wg sync.WaitGroup
+
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer func() {
+				if recovered := recover(); recovered != nil {
+					panicCh <- recovered
+				}
+			}()
+			<-start
+			handler.handleAckCommand(3, []string{"ACK", "close", "connection"})
+		}()
+	}
+
+	close(start)
+	wg.Wait()
+	close(panicCh)
+
+	for recovered := range panicCh {
+		t.Fatalf("unexpected panic while closing ack channel: %v", recovered)
+	}
+
+	select {
+	case <-handler.ackCloseReceived:
+	default:
+		t.Fatalf("expected ackCloseReceived to be closed")
+	}
+}
+
+func TestHandleSessionCommandUpdateClearsAggregateStateBeforeDirectRead(t *testing.T) {
+	resetServerLogger(t)
+
+	handler := newSessionTestHandler("session-query-reset-user")
+	readServerMessage(t, handler.serverMessages)
+
+	sawResetState := make(chan bool, 1)
+	tailCalls := 0
+	handler.commands = map[string]commandHandler{
+		"map": func(_ context.Context, _ lcontext.LContext, argc int, args []string, commandFinished func()) {
+			queryStr := strings.Join(args[1:], " ")
+			// Output is now the only aggregate (task hv0), so install a
+			// Aggregate as the real handleMapCommand does.
+			aggregate, err := maprserver.NewAggregate(queryStr, "")
+			if err != nil {
+				t.Fatalf("new output aggregate: %v", err)
+			}
+			// Use the atomic setter so this test exercises the same code path
+			// as the real handleMapCommand and avoids a direct field access race.
+			handler.setAggregate(aggregate)
+			commandFinished()
+		},
+		"tail": func(_ context.Context, _ lcontext.LContext, _ int, _ []string, commandFinished func()) {
+			tailCalls++
+			if tailCalls > 1 {
+				// Use the atomic getter; direct field access would race with
+				// concurrent reads in Shutdown/resetSessionAggregates.
+				sawResetState <- handler.getAggregate() == nil
+			}
+			commandFinished()
+		},
+	}
+
+	startPayload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app-a.log"},
+		Query: "from STATS select count(*)",
+		Regex: ".",
+	})
+	updatePayload := mustSessionPayload(t, session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app-b.log"},
+		Regex: "WARN",
+	})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", startPayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckStartOKPrefix+" 1" {
+		t.Fatalf("unexpected session start ack: %q", message)
+	}
+	// Use the atomic getter; direct field access would race with concurrent
+	// writes in handleMapCommand on another goroutine.
+	if handler.getAggregate() == nil {
+		t.Fatalf("expected query session to install aggregate state")
+	}
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "UPDATE", updatePayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckUpdateOKPrefix+" 2" {
+		t.Fatalf("unexpected session update ack: %q", message)
+	}
+
+	select {
+	case ok := <-sawResetState:
+		if !ok {
+			t.Fatalf("expected aggregate state to be cleared before direct read dispatch")
+		}
+	case <-time.After(250 * time.Millisecond):
+		t.Fatal("timed out waiting for direct read dispatch")
+	}
+}
+
+func newSessionTestHandler(userName string) *ServerHandler {
+	handler := &ServerHandler{
+		baseHandler: baseHandler{
+			done:             internal.NewDone(),
+			lines:            make(chan *line.Line, 4),
+			serverMessages:   make(chan string, 8),
+			maprMessages:     make(chan string, 4),
+			ackCloseReceived: make(chan struct{}),
+			user:             &userserver.User{Name: userName},
+			codec:            newProtocolCodec(&userserver.User{Name: userName}),
+		},
+		serverCfg: &config.ServerConfig{
+			AuthKeyEnabled: true,
+		},
+	}
+	handler.commands = map[string]commandHandler{
+		"tail": immediateNoopCommandHandler,
+		"cat":  immediateNoopCommandHandler,
+		"grep": immediateNoopCommandHandler,
+		"map":  immediateNoopCommandHandler,
+	}
+	handler.handleCommandCb = func(ctx context.Context, ltx lcontext.LContext, argc int, args []string, commandName string) {
+		if command, found := handler.commands[commandName]; found {
+			command(ctx, ltx, argc, args, func() {})
+		}
+	}
+	handler.send(handler.serverMessages, protocol.HiddenCapabilitiesPrefix+protocol.CapabilityQueryUpdateV1)
+	return handler
+}
+
+type recordedCommand struct {
+	command string
+	ctx     context.Context
+}
+
+type sessionDispatchRecorder struct {
+	starts chan recordedCommand
+	wg     sync.WaitGroup
+}
+
+func newSessionDispatchTestHandler(userName string) (*ServerHandler, *sessionDispatchRecorder) {
+	handler := newSessionTestHandler(userName)
+	recorder := &sessionDispatchRecorder{
+		starts: make(chan recordedCommand, 4),
+	}
+	handler.commands = map[string]commandHandler{
+		"tail": func(ctx context.Context, _ lcontext.LContext, argc int, args []string, commandFinished func()) {
+			recorder.starts <- recordedCommand{
+				command: strings.Join(args, " "),
+				ctx:     ctx,
+			}
+			recorder.wg.Add(1)
+			go func() {
+				defer recorder.wg.Done()
+				<-ctx.Done()
+				commandFinished()
+			}()
+		},
+	}
+	return handler, recorder
+}
+
+func newQuerySessionDispatchTestHandler(userName string) (*ServerHandler, *sessionDispatchRecorder) {
+	handler := newSessionTestHandler(userName)
+	recorder := &sessionDispatchRecorder{
+		starts: make(chan recordedCommand, 8),
+	}
+	handler.commands = map[string]commandHandler{
+		"map": func(ctx context.Context, _ lcontext.LContext, _ int, args []string, commandFinished func()) {
+			recorder.starts <- recordedCommand{
+				command: strings.Join(args, " "),
+				ctx:     ctx,
+			}
+			commandFinished()
+		},
+		"tail": func(ctx context.Context, _ lcontext.LContext, _ int, args []string, commandFinished func()) {
+			recorder.starts <- recordedCommand{
+				command: strings.Join(args, " "),
+				ctx:     ctx,
+			}
+			commandFinished()
+		},
+		"cat": func(ctx context.Context, _ lcontext.LContext, _ int, args []string, commandFinished func()) {
+			recorder.starts <- recordedCommand{
+				command: strings.Join(args, " "),
+				ctx:     ctx,
+			}
+			commandFinished()
+		},
+	}
+	return handler, recorder
+}
+
+func immediateNoopCommandHandler(_ context.Context, _ lcontext.LContext, _ int, _ []string, commandFinished func()) {
+	commandFinished()
+}
+
+func (r *sessionDispatchRecorder) waitForStart(t *testing.T) recordedCommand {
+	t.Helper()
+
+	select {
+	case started := <-r.starts:
+		return started
+	case <-time.After(250 * time.Millisecond):
+		t.Fatal("timed out waiting for dispatched session command")
+		return recordedCommand{}
+	}
+}
+
+func mustSessionPayload(t *testing.T, spec session.Spec) string {
+	t.Helper()
+
+	payload, err := json.Marshal(spec)
+	if err != nil {
+		t.Fatalf("marshal session spec: %v", err)
+	}
+	return base64.StdEncoding.EncodeToString(payload)
+}
+
+func TestParseSessionCommandWithGeneration(t *testing.T) {
+	spec := session.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+
+	action, generation, parsedSpec, err := parseSessionCommand([]string{"SESSION", "UPDATE", "7", mustSessionPayload(t, spec)}, 4)
+	if err != nil {
+		t.Fatalf("parseSessionCommand error: %v", err)
+	}
+	if action != "UPDATE" {
+		t.Fatalf("unexpected action: %s", action)
+	}
+	if generation != 7 {
+		t.Fatalf("unexpected generation: %d", generation)
+	}
+	if parsedSpec.Mode != spec.Mode {
+		t.Fatalf("unexpected parsed mode: %v", parsedSpec.Mode)
+	}
+}
+
+func TestSessionStateStoreUpdateAutoIncrementsGeneration(t *testing.T) {
+	handler := newSessionTestHandler("session-generation-user")
+	readServerMessage(t, handler.serverMessages)
+
+	startPayload := mustSessionPayload(t, session.Spec{Mode: omode.TailClient, Regex: "ERROR"})
+	updatePayload := mustSessionPayload(t, session.Spec{Mode: omode.TailClient, Regex: "WARN"})
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "START", startPayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckStartOKPrefix+" 1" {
+		t.Fatalf("unexpected session start ack: %q", message)
+	}
+
+	handler.handleSessionCommand(context.Background(), lcontext.LContext{}, 3, []string{"SESSION", "UPDATE", updatePayload}, func() {})
+	if message := readServerMessage(t, handler.serverMessages); message != sessionAckUpdateOKPrefix+" 2" {
+		t.Fatalf("unexpected session update ack: %q", message)
+	}
+}
+
+func waitForContextDone(ctx context.Context, t *testing.T) {
+	t.Helper()
+
+	select {
+	case <-ctx.Done():
+	case <-time.After(250 * time.Millisecond):
+		t.Fatal("timed out waiting for context cancellation")
+	}
+}
+
+func resetServerLogger(t *testing.T) {
+	t.Helper()
+
+	originalLogger := dlog.Server
+	dlog.Server = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Server = originalLogger
+	})
+}