summaryrefslogtreecommitdiff
path: root/internal/clients/handlers
diff options
context:
space:
mode:
Diffstat (limited to 'internal/clients/handlers')
-rw-r--r--internal/clients/handlers/basehandler.go319
-rw-r--r--internal/clients/handlers/basehandler_test.go356
-rw-r--r--internal/clients/handlers/clienthandler.go15
-rw-r--r--internal/clients/handlers/handler.go6
-rw-r--r--internal/clients/handlers/healthhandler.go13
-rw-r--r--internal/clients/handlers/maprhandler.go55
-rw-r--r--internal/clients/handlers/maprhandler_test.go278
7 files changed, 1012 insertions, 30 deletions
diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index 6f637a7..188c3aa 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -5,7 +5,10 @@ import (
"encoding/base64"
"fmt"
"io"
+ "sort"
+ "strconv"
"strings"
+ "sync"
"time"
"github.com/mimecast/dtail/internal"
@@ -19,7 +22,29 @@ type baseHandler struct {
shellStarted bool
commands chan string
receiveBuf bytes.Buffer
- status int
+
+ // pendingCommand holds the unsent tail of a dequeued command frame. A
+ // frame (e.g. a large base64-encoded regex or MapReduce query) can exceed
+ // the buffer io.Copy hands to Read (32KB), so the remainder is kept here
+ // and drained by subsequent Read calls instead of being dropped, which
+ // would corrupt the server-side command stream. Only touched by Read
+ // (single output-copy goroutine).
+ pendingCommand []byte
+ status int
+
+ capabilitiesMu sync.RWMutex
+ capabilities map[string]struct{}
+ capabilitiesCh chan struct{}
+ capabilitiesOk sync.Once
+
+ sessionAcks chan SessionAck
+}
+
+// SessionAck is a parsed hidden acknowledgement for SESSION START/UPDATE requests.
+type SessionAck struct {
+ Action string
+ Generation uint64
+ Error string
}
func (h *baseHandler) String() string {
@@ -40,6 +65,37 @@ func (h *baseHandler) Status() int {
return h.status
}
+func (h *baseHandler) Capabilities() []string {
+ h.capabilitiesMu.RLock()
+ defer h.capabilitiesMu.RUnlock()
+
+ capabilities := make([]string, 0, len(h.capabilities))
+ for capability := range h.capabilities {
+ capabilities = append(capabilities, capability)
+ }
+ sort.Strings(capabilities)
+ return capabilities
+}
+
+func (h *baseHandler) HasCapability(name string) bool {
+ h.capabilitiesMu.RLock()
+ defer h.capabilitiesMu.RUnlock()
+
+ _, ok := h.capabilities[name]
+ return ok
+}
+
+func (h *baseHandler) ReportServerError(message string) {
+ h.status = 1
+ // Route through the DIAGNOSTIC (Log) sink, not Raw: a server-error report is
+ // an audit line, not bulk payload. Via Raw it would be gated out of the
+ // client log file whenever Client.LogPayload is false (the default), silently
+ // dropping the error from the on-disk audit trail. RawLog keeps it in the
+ // file like other diagnostics while still printing it to stdout. The message
+ // carries no trailing newline; the Log sink appends one.
+ dlog.Client.RawLog(formatServerErrorMessage(h.server, message))
+}
+
// SendMessage to the server.
func (h *baseHandler) SendMessage(command string) error {
encoded := base64.StdEncoding.EncodeToString([]byte(command))
@@ -61,10 +117,8 @@ func (h *baseHandler) Write(p []byte) (n int, err error) {
for _, b := range p {
switch b {
case '\n':
- // Backwards compatible with DTail 3 (e.g. get error message from server
- // about protocol missmatch.
+ // Just add the newline to the buffer, don't treat as message delimiter
h.receiveBuf.WriteByte(b)
- fallthrough
case protocol.MessageDelimiter:
message := h.receiveBuf.String()
h.handleMessage(message)
@@ -77,14 +131,58 @@ func (h *baseHandler) Write(p []byte) (n int, err error) {
}
// Send data to the dtail server via Reader interface.
+//
+// Priority select: when Done() is closed we must still drain any pending
+// commands before returning io.EOF, because closing the connection requires
+// the '.ack close connection' message to be flushed to the server first.
+// Without this drain the server waits up to 5 seconds for the ack.
+//
+// Command frames larger than p are delivered across multiple Read calls via
+// pendingCommand (see consumeCommand); the leftover is always drained before
+// a new command is dequeued so frames are never truncated or interleaved.
func (h *baseHandler) Read(p []byte) (n int, err error) {
+ if len(h.pendingCommand) > 0 {
+ n = copy(p, h.pendingCommand)
+ h.pendingCommand = h.pendingCommand[n:]
+ if len(h.pendingCommand) == 0 {
+ // Release the backing array once the frame is fully delivered.
+ // Keeping it would retain the largest-ever frame for the handler
+ // lifetime and let the slice creep forward through its backing
+ // array on reuse; oversized frames are rare, so consumeCommand
+ // simply allocates a fresh buffer next time instead.
+ h.pendingCommand = nil
+ }
+ return
+ }
+
+ // Check for a pending command first (non-blocking), giving it priority
+ // over the Done signal so that queued acks are always delivered.
select {
case command := <-h.commands:
- n = copy(p, []byte(command))
+ return h.consumeCommand(p, command), nil
+ default:
+ }
+
+ // No command is immediately ready; block on whichever arrives first.
+ select {
+ case command := <-h.commands:
+ return h.consumeCommand(p, command), nil
case <-h.Done():
return 0, io.EOF
}
- return
+}
+
+// consumeCommand copies as much of command as fits into p and stashes the
+// remainder in pendingCommand so the next Read calls can deliver the rest of
+// the frame. pendingCommand is always empty here (Read drains and releases it
+// before dequeuing a new command), so a fresh buffer is allocated for the
+// remainder rather than reusing a long-lived one.
+func (h *baseHandler) consumeCommand(p []byte, command string) int {
+ n := copy(p, command)
+ if n < len(command) {
+ h.pendingCommand = []byte(command[n:])
+ }
+ return n
}
func (h *baseHandler) handleMessage(message string) {
@@ -92,24 +190,229 @@ func (h *baseHandler) handleMessage(message string) {
h.handleHiddenMessage(message)
return
}
+ if h.handleAuthKeyMessage(message) {
+ return
+ }
+
+ // Add newline only if the message doesn't already end with one
+ if len(message) > 0 && message[len(message)-1] == '\n' {
+ dlog.Client.Raw(message)
+ } else {
+ dlog.Client.Raw(message + "\n")
+ }
+}
+
+func (h *baseHandler) handleAuthKeyMessage(message string) bool {
+ isAuthKeyMessage, authKeyOK, authKeyDetail := parseAuthKeyMessage(message)
+ if !isAuthKeyMessage {
+ return false
+ }
+
+ if authKeyOK {
+ dlog.Client.Debug(h.server, "AUTHKEY registration accepted by server")
+ return true
+ }
+
+ if authKeyDetail == "" {
+ dlog.Client.Warn(h.server, "AUTHKEY registration failed")
+ return true
+ }
+
+ dlog.Client.Warn(h.server, "AUTHKEY registration failed", authKeyDetail)
+ return true
+}
+
+func parseAuthKeyMessage(message string) (isAuthKeyMessage bool, ok bool, detail string) {
+ if message == "" {
+ return false, false, ""
+ }
+
+ payload := strings.TrimSpace(message)
+ parts := strings.Split(payload, protocol.FieldDelimiter)
+ if len(parts) > 0 {
+ payload = strings.TrimSpace(parts[len(parts)-1])
+ }
- dlog.Client.Raw(message)
+ switch {
+ case payload == "AUTHKEY OK":
+ return true, true, ""
+ case strings.HasPrefix(payload, "AUTHKEY ERR"):
+ detail := strings.TrimSpace(strings.TrimPrefix(payload, "AUTHKEY ERR"))
+ return true, false, detail
+ default:
+ return false, false, ""
+ }
}
// Handle messages received from server which are not meant to be displayed
// to the end user.
func (h *baseHandler) handleHiddenMessage(message string) {
switch {
+ case strings.HasPrefix(message, protocol.HiddenCapabilitiesPrefix):
+ h.handleCapabilitiesMessage(message)
+ case strings.HasPrefix(message, protocol.HiddenSessionStartOKPrefix),
+ strings.HasPrefix(message, protocol.HiddenSessionUpdateOKPrefix),
+ strings.HasPrefix(message, protocol.HiddenSessionErrorPrefix):
+ h.handleSessionAckMessage(message)
case strings.HasPrefix(message, ".syn close connection"):
- go h.SendMessage(".ack close connection")
+ if err := h.SendMessage(".ack close connection"); err != nil {
+ dlog.Client.Debug(h.server, "Unable to acknowledge close connection", err)
+ }
h.Shutdown()
}
}
+func (h *baseHandler) handleCapabilitiesMessage(message string) {
+ capabilities := strings.Fields(strings.TrimPrefix(message, protocol.HiddenCapabilitiesPrefix))
+
+ h.capabilitiesMu.Lock()
+ defer h.capabilitiesMu.Unlock()
+
+ if h.capabilities == nil {
+ h.capabilities = make(map[string]struct{})
+ }
+ for _, capability := range capabilities {
+ if capability == "" {
+ continue
+ }
+ h.capabilities[capability] = struct{}{}
+ }
+
+ h.capabilitiesOk.Do(func() {
+ if h.capabilitiesCh != nil {
+ close(h.capabilitiesCh)
+ }
+ })
+}
+
func (h *baseHandler) Done() <-chan struct{} {
return h.done.Done()
}
+func (h *baseHandler) WaitForCapabilities(timeout time.Duration) bool {
+ if h.capabilitiesCh == nil {
+ return false
+ }
+
+ if timeout <= 0 {
+ select {
+ case <-h.capabilitiesCh:
+ return true
+ default:
+ return false
+ }
+ }
+
+ timer := time.NewTimer(timeout)
+ defer timer.Stop()
+
+ select {
+ case <-h.capabilitiesCh:
+ return true
+ case <-h.Done():
+ return false
+ case <-timer.C:
+ return false
+ }
+}
+
+func (h *baseHandler) WaitForSessionAck(timeout time.Duration) (SessionAck, bool) {
+ if h.sessionAcks == nil {
+ return SessionAck{}, false
+ }
+
+ if timeout <= 0 {
+ select {
+ case ack := <-h.sessionAcks:
+ return ack, true
+ default:
+ return SessionAck{}, false
+ }
+ }
+
+ timer := time.NewTimer(timeout)
+ defer timer.Stop()
+
+ select {
+ case ack := <-h.sessionAcks:
+ return ack, true
+ case <-h.Done():
+ return SessionAck{}, false
+ case <-timer.C:
+ return SessionAck{}, false
+ }
+}
+
func (h *baseHandler) Shutdown() {
h.done.Shutdown()
}
+
+func (h *baseHandler) handleSessionAckMessage(message string) {
+ ack, ok := parseSessionAckMessage(message)
+ if !ok {
+ dlog.Client.Warn(h.server, "Unable to parse session acknowledgement", message)
+ return
+ }
+ if h.sessionAcks == nil {
+ return
+ }
+
+ select {
+ case h.sessionAcks <- ack:
+ case <-h.Done():
+ default:
+ dlog.Client.Warn(h.server, "Dropping session acknowledgement because the queue is full", message)
+ }
+}
+
+func parseSessionAckMessage(message string) (SessionAck, bool) {
+ payload := strings.TrimSpace(message)
+ if payload == "" {
+ return SessionAck{}, false
+ }
+
+ switch {
+ case strings.HasPrefix(payload, protocol.HiddenSessionStartOKPrefix):
+ return parseSessionOKAck(strings.TrimPrefix(payload, protocol.HiddenSessionStartOKPrefix), "start")
+ case strings.HasPrefix(payload, protocol.HiddenSessionUpdateOKPrefix):
+ return parseSessionOKAck(strings.TrimPrefix(payload, protocol.HiddenSessionUpdateOKPrefix), "update")
+ case strings.HasPrefix(payload, protocol.HiddenSessionErrorPrefix):
+ return SessionAck{
+ Action: "error",
+ Error: strings.TrimSpace(strings.TrimPrefix(payload, protocol.HiddenSessionErrorPrefix)),
+ }, true
+ default:
+ return SessionAck{}, false
+ }
+}
+
+func parseSessionOKAck(payload string, action string) (SessionAck, bool) {
+ generationStr := strings.TrimSpace(payload)
+ if generationStr == "" {
+ return SessionAck{}, false
+ }
+
+ generation, err := strconv.ParseUint(generationStr, 10, 64)
+ if err != nil {
+ return SessionAck{}, false
+ }
+
+ return SessionAck{
+ Action: action,
+ Generation: generation,
+ }, true
+}
+
+// formatServerErrorMessage builds the "SERVER|<server>|ERROR|<message>" audit
+// line shown to the user and written to the client log file. It carries NO
+// trailing newline: it is emitted via the diagnostic (Log) sink, which appends
+// the newline itself (adding one here would produce a blank line).
+func formatServerErrorMessage(server string, message string) string {
+ return fmt.Sprintf("SERVER%s%s%sERROR%s%s",
+ protocol.FieldDelimiter,
+ server,
+ protocol.FieldDelimiter,
+ protocol.FieldDelimiter,
+ message,
+ )
+}
diff --git a/internal/clients/handlers/basehandler_test.go b/internal/clients/handlers/basehandler_test.go
new file mode 100644
index 0000000..e070e48
--- /dev/null
+++ b/internal/clients/handlers/basehandler_test.go
@@ -0,0 +1,356 @@
+package handlers
+
+import (
+ "fmt"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/mimecast/dtail/internal"
+ "github.com/mimecast/dtail/internal/io/dlog"
+ "github.com/mimecast/dtail/internal/protocol"
+)
+
+func TestParseAuthKeyMessage(t *testing.T) {
+ tests := []struct {
+ name string
+ message string
+ wantAuth bool
+ wantOK bool
+ wantInfo string
+ }{
+ {
+ name: "server formatted success",
+ message: fmt.Sprintf("SERVER%s%s%sAUTHKEY OK\n", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+ wantAuth: true,
+ wantOK: true,
+ },
+ {
+ name: "server formatted error",
+ message: fmt.Sprintf("SERVER%s%s%sAUTHKEY ERR feature disabled\n", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+ wantAuth: true,
+ wantOK: false,
+ wantInfo: "feature disabled",
+ },
+ {
+ name: "plain response success",
+ message: "AUTHKEY OK",
+ wantAuth: true,
+ wantOK: true,
+ },
+ {
+ name: "not an authkey message",
+ message: fmt.Sprintf("SERVER%s%s%ssome other message", protocol.FieldDelimiter, "host1", protocol.FieldDelimiter),
+ wantAuth: false,
+ wantOK: false,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ gotAuth, gotOK, gotInfo := parseAuthKeyMessage(tc.message)
+ if gotAuth != tc.wantAuth {
+ t.Fatalf("Unexpected auth marker: got %v want %v", gotAuth, tc.wantAuth)
+ }
+ if gotOK != tc.wantOK {
+ t.Fatalf("Unexpected ok marker: got %v want %v", gotOK, tc.wantOK)
+ }
+ if gotInfo != tc.wantInfo {
+ t.Fatalf("Unexpected info: got %q want %q", gotInfo, tc.wantInfo)
+ }
+ })
+ }
+}
+
+func TestHandleCapabilitiesMessage(t *testing.T) {
+ handler := baseHandler{
+ done: internal.NewDone(),
+ capabilities: make(map[string]struct{}),
+ capabilitiesCh: make(chan struct{}),
+ sessionAcks: make(chan SessionAck, 1),
+ }
+
+ handler.handleHiddenMessage(".syn capabilities query-update-v1 feature-two")
+
+ if !handler.HasCapability(protocol.CapabilityQueryUpdateV1) {
+ t.Fatalf("expected handler to track %q", protocol.CapabilityQueryUpdateV1)
+ }
+ if !handler.HasCapability("feature-two") {
+ t.Fatalf("expected handler to track feature-two")
+ }
+ if handler.WaitForCapabilities(10*time.Millisecond) != true {
+ t.Fatalf("expected capabilities wait to succeed")
+ }
+
+ capabilities := handler.Capabilities()
+ if len(capabilities) != 2 {
+ t.Fatalf("unexpected capabilities: %#v", capabilities)
+ }
+}
+
+func TestWaitForCapabilitiesTimeout(t *testing.T) {
+ handler := baseHandler{
+ done: internal.NewDone(),
+ capabilities: make(map[string]struct{}),
+ capabilitiesCh: make(chan struct{}),
+ sessionAcks: make(chan SessionAck, 1),
+ }
+
+ if handler.WaitForCapabilities(5 * time.Millisecond) {
+ t.Fatalf("expected capabilities wait to time out")
+ }
+}
+
+func TestFormatServerErrorMessage(t *testing.T) {
+ got := formatServerErrorMessage("srv1", "journal file targets require server capability journal-v1")
+ // No trailing newline: the message is emitted via the diagnostic (Log) sink,
+ // which appends the newline itself (see ReportServerError / RawLog).
+ want := "SERVER|srv1|ERROR|journal file targets require server capability journal-v1"
+ if got != want {
+ t.Fatalf("formatServerErrorMessage() = %q, want %q", got, want)
+ }
+}
+
+func TestParseSessionAckMessage(t *testing.T) {
+ tests := []struct {
+ name string
+ message string
+ want SessionAck
+ wantOK bool
+ }{
+ {
+ name: "start ok",
+ message: ".syn session start ok 7",
+ want: SessionAck{
+ Action: "start",
+ Generation: 7,
+ },
+ wantOK: true,
+ },
+ {
+ name: "update ok",
+ message: ".syn session update ok 8",
+ want: SessionAck{
+ Action: "update",
+ Generation: 8,
+ },
+ wantOK: true,
+ },
+ {
+ name: "error",
+ message: ".syn session err query sessions not supported yet",
+ want: SessionAck{
+ Action: "error",
+ Error: "query sessions not supported yet",
+ },
+ wantOK: true,
+ },
+ {
+ name: "invalid",
+ message: ".syn session start ok nope",
+ wantOK: false,
+ },
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ got, ok := parseSessionAckMessage(tc.message)
+ if ok != tc.wantOK {
+ t.Fatalf("unexpected ok flag: got %v want %v", ok, tc.wantOK)
+ }
+ if !tc.wantOK {
+ return
+ }
+ if got != tc.want {
+ t.Fatalf("unexpected ack: got %#v want %#v", got, tc.want)
+ }
+ })
+ }
+}
+
+func TestHandleSessionAckMessage(t *testing.T) {
+ handler := baseHandler{
+ done: internal.NewDone(),
+ sessionAcks: make(chan SessionAck, 1),
+ }
+
+ handler.handleHiddenMessage(".syn session update ok 4")
+
+ ack, ok := handler.WaitForSessionAck(10 * time.Millisecond)
+ if !ok {
+ t.Fatalf("expected session ack")
+ }
+ if ack.Action != "update" || ack.Generation != 4 {
+ t.Fatalf("unexpected session ack: %#v", ack)
+ }
+}
+
+func TestHandleCloseConnectionAcknowledgesBeforeShutdown(t *testing.T) {
+ originalLogger := dlog.Client
+ dlog.Client = &dlog.DLog{}
+ t.Cleanup(func() {
+ dlog.Client = originalLogger
+ })
+
+ handler := baseHandler{
+ done: internal.NewDone(),
+ server: "server-under-test",
+ commands: make(chan string, 1),
+ }
+
+ handler.handleHiddenMessage(".syn close connection")
+
+ select {
+ case command := <-handler.commands:
+ if command == "" {
+ t.Fatal("expected close acknowledgement command")
+ }
+ case <-time.After(10 * time.Millisecond):
+ t.Fatal("expected close acknowledgement command to be queued")
+ }
+
+ select {
+ case <-handler.Done():
+ default:
+ t.Fatal("expected handler to be shut down after close acknowledgement")
+ }
+}
+
+// TestReadDrainsAckBeforeEOF verifies that Read() always delivers the
+// '.ack close connection' command before returning io.EOF when Done() fires
+// concurrently. Without the priority-select fix, Go's non-deterministic select
+// would randomly pick the Done() case and drop the queued ack ~50% of the time.
+func TestReadDrainsAckBeforeEOF(t *testing.T) {
+ originalLogger := dlog.Client
+ dlog.Client = &dlog.DLog{}
+ t.Cleanup(func() {
+ dlog.Client = originalLogger
+ })
+
+ // Run many iterations to catch the race reliably even with -race.
+ const iterations = 500
+ for i := 0; i < iterations; i++ {
+ // Use an unbuffered commands channel so the ack write and Done() close
+ // are truly concurrent, exposing the original non-deterministic select.
+ handler := baseHandler{
+ done: internal.NewDone(),
+ server: "server-under-test",
+ commands: make(chan string, 1),
+ }
+
+ // handleHiddenMessage calls SendMessage (queues ack) then Shutdown (closes Done).
+ // After this call both handler.commands and handler.Done() are ready.
+ handler.handleHiddenMessage(".syn close connection")
+
+ buf := make([]byte, 4096)
+ n, err := handler.Read(buf)
+ if err != nil {
+ // Done() won the select — the ack was dropped. This is the bug.
+ t.Fatalf("iteration %d: Read returned io.EOF before draining the ack", i)
+ }
+ if n == 0 {
+ t.Fatalf("iteration %d: Read returned 0 bytes without an error", i)
+ }
+
+ // Second Read must now return EOF since Done() is closed and commands is empty.
+ _, err = handler.Read(buf)
+ if err == nil {
+ t.Fatalf("iteration %d: second Read should return io.EOF", i)
+ }
+ }
+}
+
+// newCommandReadTestHandler returns a handler for exercising Read directly
+// and registers a cleanup that unblocks any Read still waiting on the
+// commands channel, so a regression cannot leak a blocked goroutine.
+func newCommandReadTestHandler(t *testing.T) *baseHandler {
+ t.Helper()
+
+ handler := &baseHandler{
+ done: internal.NewDone(),
+ server: "server-under-test",
+ commands: make(chan string, 2),
+ }
+ t.Cleanup(handler.done.Shutdown)
+ return handler
+}
+
+// readCommandsWithin drains wantLen bytes from the handler on a separate
+// goroutine and fails the test when the reads do not finish within the
+// timeout. The deadline matters: under a regression that drops the remainder
+// of a partially delivered command, the next Read would block forever on the
+// empty commands channel; the timeout turns that hang into a clean assertion
+// (and shuts the handler down so the reader goroutine exits via Done).
+func readCommandsWithin(t *testing.T, handler *baseHandler, bufSize, wantLen int) []byte {
+ t.Helper()
+
+ type result struct {
+ data []byte
+ err error
+ }
+ resultCh := make(chan result, 1)
+ go func() {
+ var got []byte
+ p := make([]byte, bufSize)
+ for len(got) < wantLen {
+ n, err := handler.Read(p)
+ if err != nil {
+ resultCh <- result{got, err}
+ return
+ }
+ got = append(got, p[:n]...)
+ }
+ resultCh <- result{got, nil}
+ }()
+
+ select {
+ case res := <-resultCh:
+ if res.err != nil {
+ t.Fatalf("Read() error after %d of %d bytes: %v",
+ len(res.data), wantLen, res.err)
+ }
+ return res.data
+ case <-time.After(5 * time.Second):
+ handler.done.Shutdown()
+ t.Fatalf("Read stalled: timed out waiting for %d bytes", wantLen)
+ return nil
+ }
+}
+
+// TestBaseHandlerReadLargeCommandAcrossMultipleReads verifies that a command
+// frame larger than the caller's buffer (e.g. a huge base64-encoded regex or
+// MapReduce query) is delivered completely across multiple Read calls instead
+// of being truncated at the buffer boundary.
+func TestBaseHandlerReadLargeCommandAcrossMultipleReads(t *testing.T) {
+ handler := newCommandReadTestHandler(t)
+
+ command := "grep " + strings.Repeat("x", 1000) + ";"
+ handler.commands <- command
+
+ got := readCommandsWithin(t, handler, 32, len(command))
+ if string(got) != command {
+ t.Fatalf("large command corrupted across reads:\ngot %q\nwant %q",
+ got, command)
+ }
+}
+
+// TestBaseHandlerReadDrainsRemainderBeforeNextCommand verifies that the
+// remainder of a partially delivered command is fully drained before the next
+// queued command starts, so frames are never interleaved, and that an
+// exact-fit read leaves no stale remainder behind.
+func TestBaseHandlerReadDrainsRemainderBeforeNextCommand(t *testing.T) {
+ handler := newCommandReadTestHandler(t)
+
+ first := "first command payload;"
+ second := "second;"
+ handler.commands <- first
+ handler.commands <- second
+
+ // len(first) = 22, so first fills exactly two 11-byte reads (the second
+ // of which is the remainder), then second must start fresh.
+ got := readCommandsWithin(t, handler, 11, len(first)+len(second))
+ if string(got) != first+second {
+ t.Fatalf("commands interleaved or corrupted:\ngot %q\nwant %q",
+ got, first+second)
+ }
+}
diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go
index 27ac85e..3998e9f 100644
--- a/internal/clients/handlers/clienthandler.go
+++ b/internal/clients/handlers/clienthandler.go
@@ -10,17 +10,22 @@ type ClientHandler struct {
baseHandler
}
+var _ Handler = (*ClientHandler)(nil)
+
// NewClientHandler creates a new client handler.
func NewClientHandler(server string) *ClientHandler {
dlog.Client.Debug(server, "Creating new client handler")
return &ClientHandler{
baseHandler{
- server: server,
- shellStarted: false,
- commands: make(chan string),
- status: -1,
- done: internal.NewDone(),
+ server: server,
+ shellStarted: false,
+ commands: make(chan string),
+ status: -1,
+ done: internal.NewDone(),
+ capabilities: make(map[string]struct{}),
+ capabilitiesCh: make(chan struct{}),
+ sessionAcks: make(chan SessionAck, 4),
},
}
}
diff --git a/internal/clients/handlers/handler.go b/internal/clients/handlers/handler.go
index afa87e2..c1f7256 100644
--- a/internal/clients/handlers/handler.go
+++ b/internal/clients/handlers/handler.go
@@ -2,14 +2,20 @@ package handlers
import (
"io"
+ "time"
)
// Handler provides all methods which can be run on any client handler.
type Handler interface {
io.ReadWriter
+ Capabilities() []string
+ HasCapability(name string) bool
+ ReportServerError(message string)
SendMessage(command string) error
Server() string
Status() int
Shutdown()
Done() <-chan struct{}
+ WaitForCapabilities(timeout time.Duration) bool
+ WaitForSessionAck(timeout time.Duration) (SessionAck, bool)
}
diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
index 47b594e..763ba88 100644
--- a/internal/clients/handlers/healthhandler.go
+++ b/internal/clients/handlers/healthhandler.go
@@ -19,11 +19,14 @@ func NewHealthHandler(server string) *HealthHandler {
dlog.Client.Debug(server, "Creating new health handler")
return &HealthHandler{
baseHandler: baseHandler{
- server: server,
- shellStarted: false,
- commands: make(chan string),
- status: 2, // Assume CRITICAL status by default.
- done: internal.NewDone(),
+ server: server,
+ shellStarted: false,
+ commands: make(chan string),
+ status: 2, // Assume CRITICAL status by default.
+ done: internal.NewDone(),
+ capabilities: make(map[string]struct{}),
+ capabilitiesCh: make(chan struct{}),
+ sessionAcks: make(chan SessionAck, 4),
},
}
}
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
index 4c11470..5814fa0 100644
--- a/internal/clients/handlers/maprhandler.go
+++ b/internal/clients/handlers/maprhandler.go
@@ -5,34 +5,39 @@ import (
"github.com/mimecast/dtail/internal"
"github.com/mimecast/dtail/internal/io/dlog"
- "github.com/mimecast/dtail/internal/mapr"
"github.com/mimecast/dtail/internal/mapr/client"
"github.com/mimecast/dtail/internal/protocol"
)
+// aggregateMessagePrefix is the leading part of a mapreduce aggregate-data
+// wire message (AGGREGATE|host|data). Classifying incoming messages against
+// this full prefix, rather than just their first byte, prevents plain-mode
+// protocol acks (e.g. "AUTHKEY OK") from being fed to the aggregate parser.
+const aggregateMessagePrefix = protocol.AggregateMessageID + protocol.FieldDelimiter
+
// MaprHandler is the handler used on the client side for running mapreduce
// aggregations.
type MaprHandler struct {
baseHandler
aggregate *client.Aggregate
- query *mapr.Query
removedNl bool
}
// NewMaprHandler returns a new mapreduce client handler.
-func NewMaprHandler(server string, query *mapr.Query,
- globalGroup *mapr.GlobalGroupSet) *MaprHandler {
+func NewMaprHandler(server string, session *client.SessionState) *MaprHandler {
return &MaprHandler{
baseHandler: baseHandler{
- server: server,
- shellStarted: false,
- commands: make(chan string),
- status: -1,
- done: internal.NewDone(),
+ server: server,
+ shellStarted: false,
+ commands: make(chan string),
+ status: -1,
+ done: internal.NewDone(),
+ capabilities: make(map[string]struct{}),
+ capabilitiesCh: make(chan struct{}),
+ sessionAcks: make(chan SessionAck, 4),
},
- query: query,
- aggregate: client.NewAggregate(server, query, globalGroup),
+ aggregate: client.NewAggregate(server, session),
}
}
@@ -44,8 +49,13 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) {
h.removedNl = true
case protocol.MessageDelimiter:
message := h.baseHandler.receiveBuf.String()
+ if len(message) == 0 {
+ h.baseHandler.receiveBuf.Reset()
+ h.removedNl = false
+ continue
+ }
dlog.Client.Debug(message)
- if message[0] == 'A' {
+ if isAggregateMessage(message) {
h.handleAggregateMessage(message)
} else {
if h.removedNl {
@@ -64,6 +74,17 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) {
return len(p), nil
}
+// isAggregateMessage reports whether a wire message carries mapreduce
+// aggregate data (AGGREGATE|host|data). Only such messages may be handed to
+// the aggregate parser. Matching the full AggregateMessageID field prefix,
+// rather than just the first byte 'A', keeps plain-mode protocol acks such as
+// "AUTHKEY OK" out of the parser; those would otherwise trigger a spurious
+// "Unable to aggregate data ... expected 3 parts" error. Non-aggregate
+// messages are routed to the base handler, which recognises acks as control.
+func isAggregateMessage(message string) bool {
+ return strings.HasPrefix(message, aggregateMessagePrefix)
+}
+
// Handle a message received from server including mapr aggregation related data.
func (h *MaprHandler) handleAggregateMessage(message string) {
parts := strings.SplitN(message, protocol.FieldDelimiter, 3)
@@ -76,3 +97,13 @@ func (h *MaprHandler) handleAggregateMessage(message string) {
dlog.Client.Error("Unable to aggregate data", h.server, message, err)
}
}
+
+// Shutdown flushes any pending aggregate state before marking the handler done.
+func (h *MaprHandler) Shutdown() {
+ if h.aggregate != nil {
+ if err := h.aggregate.Flush(); err != nil {
+ dlog.Client.Error("Unable to flush aggregate data on shutdown", h.server, err)
+ }
+ }
+ h.baseHandler.Shutdown()
+}
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
+ // ta