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>

41 files changed, 9199 insertions, 335 deletions

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