merge from github.com/snonux/dtail

1 files changed, 42 insertions, 55 deletions

diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index b5045e2..b520c25 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -1,16 +1,16 @@
 package handlers
 
 import (
+	"bytes"
 	"encoding/base64"
 	"fmt"
 	"io"
-	"strconv"
 	"strings"
 	"time"
 
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/version"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
 )
 
 type baseHandler struct {
@@ -18,10 +18,20 @@ type baseHandler struct {
 	server       string
 	shellStarted bool
 	commands     chan string
-	receiveBuf   []byte
+	receiveBuf   bytes.Buffer
 	status       int
 }
 
+func (h *baseHandler) String() string {
+	return fmt.Sprintf("baseHandler(%s,server:%s,shellStarted:%v,status:%d)@%p",
+		h.done,
+		h.server,
+		h.shellStarted,
+		h.status,
+		h,
+	)
+}
+
 func (h *baseHandler) Server() string {
 	return h.server
 }
@@ -30,21 +40,13 @@ func (h *baseHandler) Status() int {
 	return h.status
 }
 
-func (h *baseHandler) Done() <-chan struct{} {
-	return h.done.Done()
-}
-
-func (h *baseHandler) Shutdown() {
-	h.done.Shutdown()
-}
-
 // SendMessage to the server.
 func (h *baseHandler) SendMessage(command string) error {
 	encoded := base64.StdEncoding.EncodeToString([]byte(command))
-	logger.Debug("Sending command", h.server, command, encoded)
+	dlog.Client.Debug("Sending command", h.server, command, encoded)
 
 	select {
-	case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded):
+	case h.commands <- fmt.Sprintf("protocol %s base64 %v;", protocol.ProtocolCompat, encoded):
 	case <-time.After(time.Second * 5):
 		return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded)
 	case <-h.Done():
@@ -57,13 +59,20 @@ func (h *baseHandler) SendMessage(command string) error {
 // Read data from the dtail server via Writer interface.
 func (h *baseHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		h.receiveBuf = append(h.receiveBuf, b)
-		if b == '\n' {
-			if len(h.receiveBuf) == 0 {
+		switch b {
+		/*
+			// NEXT: Next DTail version make it so that '\n' gets ignored. For now
+			// leave it for compatibility with older DTail server + ability to display
+			// the protocol mismatch warn message.
+			case '\n' {
 				continue
-			}
-			message := string(h.receiveBuf)
-			h.handleMessageType(message)
+		*/
+		case '\n', protocol.MessageDelimiter:
+			message := h.receiveBuf.String()
+			h.handleMessage(message)
+			h.receiveBuf.Reset()
+		default:
+			h.receiveBuf.WriteByte(b)
 		}
 	}
 
@@ -81,21 +90,13 @@ func (h *baseHandler) Read(p []byte) (n int, err error) {
 	return
 }
 
-// Handle various message types.
-func (h *baseHandler) handleMessageType(message string) {
-	if len(h.receiveBuf) == 0 {
-		return
-	}
-
-	// Hidden server commands starti with a dot "."
-	if h.receiveBuf[0] == '.' {
+func (h *baseHandler) handleMessage(message string) {
+	if len(message) > 0 && message[0] == '.' {
 		h.handleHiddenMessage(message)
-		h.receiveBuf = h.receiveBuf[:0]
 		return
 	}
 
-	logger.Raw(message)
-	h.receiveBuf = h.receiveBuf[:0]
+	dlog.Client.Raw(message)
 }
 
 // Handle messages received from server which are not meant to be displayed
@@ -103,29 +104,15 @@ func (h *baseHandler) handleMessageType(message string) {
 func (h *baseHandler) handleHiddenMessage(message string) {
 	switch {
 	case strings.HasPrefix(message, ".syn close connection"):
-		h.SendMessage(".ack close connection")
-		select {
-		case <-time.After(time.Second * 5):
-			logger.Debug("Shutting down client after timeout and sending ack to server")
-			h.Shutdown()
-		case <-h.Done():
-			return
-		}
-
-	case strings.HasPrefix(message, ".run exitstatus"):
-		splitted := strings.Split(strings.TrimSuffix(message, "\n"), " ")
-		if len(splitted) != 3 {
-			logger.Error("Unable to retrieve exitstatus", message)
-			return
-		}
-		i, err := strconv.Atoi(splitted[2])
-		if err != nil {
-			logger.Error("Unable to retrieve exitstatus", message, err)
-			return
-		}
-		logger.Debug("Retrieved exitstatus", h.status)
-		if i > h.status {
-			h.status = i
-		}
+		go h.SendMessage(".ack close connection")
+		h.Shutdown()
 	}
 }
+
+func (h *baseHandler) Done() <-chan struct{} {
+	return h.done.Done()
+}
+
+func (h *baseHandler) Shutdown() {
+	h.done.Shutdown()
+}