Introduce drun command, refactor code to use context package

1 files changed, 186 insertions, 335 deletions

diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index bed8609..3f0d6ce 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -1,17 +1,19 @@
 package handlers
 
 import (
+	"context"
+	"encoding/base64"
+	"errors"
 	"fmt"
 	"io"
 	"os"
-	"path/filepath"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/fs"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/line"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/mapr/server"
 	"github.com/mimecast/dtail/internal/omode"
 	user "github.com/mimecast/dtail/internal/user/server"
@@ -26,51 +28,33 @@ const (
 // the Bi-directional communication between SSH client and server.
 // This handler implements the handler of the SSH server.
 type ServerHandler struct {
-	// Local log file readers
-	fileReaders    []fs.FileReader
-	fileReadersMtx *sync.Mutex
-	// Channel for read lines.
-	lines chan fs.LineRead
-	// Only process log lines matching this regex.
-	regex string
-	// Server side mapr log aggregation.
-	aggregate *server.Aggregate
-	// Channel of aggregated log lines.
+	mutex              *sync.Mutex
+	lines              chan line.Line
+	regex              string
+	aggregate          *server.Aggregate
 	aggregatedMessages chan string
-	// Channel for server messages to be sent to the client.
-	serverMessages chan string
-	// Channel for hidden messages to be sent to the client.
-	hiddenMessages chan string
-	// The current payload sent to the client.
-	payload []byte
-	// The current server hostname.
-	hostname string
-	// The user connecting to dtail.
-	user *user.User
-	// To limit the server wide max amount of concurrent cats
-	catLimiter chan struct{}
-	// To limit the server wide max amount of concurrent tails
-	tailLimiter chan struct{}
-	// Server can tell handler to stop the handler.
-	stop chan struct{}
-	// Indicate that client responded to server with "ack stop connection"
-	ackStopReceived chan struct{}
-	// Stop timeout.
-	stopTimeout chan struct{}
+	serverMessages     chan string
+	payload            []byte
+	hostname           string
+	user               *user.User
+	catLimiter         chan struct{}
+	tailLimiter        chan struct{}
+	ackCloseReceived   chan struct{}
+	ctx                context.Context
+	done               chan struct{}
+	activeReaders      int
 }
 
 // NewServerHandler returns the server handler.
-func NewServerHandler(user *user.User, catLimiter chan struct{}, tailLimiter chan struct{}) *ServerHandler {
-	logger.Debug(user, "Creating tail handler")
+func NewServerHandler(ctx context.Context, user *user.User, catLimiter chan struct{}, tailLimiter chan struct{}) (*ServerHandler, <-chan struct{}) {
 	h := ServerHandler{
-		fileReadersMtx:     &sync.Mutex{},
-		lines:              make(chan fs.LineRead, 100),
+		ctx:                ctx,
+		done:               make(chan struct{}),
+		mutex:              &sync.Mutex{},
+		lines:              make(chan line.Line, 100),
 		serverMessages:     make(chan string, 10),
 		aggregatedMessages: make(chan string, 10),
-		hiddenMessages:     make(chan string, 10),
-		ackStopReceived:    make(chan struct{}),
-		stopTimeout:        make(chan struct{}),
-		stop:               make(chan struct{}),
+		ackCloseReceived:   make(chan struct{}),
 		catLimiter:         catLimiter,
 		tailLimiter:        tailLimiter,
 		regex:              ".",
@@ -85,37 +69,46 @@ func NewServerHandler(user *user.User, catLimiter chan struct{}, tailLimiter cha
 	s := strings.Split(fqdn, ".")
 	h.hostname = s[0]
 
-	return &h
+	return &h, h.done
 }
 
 // Read is to send data to the dtail client via Reader interface.
 func (h *ServerHandler) Read(p []byte) (n int, err error) {
 	for {
 		select {
+
 		case message := <-h.serverMessages:
+			if message[0] == '.' {
+				// Handle hidden message (don't display to the user, interpreted by dtail client)
+				wholePayload := []byte(fmt.Sprintf("%s\n", message))
+				n = copy(p, wholePayload)
+				return
+			}
+
+			// Handle normal server message (display to the user)
 			wholePayload := []byte(fmt.Sprintf("SERVER|%s|%s\n", h.hostname, message))
 			n = copy(p, wholePayload)
 			return
+
 		case message := <-h.aggregatedMessages:
+			// Send mapreduce-aggregated data as a message.
 			data := fmt.Sprintf("AGGREGATE|%s|%s\n", h.hostname, message)
-			//logger.Debug("Sending aggregation data", data)
 			wholePayload := []byte(data)
 			n = copy(p, wholePayload)
 			return
-		case message := <-h.hiddenMessages:
-			//logger.Debug(h.user, "Sending hidden message", message)
-			wholePayload := []byte(fmt.Sprintf(".%s\n", message))
-			n = copy(p, wholePayload)
-			return
+
 		case line := <-h.lines:
+			// Send normal file content data as a message.
 			serverInfo := []byte(fmt.Sprintf("REMOTE|%s|%3d|%v|%s|",
-				h.hostname, line.TransmittedPerc, line.Count, *line.GlobID))
+				h.hostname, line.TransmittedPerc, line.Count, line.SourceID))
 			wholePayload := append(serverInfo, line.Content[:]...)
 			n = copy(p, wholePayload)
 			return
+
 		case <-time.After(time.Second):
+			// Once in a while check whether we are done.
 			select {
-			case <-h.stop:
+			case <-h.ctx.Done():
 				return 0, io.EOF
 			default:
 			}
@@ -129,7 +122,7 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) {
 		switch c {
 		case ';':
 			commandStr := strings.TrimSpace(string(h.payload))
-			h.handleCommand(commandStr)
+			h.handleCommand(h.ctx, commandStr)
 			h.payload = nil
 		default:
 			h.payload = append(h.payload, c)
@@ -140,210 +133,167 @@ func (h *ServerHandler) Write(p []byte) (n int, err error) {
 	return
 }
 
-// Close the server handler.
-func (h *ServerHandler) Close() {
-	h.fileReadersMtx.Lock()
-	defer h.fileReadersMtx.Unlock()
+func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) {
+	logger.Debug(h.user, commandStr)
 
-	for _, reader := range h.fileReaders {
-		reader.Stop()
+	args, argc, err := h.handleProtocolVersion(strings.Split(commandStr, " "))
+	if err != nil {
+		h.send(h.serverMessages, logger.Error(h.user, err))
+		return
 	}
-	if h.aggregate != nil {
-		h.aggregate.Close()
+
+	args, argc, err = h.handleBase64(args, argc)
+	if err != nil {
+		h.send(h.serverMessages, logger.Error(h.user, err))
+		return
 	}
 
-	close(h.stop)
-}
+	if h.user.Name == config.ControlUser {
+		h.handleControlCommand(argc, args)
+		return
+	}
 
-func (h *ServerHandler) makeGlobID(path, glob string) string {
-	var idParts []string
-	pathParts := strings.Split(path, "/")
+	h.handleUserCommand(ctx, argc, args)
+}
 
-	for i, globPart := range strings.Split(glob, "/") {
-		if strings.Contains(globPart, "*") {
-			idParts = append(idParts, pathParts[i])
-		}
-	}
+func (h *ServerHandler) handleProtocolVersion(args []string) ([]string, int, error) {
+	argc := len(args)
 
-	if len(idParts) > 0 {
-		return strings.Join(idParts, "/")
+	if argc <= 2 || args[0] != "protocol" {
+		return args, argc, errors.New("unable to determine protocol version")
 	}
 
-	if len(pathParts) > 0 {
-		return pathParts[len(pathParts)-1]
+	if args[1] != version.ProtocolCompat {
+		err := fmt.Errorf("server with protool version '%s' but client with '%s', please update DTail", version.ProtocolCompat, args[1])
+		return args, argc, err
 	}
 
-	h.send(h.serverMessages, logger.Error("Empty file path given?", path, glob))
-	return ""
+	return args[2:], argc - 2, nil
 }
 
-func (h *ServerHandler) processFileGlob(mode omode.Mode, glob string, regex string) {
-	retryInterval := time.Second * 5
-	glob = filepath.Clean(glob)
-
-	errors := make(chan struct{})
-	stop := make(chan struct{})
-	defer close(stop)
+func (h *ServerHandler) handleBase64(args []string, argc int) ([]string, int, error) {
+	err := errors.New("Unable to decode client message")
 
-	go func() {
-		for {
-			select {
-			case <-errors:
-				h.send(h.serverMessages, logger.Warn(h.user, "Unable to read file(s), check server logs"))
-			case <-stop:
-				return
-			case <-h.stop:
-				return
-			}
-		}
-	}()
+	if argc != 2 || args[0] != "base64" {
+		return args, argc, err
+	}
 
-	maxRetries := 10
-	for {
-		maxRetries--
-		if maxRetries < 0 {
-			h.send(h.serverMessages, logger.Warn(h.user, "Giving up to read file(s)"))
-			h.internalClose()
-			return
-		}
+	decoded, err := base64.StdEncoding.DecodeString(args[1])
+	if err != nil {
+		return args, argc, err
+	}
+	decodedStr := string(decoded)
 
-		paths, err := filepath.Glob(glob)
-		if err != nil {
-			logger.Warn(h.user, glob, err)
-			time.Sleep(retryInterval)
-			continue
-		}
+	args = strings.Split(decodedStr, " ")
+	argc = len(decodedStr)
+	logger.Trace(h.user, "Base64 decoded received command", decodedStr, argc, args)
 
-		if numPaths := len(paths); numPaths == 0 {
-			logger.Error(h.user, "No such file(s) to read", glob)
-			select {
-			case errors <- struct{}{}:
-			case <-h.stop:
-				return
-			default:
-			}
-			time.Sleep(retryInterval)
-			continue
-		}
+	return args, argc, nil
+}
 
-		h.startReadingFiles(mode, paths, glob, regex, retryInterval, errors)
-		break
+func (h *ServerHandler) handleControlCommand(argc int, args []string) {
+	switch args[0] {
+	case "debug":
+		h.send(h.serverMessages, logger.Debug(h.user, "Receiving debug command", argc, args))
+	default:
+		logger.Warn(h.user, "Received unknown command", argc, args)
 	}
 }
 
-func (h *ServerHandler) startReadingFiles(mode omode.Mode, paths []string, glob string, regex string, retryInterval time.Duration, errors chan<- struct{}) {
-	var wg sync.WaitGroup
-	wg.Add(len(paths))
+func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []string) {
+	logger.Debug(h.user, "handleUserCommand", argc, args)
 
-	read := func(path string, wg *sync.WaitGroup) {
-		defer wg.Done()
-		globID := h.makeGlobID(path, glob)
+	switch args[0] {
+	case "grep", "cat":
+		command := newReadCommand(h, omode.CatClient)
+		h.incrementActiveReaders()
+		go func() {
+			command.Start(ctx, argc, args)
+			if h.decrementActiveReaders() == 0 {
+				h.shutdown()
+			}
+		}()
 
-		if !h.user.HasFilePermission(path) {
-			logger.Error(h.user, "No permission to read file", path, globID)
-			select {
-			case errors <- struct{}{}:
-			default:
+	case "tail":
+		command := newReadCommand(h, omode.TailClient)
+		h.incrementActiveReaders()
+		go func() {
+			command.Start(ctx, argc, args)
+			if h.decrementActiveReaders() == 0 {
+				h.shutdown()
 			}
+		}()
+
+	case "map":
+		command, aggregate, err := newMapCommand(h, argc, args)
+		if err != nil {
+			h.sendServerMessage(err.Error())
+			logger.Error(h.user, err)
 			return
 		}
 
-		h.startReadingFile(mode, path, globID, regex)
-	}
-
-	for _, path := range paths {
-		go read(path, &wg)
-	}
+		h.aggregate = aggregate
+		go func() {
+			command.Start(ctx, h.aggregatedMessages)
+			h.shutdown()
+		}()
+
+	case "run":
+		command := newRunCommand(h)
+		h.incrementActiveReaders()
+		go func() {
+			command.Start(ctx, argc, args)
+			if h.decrementActiveReaders() == 0 {
+				h.shutdown()
+			}
+		}()
 
-	wg.Wait()
-}
+	case "ack", ".ack":
+		h.handleAckCommand(argc, args)
 
-func (h *ServerHandler) startReadingFile(mode omode.Mode, path, globID, regex string) {
-	defer h.stopReadingFile(path)
-	logger.Info(h.user, "Start reading file", path, globID)
-
-	var reader fs.FileReader
-	switch mode {
-	case omode.TailClient:
-		reader = fs.NewTailFile(path, globID, h.serverMessages, h.tailLimiter)
-	case omode.GrepClient:
-		fallthrough
-	case omode.CatClient:
-		reader = fs.NewCatFile(path, globID, h.serverMessages, h.catLimiter)
 	default:
-		reader = fs.NewTailFile(path, globID, h.serverMessages, h.tailLimiter)
+		h.sendServerMessage(logger.Error(h.user, "Received unknown command", argc, args))
 	}
+}
 
-	h.fileReadersMtx.Lock()
-	h.fileReaders = append(h.fileReaders, reader)
-	h.fileReadersMtx.Unlock()
-
-	lines := h.lines
-	// Plugin mappreduce engine
-	if h.aggregate != nil {
-		lines = h.aggregate.Lines
+func (h *ServerHandler) handleAckCommand(argc int, args []string) {
+	if argc < 3 {
+		h.sendServerMessage(logger.Warn(h.user, commandParseWarning, args, argc))
+		return
 	}
-
-	for {
-		if err := reader.Start(lines, regex); err != nil {
-			logger.Error(h.user, path, globID, err)
-		}
-
-		select {
-		case <-h.stop:
-			return
-		default:
-			if !reader.Retry() {
-				return
-			}
-		}
-
-		time.Sleep(time.Second * 2)
-		logger.Info(path, globID, "Reading file again")
+	if args[1] == "close" && args[2] == "connection" {
+		close(h.ackCloseReceived)
 	}
 }
 
-func (h *ServerHandler) stopReadingFile(path string) {
-	logger.Info(h.user, "Stop reading file", path)
+func (h *ServerHandler) send(ch chan<- string, message string) {
+	select {
+	case ch <- message:
+	case <-h.ctx.Done():
+	}
+}
 
-	h.fileReadersMtx.Lock()
-	defer h.fileReadersMtx.Unlock()
+func (h *ServerHandler) sendServerMessage(message string) {
+	h.send(h.serverMessageC(), message)
+}
 
-	path = filepath.Clean(path)
-	var fileReaders []fs.FileReader
+func (h *ServerHandler) serverMessageC() chan<- string {
+	return h.serverMessages
+}
 
-	for _, reader := range h.fileReaders {
-		if reader.FilePath() == path {
-			reader.Stop()
-			continue
-		}
-		fileReaders = append(fileReaders, reader)
-	}
+func (h *ServerHandler) flush() {
+	logger.Debug(h.user, "flush()")
 
-	if len(fileReaders) == len(h.fileReaders) {
-		logger.Warn(h.user, "Didn't read file path", path)
-		return
+	if h.aggregate != nil {
+		h.aggregate.Flush()
 	}
 
-	h.fileReaders = fileReaders
-
-	if len(fileReaders) == 0 {
-		if h.aggregate != nil {
-			h.aggregate.Serialize()
-		}
-		h.allLinesSent()
+	unsentMessages := func() int {
+		return len(h.lines) + len(h.serverMessages) + len(h.aggregatedMessages)
 	}
-}
-
-func (h *ServerHandler) numUnsentMessages() int {
-	return len(h.lines) + len(h.serverMessages) + len(h.hiddenMessages) + len(h.aggregatedMessages)
-}
-
-func (h *ServerHandler) allLinesSent() {
-	defer h.internalClose()
 
 	for i := 0; i < 3; i++ {
-		if h.numUnsentMessages() == 0 {
+		if unsentMessages() == 0 {
 			logger.Debug(h.user, "All lines sent")
 			return
 		}
@@ -351,142 +301,43 @@ func (h *ServerHandler) allLinesSent() {
 		time.Sleep(time.Second)
 	}
 
-	logger.Warn(h.user, "Some lines remain unsent", h.numUnsentMessages())
+	logger.Warn(h.user, "Some lines remain unsent", unsentMessages())
 }
 
-// Handler decides to shutdown the connection, not the server itself.
-func (h *ServerHandler) internalClose() {
-	select {
-	case h.hiddenMessages <- "syn close connection":
-	case <-time.After(time.Second * 5):
-		logger.Debug(h.user, "Not waiting for ack close connection")
-		close(h.stopTimeout)
-		return
-	}
+func (h *ServerHandler) shutdown() {
+	logger.Debug(h.user, "shutdown()")
+	h.flush()
+
+	go func() {
+		select {
+		case h.serverMessageC() <- ".syn close connection":
+		case <-h.ctx.Done():
+		}
+	}()
 
 	select {
-	case <-h.Wait():
+	case <-h.ackCloseReceived:
 	case <-time.After(time.Second * 5):
-		logger.Debug(h.user, "Not waiting for ack close connection")
-		close(h.stopTimeout)
-	}
-}
-
-func (h *ServerHandler) handleCommand(commandStr string) {
-	logger.Info(h.user, commandStr)
-
-	args := strings.Split(commandStr, " ")
-	argc := len(args)
-
-	logger.Debug(h.user, "Received command", commandStr, argc, args)
-
-	if h.user.Name == config.ControlUser {
-		h.handleControlCommand(argc, args)
-		return
+		logger.Debug(h.user, "Shutdown timeout reached, enforcing shutdown")
+	case <-h.ctx.Done():
 	}
 
-	h.handleUserCommand(argc, args)
-}
-
-// Special (restricted) set of commands for anonymous ControlUser access.
-func (h *ServerHandler) handleControlCommand(argc int, args []string) {
-	switch args[0] {
-	case "ping":
-		h.send(h.hiddenMessages, "pong")
-	case "debug":
-		h.send(h.serverMessages, logger.Debug(h.user, "Receiving debug command", argc, args))
-	default:
-		logger.Warn(h.user, "Received unknown command", argc, args)
-	}
-}
-
-// Commands for authed users.
-func (h *ServerHandler) handleUserCommand(argc int, args []string) {
-	switch args[0] {
-	case "grep":
-		fallthrough
-	case "cat":
-		h.handleReadCommand(argc, args, omode.CatClient)
-	case "tail":
-		h.handleReadCommand(argc, args, omode.TailClient)
-	case "map":
-		h.handleMapCommand(argc, args)
-	case "ack":
-		h.handleAckCommand(argc, args)
-	case "ping":
-		h.send(h.hiddenMessages, "pong")
-	case "version":
-		h.send(h.serverMessages, fmt.Sprintf("Server version is "+version.String()))
-	case "debug":
-		h.send(h.serverMessages, logger.Debug(h.user, "Received debug command", argc, args))
+	select {
+	case h.done <- struct{}{}:
 	default:
-		h.send(h.serverMessages, logger.Warn(h.user, "Received unknown command", argc, args))
 	}
 }
 
-func (h *ServerHandler) handleReadCommand(argc int, args []string, mode omode.Mode) {
-	regex := "."
-	if argc >= 4 {
-		regex = args[3]
-	}
-	if argc < 3 {
-		h.send(h.serverMessages, logger.Warn(h.user, commandParseWarning, args, argc))
-		return
-	}
-	go h.processFileGlob(mode, args[1], regex)
+func (h *ServerHandler) incrementActiveReaders() {
+	// TODO: Use atomic counter variable instead, so we can get rid of the mutex
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	h.activeReaders++
 }
+func (h *ServerHandler) decrementActiveReaders() int {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+	h.activeReaders--
 
-func (h *ServerHandler) handleMapCommand(argc int, args []string) {
-	if argc < 2 {
-		h.send(h.serverMessages, logger.Warn(h.user, commandParseWarning, args, argc))
-		return
-	}
-
-	queryStr := strings.Join(args[1:], " ")
-	logger.Info(h.user, "Creating new mapr aggregator", queryStr)
-	aggregate, err := server.NewAggregate(h.aggregatedMessages, queryStr)
-
-	if err != nil {
-		h.send(h.serverMessages, logger.Error(h.user, err))
-		return
-	}
-
-	h.aggregate = aggregate
-}
-
-func (h *ServerHandler) handleAckCommand(argc int, args []string) {
-	if argc < 3 {
-		h.send(h.serverMessages, logger.Warn(h.user, commandParseWarning, args, argc))
-		return
-	}
-	if args[1] == "close" && args[2] == "connection" {
-		close(h.ackStopReceived)
-	}
-}
-
-func (h *ServerHandler) send(ch chan<- string, message string) {
-	select {
-	case ch <- message:
-	case <-h.stop:
-	}
-}
-
-// Wait (block) until server handler is closed or a timeout has exceeded.
-func (h *ServerHandler) Wait() <-chan struct{} {
-	wait := make(chan struct{})
-
-	go func() {
-		select {
-		case <-h.ackStopReceived:
-			logger.Debug(h.user, "Closing wait channel due to ACK stop received")
-			close(wait)
-		case <-h.stopTimeout:
-			logger.Debug(h.user, "Closing wait channel due to wait timeout")
-			close(wait)
-		case <-h.stop:
-			logger.Debug(h.user, "Closing wait channel due to stop")
-		}
-	}()
-
-	return wait
+	return h.activeReaders
 }