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/auth_test.go | 50 + internal/server/continuous.go | 61 +- internal/server/continuous_test.go | 202 +++ internal/server/handlers/authkeycommand_test.go | 117 ++ internal/server/handlers/basehandler.go | 547 +++++++-- internal/server/handlers/basehandler_read_test.go | 295 +++++ internal/server/handlers/commandcancel_test.go | 143 +++ internal/server/handlers/commandtimeout_test.go | 150 +++ internal/server/handlers/framesize_test.go | 202 +++ internal/server/handlers/generation_output.go | 75 ++ internal/server/handlers/generation_output_test.go | 160 +++ internal/server/handlers/healthhandler.go | 28 +- internal/server/handlers/line_writer.go | 777 ++++++++++++ internal/server/handlers/line_writer_alloc_test.go | 54 + internal/server/handlers/line_writer_test.go | 1282 ++++++++++++++++++++ internal/server/handlers/lineprocessor.go | 177 +++ internal/server/handlers/lineprocessor_test.go | 82 ++ internal/server/handlers/mapcommand.go | 19 +- .../server/handlers/mapcommand_completion_test.go | 337 +++++ internal/server/handlers/mapcommand_race_test.go | 64 + internal/server/handlers/output_manager.go | 518 ++++++++ .../server/handlers/output_manager_race_test.go | 437 +++++++ internal/server/handlers/protocol_codec.go | 87 ++ internal/server/handlers/protocol_codec_test.go | 32 + internal/server/handlers/protocol_formatter.go | 40 + internal/server/handlers/readcommand.go | 513 +++++++- .../handlers/readcommand_cancellation_test.go | 65 + .../handlers/readcommand_epoch_order_test.go | 127 ++ .../server/handlers/readcommand_glob_cap_test.go | 214 ++++ .../server/handlers/readcommand_journal_test.go | 271 +++++ .../server/handlers/readcommand_semaphore_test.go | 108 ++ .../handlers/readcommand_sendmessage_test.go | 252 ++++ internal/server/handlers/readcommand_server.go | 250 ++++ internal/server/handlers/serverhandler.go | 247 +++- internal/server/handlers/sessioncommand.go | 258 ++++ internal/server/handlers/sessioncommand_test.go | 568 +++++++++ internal/server/handlers/shutdown_coordinator.go | 108 ++ internal/server/scheduler.go | 19 +- internal/server/server.go | 329 +++-- internal/server/stats.go | 63 +- internal/server/stats_test.go | 206 ++++ 41 files changed, 9199 insertions(+), 335 deletions(-) create mode 100644 internal/server/auth_test.go create mode 100644 internal/server/continuous_test.go create mode 100644 internal/server/handlers/authkeycommand_test.go create mode 100644 internal/server/handlers/basehandler_read_test.go create mode 100644 internal/server/handlers/commandcancel_test.go create mode 100644 internal/server/handlers/commandtimeout_test.go create mode 100644 internal/server/handlers/framesize_test.go create mode 100644 internal/server/handlers/generation_output.go create mode 100644 internal/server/handlers/generation_output_test.go create mode 100644 internal/server/handlers/line_writer.go create mode 100644 internal/server/handlers/line_writer_alloc_test.go create mode 100644 internal/server/handlers/line_writer_test.go create mode 100644 internal/server/handlers/lineprocessor.go create mode 100644 internal/server/handlers/lineprocessor_test.go create mode 100644 internal/server/handlers/mapcommand_completion_test.go create mode 100644 internal/server/handlers/mapcommand_race_test.go create mode 100644 internal/server/handlers/output_manager.go create mode 100644 internal/server/handlers/output_manager_race_test.go create mode 100644 internal/server/handlers/protocol_codec.go create mode 100644 internal/server/handlers/protocol_codec_test.go create mode 100644 internal/server/handlers/protocol_formatter.go create mode 100644 internal/server/handlers/readcommand_cancellation_test.go create mode 100644 internal/server/handlers/readcommand_epoch_order_test.go create mode 100644 internal/server/handlers/readcommand_glob_cap_test.go create mode 100644 internal/server/handlers/readcommand_journal_test.go create mode 100644 internal/server/handlers/readcommand_semaphore_test.go create mode 100644 internal/server/handlers/readcommand_sendmessage_test.go create mode 100644 internal/server/handlers/readcommand_server.go create mode 100644 internal/server/handlers/sessioncommand.go create mode 100644 internal/server/handlers/sessioncommand_test.go create mode 100644 internal/server/handlers/shutdown_coordinator.go create mode 100644 internal/server/stats_test.go (limited to 'internal/server') diff --git a/internal/server/auth_test.go b/internal/server/auth_test.go new file mode 100644 index 0000000..ecd3c08 --- /dev/null +++ b/internal/server/auth_test.go @@ -0,0 +1,50 @@ +package server + +import ( + "crypto/subtle" + "testing" +) + +// TestConstantTimePasswordCompare verifies that password comparisons use +// constant-time comparison to prevent timing side-channel attacks. +// This is a regression test for the bug where `!=` was used directly, +// leaking timing information to an attacker who can measure response latency. +func TestConstantTimePasswordCompare(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + a string + b string + wantSame bool + }{ + {"equal passwords", "secret123", "secret123", true}, + {"different passwords", "secret123", "wrongpass", false}, + {"empty vs non-empty", "", "secret", false}, + {"both empty", "", "", true}, + {"prefix match only", "secret", "secretXYZ", false}, + {"suffix match only", "XYZsecret", "secret", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // secretsEqual is the function under test — it must use + // crypto/subtle.ConstantTimeCompare, not plain ==. + got := secretsEqual(tt.a, tt.b) + if got != tt.wantSame { + t.Errorf("secretsEqual(%q, %q) = %v, want %v", + tt.a, tt.b, got, tt.wantSame) + } + + // Cross-check with the reference implementation to ensure our + // helper is semantically correct, not just timing-safe. + reference := subtle.ConstantTimeCompare([]byte(tt.a), []byte(tt.b)) == 1 + if got != reference { + t.Errorf("secretsEqual(%q, %q) = %v but crypto/subtle gives %v", + tt.a, tt.b, got, reference) + } + }) + } +} diff --git a/internal/server/continuous.go b/internal/server/continuous.go index ac5c686..f3ee4fa 100644 --- a/internal/server/continuous.go +++ b/internal/server/continuous.go @@ -13,10 +13,32 @@ import ( gossh "golang.org/x/crypto/ssh" ) -type continuous struct{} +type continuousClient interface { + Start(context.Context, <-chan string) int +} -func newContinuous() *continuous { - return &continuous{} +type continuous struct { + cfg config.RuntimeConfig + newMaprClient func(config.Args, clients.MaprClientMode) (continuousClient, error) + dayChangeWatcher func(context.Context) bool + retryInterval time.Duration + now func() time.Time + newTicker func(time.Duration) (<-chan time.Time, func()) +} + +func newContinuous(cfg config.RuntimeConfig) *continuous { + c := &continuous{cfg: cfg} + c.retryInterval = time.Minute + c.now = time.Now + c.newTicker = func(d time.Duration) (<-chan time.Time, func()) { + ticker := time.NewTicker(d) + return ticker.C, ticker.Stop + } + c.newMaprClient = func(args config.Args, mode clients.MaprClientMode) (continuousClient, error) { + return clients.NewMaprClient(args, mode) + } + c.dayChangeWatcher = c.waitForDayChange + return c } func (c *continuous) start(ctx context.Context) { @@ -26,17 +48,20 @@ func (c *continuous) start(ctx context.Context) { } func (c *continuous) runJobs(ctx context.Context) { - for _, job := range config.Server.Continuous { + for i := range c.cfg.Server.Continuous { + job := &c.cfg.Server.Continuous[i] if !job.Enable { dlog.Server.Debug(job.Name, "Not running job as not enabled") continue } - go func(job config.Continuous) { + go func(job *config.Continuous) { c.runJob(ctx, job) + retryTicker := time.NewTicker(c.retryInterval) + defer retryTicker.Stop() for { select { - // Retry after a minute - case <-time.After(time.Minute): + // Retry after the configured interval. + case <-retryTicker.C: c.runJob(ctx, job) case <-ctx.Done(): return @@ -46,14 +71,14 @@ func (c *continuous) runJobs(ctx context.Context) { } } -func (c *continuous) runJob(ctx context.Context, job config.Continuous) { +func (c *continuous) runJob(ctx context.Context, job *config.Continuous) { dlog.Server.Debug(job.Name, "Processing job") files := fillDates(job.Files) outfile := fillDates(job.Outfile) servers := strings.Join(job.Servers, ",") if servers == "" { - servers = config.Server.SSHBindAddress + servers = c.cfg.Server.SSHBindAddress } args := config.Args{ @@ -67,7 +92,7 @@ func (c *continuous) runJob(ctx context.Context, job config.Continuous) { args.SSHAuthMethods = append(args.SSHAuthMethods, gossh.Password(job.Name)) args.QueryStr = fmt.Sprintf("%s outfile %s", job.Query, outfile) - client, err := clients.NewMaprClient(args, clients.NonCumulativeMode) + client, err := c.newMaprClient(args, clients.NonCumulativeMode) if err != nil { dlog.Server.Error(fmt.Sprintf("Unable to create job %s", job.Name), err) return @@ -77,7 +102,7 @@ func (c *continuous) runJob(ctx context.Context, job config.Continuous) { defer cancel() if job.RestartOnDayChange { go func() { - if c.waitForDayChange(ctx) { + if c.dayChangeWatcher(jobCtx) { dlog.Server.Info(fmt.Sprintf("Canceling job %s due to day change", job.Name)) cancel() } @@ -95,11 +120,13 @@ func (c *continuous) runJob(ctx context.Context, job config.Continuous) { } func (c *continuous) waitForDayChange(ctx context.Context) bool { - startTime := time.Now() + startTime := c.now() + tickCh, stop := c.newTicker(time.Second) + defer stop() for { select { - case <-time.After(time.Second): - if time.Now().Day() != startTime.Day() { + case <-tickCh: + if !sameCalendarDay(c.now(), startTime) { return true } case <-ctx.Done(): @@ -107,3 +134,9 @@ func (c *continuous) waitForDayChange(ctx context.Context) bool { } } } + +func sameCalendarDay(a, b time.Time) bool { + ay, am, ad := a.Date() + by, bm, bd := b.Date() + return ay == by && am == bm && ad == bd +} diff --git a/internal/server/continuous_test.go b/internal/server/continuous_test.go new file mode 100644 index 0000000..12d9f18 --- /dev/null +++ b/internal/server/continuous_test.go @@ -0,0 +1,202 @@ +package server + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/mimecast/dtail/internal/clients" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/dlog" +) + +func TestSameCalendarDay(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + a time.Time + b time.Time + want bool + }{ + { + name: "same day", + a: time.Date(2026, time.January, 15, 10, 0, 0, 0, time.UTC), + b: time.Date(2026, time.January, 15, 23, 59, 59, 0, time.UTC), + want: true, + }, + { + name: "same day-of-month in different months", + a: time.Date(2026, time.January, 15, 10, 0, 0, 0, time.UTC), + b: time.Date(2026, time.February, 15, 10, 0, 0, 0, time.UTC), + want: false, + }, + { + name: "same day-of-month across years", + a: time.Date(2025, time.December, 31, 10, 0, 0, 0, time.UTC), + b: time.Date(2026, time.January, 31, 10, 0, 0, 0, time.UTC), + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + if got := sameCalendarDay(tt.a, tt.b); got != tt.want { + t.Fatalf("sameCalendarDay(%v, %v) = %v, want %v", tt.a, tt.b, got, tt.want) + } + }) + } +} + +func TestContinuousRunJobsReleasesDayChangeWatcherAcrossRetries(t *testing.T) { + dlog.Server = &dlog.DLog{} + + c := newContinuous(config.RuntimeConfig{ + Server: &config.ServerConfig{ + SSHBindAddress: "127.0.0.1", + }, + }) + c.retryInterval = 25 * time.Millisecond + + var watcherStarts int32 + var watcherExits int32 + started := make(chan struct{}, 1) + release := make(chan struct{}, 1) + c.newMaprClient = func(args config.Args, mode clients.MaprClientMode) (continuousClient, error) { + return blockingContinuousClient{ + started: started, + release: release, + }, nil + } + c.dayChangeWatcher = func(ctx context.Context) bool { + atomic.AddInt32(&watcherStarts, 1) + defer atomic.AddInt32(&watcherExits, 1) + return c.waitForDayChange(ctx) + } + + job := config.Continuous{} + job.Enable = true + job.RestartOnDayChange = true + c.cfg.Server.Continuous = []config.Continuous{job} + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + c.runJobs(ctx) + close(done) + }() + + for i := int32(1); i <= 5; i++ { + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for retry %d to start", i) + } + + waitForCounterAtLeast(t, func() int32 { + return atomic.LoadInt32(&watcherStarts) + }, i) + + release <- struct{}{} + + waitForCounterAtLeast(t, func() int32 { + return atomic.LoadInt32(&watcherExits) + }, i) + } + + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("continuous job runner did not stop after cancellation") + } +} + +func TestContinuousWaitForDayChangeDetectsMonthBoundary(t *testing.T) { + dlog.Server = &dlog.DLog{} + + c := newContinuous(config.RuntimeConfig{}) + + start := time.Date(2026, time.January, 31, 23, 59, 59, 0, time.UTC) + sameDay := time.Date(2026, time.January, 31, 23, 59, 59, 500_000_000, time.UTC) + nextDay := time.Date(2026, time.February, 1, 0, 0, 0, 0, time.UTC) + + var nowCalls int32 + c.now = func() time.Time { + switch atomic.AddInt32(&nowCalls, 1) { + case 1: + return start + case 2: + return sameDay + default: + return nextDay + } + } + + tickCh := make(chan time.Time, 2) + c.newTicker = func(time.Duration) (<-chan time.Time, func()) { + return tickCh, func() {} + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := make(chan bool, 1) + go func() { + result <- c.waitForDayChange(ctx) + }() + + tickCh <- start + select { + case got := <-result: + t.Fatalf("waitForDayChange returned after same-day tick: %v", got) + case <-time.After(100 * time.Millisecond): + } + + tickCh <- nextDay + select { + case got := <-result: + if !got { + t.Fatal("waitForDayChange returned false after the month boundary tick") + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for waitForDayChange to detect the month boundary") + } +} + +type blockingContinuousClient struct { + started chan<- struct{} + release <-chan struct{} +} + +func (f blockingContinuousClient) Start(context.Context, <-chan string) int { + f.started <- struct{}{} + <-f.release + return 0 +} + +func waitForCounterAtLeast(t *testing.T, current func() int32, min int32) { + t.Helper() + + deadline := time.NewTimer(2 * time.Second) + defer deadline.Stop() + + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + + for { + if current() >= min { + return + } + + select { + case <-deadline.C: + t.Fatalf("timed out waiting for counter to reach %d, got %d", min, current()) + case <-ticker.C: + } + } +} diff --git a/internal/server/handlers/authkeycommand_test.go b/internal/server/handlers/authkeycommand_test.go new file mode 100644 index 0000000..a454e94 --- /dev/null +++ b/internal/server/handlers/authkeycommand_test.go @@ -0,0 +1,117 @@ +package handlers + +import ( + "context" + "crypto/ed25519" + "encoding/base64" + "testing" + "time" + + "github.com/mimecast/dtail/internal" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/lcontext" + sshserver "github.com/mimecast/dtail/internal/ssh/server" + userserver "github.com/mimecast/dtail/internal/user/server" + + gossh "golang.org/x/crypto/ssh" +) + +func TestHandleAuthKeyCommandSuccess(t *testing.T) { + handler := newAuthKeyTestHandler("authkey-success-user", true) + key := handlerTestPublicKey(t, 31) + keyArg := base64.StdEncoding.EncodeToString(key.Marshal()) + + commandFinished := false + handler.handleAuthKeyCommand(context.Background(), lcontext.LContext{}, 2, + []string{"AUTHKEY", keyArg}, func() { + commandFinished = true + }) + + if !commandFinished { + t.Fatalf("Expected commandFinished callback to be called") + } + if message := readServerMessage(t, handler.serverMessages); message != "AUTHKEY OK\n" { + t.Fatalf("Unexpected response: %q", message) + } + if !handler.authKeyStore.Has(handler.user.Name, key) { + t.Fatalf("Expected key to be stored for user") + } + handler.authKeyStore.Remove(handler.user.Name, key) +} + +func TestHandleAuthKeyCommandFeatureDisabled(t *testing.T) { + handler := newAuthKeyTestHandler("authkey-disabled-user", false) + key := handlerTestPublicKey(t, 32) + keyArg := base64.StdEncoding.EncodeToString(key.Marshal()) + + handler.handleAuthKeyCommand(context.Background(), lcontext.LContext{}, 2, + []string{"AUTHKEY", keyArg}, func() {}) + + if message := readServerMessage(t, handler.serverMessages); message != "AUTHKEY ERR feature disabled\n" { + t.Fatalf("Unexpected response: %q", message) + } + if handler.authKeyStore.Has(handler.user.Name, key) { + t.Fatalf("Expected no key to be stored while feature is disabled") + } +} + +func TestHandleAuthKeyCommandInvalidPayload(t *testing.T) { + handler := newAuthKeyTestHandler("authkey-invalid-user", true) + + handler.handleAuthKeyCommand(context.Background(), lcontext.LContext{}, 2, + []string{"AUTHKEY", "not-base64"}, func() {}) + + if message := readServerMessage(t, handler.serverMessages); message != "AUTHKEY ERR invalid base64\n" { + t.Fatalf("Unexpected response for invalid base64: %q", message) + } + + validButNonSSH := base64.StdEncoding.EncodeToString([]byte("not-an-ssh-key")) + handler.handleAuthKeyCommand(context.Background(), lcontext.LContext{}, 2, + []string{"AUTHKEY", validButNonSSH}, func() {}) + if message := readServerMessage(t, handler.serverMessages); message != "AUTHKEY ERR invalid public key\n" { + t.Fatalf("Unexpected response for invalid key bytes: %q", message) + } +} + +func newAuthKeyTestHandler(userName string, authKeyEnabled bool) *ServerHandler { + return &ServerHandler{ + baseHandler: baseHandler{ + done: internal.NewDone(), + serverMessages: make(chan string, 4), + user: &userserver.User{Name: userName}, + }, + serverCfg: &config.ServerConfig{ + AuthKeyEnabled: authKeyEnabled, + }, + authKeyStore: sshserver.NewAuthKeyStore(time.Hour, 5), + } +} + +func handlerTestPublicKey(t *testing.T, seedByte byte) gossh.PublicKey { + t.Helper() + + seed := make([]byte, ed25519.SeedSize) + for i := range seed { + seed[i] = seedByte + } + + privateKey := ed25519.NewKeyFromSeed(seed) + publicKey, err := gossh.NewPublicKey(privateKey.Public()) + if err != nil { + t.Fatalf("Unable to build ssh public key: %s", err.Error()) + } + + return publicKey +} + +func readServerMessage(t *testing.T, messages <-chan string) string { + t.Helper() + + select { + case message := <-messages: + return message + case <-time.After(time.Second): + t.Fatalf("Timed out waiting for server message") + return "" + } +} diff --git a/internal/server/handlers/basehandler.go b/internal/server/handlers/basehandler.go index f6ab3db..1ac159f 100644 --- a/internal/server/handlers/basehandler.go +++ b/internal/server/handlers/basehandler.go @@ -3,8 +3,6 @@ package handlers import ( "bytes" "context" - "encoding/base64" - "errors" "fmt" "io" "strconv" @@ -19,26 +17,76 @@ import ( "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/io/pool" "github.com/mimecast/dtail/internal/lcontext" - "github.com/mimecast/dtail/internal/mapr/server" + maprserver "github.com/mimecast/dtail/internal/mapr/server" "github.com/mimecast/dtail/internal/protocol" user "github.com/mimecast/dtail/internal/user/server" ) type handleCommandCb func(context.Context, lcontext.LContext, int, []string, string) +// commandCancelKeyType is a private key type for stashing a per-command +// context.CancelFunc inside a context.Context. It is used to hand the cancel +// ownership from handleCommand (which creates the context) to the command +// completion callback (handleUserCommand.commandFinished), which is the only +// place that knows when the asynchronous command is actually done. +type commandCancelKeyType struct{} + +var commandCancelKey commandCancelKeyType + +// withCommandCancel returns a derived context that carries the per-command +// cancel func. See cancelCommandContext for the matching consumer. +func withCommandCancel(ctx context.Context, cancel context.CancelFunc) context.Context { + if cancel == nil { + return ctx + } + return context.WithValue(ctx, commandCancelKey, cancel) +} + +// cancelCommandContext invokes the per-command cancel func stashed on ctx (if +// any) exactly once. It is a no-op when ctx carries no cancel (for example in +// the session-command path where the session state owns the cancel). +func cancelCommandContext(ctx context.Context) { + cancel, ok := ctx.Value(commandCancelKey).(context.CancelFunc) + if !ok || cancel == nil { + return + } + cancel() +} + type baseHandler struct { - done *internal.Done - handleCommandCb handleCommandCb - lines chan *line.Line - aggregate *server.Aggregate + done *internal.Done + handleCommandCb handleCommandCb + lines chan *line.Line + + // aggregate is written by handleMapCommand on the command-dispatch + // goroutine and read concurrently by Shutdown, Aggregate, and + // resetSessionAggregates. Using atomic.Pointer eliminates the data race + // without requiring h.mutex to be held around every access site. + aggregate atomic.Pointer[maprserver.Aggregate] + maprMessages chan string serverMessages chan string hostname string user *user.User ackCloseReceived chan struct{} + ackCloseOnce sync.Once activeCommands int32 - readBuf bytes.Buffer - writeBuf bytes.Buffer + codec protocolCodec + + // readBuf holds the formatted protocol message currently being sent to + // the client. It is only touched by Read (single session output + // goroutine) and retains any bytes that did not fit into the caller's + // buffer, so messages larger than one Read are delivered across multiple + // calls instead of being truncated (see Read/drainReadBuf). + readBuf bytes.Buffer + writeBuf bytes.Buffer + + // maxCommandFrameSize is the maximum number of bytes that may be buffered + // between two ';' delimiters. When a frame grows beyond this limit the + // Write method closes the session immediately to prevent a malicious or + // misbehaving client from exhausting server memory. The value is set at + // construction time from ServerConfig.MaxCommandFrameSize. + maxCommandFrameSize int // Some global options + sync primitives required. once sync.Once @@ -46,10 +94,30 @@ type baseHandler struct { quiet bool plain bool serverless bool + + output outputManager + + activeGeneration func() uint64 +} + +// getAggregate returns the current output MapReduce aggregate atomically. +func (h *baseHandler) getAggregate() *maprserver.Aggregate { + return h.aggregate.Load() } -// Shutdown the handler. +// setAggregate stores a output MapReduce aggregate atomically. +func (h *baseHandler) setAggregate(ta *maprserver.Aggregate) { + h.aggregate.Store(ta) +} + +// Shutdown the handler. Uses atomic accessors to read aggregate pointers so +// the reads are race-free with concurrent writes from handleMapCommand. func (h *baseHandler) Shutdown() { + // Shutdown output aggregate if present. + if ta := h.getAggregate(); ta != nil { + dlog.Server.Info(h.user, "Shutting down output aggregate") + ta.Shutdown() + } h.done.Shutdown() } @@ -59,73 +127,143 @@ func (h *baseHandler) Done() <-chan struct{} { } // Read is to send data to the dtail client via Reader interface. +// +// A formatted protocol message can be larger than p (io.Copy drives this +// reader with a 32KB buffer while MaxLineLength allows lines up to 1MB), so +// each Read drains any bytes left over from a previous call first and every +// message path keeps its unsent remainder in readBuf across calls. Dropping +// the remainder would truncate long lines and lose the trailing message +// delimiter, desyncing the client-side parser. This mirrors the remainder +// buffer used by the output path (outputManager.tryRead). func (h *baseHandler) Read(p []byte) (n int, err error) { - defer h.readBuf.Reset() + if h.readBuf.Len() > 0 { + return h.drainReadBuf(p), nil + } - select { - case message := <-h.serverMessages: - if len(message) > 0 && message[0] == '.' { - // Handle hidden message (don't display to the user) - h.readBuf.WriteString(message) - h.readBuf.WriteByte(protocol.MessageDelimiter) - n = copy(p, h.readBuf.Bytes()) - return + for { + if n, handled := h.output.tryRead(p, h.user, h.shouldDropGeneration); handled { + if n == 0 { + continue + } + return n, nil } - if h.serverless { - return + pollInterval := time.Second + if h.output.enabled() { + // Output reads require tighter wake-ups so we can continue draining the output channel. + pollInterval = h.output.resolvedReadRetryInterval() } + poll := time.After(pollInterval) - // Handle normal server message (display to the user) - h.readBuf.WriteString("SERVER") - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(h.hostname) - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(message) - h.readBuf.WriteByte(protocol.MessageDelimiter) - n = copy(p, h.readBuf.Bytes()) - - case message := <-h.maprMessages: - // Send mapreduce-aggregated data as a message. - h.readBuf.WriteString("AGGREGATE") - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(h.hostname) - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(message) - h.readBuf.WriteByte(protocol.MessageDelimiter) - n = copy(p, h.readBuf.Bytes()) - - case line := <-h.lines: - if !h.plain { - h.readBuf.WriteString("REMOTE") + select { + case message := <-h.serverMessages: + generation, decodedMessage := decodeGeneratedMessage(message) + if h.shouldDropGeneration(generation) { + continue + } + message = decodedMessage + if len(message) > 0 && message[0] == '.' { + // Handle hidden message (don't display to the user) + h.readBuf.WriteString(message) + h.readBuf.WriteByte(protocol.MessageDelimiter) + n = h.drainReadBuf(p) + return + } + + if h.serverless { + return + } + + // Skip empty server messages when in plain mode + if h.plain && (message == "" || message == "\n") { + return + } + + // Handle normal server message (display to the user). + formatServerMessage(&h.readBuf, h.hostname, message, h.plain) + n = h.drainReadBuf(p) + return + + case message := <-h.maprMessages: + generation, decodedMessage := decodeGeneratedMessage(message) + if h.shouldDropGeneration(generation) { + continue + } + message = decodedMessage + // Send mapreduce-aggregated data as a message. The leading + // AggregateMessageID field lets the mapr client tell aggregate + // data apart from plain server acks that happen to start with 'A'. + h.readBuf.WriteString(protocol.AggregateMessageID) h.readBuf.WriteString(protocol.FieldDelimiter) h.readBuf.WriteString(h.hostname) h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(fmt.Sprintf("%3d", line.TransmittedPerc)) - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(fmt.Sprintf("%v", line.Count)) - h.readBuf.WriteString(protocol.FieldDelimiter) - h.readBuf.WriteString(line.SourceID) - h.readBuf.WriteString(protocol.FieldDelimiter) - } - h.readBuf.WriteString(line.Content.String()) - h.readBuf.WriteByte(protocol.MessageDelimiter) - n = copy(p, h.readBuf.Bytes()) - pool.RecycleBytesBuffer(line.Content) - line.Recycle() + h.readBuf.WriteString(message) + h.readBuf.WriteByte(protocol.MessageDelimiter) + n = h.drainReadBuf(p) + return + + case line := <-h.lines: + if line == nil { + continue + } + if h.shouldDropGeneration(line.Generation) { + pool.RecycleBytesBuffer(line.Content) + line.Recycle() + continue + } + if h.plain { + h.readBuf.Write(line.Content.Bytes()) + h.readBuf.WriteByte(protocol.MessageDelimiter) + } else { + formatRemoteLine( + &h.readBuf, + h.hostname, + fmt.Sprintf("%3d", line.TransmittedPerc), + line.Count, + line.SourceID, + line.Content.Bytes(), + ) + } + n = h.drainReadBuf(p) + pool.RecycleBytesBuffer(line.Content) + line.Recycle() + return - case <-time.After(time.Second): - select { case <-h.done.Done(): err = io.EOF return - default: + + case <-poll: + // Wake periodically so output mode transitions don't leave this read blocked forever. + select { + case <-h.done.Done(): + err = io.EOF + return + default: + } + return } } - return +} + +// drainReadBuf copies as many buffered message bytes as fit into p and keeps +// the remainder in readBuf for subsequent Read calls. bytes.Buffer.Read +// consumes exactly the bytes it returns, so nothing is ever discarded; its +// io.EOF (only possible on an empty buffer) is deliberately not propagated +// because an empty buffer here simply means there is nothing left to drain. +func (h *baseHandler) drainReadBuf(p []byte) int { + n, _ := h.readBuf.Read(p) + return n } // Write is to receive data from the dtail client via Writer interface. +// Each byte is accumulated in writeBuf until a ';' delimiter arrives, at which +// point the buffered frame is dispatched as a command and the buffer is reset. +// +// To prevent a client from exhausting server memory with an unterminated frame, +// the buffer length is checked against maxCommandFrameSize on every append. When +// the limit is exceeded the session is shut down and io.ErrClosedPipe is returned +// so the SSH layer tears down the connection. func (h *baseHandler) Write(p []byte) (n int, err error) { for _, b := range p { switch b { @@ -134,6 +272,19 @@ func (h *baseHandler) Write(p []byte) (n int, err error) { h.writeBuf.Reset() default: h.writeBuf.WriteByte(b) + // Guard against unbounded frame growth: a client could send bytes + // without ever emitting a ';' delimiter and grow the buffer + // indefinitely. Reject and close when the configurable limit is hit. + if h.maxCommandFrameSize > 0 && h.writeBuf.Len() > h.maxCommandFrameSize { + dlog.Server.Error(h.user, + "command frame exceeds maximum size, closing session", + "frameSize", h.writeBuf.Len(), + "limit", h.maxCommandFrameSize, + ) + h.writeBuf.Reset() + h.done.Shutdown() + return len(p), io.ErrClosedPipe + } } } n = len(p) @@ -153,79 +304,145 @@ func (h *baseHandler) handleCommand(commandStr string) { h.sendln(h.serverMessages, dlog.Server.Error(h.user, err)) return } - ctx, cancel := context.WithCancel(context.Background()) - go func() { - <-h.done.Done() + ctx, cancel := h.newCommandContext(context.Background()) + // Cancel ownership is transferred to the command completion callback + // (see cancelCommandContext + handleUserCommand.commandFinished) so the + // per-command context and its watcher goroutine are released once the + // (possibly asynchronous) command has finished. If dispatch fails before + // the callback is ever invoked we must cancel here to avoid a leak. + ctx = withCommandCancel(ctx, cancel) + + if err := h.dispatchCommand(ctx, args, argc); err != nil { cancel() - }() + h.sendln(h.serverMessages, dlog.Server.Error(h.user, err)) + } +} + +func (h *baseHandler) dispatchCommand(ctx context.Context, args []string, argc int) error { + // Strip and apply a leading "timeout N ..." prefix. The client emits + // this when --timeout>0 (see internal/session/spec.go queryCommands); it + // caps how long the server collects data for that read command before its + // context is canceled. Handling it here covers both the legacy command + // stream and the SESSION dispatch path, which both funnel through here. + ctx, args, argc, err := applyCommandTimeout(ctx, args, argc) + if err != nil { + return err + } - parts := strings.Split(args[0], ":") + parts := strings.SplitN(args[0], ":", 2) commandName := parts[0] // Either no options or empty options provided. if len(parts) == 1 || len(parts[1]) == 0 { h.handleCommandCb(ctx, lcontext.LContext{}, argc, args, commandName) - return + return nil } - options, ltx, err := config.DeserializeOptions(parts[1:]) + options, ltx, err := config.DeserializeOptions([]string{parts[1]}) if err != nil { - h.sendln(h.serverMessages, dlog.Server.Error(h.user, err)) - return + return err } h.handleOptions(options) h.handleCommandCb(ctx, ltx, argc, args, commandName) + return nil } -func (h *baseHandler) handleProtocolVersion(args []string) ([]string, int, string, error) { - argc := len(args) - var add string - - if argc <= 2 || args[0] != "protocol" { - return args, argc, add, errors.New("unable to determine protocol version") +// maxCommandTimeoutSeconds caps the "timeout N " prefix value. 24h is far +// beyond any realistic collection window yet nowhere near the int64 overflow +// point of time.Duration (~292 years in nanoseconds), so it doubles as an +// overflow guard for the multiplication in applyCommandTimeout. +const maxCommandTimeoutSeconds = 24 * 60 * 60 + +// applyCommandTimeout detects a leading "timeout N ..." command prefix (as +// emitted by the client when --timeout>0) and returns a context that is +// canceled after N seconds together with the remaining command (the prefix +// stripped). This restores the original server-side deadline semantics: "Max +// time dtail server will collect data until disconnection". When no timeout +// prefix is present, or N<=0, the context and args are returned unchanged so +// the --timeout 0 / unset case behaves exactly as before. +// +// The timeout child cancel is chained onto the per-command cancel already +// stashed on ctx (if any) so cancelCommandContext, invoked once the command +// finishes, releases both the parent cancel and the timeout timer. In the +// session-dispatch path ctx carries no per-command cancel, so the returned +// context is the sole owner and its cancel still fires on command completion. +func applyCommandTimeout(ctx context.Context, args []string, argc int) (context.Context, []string, int, error) { + if argc < 3 || args[0] != "timeout" { + return ctx, args, argc, nil } - if args[1] != protocol.ProtocolCompat { - clientCompat, _ := strconv.Atoi(args[1]) - serverCompat, _ := strconv.Atoi(protocol.ProtocolCompat) - if clientCompat <= 3 { - // Protocol version 3 or lower expect a newline as message separator - // One day (after 2 major versions) this exception may be removed! - add = "\n" - } + seconds, err := strconv.Atoi(args[1]) + if err != nil { + return ctx, args, argc, fmt.Errorf("invalid timeout value %q: %w", args[1], err) + } + // Reject absurd values rather than clamp: an out-of-range N is a client + // mistake, and erroring (like the non-numeric case above) surfaces it + // instead of silently substituting a different deadline. This also guards + // against int64 overflow in time.Duration(seconds)*time.Second below, which + // for a huge N would wrap to a negative (already-elapsed) deadline and + // cancel the read immediately. + if seconds > maxCommandTimeoutSeconds { + return ctx, args, argc, fmt.Errorf("timeout value %d exceeds maximum of %d seconds", + seconds, maxCommandTimeoutSeconds) + } + if seconds <= 0 { + return ctx, args[2:], argc - 2, nil + } - toUpdate := "client" - if clientCompat > serverCompat { - toUpdate = "server" + timeoutCtx, cancel := context.WithTimeout(ctx, time.Duration(seconds)*time.Second) + parentCancel, _ := ctx.Value(commandCancelKey).(context.CancelFunc) + combined := func() { + cancel() + if parentCancel != nil { + parentCancel() } - err := fmt.Errorf("the DTail server protocol version '%s' does not match "+ - "client protocol version '%s', please update DTail %s", - protocol.ProtocolCompat, args[1], toUpdate) - return args, argc, add, err } - return args[2:], argc - 2, add, nil + return withCommandCancel(timeoutCtx, combined), args[2:], argc - 2, nil +} + +func (h *baseHandler) handleProtocolVersion(args []string) ([]string, int, string, error) { + return h.codec.handleProtocolVersion(args) } func (h *baseHandler) handleBase64(args []string, argc int) ([]string, int, error) { - err := errors.New("unable to decode client message, DTail server and client " + - "versions may not be compatible") - if argc != 2 || args[0] != "base64" { - return args, argc, err - } + return h.codec.handleBase64(args, argc) +} - decoded, err := base64.StdEncoding.DecodeString(args[1]) - if err != nil { - return args, argc, err +func (h *baseHandler) handleRawCommand(ctx context.Context, command string) error { + args := strings.Fields(command) + if len(args) == 0 { + return fmt.Errorf("empty command") } - decodedStr := string(decoded) + return h.dispatchCommand(ctx, args, len(args)) +} - args = strings.Split(decodedStr, " ") - argc = len(decodedStr) - dlog.Server.Trace(h.user, "Base64 decoded received command", - decodedStr, argc, args) +// newCommandContext creates a cancellable context for a single command +// invocation. The caller owns the returned cancel func and MUST invoke it +// exactly once (typically via defer or through the per-command cancel +// stashed on the context, see withCommandCancel/cancelCommandContext). +// Failing to cancel leaks both the context and the watcher goroutine +// spawned below, because the watcher only returns when the handler is shut +// down; on long-lived sessions (:reload, continuous/scheduled workloads) +// those leaks accumulate per command. +// +// The watcher goroutine doubles as a defensive safety net: even if a +// caller forgets to cancel, handler shutdown still drains it by cancelling +// the context via <-h.done.Done(). +func (h *baseHandler) newCommandContext(parent context.Context) (context.Context, context.CancelFunc) { + if parent == nil { + parent = context.Background() + } - return args, argc, nil + ctx, cancel := context.WithCancel(parent) + go func() { + select { + case <-h.done.Done(): + cancel() + case <-ctx.Done(): + } + }() + return ctx, cancel } func (h *baseHandler) handleAckCommand(argc int, args []string) { @@ -237,11 +454,9 @@ func (h *baseHandler) handleAckCommand(argc int, args []string) { return } if args[1] == "close" && args[2] == "connection" { - select { - case <-h.ackCloseReceived: - default: + h.ackCloseOnce.Do(func() { close(h.ackCloseReceived) - } + }) } } @@ -278,24 +493,76 @@ func (h *baseHandler) sendln(ch chan<- string, message string) { h.send(ch, message+"\n") } +func (h *baseHandler) shouldDropGeneration(generation uint64) bool { + if generation == 0 || h.activeGeneration == nil { + return false + } + + activeGeneration := h.activeGeneration() + if activeGeneration == 0 { + return false + } + + return activeGeneration != generation +} + func (h *baseHandler) flush() { dlog.Server.Trace(h.user, "flush()") numUnsentMessages := func() int { - return len(h.lines) + len(h.serverMessages) + len(h.maprMessages) + lineCount := len(h.lines) + serverCount := len(h.serverMessages) + maprCount := len(h.maprMessages) + outputCount := h.output.channelLen() + dlog.Server.Trace(h.user, "flush", "lines", lineCount, "server", serverCount, "mapr", maprCount, "output", outputCount) + return lineCount + serverCount + maprCount + outputCount + } + + // Use atomic accessors to avoid a data race with handleMapCommand, which + // may be concurrently writing aggregate pointers on another goroutine. + maxWait := time.Second + if h.output.enabled() || h.getAggregate() != nil { + maxWait = 3 * time.Second + } + if h.serverless && maxWait < 5*time.Second { + maxWait = 5 * time.Second } - for i := 0; i < 10; i++ { - if numUnsentMessages() == 0 { + + deadline := time.Now().Add(maxWait) + for i := 0; ; i++ { + unsent := numUnsentMessages() + if unsent == 0 { dlog.Server.Debug(h.user, "ALL lines sent", fmt.Sprintf("%p", h)) return } - dlog.Server.Debug(h.user, "Still lines to be sent") + if time.Now().After(deadline) { + dlog.Server.Warn(h.user, "Some lines remain unsent", unsent) + return + } + dlog.Server.Debug(h.user, "Still lines to be sent", "iteration", i, "unsent", unsent, "deadline", deadline.Sub(time.Now())) time.Sleep(time.Millisecond * 10) } - dlog.Server.Warn(h.user, "Some lines remain unsent", numUnsentMessages()) } func (h *baseHandler) shutdown() { - dlog.Server.Debug(h.user, "shutdown()") + // Log current state at shutdown + activeCommands := atomic.LoadInt32(&h.activeCommands) + dlog.Server.Info(h.user, "shutdown() called", "activeCommands", activeCommands, "outputMode", h.output.enabled()) + + // In output mode, ensure all data is flushed before shutdown + if h.output.enabled() { + h.flushOutput() + } + + // Shutdown the aggregate BEFORE flush to ensure MapReduce data is available. + // Use the atomic accessor to avoid a data race with handleMapCommand which + // may be concurrently storing the aggregate pointer on another goroutine. + if ta := h.getAggregate(); ta != nil { + dlog.Server.Info(h.user, "Shutting down output aggregate in shutdown()") + ta.Shutdown() + // Give time for serialization to complete. + time.Sleep(100 * time.Millisecond) + } + h.flush() go func() { @@ -322,3 +589,55 @@ func (h *baseHandler) decrementActiveCommands() int32 { atomic.AddInt32(&h.activeCommands, -1) return atomic.LoadInt32(&h.activeCommands) } + +// EnableDirectOutput enables output mode for direct line processing. It is an +// atomic check-and-enable: the return value is true when this call switched +// output mode on and false when it was already active (in which case the +// existing output state is left untouched). +func (h *baseHandler) EnableDirectOutput() bool { + return h.output.enable() +} + +// DirectOutputActive returns true if output mode is enabled +func (h *baseHandler) DirectOutputActive() bool { + return h.output.enabled() +} + +// HasOutputEOF returns true when a output EOF channel exists. +func (h *baseHandler) HasOutputEOF() bool { + return h.output.hasEOF() +} + +// OutputEpoch returns the current output handshake epoch. Capture it before +// checking the pending-work count and pass it to SignalOutputEOF so a stale +// "batch over" decision cannot EOF a batch that joined in between. +func (h *baseHandler) OutputEpoch() uint64 { + return h.output.currentEpoch() +} + +// SignalOutputEOF closes the output EOF channel once, unless the handshake +// epoch has advanced past the given captured value (i.e. another command +// joined the output session since), in which case the stale signal is dropped. +func (h *baseHandler) SignalOutputEOF(epoch uint64) { + h.output.signalEOF(epoch) +} + +// flushOutput ensures all output channel data is processed +func (h *baseHandler) flushOutput() { + h.output.flush(h.user) +} + +// GetOutputChannel returns the output lines channel for direct writing +func (h *baseHandler) GetOutputChannel() chan []byte { + return h.output.channel() +} + +// OutputChannelLen returns current output channel buffered size. +func (h *baseHandler) OutputChannelLen() int { + return h.output.channelLen() +} + +// WaitForOutputEOFAck waits until output reader acknowledges EOF or timeout. +func (h *baseHandler) WaitForOutputEOFAck(timeout time.Duration) bool { + return h.output.waitForEOFAck(timeout) +} diff --git a/internal/server/handlers/basehandler_read_test.go b/internal/server/handlers/basehandler_read_test.go new file mode 100644 index 0000000..e228ddb --- /dev/null +++ b/internal/server/handlers/basehandler_read_test.go @@ -0,0 +1,295 @@ +package handlers + +import ( + "bytes" + "fmt" + "testing" + "time" + + "github.com/mimecast/dtail/internal" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/io/line" + "github.com/mimecast/dtail/internal/protocol" +) + +// newReadTestHandler returns a baseHandler suitable for exercising Read +// directly, without any output or generation scoping involved. +func newReadTestHandler() baseHandler { + return baseHandler{ + done: internal.NewDone(), + lines: make(chan *line.Line, 4), + serverMessages: make(chan string, 4), + maprMessages: make(chan string, 4), + hostname: "testhost", + } +} + +// readExactly drains wantLen bytes from the handler using a buffer of bufSize +// bytes per Read call, failing the test when Read errors or stalls. Stalls are +// detected fast via consecutive empty reads and a wall-clock deadline: an +// empty Read means the handler blocked ~1s in its poll select without data, +// so two in a row indicate the remaining message bytes were dropped (e.g. a +// regression back to resetting readBuf) and we fail with a diagnostic instead +// of running into the go test timeout. +func readExactly(t *testing.T, handler *baseHandler, bufSize, wantLen int) []byte { + t.Helper() + + var got []byte + p := make([]byte, bufSize) + deadline := time.Now().Add(10 * time.Second) + emptyReads := 0 + for len(got) < wantLen { + if time.Now().After(deadline) { + t.Fatalf("Read stalled: got %d of %d bytes before deadline", + len(got), wantLen) + } + n, err := handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if n == 0 { + emptyReads++ + if emptyReads >= 2 { + t.Fatalf("Read stalled: got %d of %d bytes, %d consecutive empty reads", + len(got), wantLen, emptyReads) + } + continue + } + emptyReads = 0 + got = append(got, p[:n]...) + } + return got +} + +// expectedRemoteLine renders the protocol message Read is expected to emit +// for a line delivered via the lines channel in non-plain mode. +func expectedRemoteLine(content []byte, count uint64, sourceID string) []byte { + var want bytes.Buffer + formatRemoteLine(&want, "testhost", fmt.Sprintf("%3d", 100), count, + sourceID, content) + return want.Bytes() +} + +// TestBaseHandlerReadLargeLineAcrossMultipleReads reproduces the original +// bug: a line larger than the caller's buffer must arrive completely, +// including the trailing message delimiter, across multiple Read calls. +func TestBaseHandlerReadLargeLineAcrossMultipleReads(t *testing.T) { + handler := newReadTestHandler() + + content := bytes.Repeat([]byte("x"), 1000) + handler.lines <- line.New(bytes.NewBuffer(append([]byte{}, content...)), + 1, 100, "test.log") + + want := expectedRemoteLine(content, 1, "test.log") + got := readExactly(t, &handler, 32, len(want)) + + if !bytes.Equal(got, want) { + t.Fatalf("large line corrupted across reads:\ngot %q\nwant %q", got, want) + } + if got[len(got)-1] != protocol.MessageDelimiter { + t.Fatalf("message delimiter lost, last byte = %q", got[len(got)-1]) + } +} + +// TestBaseHandlerReadLargeServerMessageAcrossMultipleReads verifies the +// serverMessages path keeps its remainder across Read calls as well. +func TestBaseHandlerReadLargeServerMessageAcrossMultipleReads(t *testing.T) { + handler := newReadTestHandler() + + message := "server says: " + string(bytes.Repeat([]byte("y"), 500)) + handler.serverMessages <- message + + var want bytes.Buffer + formatServerMessage(&want, "testhost", message, false) + got := readExactly(t, &handler, 32, want.Len()) + + if !bytes.Equal(got, want.Bytes()) { + t.Fatalf("large server message corrupted across reads:\ngot %q\nwant %q", + got, want.Bytes()) + } +} + +// TestBaseHandlerReadLargeMaprMessageAcrossMultipleReads verifies the +// maprMessages AGGREGATE path keeps its remainder across Read calls. +func TestBaseHandlerReadLargeMaprMessageAcrossMultipleReads(t *testing.T) { + handler := newReadTestHandler() + + message := "aggregated " + string(bytes.Repeat([]byte("m"), 500)) + handler.maprMessages <- message + + var want bytes.Buffer + want.WriteString("AGGREGATE") + want.WriteString(protocol.FieldDelimiter) + want.WriteString("testhost") + want.WriteString(protocol.FieldDelimiter) + want.WriteString(message) + want.WriteByte(protocol.MessageDelimiter) + + got := readExactly(t, &handler, 32, want.Len()) + if !bytes.Equal(got, want.Bytes()) { + t.Fatalf("large mapr message corrupted across reads:\ngot %q\nwant %q", + got, want.Bytes()) + } +} + +// TestBaseHandlerReadLargeHiddenMessageAcrossMultipleReads verifies the +// hidden-message path (messages starting with '.') keeps its remainder across +// Read calls: hidden messages are forwarded verbatim plus the delimiter. +func TestBaseHandlerReadLargeHiddenMessageAcrossMultipleReads(t *testing.T) { + handler := newReadTestHandler() + + message := ".hidden " + string(bytes.Repeat([]byte("h"), 500)) + handler.serverMessages <- message + + want := append([]byte(message), protocol.MessageDelimiter) + got := readExactly(t, &handler, 32, len(want)) + if !bytes.Equal(got, want) { + t.Fatalf("large hidden message corrupted across reads:\ngot %q\nwant %q", + got, want) + } +} + +// TestBaseHandlerReadDrainsRemainderBeforeOutputData verifies that a pending +// readBuf remainder is fully delivered before any output payload when output +// mode is toggled on between Reads: the remainder belongs to a message that +// was already accepted for delivery, so output output must not preempt it. +func TestBaseHandlerReadDrainsRemainderBeforeOutputData(t *testing.T) { + // The output read path logs via dlog.Server, which is nil in unit tests; + // stub it out like the other output tests in this package do. + originalLogger := dlog.Server + dlog.Server = &dlog.DLog{} + t.Cleanup(func() { dlog.Server = originalLogger }) + + handler := newReadTestHandler() + + message := "regular " + string(bytes.Repeat([]byte("r"), 200)) + handler.serverMessages <- message + var want bytes.Buffer + formatServerMessage(&want, "testhost", message, false) + + // First small Read leaves the rest of the message in readBuf. + p := make([]byte, 16) + n, err := handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + got := append([]byte{}, p[:n]...) + + // Toggle output mode on mid-message and queue a output payload. + if !handler.output.enable() { + t.Fatal("expected output enable to switch output mode on") + } + outputPayload := []byte("output payload data") + handler.output.channel() <- outputPayload + + total := want.Len() + len(outputPayload) + got = append(got, readExactly(t, &handler, 16, total-len(got))...) + + if !bytes.Equal(got[:want.Len()], want.Bytes()) { + t.Fatalf("output data preempted pending remainder:\ngot %q\nwant %q", + got[:want.Len()], want.Bytes()) + } + if !bytes.Equal(got[want.Len():], outputPayload) { + t.Fatalf("output payload corrupted:\ngot %q\nwant %q", + got[want.Len():], outputPayload) + } +} + +// TestBaseHandlerReadDelimiterAloneInFinalRead verifies the boundary where +// the caller's buffer holds everything except the trailing delimiter, so the +// final Read must deliver the delimiter as its sole byte. +func TestBaseHandlerReadDelimiterAloneInFinalRead(t *testing.T) { + handler := newReadTestHandler() + + message := "delimiter boundary" + var want bytes.Buffer + formatServerMessage(&want, "testhost", message, false) + handler.serverMessages <- message + + p := make([]byte, want.Len()-1) + n, err := handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if n != want.Len()-1 || !bytes.Equal(p[:n], want.Bytes()[:n]) { + t.Fatalf("first read mismatch: n = %d, want %d", n, want.Len()-1) + } + + n, err = handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if n != 1 || p[0] != protocol.MessageDelimiter { + t.Fatalf("expected lone delimiter in final read, got %q", p[:n]) + } +} + +// TestBaseHandlerReadExactFitBuffer verifies that a message exactly filling +// the caller's buffer leaves no stale remainder behind: the next message must +// start fresh instead of being glued to leftover bytes. +func TestBaseHandlerReadExactFitBuffer(t *testing.T) { + handler := newReadTestHandler() + + first := "exact fit" + var firstWant bytes.Buffer + formatServerMessage(&firstWant, "testhost", first, false) + + handler.serverMessages <- first + p := make([]byte, firstWant.Len()) + n, err := handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if n != firstWant.Len() || !bytes.Equal(p[:n], firstWant.Bytes()) { + t.Fatalf("exact-fit read mismatch:\ngot %q\nwant %q", p[:n], firstWant.Bytes()) + } + + second := "next message" + var secondWant bytes.Buffer + formatServerMessage(&secondWant, "testhost", second, false) + + handler.serverMessages <- second + got := readExactly(t, &handler, firstWant.Len(), secondWant.Len()) + if !bytes.Equal(got, secondWant.Bytes()) { + t.Fatalf("stale remainder leaked into next message:\ngot %q\nwant %q", + got, secondWant.Bytes()) + } +} + +// TestBaseHandlerReadMultipleQueuedMessages verifies that several queued +// messages, each larger than the read buffer, arrive back to back in order +// and without corruption. +func TestBaseHandlerReadMultipleQueuedMessages(t *testing.T) { + handler := newReadTestHandler() + + var want bytes.Buffer + for i := 0; i < 3; i++ { + content := bytes.Repeat([]byte{byte('a' + i)}, 100) + handler.lines <- line.New(bytes.NewBuffer(append([]byte{}, content...)), + uint64(i+1), 100, "queued.log") + want.Write(expectedRemoteLine(content, uint64(i+1), "queued.log")) + } + + got := readExactly(t, &handler, 16, want.Len()) + if !bytes.Equal(got, want.Bytes()) { + t.Fatalf("queued messages corrupted:\ngot %q\nwant %q", got, want.Bytes()) + } +} + +// TestBaseHandlerReadPlainEmptyLine verifies the edge case of an empty line +// in plain mode: the message consists of the delimiter only. +func TestBaseHandlerReadPlainEmptyLine(t *testing.T) { + handler := newReadTestHandler() + handler.plain = true + + handler.lines <- line.New(&bytes.Buffer{}, 1, 100, "empty.log") + + p := make([]byte, 8) + n, err := handler.Read(p) + if err != nil { + t.Fatalf("Read() error = %v", err) + } + if n != 1 || p[0] != protocol.MessageDelimiter { + t.Fatalf("expected single delimiter byte, got %q", p[:n]) + } +} diff --git a/internal/server/handlers/commandcancel_test.go b/internal/server/handlers/commandcancel_test.go new file mode 100644 index 0000000..cac1475 --- /dev/null +++ b/internal/server/handlers/commandcancel_test.go @@ -0,0 +1,143 @@ +package handlers + +import ( + "context" + "encoding/base64" + "runtime" + "testing" + "time" + + "github.com/mimecast/dtail/internal" + "github.com/mimecast/dtail/internal/lcontext" + "github.com/mimecast/dtail/internal/protocol" +) + +// TestHandleCommandCancelsContextAfterCommandFinished verifies that +// baseHandler.handleCommand no longer discards the cancel func returned by +// newCommandContext. Pre-fix the cancel was dropped, so the per-command +// context (and the watcher goroutine spawned by newCommandContext) leaked +// for the lifetime of the SSH session. The cancel must fire when +// commandFinished is invoked. +func TestHandleCommandCancelsContextAfterCommandFinished(t *testing.T) { + resetServerLogger(t) + + handler := newSessionTestHandler("handle-command-cancel-user") + readServerMessage(t, handler.serverMessages) + handler.handleCommandCb = handler.handleUserCommand + + type captured struct { + ctx context.Context + finish func() + } + ch := make(chan captured, 1) + handler.commands = map[string]commandHandler{ + // AUTHKEY is a side-effect command so commandFinished does not + // trigger handler shutdown, keeping the test focused on the + // per-command cancel contract. + "AUTHKEY": func(ctx context.Context, _ lcontext.LContext, _ int, _ []string, commandFinished func()) { + ch <- captured{ctx: ctx, finish: commandFinished} + }, + } + + encoded := base64.StdEncoding.EncodeToString([]byte("AUTHKEY dummy")) + handler.handleCommand("protocol " + protocol.ProtocolCompat + " base64 " + encoded) + + var got captured + select { + case got = <-ch: + case <-time.After(500 * time.Millisecond): + t.Fatal("AUTHKEY command was not dispatched") + } + + select { + case <-got.ctx.Done(): + t.Fatal("per-command context cancelled before commandFinished ran") + default: + } + + got.finish() + + select { + case <-got.ctx.Done(): + case <-time.After(500 * time.Millisecond): + t.Fatal("per-command context was not cancelled after commandFinished ran; cancel was discarded") + } +} + +// TestNewCommandContextReleasesWatcherGoroutine ensures the watcher +// goroutine spawned by newCommandContext exits promptly once either the +// per-command cancel fires or the handler is shut down. This is the +// defensive safety net that keeps a leak from accumulating even if a +// future caller forgets to invoke cancel. +func TestNewCommandContextReleasesWatcherGoroutine(t *testing.T) { + h := &baseHandler{done: internal.NewDone()} + t.Cleanup(h.done.Shutdown) + + // Warm up so any lazily-started runtime goroutines are already up. + _, cancel := h.newCommandContext(context.Background()) + cancel() + time.Sleep(20 * time.Millisecond) + + baseline := runtime.NumGoroutine() + + const N = 100 + for i := 0; i < N; i++ { + _, cancel := h.newCommandContext(context.Background()) + cancel() + } + + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if delta := runtime.NumGoroutine() - baseline; delta <= 4 { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("watcher goroutines leaked: delta=%d (expected <= 4)", runtime.NumGoroutine()-baseline) +} + +// TestNewCommandContextHandlerShutdownReleasesWatcher verifies the +// defensive safety net: if a caller forgets to cancel a per-command +// context, shutting down the handler still drains the watcher goroutine +// rather than leaving it blocked until process exit. +func TestNewCommandContextHandlerShutdownReleasesWatcher(t *testing.T) { + h := &baseHandler{done: internal.NewDone()} + + // Warm up. + _, cancel := h.newCommandContext(context.Background()) + cancel() + time.Sleep(20 * time.Millisecond) + + baseline := runtime.NumGoroutine() + + const N = 50 + ctxs := make([]context.Context, 0, N) + for i := 0; i < N; i++ { + ctx, _ := h.newCommandContext(context.Background()) + ctxs = append(ctxs, ctx) + } + + if delta := runtime.NumGoroutine() - baseline; delta < N/2 { + t.Fatalf("expected goroutines to accumulate before shutdown, delta=%d", delta) + } + + h.done.Shutdown() + + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if delta := runtime.NumGoroutine() - baseline; delta <= 4 { + // All watcher goroutines should have observed the contexts + // being cancelled via the defensive done.Done() branch. + for _, ctx := range ctxs { + select { + case <-ctx.Done(): + default: + t.Fatalf("context not cancelled by handler shutdown") + } + } + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("watcher goroutines leaked past shutdown: delta=%d", runtime.NumGoroutine()-baseline) +} diff --git a/internal/server/handlers/commandtimeout_test.go b/internal/server/handlers/commandtimeout_test.go new file mode 100644 index 0000000..3a31c17 --- /dev/null +++ b/internal/server/handlers/commandtimeout_test.go @@ -0,0 +1,150 @@ +package handlers + +import ( + "context" + "reflect" + "testing" + "time" + + "github.com/mimecast/dtail/internal/lcontext" +) + +// TestApplyCommandTimeout exercises the pure prefix-stripping helper that +// restores server-side handling of the "timeout N ..." command emitted by +// the client when --timeout>0 (see internal/session/spec.go). The server-side +// parser was removed in 2020 while the client kept emitting the prefix, so the +// command reached the dispatcher as an unknown command "timeout". Each case +// checks the returned args/argc and whether a context deadline was applied. +func TestApplyCommandTimeout(t *testing.T) { + tests := []struct { + name string + args []string + wantArgs []string + wantErr bool + wantDeadline bool // expect a context deadline on the returned ctx + }{ + { + name: "no timeout prefix passes through unchanged", + args: []string{"cat", "file", "regex:noop"}, + wantArgs: []string{"cat", "file", "regex:noop"}, + }, + { + name: "positive timeout strips prefix and sets deadline", + args: []string{"timeout", "5", "cat", "file", "regex:noop"}, + wantArgs: []string{"cat", "file", "regex:noop"}, + wantDeadline: true, + }, + { + name: "tail read command timeout strips prefix", + args: []string{"timeout", "30", "tail", "file", "regex:noop"}, + wantArgs: []string{"tail", "file", "regex:noop"}, + wantDeadline: true, + }, + { + name: "zero timeout strips prefix without a deadline", + args: []string{"timeout", "0", "cat", "file"}, + wantArgs: []string{"cat", "file"}, + }, + { + name: "negative timeout strips prefix without a deadline", + args: []string{"timeout", "-5", "cat", "file"}, + wantArgs: []string{"cat", "file"}, + }, + { + name: "non-numeric timeout is rejected", + args: []string{"timeout", "abc", "cat", "file"}, + wantErr: true, + }, + { + // A huge N would overflow time.Duration(seconds)*time.Second into a + // negative (already-elapsed) deadline; the max-seconds guard rejects + // it instead of cancelling the read immediately. + name: "out-of-range timeout is rejected (overflow guard)", + args: []string{"timeout", "9223372036854775807", "cat", "file"}, + wantErr: true, + }, + { + name: "bare timeout with too few args is not treated as a prefix", + args: []string{"timeout", "5"}, + wantArgs: []string{"timeout", "5"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx, gotArgs, gotArgc, err := applyCommandTimeout(context.Background(), tc.args, len(tc.args)) + if tc.wantErr { + if err == nil { + t.Fatalf("expected error, got nil (args=%v)", gotArgs) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !reflect.DeepEqual(gotArgs, tc.wantArgs) { + t.Fatalf("args = %v, want %v", gotArgs, tc.wantArgs) + } + if gotArgc != len(tc.wantArgs) { + t.Fatalf("argc = %d, want %d", gotArgc, len(tc.wantArgs)) + } + _, hasDeadline := ctx.Deadline() + if hasDeadline != tc.wantDeadline { + t.Fatalf("ctx deadline present = %v, wan