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, 205 insertions, 42 deletions

diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 69bebc4..f957af7 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -2,7 +2,11 @@ package handlers
 
 import (
 	"context"
+	"encoding/base64"
+	"os/exec"
+	"runtime"
 	"strings"
+	"sync/atomic"
 
 	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/config"
@@ -10,7 +14,11 @@ import (
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/lcontext"
 	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/protocol"
+	sshserver "github.com/mimecast/dtail/internal/ssh/server"
 	user "github.com/mimecast/dtail/internal/user/server"
+
+	gossh "golang.org/x/crypto/ssh"
 )
 
 // ServerHandler implements the Reader and Writer interfaces to handle
@@ -18,30 +26,59 @@ import (
 // This handler implements the handler of the SSH server.
 type ServerHandler struct {
 	baseHandler
-	catLimiter  chan struct{}
-	tailLimiter chan struct{}
-	regex       string
+	catLimiter   chan struct{}
+	tailLimiter  chan struct{}
+	serverCfg    *config.ServerConfig
+	authKeyStore *sshserver.AuthKeyStore
+	regex        string
+	commands     map[string]commandHandler
+	sessionState sessionCommandState
+	// Track pending files waiting for limiter slots
+	pendingFiles int32
 }
 
+type commandHandler func(context.Context, lcontext.LContext, int, []string, func())
+
+var _ Handler = (*ServerHandler)(nil)
+
+var serverJournalCapabilityAvailable = detectJournalCapabilityAvailable()
+var advertisedServerCapabilities = serverCapabilities(runtime.GOOS, serverJournalCapabilityAvailable)
+
 // NewServerHandler returns the server handler.
 func NewServerHandler(user *user.User, catLimiter,
-	tailLimiter chan struct{}) *ServerHandler {
+	tailLimiter chan struct{}, serverCfg *config.ServerConfig,
+	authKeyStore *sshserver.AuthKeyStore) *ServerHandler {
 
 	dlog.Server.Debug(user, "Creating new server handler")
+	if serverCfg == nil {
+		dlog.Server.FatalPanic("Missing server config in NewServerHandler")
+	}
+
+	if authKeyStore == nil {
+		dlog.Server.FatalPanic("NewServerHandler: authKeyStore must not be nil")
+	}
+
 	h := ServerHandler{
 		baseHandler: baseHandler{
-			done:             internal.NewDone(),
-			lines:            make(chan *line.Line, 100),
-			serverMessages:   make(chan string, 10),
-			maprMessages:     make(chan string, 10),
-			ackCloseReceived: make(chan struct{}),
-			user:             user,
+			done:                internal.NewDone(),
+			lines:               make(chan *line.Line, 100),
+			serverMessages:      make(chan string, 10),
+			maprMessages:        make(chan string, 10),
+			ackCloseReceived:    make(chan struct{}),
+			user:                user,
+			codec:               newProtocolCodec(user),
+			maxCommandFrameSize: serverCfg.MaxCommandFrameSize,
 		},
-		catLimiter:  catLimiter,
-		tailLimiter: tailLimiter,
-		regex:       ".",
+		catLimiter:   catLimiter,
+		tailLimiter:  tailLimiter,
+		serverCfg:    serverCfg,
+		authKeyStore: authKeyStore,
+		regex:        ".",
 	}
 	h.handleCommandCb = h.handleUserCommand
+	h.commands = h.newCommandRegistry()
+	h.output.configure(h.outputManagerConfig())
+	h.baseHandler.activeGeneration = h.sessionState.currentGeneration
 
 	fqdn, err := config.Hostname()
 	if err != nil {
@@ -50,53 +87,179 @@ func NewServerHandler(user *user.User, catLimiter,
 
 	s := strings.Split(fqdn, ".")
 	h.hostname = s[0]
+	h.send(h.serverMessages, protocol.HiddenCapabilitiesPrefix+advertisedServerCapabilities)
 
 	return &h
 }
 
+func detectJournalCapabilityAvailable() bool {
+	if runtime.GOOS != "linux" {
+		return false
+	}
+	_, err := exec.LookPath("journalctl")
+	return err == nil
+}
+
+func serverCapabilities(goos string, journalctlAvailable bool) string {
+	capabilities := []string{protocol.CapabilityQueryUpdateV1}
+	if goos == "linux" && journalctlAvailable {
+		capabilities = append(capabilities, protocol.CapabilityJournalV1)
+	}
+	return strings.Join(capabilities, " ")
+}
+
 func (h *ServerHandler) handleUserCommand(ctx context.Context, ltx lcontext.LContext,
 	argc int, args []string, commandName string) {
 
 	dlog.Server.Debug(h.user, "Handling user command", argc, args)
+	shutdownOnCompletion := shouldShutdownOnCommandCompletion(commandName)
 	h.incrementActiveCommands()
 	commandFinished := func() {
-		if h.decrementActiveCommands() == 0 {
+		activeCommands := h.decrementActiveCommands()
+		pendingFiles := atomic.LoadInt32(&h.pendingFiles)
+		dlog.Server.Debug(h.user, "Command finished", "activeCommands", activeCommands, "pendingFiles", pendingFiles)
+
+		// Release the per-command context + watcher goroutine created for
+		// this invocation (see baseHandler.handleCommand). In the session
+		// dispatch path ctx carries no command cancel and this is a no-op;
+		// the session state owns cancellation there.
+		cancelCommandContext(ctx)
+
+		// Only shutdown if no active commands AND no pending files.
+		// AUTHKEY is a session-side effect command and should not terminate the shell
+		// because user commands may still follow in the same session.
+		if shutdownOnCompletion && activeCommands == 0 && pendingFiles == 0 && !h.sessionState.keepAlive() {
 			h.shutdown()
 		}
 	}
 
-	switch commandName {
-	case "grep", "cat":
-		command := newReadCommand(h, omode.CatClient)
-		go func() {
-			command.Start(ctx, ltx, argc, args, 1)
-			commandFinished()
-		}()
-	case "tail":
-		command := newReadCommand(h, omode.TailClient)
-		go func() {
-			command.Start(ctx, ltx, argc, args, 10)
-			commandFinished()
-		}()
-	case "map":
-		command, aggregate, err := newMapCommand(h, argc, args)
-		if err != nil {
-			h.sendln(h.serverMessages, err.Error())
-			dlog.Server.Error(h.user, err)
-			commandFinished()
-			return
-		}
-		h.aggregate = aggregate
+	handler, found := h.commands[commandName]
+	if !found {
+		h.sendln(h.serverMessages, dlog.Server.Error(h.user,
+			"Received unknown user command", commandName, argc, args))
+		commandFinished()
+		return
+	}
+
+	handler(ctx, ltx, argc, args, commandFinished)
+}
+
+func shouldShutdownOnCommandCompletion(commandName string) bool {
+	switch {
+	case strings.EqualFold(commandName, "AUTHKEY"):
+		return false
+	case strings.EqualFold(commandName, "SESSION"):
+		return false
+	default:
+		return true
+	}
+}
+
+func (h *ServerHandler) newCommandRegistry() map[string]commandHandler {
+	return map[string]commandHandler{
+		"grep":    h.makeReadCommandHandler(omode.GrepClient, 1),
+		"cat":     h.makeReadCommandHandler(omode.CatClient, 1),
+		"tail":    h.makeReadCommandHandler(omode.TailClient, 10),
+		"map":     h.handleMapCommand,
+		".ack":    h.handleAckUserCommand,
+		"AUTHKEY": h.handleAuthKeyCommand,
+		"SESSION": h.handleSessionCommand,
+		"authkey": h.handleAuthKeyCommand,
+		"session": h.handleSessionCommand,
+	}
+}
+
+func (h *ServerHandler) makeReadCommandHandler(mode omode.Mode, tailBackoff int) commandHandler {
+	return func(ctx context.Context, ltx lcontext.LContext, argc int, args []string, commandFinished func()) {
+		command := newReadCommand(h, mode)
 		go func() {
-			command.Start(ctx, h.maprMessages)
+			command.Start(ctx, ltx, argc, args, tailBackoff)
 			commandFinished()
 		}()
-	case ".ack":
-		h.handleAckCommand(argc, args)
+	}
+}
+
+func (h *ServerHandler) handleMapCommand(ctx context.Context, _ lcontext.LContext, argc int, args []string, commandFinished func()) {
+	command, aggregate, err := newMapCommand(h, argc, args)
+	if err != nil {
+		h.sendln(h.serverMessages, err.Error())
+		dlog.Server.Error(h.user, err)
 		commandFinished()
-	default:
-		h.sendln(h.serverMessages, dlog.Server.Error(h.user,
-			"Received unknown user command", commandName, argc, args))
+		return
+	}
+
+	// Use the atomic setter so concurrent reads from Shutdown, Aggregate,
+	// and resetSessionAggregates are race-free.
+	h.setAggregate(aggregate)
+	maprMessages, closeMaprMessages := h.newGeneratedMaprMessagesChannel(ctx, sessionGenerationFromContext(ctx))
+	go func() {
+		command.Start(ctx, maprMessages)
+		closeMaprMessages()
 		commandFinished()
+	}()
+}
+
+func (h *ServerHandler) handleAckUserCommand(_ context.Context, _ lcontext.LContext, argc int, args []string, commandFinished func()) {
+	h.handleAckCommand(argc, args)
+	commandFinished()
+}
+
+func (h *ServerHandler) handleAuthKeyCommand(_ context.Context, _ lcontext.LContext,
+	argc int, args []string, commandFinished func()) {
+
+	defer commandFinished()
+
+	if !h.serverCfg.AuthKeyEnabled {
+		h.sendln(h.serverMessages, "AUTHKEY ERR feature disabled")
+		return
+	}
+
+	if argc < 2 || strings.TrimSpace(args[1]) == "" {
+		h.sendln(h.serverMessages, "AUTHKEY ERR missing public key")
+		return
+	}
+
+	decodedPubKey, err := base64.StdEncoding.DecodeString(args[1])
+	if err != nil {
+		h.sendln(h.serverMessages, "AUTHKEY ERR invalid base64")
+		return
+	}
+
+	pubKey, err := gossh.ParsePublicKey(decodedPubKey)
+	if err != nil {
+		h.sendln(h.serverMessages, "AUTHKEY ERR invalid public key")
+		return
+	}
+
+	if h.authKeyStore == nil {
+		h.sendln(h.serverMessages, "AUTHKEY ERR internal key store unavailable")
+		return
+	}
+	h.authKeyStore.Add(h.user.Name, pubKey)
+	h.sendln(h.serverMessages, "AUTHKEY OK")
+}
+
+func (h *ServerHandler) newGeneratedMaprMessagesChannel(ctx context.Context, generation uint64) (chan string, func()) {
+	maprMessages := make(chan string, 16)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for {
+			select {
+			case message, ok := <-maprMessages:
+				if !ok {
+					return
+				}
+				h.send(h.maprMessages, encodeGeneratedMessage(generation, message))
+			case <-ctx.Done():
+				return
+			case <-h.done.Done():
+				return
+			}
+		}
+	}()
+	return maprMessages, func() {
+		close(maprMessages)
+		<-done
 	}
 }