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, 218 insertions, 30 deletions

diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
index 5c3d455..b98f815 100644
--- a/internal/clients/connectors/serverconnection.go
+++ b/internal/clients/connectors/serverconnection.go
@@ -2,20 +2,38 @@ package connectors
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"io"
+	"net"
+	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
-	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
+	sessionspec "github.com/mimecast/dtail/internal/session"
 	"github.com/mimecast/dtail/internal/ssh/client"
 
 	"golang.org/x/crypto/ssh"
 )
 
+// SSHSettings provides the connection settings needed by ServerConnection.
+type SSHSettings interface {
+	SSHPort() int
+	SSHConnectTimeout() time.Duration
+}
+
+const (
+	defaultSSHConnectTimeout = 2 * time.Second
+	defaultSSHPort           = 2222
+	defaultCapabilityWait    = 250 * time.Millisecond
+)
+
 // ServerConnection represents a connection to a single remote dtail server via
 // SSH protocol.
 type ServerConnection struct {
@@ -28,30 +46,61 @@ type ServerConnection struct {
 	config          *ssh.ClientConfig
 	handler         handlers.Handler
 	commands        []string
+	sessionSpec     sessionspec.Spec
+	sessionState    committedSessionState
+	interactive     bool
+	authKeyPath     string
+	authKeyDisabled bool
 	hostKeyCallback client.HostKeyCallback
-	throttlingDone  bool
+	// throttleReleased ensures the throttle slot is returned to throttleCh
+	// exactly once, even if both the early-release path in handle() and the
+	// deferred cleanup in Start() execute concurrently or the same goroutine
+	// hits both paths.  sync.Once is safe for concurrent callers whereas the
+	// previous bool guard was not synchronized.
+	throttleReleased sync.Once
 }
 
+var _ Connector = (*ServerConnection)(nil)
+
 // NewServerConnection returns a new DTail SSH server connection.
 func NewServerConnection(server string, userName string,
 	authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback,
-	handler handlers.Handler, commands []string) *ServerConnection {
+	handler handlers.Handler, commands []string, sessionSpec sessionspec.Spec,
+	interactive bool, authKeyPath string, authKeyDisabled bool, settings SSHSettings) *ServerConnection {
 
 	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
+	sshConnectTimeout := defaultSSHConnectTimeout
+	defaultPort := defaultSSHPort
+	if settings != nil {
+		sshConnectTimeout = settings.SSHConnectTimeout()
+		defaultPort = settings.SSHPort()
+	}
+	if sshConnectTimeout <= 0 {
+		sshConnectTimeout = defaultSSHConnectTimeout
+	}
+	if defaultPort <= 0 {
+		defaultPort = defaultSSHPort
+	}
+
 	c := ServerConnection{
 		hostKeyCallback: hostKeyCallback,
 		server:          server,
 		handler:         handler,
 		commands:        commands,
+		sessionSpec:     sessionSpec,
+		interactive:     interactive,
+		authKeyPath:     resolveAuthKeyPath(authKeyPath),
+		authKeyDisabled: authKeyDisabled,
 		config: &ssh.ClientConfig{
-			User:            userName,
-			Auth:            authMethods,
-			HostKeyCallback: hostKeyCallback.Wrap(),
-			Timeout:         time.Second * 2,
+			User:    userName,
+			Auth:    authMethods,
+			Timeout: sshConnectTimeout,
+			// HostKeyCallback is assigned per-handshake in dial() so the
+			// callback can honour the handshake's context (see dial()).
 		},
 	}
 
-	c.initServerPort()
+	c.initServerPort(defaultPort)
 	return &c
 }
 
@@ -61,12 +110,42 @@ func (c *ServerConnection) Server() string { return c.server }
 // Handler returns the handler used for the connection.
 func (c *ServerConnection) Handler() handlers.Handler { return c.handler }
 
+// SupportsQueryUpdates reports whether the remote server advertised the
+// runtime query replacement capability. Older servers simply time out and
+// return false here without affecting the legacy command path.
+func (c *ServerConnection) SupportsQueryUpdates(timeout time.Duration) bool {
+	return supportsQueryUpdates(c.handler, timeout)
+}
+
+// ApplySessionSpec starts or updates the interactive session state on the
+// existing SSH connection when runtime query updates are supported.
+func (c *ServerConnection) ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error {
+	return applySessionSpec(c.server, c.handler, &c.sessionState, spec, timeout)
+}
+
+// ApplySessionSpecWithGeneration starts or updates the interactive session
+// state using an explicit committed generation as the base for the update.
+func (c *ServerConnection) ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error {
+	return applySessionSpecWithGeneration(c.server, c.handler, &c.sessionState, spec, generation, false, timeout)
+}
+
+// CommittedSession returns the last server-acknowledged session state.
+func (c *ServerConnection) CommittedSession() (sessionspec.Spec, uint64, bool) {
+	return c.sessionState.snapshot()
+}
+
+// RestoreCommittedSession resets the local session snapshot without advancing
+// the generation.
+func (c *ServerConnection) RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) {
+	c.sessionState.restore(spec, generation, committed)
+}
+
 // Attempt to parse the server port address from the provided server FQDN.
-func (c *ServerConnection) initServerPort() {
+func (c *ServerConnection) initServerPort(defaultPort int) {
 	parts := strings.Split(c.server, ":")
 	if len(parts) == 1 {
 		c.hostname = c.server
-		c.port = config.Common.SSHPort
+		c.port = defaultPort
 		return
 	}
 
@@ -99,12 +178,15 @@ func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
 
 	go func() {
 		defer func() {
-			if !c.throttlingDone {
-				dlog.Client.Debug(c.server, "Unthrottling connection (1)",
+			// Release the throttle slot on the way out regardless of which
+			// code path already attempted it.  throttleReleased.Do guarantees
+			// the drain happens exactly once even if handle() already released
+			// the slot early (the fast path when the session is fully up).
+			c.throttleReleased.Do(func() {
+				dlog.Client.Debug(c.server, "Unthrottling connection (cleanup)",
 					len(throttleCh), cap(throttleCh))
-				c.throttlingDone = true
 				<-throttleCh
-			}
+			})
 			cancel()
 		}()
 
@@ -133,10 +215,33 @@ func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc,
 	address := fmt.Sprintf("%s:%d", c.hostname, c.port)
 	dlog.Client.Debug(c.server, "Dialing into the connection", address)
 
-	client, err := ssh.Dial("tcp", address, c.config)
+	// Use context-aware dialing to enable proper cancellation during connection establishment.
+	// TCP KeepAlive (30s) prevents silent connection failures on long-lived connections.
+	dialer := &net.Dialer{
+		Timeout:   c.config.Timeout, // Use the SSH config timeout (2 seconds)
+		KeepAlive: 30 * time.Second, // Standard Go default for connection health monitoring
+	}
+
+	// Establish TCP connection with context support for cancellation
+	conn, err := dialer.DialContext(ctx, "tcp", address)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to dial TCP connection to %s: %w", address, err)
 	}
+
+	// Perform SSH handshake over the established TCP connection. Build a
+	// per-handshake ssh.ClientConfig so the host-key callback is bound to
+	// ctx and unblocks cleanly if the handshake is cancelled (e.g. when the
+	// user aborts before responding to the unknown-host prompt).
+	handshakeConfig := *c.config
+	handshakeConfig.HostKeyCallback = c.hostKeyCallback.Wrap(ctx)
+	sshConn, chans, reqs, err := ssh.NewClientConn(conn, address, &handshakeConfig)
+	if err != nil {
+		conn.Close()
+		return fmt.Errorf("SSH handshake failed for %s: %w", address, err)
+	}
+
+	// Create SSH client from the connection components
+	client := ssh.NewClient(sshConn, chans, reqs)
 	defer client.Close()
 
 	return c.session(ctx, cancel, client, throttleCh)
@@ -149,7 +254,7 @@ func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFun
 	dlog.Client.Debug(c.server, "Creating SSH session")
 	session, err := client.NewSession()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to create SSH session for %s: %w", c.server, err)
 	}
 	defer session.Close()
 	return c.handle(ctx, cancel, session, throttleCh)
@@ -161,14 +266,14 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 	dlog.Client.Debug(c.server, "Creating handler for SSH session")
 	stdinPipe, err := session.StdinPipe()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get SSH session stdin pipe for %s: %w", c.server, err)
 	}
 	stdoutPipe, err := session.StdoutPipe()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get SSH session stdout pipe for %s: %w", c.server, err)
 	}
 	if err := session.Shell(); err != nil {
-		return err
+		return fmt.Errorf("failed to start SSH shell for %s: %w", c.server, err)
 	}
 
 	go func() {
@@ -191,22 +296,105 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 		}
 	}()
 
-	// Send all commands to client.
-	for _, command := range c.commands {
-		dlog.Client.Debug(command)
-		if err := c.handler.SendMessage(command); err != nil {
-			dlog.Client.Debug(err)
-		}
+	if c.authKeyDisabled {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration because auth-key is disabled")
+	} else {
+		c.sendAuthKeyRegistrationCommand()
+	}
+
+	if err := dispatchInitialCommands(c.server, c.handler, c.commands, c.interactive, c.sessionSpec, &c.sessionState); err != nil {
+		c.handler.Shutdown()
+		return err
 	}
 
-	if !c.throttlingDone {
-		dlog.Client.Debug(c.server, "Unthrottling connection (2)",
+	// Release the throttle slot as soon as the session is fully established so
+	// the next pending connection can proceed without waiting for this session
+	// to finish.  throttleReleased.Do is idempotent: if the deferred cleanup
+	// in Start() fires first (e.g. on a dial error path that never reaches
+	// here), the slot is still returned exactly once.
+	c.throttleReleased.Do(func() {
+		dlog.Client.Debug(c.server, "Unthrottling connection (session up)",
 			len(throttleCh), cap(throttleCh))
-		c.throttlingDone = true
 		<-throttleCh
-	}
+	})
 
 	<-ctx.Done()
 	c.handler.Shutdown()
 	return nil
 }
+
+// resolveAuthKeyPath returns the effective auth-key path. When the provided
+// path is non-empty it is used as-is. Otherwise the function falls back to
+// $HOME/.ssh/id_rsa. If HOME is also empty it returns "" so that the AUTHKEY
+// registration step (sendAuthKeyRegistrationCommand) will skip gracefully
+// instead of trying to open a path that the SSH library cannot expand.
+func resolveAuthKeyPath(authKeyPath string) string {
+	if strings.TrimSpace(authKeyPath) != "" {
+		return authKeyPath
+	}
+	homeDir := os.Getenv("HOME")
+	if homeDir == "" {
+		return ""
+	}
+	return filepath.Join(homeDir, ".ssh", "id_rsa")
+}
+
+func (c *ServerConnection) sendAuthKeyRegistrationCommand() {
+	authKeyPubPath := c.authKeyPath + ".pub"
+	authKeyPubBytes, err := os.ReadFile(authKeyPubPath)
+	if err != nil {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, unable to read public key", authKeyPubPath, err)
+		return
+	}
+
+	authKeyBase64, err := extractAuthKeyBase64(authKeyPubBytes)
+	if err != nil {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, invalid public key file", authKeyPubPath, err)
+		return
+	}
+
+	if err := c.handler.SendMessage("AUTHKEY " + authKeyBase64); err != nil {
+		dlog.Client.Debug(c.server, "Unable to send AUTHKEY registration command", err)
+		return
+	}
+	dlog.Client.Debug(c.server, "Sent AUTHKEY registration command", authKeyPubPath)
+}
+
+func extractAuthKeyBase64(authKeyPubBytes []byte) (string, error) {
+	authKeyPubContent := string(authKeyPubBytes)
+	for _, line := range strings.Split(authKeyPubContent, "\n") {
+		trimmedLine := strings.TrimSpace(line)
+		if trimmedLine == "" || strings.HasPrefix(trimmedLine, "#") {
+			continue
+		}
+
+		fields := strings.Fields(trimmedLine)
+		if len(fields) < 2 {
+			return "", fmt.Errorf("expected authorized key format '<type> <base64-key> [comment]'")
+		}
+
+		authKeyBase64 := strings.TrimSpace(fields[1])
+		if _, err := base64.StdEncoding.DecodeString(authKeyBase64); err != nil {
+			return "", fmt.Errorf("invalid base64 public key: %w", err)
+		}
+
+		return authKeyBase64, nil
+	}
+
+	return "", fmt.Errorf("no public key found")
+}
+
+func supportsQueryUpdates(handler handlers.Handler, timeout time.Duration) bool {
+	if handler == nil {
+		return false
+	}
+
+	if timeout <= 0 {
+		timeout = defaultCapabilityWait
+	}
+	if !handler.WaitForCapabilities(timeout) {
+		return false
+	}
+
+	return handler.HasCapability(protocol.CapabilityQueryUpdateV1)
+}