From 849951be1d1a7ee9f9302006ccb187bf5b4e36f3 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Wed, 22 Jul 2026 23:51:18 +0300 Subject: =?UTF-8?q?feat:=20DTail=20fork=20=E2=80=94=20server/client=20feat?= =?UTF-8?q?ure=20development?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- internal/server/handlers/serverhandler.go | 247 +++++++++++++++++++++++++----- 1 file changed, 205 insertions(+), 42 deletions(-) (limited to 'internal/server/handlers/serverhandler.go') 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 } } -- cgit v1.2.3