merge

6 files changed, 37 insertions, 28 deletions

diff --git a/internal/server/continuous.go b/internal/server/continuous.go
index 583d136..f75c732 100644
--- a/internal/server/continuous.go
+++ b/internal/server/continuous.go
@@ -92,7 +92,7 @@ func (c *continuous) runJob(ctx context.Context, job config.Continuous) {
 	}
 
 	logger.Info(fmt.Sprintf("Starting job %s", job.Name))
-	status := client.Start(jobCtx, make(chan struct{}))
+	status := client.Start(jobCtx, make(chan string))
 	logMessage := fmt.Sprintf("Job exited with status %d", status)
 
 	if status != 0 {
diff --git a/internal/server/handlers/controlhandler.go b/internal/server/handlers/controlhandler.go
index 9a8eb75..1e17c78 100644
--- a/internal/server/handlers/controlhandler.go
+++ b/internal/server/handlers/controlhandler.go
@@ -41,10 +41,12 @@ func NewControlHandler(user *user.User) *ControlHandler {
 	return &h
 }
 
+// Shutdown the handler.
 func (h *ControlHandler) Shutdown() {
 	h.done.Shutdown()
 }
 
+// Done channel of the handler.
 func (h *ControlHandler) Done() <-chan struct{} {
 	return h.done.Done()
 }
@@ -90,9 +92,7 @@ func (h *ControlHandler) handleCommand(command string) {
 	case "health":
 		h.serverMessages <- "OK: DTail SSH Server seems fine"
 		h.serverMessages <- "done;"
-	case "debug":
-		h.serverMessages <- logger.Debug(h.user, "Receiving debug command", command, s)
 	default:
-		h.serverMessages <- logger.Warn(h.user, "Received unknown control command", command, s)
+		h.serverMessages <- logger.Error(h.user, "Received unknown control command", command, s)
 	}
 }
diff --git a/internal/server/handlers/readcommand.go b/internal/server/handlers/readcommand.go
index 0f9207d..5bab26f 100644
--- a/internal/server/handlers/readcommand.go
+++ b/internal/server/handlers/readcommand.go
@@ -25,37 +25,29 @@ func newReadCommand(server *ServerHandler, mode omode.Mode) *readCommand {
 	}
 }
 
-func (r *readCommand) Start(ctx context.Context, argc int, args []string) {
+func (r *readCommand) Start(ctx context.Context, argc int, args []string, retries int) {
 	re := regex.NewNoop()
 
 	if argc >= 4 {
 		deserializedRegex, err := regex.Deserialize(strings.Join(args[2:], " "))
 		if err != nil {
-			logger.Error(err)
 			r.server.sendServerMessage(logger.Error(r.server.user, commandParseWarning, err))
 			return
 		}
 		re = deserializedRegex
 	}
 	if argc < 3 {
-		r.server.sendServerMessage(logger.Warn(r.server.user, commandParseWarning, args, argc))
+		r.server.sendServerWarnMessage(logger.Warn(r.server.user, commandParseWarning, args, argc))
 		return
 	}
-	r.readGlob(ctx, args[1], re)
+	r.readGlob(ctx, args[1], re, retries)
 }
 
-func (r *readCommand) readGlob(ctx context.Context, glob string, re regex.Regex) {
+func (r *readCommand) readGlob(ctx context.Context, glob string, re regex.Regex, retries int) {
 	retryInterval := time.Second * 5
 	glob = filepath.Clean(glob)
 
-	maxRetries := 10
-	for {
-		maxRetries--
-		if maxRetries < 0 {
-			r.server.sendServerMessage(logger.Warn(r.server.user, "Giving up to read file(s)"))
-			return
-		}
-
+	for retryCount := 0; retryCount < retries; retryCount++ {
 		paths, err := filepath.Glob(glob)
 		if err != nil {
 			logger.Warn(r.server.user, glob, err)
@@ -65,7 +57,7 @@ func (r *readCommand) readGlob(ctx context.Context, glob string, re regex.Regex)
 
 		if numPaths := len(paths); numPaths == 0 {
 			logger.Error(r.server.user, "No such file(s) to read", glob)
-			r.server.sendServerMessage(logger.Warn(r.server.user, "Unable to read file(s), check server logs"))
+			r.server.sendServerWarnMessage(logger.Warn(r.server.user, "Unable to read file(s), check server logs"))
 			select {
 			case <-ctx.Done():
 				return
@@ -76,8 +68,11 @@ func (r *readCommand) readGlob(ctx context.Context, glob string, re regex.Regex)
 		}
 
 		r.readFiles(ctx, paths, glob, re, retryInterval)
-		break
+		return
 	}
+
+	r.server.sendServerWarnMessage(logger.Warn(r.server.user, "Giving up to read file(s)"))
+	return
 }
 
 func (r *readCommand) readFiles(ctx context.Context, paths []string, glob string, re regex.Regex, retryInterval time.Duration) {
@@ -97,7 +92,7 @@ func (r *readCommand) readFileIfPermissions(ctx context.Context, wg *sync.WaitGr
 
 	if !r.server.user.HasFilePermission(path, "readfiles") {
 		logger.Error(r.server.user, "No permission to read file", path, globID)
-		r.server.sendServerMessage(logger.Warn(r.server.user, "Unable to read file(s), check server logs"))
+		r.server.sendServerWarnMessage(logger.Warn(r.server.user, "Unable to read file(s), check server logs"))
 		return
 	}
 
@@ -161,6 +156,6 @@ func (r *readCommand) makeGlobID(path, glob string) string {
 		return pathParts[len(pathParts)-1]
 	}
 
-	r.server.sendServerMessage(logger.Error("Empty file path given?", path, glob))
+	r.server.sendServerWarnMessage(logger.Warn("Empty file path given?", path, glob))
 	return ""
 }
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index db917bd..185e7c2 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -43,6 +43,7 @@ type ServerHandler struct {
 	ackCloseReceived   chan struct{}
 	activeCommands     int32
 	activeReaders      int32
+	quiet              bool
 }
 
 // NewServerHandler returns the server handler.
@@ -70,10 +71,12 @@ func NewServerHandler(user *user.User, catLimiter, tailLimiter chan struct{}) *S
 	return &h
 }
 
+// Shutdown the handler.
 func (h *ServerHandler) Shutdown() {
 	h.done.Shutdown()
 }
 
+// Done channel of the handler.
 func (h *ServerHandler) Done() <-chan struct{} {
 	return h.done.Done()
 }
@@ -243,13 +246,19 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		commandFinished()
 		return
 	}
+	if quiet, ok := options["quiet"]; ok {
+		if quiet == "true" {
+			logger.Debug(h.user, "Enabling quiet mode")
+			h.quiet = true
+		}
+	}
 
 	switch commandName {
 	case "grep", "cat":
 		command := newReadCommand(h, omode.CatClient)
 		go func() {
 			h.incrementActiveReaders()
-			command.Start(ctx, argc, args)
+			command.Start(ctx, argc, args, 1)
 			readerFinished()
 			commandFinished()
 		}()
@@ -258,7 +267,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 		command := newReadCommand(h, omode.TailClient)
 		go func() {
 			h.incrementActiveReaders()
-			command.Start(ctx, argc, args)
+			command.Start(ctx, argc, args, 10)
 			readerFinished()
 			commandFinished()
 		}()
@@ -290,7 +299,7 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 
 func (h *ServerHandler) handleAckCommand(argc int, args []string) {
 	if argc < 3 {
-		h.sendServerMessage(logger.Warn(h.user, commandParseWarning, args, argc))
+		h.sendServerWarnMessage(logger.Warn(h.user, commandParseWarning, args, argc))
 		return
 	}
 	if args[1] == "close" && args[2] == "connection" {
@@ -309,6 +318,13 @@ func (h *ServerHandler) sendServerMessage(message string) {
 	h.send(h.serverMessageC(), message)
 }
 
+func (h *ServerHandler) sendServerWarnMessage(message string) {
+	if h.quiet {
+		return
+	}
+	h.send(h.serverMessageC(), message)
+}
+
 func (h *ServerHandler) serverMessageC() chan<- string {
 	return h.serverMessages
 }
diff --git a/internal/server/scheduler.go b/internal/server/scheduler.go
index 9d76a3b..a1e9e36 100644
--- a/internal/server/scheduler.go
+++ b/internal/server/scheduler.go
@@ -93,7 +93,7 @@ func (s *scheduler) runJobs(ctx context.Context) {
 		defer cancel()
 
 		logger.Info(fmt.Sprintf("Starting job %s", job.Name))
-		status := client.Start(jobCtx, make(chan struct{}))
+		status := client.Start(jobCtx, make(chan string))
 		logMessage := fmt.Sprintf("Job exited with status %d", status)
 
 		if status != 0 {
diff --git a/internal/server/server.go b/internal/server/server.go
index d8d43c9..a20737e 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -75,9 +75,7 @@ func (s *Server) Start(ctx context.Context) int {
 	go s.cont.start(ctx)
 	go s.listenerLoop(ctx, listener)
 
-	select {
-	case <-ctx.Done():
-	}
+	<-ctx.Done()
 
 	// For future use.
 	return 0