Merge branch 'develop' of gitlab.devuk.mimecast.lan:Storage/dtail into develop

2 files changed, 16 insertions, 10 deletions

diff --git a/internal/server/handlers/controlhandler.go b/internal/server/handlers/controlhandler.go
index a33a78b..daa9835 100644
--- a/internal/server/handlers/controlhandler.go
+++ b/internal/server/handlers/controlhandler.go
@@ -87,6 +87,6 @@ func (h *ControlHandler) handleCommand(ctx context.Context, command string) {
 	case "debug":
 		h.serverMessages <- logger.Debug(h.user, "Receiving debug command", command, s)
 	default:
-		h.serverMessages <- logger.Warn(h.user, "Received unknown command", command, s)
+		h.serverMessages <- logger.Warn(h.user, "Received unknown control command", command, s)
 	}
 }
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index b638e69..739696c 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -176,11 +176,11 @@ func (h *ServerHandler) handleCommand(ctx context.Context, commandStr string) {
 			<-commandCtx.Done()
 			cancel()
 		}()
-		h.handleUserCommand(commandCtx, argc, args)
+		h.handleUserCommand(commandCtx, argc, args, timeout)
 		return
 	}
 
-	h.handleUserCommand(ctx, argc, args)
+	h.handleUserCommand(ctx, argc, args, timeout)
 }
 
 func (h *ServerHandler) handleProtocolVersion(args []string) ([]string, int, error) {
@@ -233,11 +233,11 @@ func (h *ServerHandler) handleControlCommand(argc int, args []string) {
 	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)
+		logger.Warn(h.user, "Received unknown control command", argc, args)
 	}
 }
 
-func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []string) {
+func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []string, timeout time.Duration) {
 	logger.Debug(h.user, "handleUserCommand", argc, args)
 
 	h.incrementActiveCommands()
@@ -248,7 +248,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 	}
 
 	splitted := strings.Split(args[0], ":")
-	command := splitted[0]
+	commandName := splitted[0]
 
 	options, err := readOptions(splitted[1:])
 	if err != nil {
@@ -257,7 +257,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		return
 	}
 
-	switch command {
+	switch commandName {
 	case "grep", "cat":
 		command := newReadCommand(h, omode.CatClient)
 		h.incrementActiveCommands()
@@ -295,7 +295,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		jobName, _ := options["jobName"]
 		logger.Debug(h.user, "run", options)
 
-		if val, ok := options["background"]; ok && val == "cancel" {
+		if val, ok := options["background"]; ok && (val == "cancel" || val == "stop") {
 			if err := h.background.Cancel(h.user.Name, jobName); err != nil {
 				h.sendServerMessage(logger.Error(h.user, err, jobName, args))
 			} else {
@@ -329,7 +329,13 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		wg.Add(1)
 
 		if background {
-			commandCtx, cancel := context.WithCancel(h.serverCtx)
+			if timeout == 0 {
+				// Set default background timeout.
+				timeout = time.Hour * 1
+			}
+			// Use a new context based on the server context, so that background job does not get
+			// terminated when handler/SSH connection terminates.
+			commandCtx, cancel := context.WithTimeout(h.serverCtx, timeout)
 
 			if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil {
 				h.sendServerMessage(logger.Error(h.user, err, jobName, args))
@@ -367,7 +373,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		finished()
 
 	default:
-		h.sendServerMessage(logger.Error(h.user, "Received unknown command", argc, args))
+		h.sendServerMessage(logger.Error(h.user, "Received unknown user command", commandName, argc, args, options))
 		finished()
 	}
 }