diff options
| author | Paul Buetow <paul@buetow.org> | 2026-07-22 23:51:18 +0300 |
|---|---|---|
| committer | Paul Buetow <paul@buetow.org> | 2026-07-22 23:51:18 +0300 |
| commit | 849951be1d1a7ee9f9302006ccb187bf5b4e36f3 (patch) | |
| tree | 496c924a03a9ea6212e29bb4699e268066ebad81 /internal/clients/handlers/maprhandler_test.go | |
| parent | bf78b3abffee6d49c08ca2980156afc455994969 (diff) | |
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>
Diffstat (limited to 'internal/clients/handlers/maprhandler_test.go')
| -rw-r--r-- | internal/clients/handlers/maprhandler_test.go | 278 |
1 files changed, 278 insertions, 0 deletions
diff --git a/internal/clients/handlers/maprhandler_test.go b/internal/clients/handlers/maprhandler_test.go new file mode 100644 index 0000000..89df43a --- /dev/null +++ b/internal/clients/handlers/maprhandler_test.go @@ -0,0 +1,278 @@ +package handlers + +import ( + "bytes" + "context" + "io" + "os" + "strings" + "sync" + "testing" + + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/mapr" + maprclient "github.com/mimecast/dtail/internal/mapr/client" + "github.com/mimecast/dtail/internal/protocol" + "github.com/mimecast/dtail/internal/source" +) + +func TestMaprHandlerShutdownFlushesPendingAggregateState(t *testing.T) { + query, err := mapr.NewQuery("select status,count(status) from stats group by status") + if err != nil { + t.Fatalf("NewQuery() error = %v", err) + } + + session := maprclient.NewSessionState(query) + handler := NewMaprHandler("srv1", session) + countStorage := handlerCountStorage(t, query) + + message := strings.Join([]string{ + "ERROR", + "2", + countStorage + protocol.AggregateKVDelimiter + "2", + "", + }, protocol.AggregateDelimiter) + if err := handler.aggregate.Aggregate(message); err != nil { + t.Fatalf("Aggregate() error = %v", err) + } + + handler.Shutdown() + + result, numRows, err := session.Snapshot().GlobalGroup.Result(query, 10, nil) + if err != nil { + t.Fatalf("Result() error = %v", err) + } + if numRows != 1 { + t.Fatalf("numRows = %d, want 1", numRows) + } + if !strings.Contains(result, "2") { + t.Fatalf("expected flushed aggregate row, got %q", result) + } +} + +func TestMaprHandlerWriteEmptyMessageBetweenDelimiters(t *testing.T) { + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) + + query, err := mapr.NewQuery("select status,count(status) from stats group by status") + if err != nil { + t.Fatalf("NewQuery() error = %v", err) + } + + session := maprclient.NewSessionState(query) + handler := NewMaprHandler("srv1", session) + + defer func() { + if r := recover(); r != nil { + t.Fatalf("MaprHandler.Write panicked on empty protocol message: %v", r) + } + }() + + // Two consecutive MessageDelimiter bytes produce an empty message + // between them. A leading delimiter yields an empty message too. + // Both must be tolerated without panicking. + input := []byte{ + protocol.MessageDelimiter, + protocol.MessageDelimiter, + } + if _, err := handler.Write(input); err != nil { + t.Fatalf("Write() error = %v", err) + } +} + +// TestMaprHandlerClassifiesAuthKeyAckAsControl is a regression test for the +// dmap client feeding the server's "AUTHKEY OK" acknowledgement into the +// aggregate parser. In plain output mode the ack arrives on the wire verbatim +// (no SERVER|host| prefix), so it begins with the letter 'A' just like a real +// AGGREGATE|host|data message. Classifying on the full AggregateMessageID +// field prefix, instead of only the first byte, keeps such acks (and any +// sibling control message that merely starts with 'A') out of the aggregate +// parser, which previously logged a spurious +// "Unable to aggregate data ... expected 3 parts" error. +func TestMaprHandlerClassifiesAuthKeyAckAsControl(t *testing.T) { + aggregate := protocol.AggregateMessageID + protocol.FieldDelimiter + + "host1" + protocol.FieldDelimiter + "payload" + + tests := []struct { + name string + message string + wantAggregate bool + }{ + { + name: "genuine aggregate data", + message: aggregate, + wantAggregate: true, + }, + { + name: "plain-mode authkey ack", + message: "AUTHKEY OK", + wantAggregate: false, + }, + { + name: "server-prefixed authkey ack", + message: "SERVER" + protocol.FieldDelimiter + "host1" + protocol.FieldDelimiter + "AUTHKEY OK", + wantAggregate: false, + }, + { + name: "unrelated message starting with A", + message: "Application ready", + wantAggregate: false, + }, + { + // Adversarial: the AGGREGATE| tag appears, but embedded in a + // later field rather than as the leading field. Only the leading + // tag may classify a message as aggregate data. + message: "SERVER" + protocol.FieldDelimiter + "host1" + protocol.FieldDelimiter + aggregate, + name: "embedded aggregate tag is not the leading field", + wantAggregate: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := isAggregateMessage(tc.message); got != tc.wantAggregate { + t.Fatalf("isAggregateMessage(%q) = %v, want %v", + tc.message, got, tc.wantAggregate) + } + }) + } +} + +// TestMaprHandlerWriteAuthKeyAckEmitsNoAggregateError feeds a plain-mode +// message stream (an "AUTHKEY OK" ack followed by a genuine aggregate message) +// through Write and asserts two things by inspecting the captured client log: +// +// 1. no spurious "Unable to aggregate data ... expected 3 parts" error is +// emitted for the ack (the exact regression symptom); and +// 2. the genuine AGGREGATE|host|data message still produces its aggregate row. +// +// The log assertion is what makes this a real regression guard: with the old +// first-byte (message[0] == 'A') classifier the ack was fed to the aggregate +// parser and error-logged, so this test fails red against that code and green +// against the current prefix-based classifier. +func TestMaprHandlerWriteAuthKeyAckEmitsNoAggregateError(t *testing.T) { + ensureClientStdoutLogger(t) + + query, err := mapr.NewQuery("select status,count(status) from stats group by status") + if err != nil { + t.Fatalf("NewQuery() error = %v", err) + } + + session := maprclient.NewSessionState(query) + handler := NewMaprHandler("srv1", session) + countStorage := handlerCountStorage(t, query) + + // A genuine aggregate wire message: AGGREGATE|host|<serialized set>. + serialized := strings.Join([]string{ + "ERROR", + "2", + countStorage + protocol.AggregateKVDelimiter + "2", + "", + }, protocol.AggregateDelimiter) + aggregate := protocol.AggregateMessageID + protocol.FieldDelimiter + + "host1" + protocol.FieldDelimiter + serialized + + // Plain-mode ack first, then the genuine aggregate message, each + // terminated by the protocol message delimiter. + var input []byte + input = append(input, []byte("AUTHKEY OK")...) + input = append(input, protocol.MessageDelimiter) + input = append(input, []byte(aggregate)...) + input = append(input, protocol.MessageDelimiter) + + logOutput := captureStdout(t, func() { + if _, err := handler.Write(input); err != nil { + t.Fatalf("Write() error = %v", err) + } + handler.Shutdown() + }) + + if strings.Contains(logOutput, "Unable to aggregate data") || + strings.Contains(logOutput, "expected 3 parts") { + t.Fatalf("AUTHKEY OK ack was fed to the aggregate parser; "+ + "captured client log:\n%s", logOutput) + } + + result, numRows, err := session.Snapshot().GlobalGroup.Result(query, 10, nil) + if err != nil { + t.Fatalf("Result() error = %v", err) + } + if numRows != 1 { + t.Fatalf("numRows = %d, want 1 (only the genuine aggregate message)", numRows) + } + if !strings.Contains(result, "2") { + t.Fatalf("expected the genuine aggregate row, got %q", result) + } +} + +// ensureClientStdoutLogger initialises the global client logger so that it +// writes plain (uncolored) lines to os.Stdout at error level. This lets tests +// capture emitted log lines via captureStdout. It is idempotent and safe under +// -count>1: dlog.Start runs at most once per process (guarded on dlog.Client), +// matching the pattern used by the mapr server tests. +func ensureClientStdoutLogger(t *testing.T) { + t.Helper() + if config.Common == nil { + config.Common = &config.CommonConfig{Logger: "stdout", LogLevel: "error"} + } + if config.Client == nil { + config.Client = &config.ClientConfig{} + } + // Force the plain log path so captured output has no color escape codes. + config.Client.TermColorsEnable = false + if dlog.Client == nil { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + var wg sync.WaitGroup + wg.Add(1) + dlog.Start(ctx, &wg, source.Client) + } +} + +// captureStdout redirects os.Stdout to a pipe for the duration of fn and +// returns everything written to it. The stdout logger writes synchronously via +// fmt.Println, so all log lines emitted by fn are captured once fn returns. +func captureStdout(t *testing.T, fn func()) string { + t.Helper() + + orig := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe() error = %v", err) + } + os.Stdout = w + + collected := make(chan string, 1) + go func() { + var buf bytes.Buffer + _, _ = io.Copy(&buf, r) + collected <- buf.String() + }() + + fn() + + os.Stdout = orig + if err := w.Close(); err != nil { + t.Fatalf("closing stdout pipe: %v", err) + } + out := <-collected + if err := r.Close(); err != nil { + t.Fatalf("closing stdout pipe reader: %v", err) + } + return out +} + +func handlerCountStorage(t *testing.T, query *mapr.Query) string { + t.Helper() + for _, selectCondition := range query.Select { + if selectCondition.Operation == mapr.Count { + return selectCondition.FieldStorage + } + } + t.Fatalf("query %q does not contain count() storage", query.RawQuery) + return "" +} |
