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, 880 insertions, 0 deletions

diff --git a/internal/clients/connectors/serverconnection_test.go b/internal/clients/connectors/serverconnection_test.go
new file mode 100644
index 0000000..4e22dc7
--- /dev/null
+++ b/internal/clients/connectors/serverconnection_test.go
@@ -0,0 +1,880 @@
+package connectors
+
+import (
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/protocol"
+	sessionspec "github.com/mimecast/dtail/internal/session"
+
+	"golang.org/x/crypto/ssh"
+)
+
+// TestResolveAuthKeyPathNoLiteralPath is a regression test for the bug
+// described in task l6: when authKeyPath is empty and HOME is also unset,
+// resolveAuthKeyPath must return "" instead of a mangled path like
+// "/.ssh/id_rsa" or the literal "~/.ssh/id_rsa" that the SSH stack cannot use.
+func TestResolveAuthKeyPathNoLiteralPath(t *testing.T) {
+	// Unset HOME so the environment fallback is also empty.
+	t.Setenv("HOME", "")
+
+	got := resolveAuthKeyPath("")
+	if got != "" {
+		t.Fatalf("resolveAuthKeyPath(\"\") with empty HOME = %q; want \"\"", got)
+	}
+}
+
+// TestResolveAuthKeyPathExplicitPathPassedThrough verifies that a non-empty
+// explicit auth key path is returned unchanged.
+func TestResolveAuthKeyPathExplicitPathPassedThrough(t *testing.T) {
+	got := resolveAuthKeyPath("/custom/key")
+	if got != "/custom/key" {
+		t.Fatalf("resolveAuthKeyPath(\"/custom/key\") = %q; want \"/custom/key\"", got)
+	}
+}
+
+// TestResolveAuthKeyPathFallsBackToHome verifies that when authKeyPath is empty
+// but HOME is set, the function returns the expected default path.
+func TestResolveAuthKeyPathFallsBackToHome(t *testing.T) {
+	t.Setenv("HOME", "/home/testuser")
+
+	got := resolveAuthKeyPath("")
+	want := "/home/testuser/.ssh/id_rsa"
+	if got != want {
+		t.Fatalf("resolveAuthKeyPath(\"\") = %q; want %q", got, want)
+	}
+}
+
+func TestExtractAuthKeyBase64(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	t.Run("valid authorized key line", func(t *testing.T) {
+		pubKey := []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n")
+
+		got, err := extractAuthKeyBase64(pubKey)
+		if err != nil {
+			t.Fatalf("Expected valid key, got error: %v", err)
+		}
+		if got != "AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" {
+			t.Fatalf("Unexpected base64 payload: %s", got)
+		}
+	})
+
+	t.Run("invalid key format", func(t *testing.T) {
+		_, err := extractAuthKeyBase64([]byte("not-a-valid-authorized-key-line"))
+		if err == nil {
+			t.Fatalf("Expected parse error for invalid key format")
+		}
+	})
+
+	t.Run("invalid base64 payload", func(t *testing.T) {
+		_, err := extractAuthKeyBase64([]byte("ssh-ed25519 !!! not-valid\n"))
+		if err == nil {
+			t.Fatalf("Expected error for invalid base64 payload")
+		}
+	})
+}
+
+func TestSendAuthKeyRegistrationCommand(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	tempDir := t.TempDir()
+	privateKeyPath := filepath.Join(tempDir, "id_rsa")
+	publicKeyPath := privateKeyPath + ".pub"
+	if err := os.WriteFile(publicKeyPath,
+		[]byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n"), 0600); err != nil {
+		t.Fatalf("Unable to write public key test file: %v", err)
+	}
+
+	handler := &mockHandler{}
+	conn := &ServerConnection{
+		server:      "srv1",
+		handler:     handler,
+		authKeyPath: privateKeyPath,
+	}
+
+	conn.sendAuthKeyRegistrationCommand()
+
+	if len(handler.commands) != 1 {
+		t.Fatalf("Expected one AUTHKEY command, got %d", len(handler.commands))
+	}
+	expected := "AUTHKEY AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+	if handler.commands[0] != expected {
+		t.Fatalf("Unexpected AUTHKEY command.\nexpected: %s\ngot:      %s", expected, handler.commands[0])
+	}
+}
+
+func TestNewServerConnectionUsesInjectedSettings(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := NewServerConnection(
+		"srv1",
+		"user",
+		nil,
+		testHostKeyCallback{},
+		&mockHandler{},
+		nil,
+		sessionspec.Spec{},
+		false,
+		"",
+		false,
+		testSSHSettings{port: 3022, timeout: 5 * time.Second},
+	)
+
+	if conn.hostname != "srv1" {
+		t.Fatalf("Expected hostname srv1, got %q", conn.hostname)
+	}
+	if conn.port != 3022 {
+		t.Fatalf("Expected injected port 3022, got %d", conn.port)
+	}
+	if conn.config.Timeout != 5*time.Second {
+		t.Fatalf("Expected injected timeout 5s, got %v", conn.config.Timeout)
+	}
+}
+
+func TestNewServerConnectionFallsBackToDefaults(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := NewServerConnection(
+		"srv1",
+		"user",
+		nil,
+		testHostKeyCallback{},
+		&mockHandler{},
+		nil,
+		sessionspec.Spec{},
+		false,
+		"",
+		false,
+		testSSHSettings{},
+	)
+
+	if conn.port != defaultSSHPort {
+		t.Fatalf("Expected default port %d, got %d", defaultSSHPort, conn.port)
+	}
+	if conn.config.Timeout != defaultSSHConnectTimeout {
+		t.Fatalf("Expected default timeout %v, got %v", defaultSSHConnectTimeout, conn.config.Timeout)
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdates(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+		},
+	}
+
+	if !conn.SupportsQueryUpdates(10 * time.Millisecond) {
+		t.Fatalf("expected query-update capability to be detected")
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdatesFallsBackForOlderServers(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{},
+	}
+
+	if conn.SupportsQueryUpdates(5 * time.Millisecond) {
+		t.Fatalf("expected old-server fallback when no capability is advertised")
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdatesRequiresCapabilityFlag(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{
+			waitForCapabilities: true,
+		},
+	}
+
+	if conn.SupportsQueryUpdates(10 * time.Millisecond) {
+		t.Fatalf("expected capability wait success alone to be insufficient")
+	}
+}
+
+func TestServerConnectionApplySessionSpecStart(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		server: "srv1",
+		handler: &mockHandler{
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+			sessionAcks: []handlers.SessionAck{{
+				Action:     "start",
+				Generation: 1,
+			}},
+		},
+	}
+
+	spec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	if err := conn.ApplySessionSpec(spec, 10*time.Millisecond); err != nil {
+		t.Fatalf("ApplySessionSpec() error = %v", err)
+	}
+
+	mock := conn.handler.(*mockHandler)
+	if len(mock.commands) != 1 {
+		t.Fatalf("expected one session command, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 1 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecUpdateUsesNextGeneration(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 4},
+			{Action: "update", Generation: 5},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	updateSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("update ApplySessionSpec() error = %v", err)
+	}
+	if len(mock.commands) != 2 {
+		t.Fatalf("expected two session commands, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 5 || committedSpec.Regex != "WARN" {
+		t.Fatalf("unexpected committed session after update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecReappliesPreviousSpecForRollback(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 4},
+			{Action: "update", Generation: 5},
+			{Action: "update", Generation: 6},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	updateSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("update ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("rollback ApplySessionSpec() error = %v", err)
+	}
+	if len(mock.commands) != 3 {
+		t.Fatalf("expected three session commands, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 6 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session after rollback: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecFallsBackForUnsupportedServer(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{},
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"}, 5*time.Millisecond)
+	if !errors.Is(err, ErrSessionUnsupported) {
+		t.Fatalf("expected ErrSessionUnsupported, got %v", err)
+	}
+}
+
+func TestRequireJournalCapability(t *testing.T) {
+	tests := []struct {
+		name                string
+		spec                sessionspec.Spec
+		waitForCapabilities bool
+		capabilities        map[string]bool
+		wantErr             error
+		wantServerError     bool
+	}{
+		{
+			name: "journal file with journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityJournalV1: true,
+			},
+		},
+		{
+			name: "journal file without journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+			wantErr:         ErrJournalUnsupported,
+			wantServerError: true,
+		},
+		{
+			name: "journal file without capabilities advertisement",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			wantErr:         ErrJournalUnsupported,
+			wantServerError: true,
+		},
+		{
+			name: "regular file without journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"/var/log/app.log"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			handler := &mockHandler{
+				waitForCapabilities: tc.waitForCapabilities,
+				capabilities:        tc.capabilities,
+			}
+
+			err := requireJournalCapability("srv1", handler, tc.spec, 10*time.Millisecond)
+			if !errors.Is(err, tc.wantErr) {
+				t.Fatalf("requireJournalCapability() error = %v, want %v", err, tc.wantErr)
+			}
+			if got := handler.serverError != ""; got != tc.wantServerError {
+				t.Fatalf("server error recorded = %v, want %v", got, tc.wantServerError)
+			}
+			if tc.wantServerError && !strings.Contains(handler.serverError, protocol.CapabilityJournalV1) {
+				t.Fatalf("server error %q does not mention %s", handler.serverError, protocol.CapabilityJournalV1)
+			}
+		})
+	}
+}
+
+func TestDispatchInitialCommandsRejectsJournalWithoutCapability(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	spec := sessionspec.Spec{
+		Mode:  omode.CatClient,
+		Files: []string{"journal:ssh.service"},
+	}
+
+	err := dispatchInitialCommands("srv1", handler, []string{"cat: journal:ssh.service ."}, false, spec, &committedSessionState{})
+	if !errors.Is(err, ErrJournalUnsupported) {
+		t.Fatalf("expected ErrJournalUnsupported, got %v", err)
+	}
+	if len(handler.commands) != 0 {
+		t.Fatalf("expected no commands to be sent, got %#v", handler.commands)
+	}
+	if handler.Status() != 1 {
+		t.Fatalf("handler status = %d, want 1", handler.Status())
+	}
+}
+
+func TestDispatchInitialCommandsRejectsInteractiveJournalWithoutCapability(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	spec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"journal:ssh.service"},
+	}
+
+	err := dispatchInitialCommands("srv1", handler, []string{"tail: journal:ssh.service ."}, true, spec, &committedSessionState{})
+	if !errors.Is(err, ErrJournalUnsupported) {
+		t.Fatalf("expected ErrJournalUnsupported, got %v", err)
+	}
+	if len(handler.commands) != 0 {
+		t.Fatalf("expected no commands to be sent, got %#v", handler.commands)
+	}
+	if handler.Status() != 1 {
+		t.Fatalf("handler status = %d, want 1", handler.Status())
+	}
+}
+
+func TestServerConnectionApplySessionSpecPreservesCommittedStateOnRejectedUpdate(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 2},
+			{Action: "error", Error: "bad reload"},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"}
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "WARN"}, 10*time.Millisecond)
+	if !errors.Is(err, ErrSessionRejected) {
+		t.Fatalf("expected ErrSessionRejected, got %v", err)
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 2 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session after rejected update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecRejectsUnexpectedAck(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "update", Generation: 1},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}, 10*time.Millisecond)
+	if !errors.Is(err, ErrUnexpectedSessionAck) {
+		t.Fatalf("expected ErrUnexpectedSessionAck, got %v", err)
+	}
+	if _, _, ok := conn.CommittedSession(); ok {
+		t.Fatalf("unexpected committed session after mismatched ack")
+	}
+}
+
+func TestServerConnectionApplySessionSpecTimesOutWaitingForAck(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}, 10*time.Millisecond)
+	if !errors.Is(err, ErrSessionAckTimeout) {
+		t.Fatalf("expected ErrSessionAckTimeout, got %v", err)
+	}
+	if len(mock.commands) != 1 {
+		t.Fatalf("expected session command to be sent before timeout, got %d", len(mock.commands))
+	}
+	if _, _, ok := conn.CommittedSession(); ok {
+		t.Fatalf("unexpected committed session after missing ack")
+	}
+}
+
+func TestApplySessionSpecSerializesConcurrentBootstrapAndReload(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := newBlockingSessionHandler()
+	state := &committedSessionState{}
+
+	initialSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	reloadSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	initialErrCh := make(chan error, 1)
+	go func() {
+		initialErrCh <- dispatchInitialCommands("srv1", handler, nil, true, initialSpec, state)
+	}()
+
+	firstCommand := <-handler.commandsCh
+	if !strings.HasPrefix(firstCommand, "SESSION START ") {
+		t.Fatalf("expected initial SESSION START command, got %q", firstCommand)
+	}
+
+	reloadErrCh := make(chan error, 1)
+	go func() {
+		reloadErrCh <- applySessionSpec("srv1", handler, state, reloadSpec, 50*time.Millisecond)
+	}()
+
+	select {
+	case command := <-handler.commandsCh:
+		t.Fatalf("unexpected concurrent session command before bootstrap ack: %q", command)
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	handler.ackCh <- handlers.SessionAck{Action: "start", Generation: 1}
+	if err := <-initialErrCh; err != nil {
+		t.Fatalf("dispatchInitialCommands() error = %v", err)
+	}
+
+	secondCommand := <-handler.commandsCh
+	if !strings.HasPrefix(secondCommand, "SESSION UPDATE 2 ") {
+		t.Fatalf("expected reload to send SESSION UPDATE after bootstrap, got %q", secondCommand)
+	}
+
+	handler.ackCh <- handlers.SessionAck{Action: "update", Generation: 2}
+	if err := <-reloadErrCh; err != nil {
+		t.Fatalf("applySessionSpec() error = %v", err)
+	}
+
+	committedSpec, generation, ok := state.snapshot()
+	if !ok || generation != 2 || committedSpec.Regex != "WARN" {
+		t.Fatalf("unexpected committed session after reload: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+// TestThrottleReleasedIsIdempotent verifies that calling the throttle-release
+// logic from two concurrent goroutines drains throttleCh exactly once and does
+// not panic or block.  This is a regression test for the data race that existed
+// when the old bool guard (throttlingDone) was read and written without
+// synchronization: under -race two goroutines could both observe the bool as
+// false and both attempt to drain the channel, stealing an extra slot.
+func TestThrottleReleasedIsIdempotent(t *testing.T) {
+	t.Parallel()
+
+	// throttleCh is buffered with 1 slot, as in the real Start() path.
+	throttleCh := make(chan struct{}, 1)
+	throttleCh <- struct{}{} // occupy the one slot
+
+	conn := &ServerConnection{}
+
+	const workers = 64
+	var wg sync.WaitGroup
+	wg.Add(workers)
+
+	// Simulate workers racing to release the throttle slot (e.g. handle()
+	// early-release and the defer cleanup in Start() firing around the same
+	// time).  Only one drain must succeed; the rest must be no-ops.
+	for range workers {
+		go func() {
+			defer wg.Done()
+			conn.throttleReleased.Do(func() {
+				<-throttleCh
+			})
+		}()
+	}
+
+	wg.Wait()
+
+	// throttleCh must be empty: exactly one goroutine drained it.
+	if len(throttleCh) != 0 {
+		t.Fatalf("throttleCh length = %d, want 0 (slot was not released)", len(throttleCh))
+	}
+
+	// Confirm a second occupant can now be added, proving the slot is free.
+	select {
+	case throttleCh <- struct{}{}:
+		// expected: slot was freed exactly once
+	default:
+		t.Fatal("throttleCh full after release, expected one free slot")
+	}
+}
+
+type testSSHSettings struct {
+	port    int
+	timeout time.Duration
+}
+
+func (s testSSHSettings) SSHPort() int {
+	return s.port
+}
+
+func (s testSSHSettings) SSHConnectTimeout() time.Duration {
+	return s.timeout
+}
+
+type testHostKeyCallback struct{}
+
+func (testHostKeyCallback) Wrap(context.Context) ssh.HostKeyCallback {
+	return ssh.InsecureIgnoreHostKey()
+}
+
+func (testHostKeyCallback) Untrusted(string) bool {
+	return false
+}
+
+func (testHostKeyCallback) PromptAddHosts(context.Context) {}
+
+func resetClientLogger(t *testing.T) {
+	t.Helper()
+
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+}
+
+type mockHandler struct {
+	commands            []string
+	capabilities        map[string]bool
+	waitForCapabilities bool
+	sessionAcks         []handlers.SessionAck
+	serverError         string
+	status              int
+}
+
+var _ handlers.Handler = (*mockHandler)(nil)
+
+func (m *mockHandler) SendMessage(command string) error {
+	m.commands = append(m.commands, command)
+	return nil
+}
+
+func (m *mockHandler) Capabilities() []string {
+	var capabilities []string
+	for capability := range m.capabilities {
+		capabilities = append(capabilities, capability)
+	}
+	return capabilities
+}
+
+func (m *mockHandler) HasCapability(name string) bool {
+	return m.capabilities[name]
+}
+
+func (m *mockHandler) ReportServerError(message string) {
+	m.serverError = message
+	m.status = 1
+}
+
+func (m *mockHandler) Server() string {
+	return "mock"
+}
+
+func (m *mockHandler) Status() int {
+	return m.status
+}
+
+func (m *mockHandler) Shutdown() {}
+
+func (m *mockHandler) Done() <-chan struct{} {
+	ch := make(chan struct{})
+	close(ch)
+	return ch
+}
+
+func (m *mockHandler) WaitForCapabilities(timeout time.Duration) bool {
+	return m.waitForCapabilities
+}
+
+func (m *mockHandler) WaitForSessionAck(timeout time.Duration) (handlers.SessionAck, bool) {
+	if timeout <= 0 {
+		return handlers.SessionAck{}, false
+	}
+	if len(m.sessionAcks) == 0 {
+		return handlers.SessionAck{}, false
+	}
+
+	ack := m.sessionAcks[0]
+	m.sessionAcks = m.sessionAcks[1:]
+	return ack, true
+}
+
+func (m *mockHandler) Read(_ []byte) (int, error) {
+	return 0, nil
+}
+
+func (m *mockHandler) Write(p []byte) (int, error) {
+	return len(p), nil
+}
+
+type blockingSessionHandler struct {
+	mu           sync.Mutex
+	commands     []string
+	commandsCh   chan string
+	ackCh        chan handlers.SessionAck
+	capabilities map[string]bool
+}
+
+func newBlockingSessionHandler() *blockingSessionHandler {
+	return &blockingSessionHandler{
+		commandsCh: make(chan string, 8),
+		ackCh:      make(chan handlers.SessionAck, 8),
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+}
+
+var _ handlers.Handler = (*blockingSessionHandler)(nil)
+
+func (h *blockingSessionHandler) SendMessage(command string) error {
+	h.mu.Lock()
+	h.commands = append(h.commands, command)
+	h.mu.Unlock()
+	h.commandsCh <- command
+	return nil
+}
+
+func (h *blockingSessionHandler) Capabilities() []string {
+	capabilities := make([]string, 0, len(h.capabilities))
+	for capability := range h.capabilities {
+		capabilities = append(capabilities, capability)
+	}
+	return capabilities
+}
+
+func (h *blockingSessionHandler) HasCapability(name string) bool {
+	return h.capabilities[name]
+}
+
+func (*blockingSessionHandler) ReportServerError(string) {}
+
+func (*blockingSessionHandler) Server() string {
+	return "mock"
+}
+
+func (*blockingSessionHandler) Status() int {
+	return 0
+}
+
+func (*blockingSessionHandler) Shutdown() {}
+
+func (*blockingSessionHandler) Done() <-chan struct{} {
+	return make(chan struct{})
+}
+
+func (*blockingSessionHandler) WaitForCapabilities(time.Duration) bool {
+	return true
+}
+
+func (h *blockingSessionHandler) WaitForSessionAck(timeout time.Duration) (handlers.SessionAck, bool) {
+	if timeout <= 0 {
+		select {
+		case ack := <-h.ackCh:
+			return ack, true
+		default:
+			return handlers.SessionAck{}, false
+		}
+	}
+
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	select {
+	case ack := <-h.ackCh:
+		return ack, true
+	case <-timer.C:
+		return handlers.SessionAck{}, false
+	}
+}
+
+func (*blockingSessionHandler) Read(_ []byte) (int, error) {
+	return 0, nil
+}
+
+func (*blockingSessionHandler) Write(p []byte) (int, error) {
+	return len(p), nil
+}