merge develop

16 files changed, 686 insertions, 575 deletions

diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index f83fcfd..4a7bd84 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -2,15 +2,13 @@ package clients
 
 import (
 	"context"
-	"fmt"
-	"strings"
 	"sync"
 	"time"
 
-	"github.com/mimecast/dtail/internal/clients/remote"
+	"github.com/mimecast/dtail/internal/clients/connectors"
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/discovery"
-	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/regex"
 	"github.com/mimecast/dtail/internal/ssh/client"
 
@@ -19,13 +17,13 @@ import (
 
 // This is the main client data structure.
 type baseClient struct {
-	Args
+	config.Args
 	// To display client side stats
 	stats *stats
 	// List of remote servers to connect to.
 	servers []string
 	// We have one connection per remote server.
-	connections []*remote.Connection
+	connections []connectors.Connector
 	// SSH auth methods to use to connect to the remote servers.
 	sshAuthMethods []gossh.AuthMethod
 	// To deal with SSH host keys
@@ -41,7 +39,7 @@ type baseClient struct {
 }
 
 func (c *baseClient) init() {
-	logger.Debug("Initiating base client")
+	dlog.Client.Debug("Initiating base client", c.Args.String())
 
 	flag := regex.Default
 	if c.Args.RegexInvert {
@@ -49,12 +47,16 @@ func (c *baseClient) init() {
 	}
 	regex, err := regex.New(c.Args.RegexStr, flag)
 	if err != nil {
-		logger.FatalExit(c.Regex, "invalid regex!", err, regex)
+		dlog.Client.FatalPanic(c.Regex, "Invalid regex!", err, regex)
 	}
 	c.Regex = regex
-	logger.Debug("Regex", c.Regex)
 
-	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, c.throttleCh, c.Args.PrivateKeyPathFile)
+	if c.Args.Serverless {
+		return
+	}
+	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(
+		c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts,
+		c.throttleCh, c.Args.PrivateKeyPathFile)
 }
 
 func (c *baseClient) makeConnections(maker maker) {
@@ -62,26 +64,31 @@ func (c *baseClient) makeConnections(maker maker) {
 
 	discoveryService := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle)
 	for _, server := range discoveryService.ServerList() {
-		c.connections = append(c.connections, c.makeConnection(server, c.sshAuthMethods, c.hostKeyCallback))
+		c.connections = append(c.connections, c.makeConnection(server,
+			c.sshAuthMethods, c.hostKeyCallback))
 	}
 
 	c.stats = newTailStats(len(c.connections))
 }
 
 func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status int) {
-	// Periodically check for unknown hosts, and ask the user whether to trust them or not.
-	go c.hostKeyCallback.PromptAddHosts(ctx)
+	dlog.Client.Trace("Starting base client")
+	// Can be nil when serverless.
+	if c.hostKeyCallback != nil {
+		// Periodically check for unknown hosts, and ask the user whether to trust them or not.
+		go c.hostKeyCallback.PromptAddHosts(ctx)
+	}
 	// Print client stats every time something on statsCh is recieved.
 	go c.stats.Start(ctx, c.throttleCh, statsCh, c.Args.Quiet)
-	// Keep count of active connections
-	active := make(chan struct{}, len(c.connections))
 
+	var wg sync.WaitGroup
+	wg.Add(len(c.connections))
 	var mutex sync.Mutex
-	for i, conn := range c.connections {
-		go func(i int, conn *remote.Connection) {
-			connStatus := c.start(ctx, active, i, conn)
 
-			// Update global status.
+	for i, conn := range c.connections {
+		go func(i int, conn connectors.Connector) {
+			defer wg.Done()
+			connStatus := c.startConnection(ctx, i, conn)
 			mutex.Lock()
 			defer mutex.Unlock()
 			if connStatus > status {
@@ -90,15 +97,12 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i
 		}(i, conn)
 	}
 
-	c.waitUntilDone(ctx, active)
+	wg.Wait()
 	return
 }
 
-func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, conn *remote.Connection) (status int) {
-	// Increment connection count
-	active <- struct{}{}
-	// Derement connection count
-	defer func() { <-active }()
+func (c *baseClient) startConnection(ctx context.Context, i int,
+	conn connectors.Connector) (status int) {
 
 	for {
 		connCtx, cancel := context.WithCancel(ctx)
@@ -106,80 +110,25 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con
 
 		conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh)
 		// Retrieve status code from handler (dtail client will exit with that status)
-		status = conn.Handler.Status()
+		status = conn.Handler().Status()
 
 		if !c.retry {
 			return
 		}
 
 		time.Sleep(time.Second * 2)
-		logger.Debug(conn.Server, "Reconnecting")
-
-		conn = c.makeConnection(conn.Server, c.sshAuthMethods, c.hostKeyCallback)
+		dlog.Client.Debug(conn.Server(), "Reconnecting")
+		conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback)
 		c.connections[i] = conn
 	}
 }
 
-func (c *baseClient) makeCommandOptions() map[string]string {
-	options := make(map[string]string)
-
-	if c.Args.Quiet {
-		options["quiet"] = fmt.Sprintf("%v", c.Args.Quiet)
-	}
-	if c.Args.LContext.MaxCount != 0 {
-		options["max"] = fmt.Sprintf("%d", c.Args.LContext.MaxCount)
-	}
-	if c.Args.LContext.BeforeContext != 0 {
-		options["before"] = fmt.Sprintf("%d", c.Args.LContext.BeforeContext)
-	}
-	if c.Args.LContext.AfterContext != 0 {
-		options["after"] = fmt.Sprintf("%d", c.Args.LContext.AfterContext)
-	}
-
-	return options
-}
-
-func (c *baseClient) commandOptionsToString(options map[string]string) string {
-	var sb strings.Builder
-
-	count := 0
-	for k, v := range options {
-		if count > 0 {
-			sb.WriteString(":")
-		}
-		sb.WriteString(fmt.Sprintf("%s=%s", k, v))
-		count++
-	}
-
-	return sb.String()
-}
-
-func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) *remote.Connection {
-	conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback)
-	conn.Handler = c.maker.makeHandler(server)
-	conn.Commands = c.maker.makeCommands(c.makeCommandOptions())
-
-	return conn
-}
-
-func (c *baseClient) waitUntilDone(ctx context.Context, active chan struct{}) {
-	defer logger.Debug("Terminated connection")
-
-	// We want to have at least one active connection
-	<-active
-	// Put it back on the channel
-	active <- struct{}{}
-
-	if c.Mode == omode.TailClient && c.retry {
-		<-ctx.Done()
-	}
-
-	for {
-		numActive := len(active)
-		if numActive == 0 {
-			return
-		}
-		logger.Debug("Active connections", numActive)
-		time.Sleep(time.Second)
+func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod,
+	hostKeyCallback client.HostKeyCallback) connectors.Connector {
+	if c.Args.Serverless {
+		return connectors.NewServerless(c.UserName, c.maker.makeHandler(server),
+			c.maker.makeCommands())
 	}
+	return connectors.NewServerConnection(server, c.UserName, sshAuthMethods,
+		hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands())
 }
diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go
index db892f1..bd65560 100644
--- a/internal/clients/catclient.go
+++ b/internal/clients/catclient.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/omode"
 )
 
@@ -16,11 +18,10 @@ type CatClient struct {
 }
 
 // NewCatClient returns a new cat client.
-func NewCatClient(args Args) (*CatClient, error) {
+func NewCatClient(args config.Args) (*CatClient, error) {
 	if args.RegexStr != "" {
 		return nil, errors.New("Can't use regex with 'cat' operating mode")
 	}
-
 	args.Mode = omode.CatClient
 
 	c := CatClient{
@@ -33,7 +34,6 @@ func NewCatClient(args Args) (*CatClient, error) {
 
 	c.init()
 	c.makeConnections(c)
-
 	return &c, nil
 }
 
@@ -41,10 +41,14 @@ func (c CatClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c CatClient) makeCommands(options map[string]string) (commands []string) {
-	optionsStr := c.commandOptionsToString(options)
+func (c CatClient) makeCommands() (commands []string) {
+	regex, err := c.Regex.Serialize()
+	if err != nil {
+		dlog.Client.FatalPanic(err)
+	}
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s",
+			c.Mode.String(), c.Args.SerializeOptions(), file, regex))
 	}
 	return
 }
diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go
new file mode 100644
index 0000000..3ab6a08
--- /dev/null
+++ b/internal/clients/connectors/connector.go
@@ -0,0 +1,17 @@
+package connectors
+
+import (
+	"context"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+)
+
+// Connector interface.
+type Connector interface {
+	// Start the connection.
+	Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{})
+	// Server hostname.
+	Server() string
+	// Handler for the connection.
+	Handler() handlers.Handler
+}
diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
new file mode 100644
index 0000000..aeb2a41
--- /dev/null
+++ b/internal/clients/connectors/serverconnection.go
@@ -0,0 +1,206 @@
+package connectors
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	"golang.org/x/crypto/ssh"
+)
+
+// ServerConnection represents a connection to a single remote dtail server via
+// SSH protocol.
+type ServerConnection struct {
+	// The full server string as received from the server discovery (can be with port number)
+	server string
+	// Only the hostname or FQDN (without the port number)
+	hostname string
+	// Only the port number.
+	port            int
+	config          *ssh.ClientConfig
+	handler         handlers.Handler
+	commands        []string
+	hostKeyCallback client.HostKeyCallback
+	throttlingDone  bool
+}
+
+// NewServerConnection returns a new DTail SSH server connection.
+func NewServerConnection(server string, userName string,
+	authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback,
+	handler handlers.Handler, commands []string) *ServerConnection {
+
+	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
+	c := ServerConnection{
+		hostKeyCallback: hostKeyCallback,
+		server:          server,
+		handler:         handler,
+		commands:        commands,
+		config: &ssh.ClientConfig{
+			User:            userName,
+			Auth:            authMethods,
+			HostKeyCallback: hostKeyCallback.Wrap(),
+			Timeout:         time.Second * 2,
+		},
+	}
+
+	c.initServerPort()
+	return &c
+}
+
+// Server returns the server hostname connected to.
+func (c *ServerConnection) Server() string { return c.server }
+
+// Handler returns the handler used for the connection.
+func (c *ServerConnection) Handler() handlers.Handler { return c.handler }
+
+// Attempt to parse the server port address from the provided server FQDN.
+func (c *ServerConnection) initServerPort() {
+	parts := strings.Split(c.server, ":")
+	if len(parts) == 1 {
+		c.hostname = c.server
+		c.port = config.Common.SSHPort
+		return
+	}
+
+	dlog.Client.Debug("Parsing port from hostname", parts)
+	port, err := strconv.Atoi(parts[1])
+	if err != nil {
+		dlog.Client.FatalPanic("Unable to parse client port", c.server, parts, err)
+	}
+	c.hostname = parts[0]
+	c.port = port
+}
+
+// Start the connection to the server.
+func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) {
+
+	// Throttle how many connections can be established concurrently (based on ch length)
+	dlog.Client.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh))
+
+	select {
+	case throttleCh <- struct{}{}:
+	case <-ctx.Done():
+		dlog.Client.Debug(c.server, "Not establishing connection as context is done",
+			len(throttleCh), cap(throttleCh))
+		return
+	}
+
+	dlog.Client.Debug(c.server, "Throttling says that the connection can be established",
+		len(throttleCh), cap(throttleCh))
+
+	go func() {
+		defer func() {
+			if !c.throttlingDone {
+				dlog.Client.Debug(c.server, "Unthrottling connection (1)",
+					len(throttleCh), cap(throttleCh))
+				c.throttlingDone = true
+				<-throttleCh
+			}
+			cancel()
+		}()
+
+		if err := c.dial(ctx, cancel, throttleCh, statsCh); err != nil {
+			dlog.Client.Warn(c.server, err)
+			if c.hostKeyCallback.Untrusted(c.server) {
+				dlog.Client.Debug(c.server, "Not trusting host")
+			}
+		}
+	}()
+
+	<-ctx.Done()
+}
+
+// Dail into a new SSH connection. Close connection in case of an error.
+func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Incrementing connection stats")
+	statsCh <- struct{}{}
+	defer func() {
+		dlog.Client.Debug(c.server, "Decrementing connection stats")
+		<-statsCh
+	}()
+
+	address := fmt.Sprintf("%s:%d", c.hostname, c.port)
+	dlog.Client.Debug(c.server, "Dialing into the connection", address)
+
+	client, err := ssh.Dial("tcp", address, c.config)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	return c.session(ctx, cancel, client, throttleCh)
+}
+
+// Create the SSH session. Close the session in case of an error.
+func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFunc,
+	client *ssh.Client, throttleCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Creating SSH session")
+	session, err := client.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+	return c.handle(ctx, cancel, session, throttleCh)
+}
+
+func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc,
+	session *ssh.Session, throttleCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Creating handler for SSH session")
+	stdinPipe, err := session.StdinPipe()
+	if err != nil {
+		return err
+	}
+	stdoutPipe, err := session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	if err := session.Shell(); err != nil {
+		return err
+	}
+
+	go func() {
+		io.Copy(stdinPipe, c.handler)
+		cancel()
+	}()
+	go func() {
+		io.Copy(c.handler, stdoutPipe)
+		cancel()
+	}()
+	go func() {
+		select {
+		case <-c.handler.Done():
+		case <-ctx.Done():
+		}
+		cancel()
+	}()
+
+	// Send all commands to client.
+	for _, command := range c.commands {
+		dlog.Client.Debug(command)
+		c.handler.SendMessage(command)
+	}
+
+	if !c.throttlingDone {
+		dlog.Client.Debug(c.server, "Unthrottling connection (2)",
+			len(throttleCh), cap(throttleCh))
+		c.throttlingDone = true
+		<-throttleCh
+	}
+
+	<-ctx.Done()
+	c.handler.Shutdown()
+	return nil
+}
diff --git a/internal/clients/connectors/serverless.go b/internal/clients/connectors/serverless.go
new file mode 100644
index 0000000..2ff490a
--- /dev/null
+++ b/internal/clients/connectors/serverless.go
@@ -0,0 +1,116 @@
+package connectors
+
+import (
+	"context"
+	"io"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	serverHandlers "github.com/mimecast/dtail/internal/server/handlers"
+	user "github.com/mimecast/dtail/internal/user/server"
+)
+
+// Serverless creates a server object directly without TCP.
+type Serverless struct {
+	handler  handlers.Handler
+	commands []string
+	userName string
+}
+
+// NewServerless starts a new serverless session.
+func NewServerless(userName string, handler handlers.Handler,
+	commands []string) *Serverless {
+
+	dlog.Client.Debug("Creating new serverless connector", handler, commands)
+	return &Serverless{
+		userName: userName,
+		handler:  handler,
+		commands: commands,
+	}
+}
+
+// Server returns serverless server indicator.
+func (s *Serverless) Server() string {
+	return "local(serverless)"
+}
+
+// Handler returns the handler used for the serverless connection.
+func (s *Serverless) Handler() handlers.Handler {
+	return s.handler
+}
+
+// Start the serverless connection.
+func (s *Serverless) Start(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) {
+
+	dlog.Client.Debug("Starting serverless connector")
+	go func() {
+		defer cancel()
+
+		if err := s.handle(ctx, cancel); err != nil {
+			dlog.Client.Warn(err)
+		}
+	}()
+	<-ctx.Done()
+}
+
+func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) error {
+	dlog.Client.Debug("Creating server handler for a serverless session")
+
+	user, err := user.New(s.userName, s.Server())
+	if err != nil {
+		return err
+	}
+
+	var serverHandler serverHandlers.Handler
+	switch s.userName {
+	case config.HealthUser:
+		dlog.Client.Debug("Creating serverless health handler")
+		serverHandler = serverHandlers.NewHealthHandler(user)
+	default:
+		dlog.Client.Debug("Creating serverless server handler")
+		serverHandler = serverHandlers.NewServerHandler(
+			user,
+			make(chan struct{}, config.Server.MaxConcurrentCats),
+			make(chan struct{}, config.Server.MaxConcurrentTails),
+		)
+	}
+
+	terminate := func() {
+		dlog.Client.Debug("Terminating serverless connection")
+		serverHandler.Shutdown()
+		cancel()
+	}
+
+	go func() {
+		io.Copy(serverHandler, s.handler)
+		dlog.Client.Trace("io.Copy(serverHandler, s.handler) => done")
+		terminate()
+	}()
+	go func() {
+		io.Copy(s.handler, serverHandler)
+		dlog.Client.Trace("io.Copy(s.handler, serverHandler) => done")
+		terminate()
+	}()
+	go func() {
+		select {
+		case <-s.handler.Done():
+			dlog.Client.Trace("<-s.handler.Done()")
+		case <-ctx.Done():
+			dlog.Client.Trace("<-ctx.Done()")
+		}
+		terminate()
+	}()
+
+	// Send all commands to client.
+	for _, command := range s.commands {
+		dlog.Client.Debug("Sending command to serverless server", command)
+		s.handler.SendMessage(command)
+	}
+
+	<-ctx.Done()
+	dlog.Client.Trace("s.handler.Shutdown()")
+	s.handler.Shutdown()
+	return nil
+}
diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go
index 567193a..7521c67 100644
--- a/internal/clients/grepclient.go
+++ b/internal/clients/grepclient.go
@@ -7,16 +7,19 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/omode"
 )
 
-// GrepClient searches a remote file for all lines matching a regular expression. Only the matching lines are displayed.
+// GrepClient searches a remote file for all lines matching a regular
+// expression. Only the matching lines are displayed.
 type GrepClient struct {
 	baseClient
 }
 
 // NewGrepClient creates a new grep client.
-func NewGrepClient(args Args) (*GrepClient, error) {
+func NewGrepClient(args config.Args) (*GrepClient, error) {
 	if args.RegexStr == "" {
 		return nil, errors.New("No regex specified, use '-regex' flag")
 	}
@@ -32,7 +35,6 @@ func NewGrepClient(args Args) (*GrepClient, error) {
 
 	c.init()
 	c.makeConnections(c)
-
 	return &c, nil
 }
 
@@ -40,11 +42,14 @@ func (c GrepClient) makeHandler(server string) handlers.Handler {
 	return handlers.NewClientHandler(server)
 }
 
-func (c GrepClient) makeCommands(options map[string]string) (commands []string) {
-	optionsStr := c.commandOptionsToString(options)
+func (c GrepClient) makeCommands() (commands []string) {
+	regex, err := c.Regex.Serialize()
+	if err != nil {
+		dlog.Client.FatalPanic(err)
+	}
 	for _, file := range strings.Split(c.What, ",") {
-		commands = append(commands, fmt.Sprintf("%s:%s %s %s", c.Mode.String(), optionsStr, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s:%s %s %s",
+			c.Mode.String(), c.Args.SerializeOptions(), file, regex))
 	}
-
 	return
 }
diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index 602a7ac..b520c25 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"bytes"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -8,8 +9,8 @@ import (
 	"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 {
@@ -17,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
 }
@@ -29,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():
@@ -56,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)
 		}
 	}
 
@@ -77,31 +87,32 @@ func (h *baseHandler) Read(p []byte) (n int, err error) {
 	case <-h.Done():
 		return 0, io.EOF
 	}
-
 	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
 // to the end user.
 func (h *baseHandler) handleHiddenMessage(message string) {
-	if strings.HasPrefix(message, ".syn close connection") {
-		h.SendMessage(".ack close connection")
+	switch {
+	case strings.HasPrefix(message, ".syn close connection"):
+		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()
+}
diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go
index 2bcb038..27ac85e 100644
--- a/internal/clients/handlers/clienthandler.go
+++ b/internal/clients/handlers/clienthandler.go
@@ -2,7 +2,7 @@ package handlers
 
 import (
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 )
 
 // ClientHandler is the basic client handler interface.
@@ -12,7 +12,7 @@ type ClientHandler struct {
 
 // NewClientHandler creates a new client handler.
 func NewClientHandler(server string) *ClientHandler {
-	logger.Debug(server, "Creating new client handler")
+	dlog.Client.Debug(server, "Creating new client handler")
 
 	return &ClientHandler{
 		baseHandler{
diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
index 0440706..47b594e 100644
--- a/internal/clients/handlers/healthhandler.go
+++ b/internal/clients/handlers/healthhandler.go
@@ -1,88 +1,56 @@
 package handlers
 
 import (
-	"errors"
-	"fmt"
-	"time"
+	"strings"
 
 	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
 )
 
-// HealthHandler implements the handler required for health checks.
+// HealthHandler is the handler used on the client side for running mapreduce
+// aggregations.
 type HealthHandler struct {
-	done *internal.Done
-	// Buffer of incoming data from server.
-	receiveBuf []byte
-	// To send commands to the server.
-	commands chan string
-	// To receive messages from the server.
-	receive chan<- string
-	// The remote server address
-	server string
-	// The return status.
-	status int
-}
-
-// NewHealthHandler returns a new health check handler.
-func NewHealthHandler(server string, receive chan<- string) *HealthHandler {
-	h := HealthHandler{
-		server:   server,
-		receive:  receive,
-		commands: make(chan string),
-		status:   -1,
-		done:     internal.NewDone(),
-	}
-
-	return &h
-}
-
-// Server returns the remote server name.
-func (h *HealthHandler) Server() string {
-	return h.server
-}
-
-// Status of the handler.
-func (h *HealthHandler) Status() int {
-	return h.status
-}
-
-// Done returns done channel of the handler.
-func (h *HealthHandler) Done() <-chan struct{} {
-	return h.done.Done()
-}
-
-// Shutdown the handler.
-func (h *HealthHandler) Shutdown() {
-	h.done.Shutdown()
-}
-
-// SendMessage sends a DTail command to the server.
-func (h *HealthHandler) SendMessage(command string) error {
-	select {
-	case h.commands <- fmt.Sprintf("%s;", command):
-	case <-time.NewTimer(time.Second * 10).C:
-		return errors.New("Timed out sending command " + command)
-	case <-h.Done():
+	baseHandler
+}
+
+// NewHealthHandler returns a new health client handler.
+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(),
+		},
 	}
-
-	return nil
 }
 
-// Server writes byte stream to client.
+// Read data from the dtail server via Writer interface.
 func (h *HealthHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		h.receiveBuf = append(h.receiveBuf, b)
-		if b == '\n' {
-			h.receive <- string(h.receiveBuf)
-			h.receiveBuf = h.receiveBuf[:0]
+		switch b {
+		case '\n', protocol.MessageDelimiter:
+			message := h.baseHandler.receiveBuf.String()
+			h.handleMessage(message)
+			h.baseHandler.receiveBuf.Reset()
+		default:
+			h.baseHandler.receiveBuf.WriteByte(b)
 		}
 	}
-
 	return len(p), nil
 }
 
-// Server reads byte stream from client.
-func (h *HealthHandler) Read(p []byte) (n int, err error) {
-	n = copy(p, []byte(<-h.commands))
-	return
+func (h *HealthHandler) handleMessage(message string) {
+	if len(message) > 0 && message[0] == '.' {
+		h.baseHandler.handleHiddenMessage(message)
+		return
+	}
+	s := strings.Split(message, protocol.FieldDelimiter)
+	message = s[len(s)-1]
+	if message == "OK" {
+		h.baseHandler.status = 0
+	}
 }
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
index fb71c8f..8718b35 100644
--- a/internal/clients/handlers/maprhandler.go
+++ b/internal/clients/handlers/maprhandler.go
@@ -4,21 +4,24 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"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"
 )
 
-// MaprHandler is the handler used on the client side for running mapreduce aggregations.
+// MaprHandler is the handler used on the client side for running mapreduce
+// aggregations.
 type MaprHandler struct {
 	baseHandler