code cleanup and minor refactorings

3 files changed, 26 insertions, 159 deletions

diff --git a/internal/server/handlers/runcommand.go b/internal/server/handlers/runcommand.go
deleted file mode 100644
index 8e5895b..0000000
--- a/internal/server/handlers/runcommand.go
+++ /dev/null
@@ -1,111 +0,0 @@
-package handlers
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"io/ioutil"
-	"os"
-	"os/exec"
-	"strings"
-	"sync"
-	"time"
-
-	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/io/run"
-)
-
-type runCommand struct {
-	server *ServerHandler
-	run    run.Run
-}
-
-func newRunCommand(server *ServerHandler) runCommand {
-	return runCommand{
-		server: server,
-	}
-}
-
-func (r runCommand) StartBackground(ctx context.Context, wg *sync.WaitGroup, argc int, args, outerArgs []string) error {
-	if argc < 2 {
-		return fmt.Errorf("%s: args:%v argc:%d", commandParseWarning, args, argc)
-	}
-
-	ec := make(chan int, 1)
-	var pid int
-	var err error
-
-	command := strings.Join(args[1:], " ")
-	if strings.Contains(command, ";") || strings.Contains(command, "\n") {
-		if pid, err = r.startScript(ctx, wg, ec, command, outerArgs); err != nil {
-			r.server.sendServerMessage(".run exitstatus 255")
-			return err
-		}
-		return nil
-	}
-
-	if pid, err = r.start(ctx, wg, ec, strings.TrimSpace(command), outerArgs); err != nil {
-		r.server.sendServerMessage(".run exitstatus 255")
-		return err
-	}
-
-	exitCode := <-ec
-	r.server.sendServerMessage(fmt.Sprintf(".run exitstatus %d", exitCode))
-	r.server.sendServerMessage(logger.Info(fmt.Sprintf("Process %d exited with status %d", pid, exitCode)))
-
-	return nil
-}
-
-func (r runCommand) startScript(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, script string, outerArgs []string) (int, error) {
-	if _, err := os.Stat(config.Common.TmpDir); os.IsNotExist(err) {
-		return -1, err
-	}
-
-	timestamp := time.Now().UnixNano()
-	scriptPath := fmt.Sprintf("%s/%s_%v.sh", config.Common.TmpDir, r.server.user.Name, timestamp)
-
-	// TODO: On dserver startup delete all previously written scripts (there might be left overs due to a crash or so)
-	logger.Debug(r.server.user, "Writing temp script", scriptPath)
-
-	script = fmt.Sprintf("#!/bin/sh\n%s", script)
-	if err := ioutil.WriteFile(scriptPath, []byte(script), 0700); err != nil {
-		return -1, err
-	}
-
-	pid, err := r.start(ctx, wg, ec, scriptPath, outerArgs)
-	go func() {
-		wg.Wait()
-		logger.Debug("Deleting script", scriptPath)
-		os.Remove(scriptPath)
-	}()
-
-	return pid, err
-}
-
-func (r runCommand) start(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, command string, outerArgs []string) (int, error) {
-	if len(command) == 0 {
-		return -1, errors.New("Empty command provided")
-	}
-
-	splitted := strings.Split(command, " ")
-	path := splitted[0]
-	args := splitted[1:]
-	args = append(args, outerArgs...)
-
-	qualifiedPath, err := exec.LookPath(path)
-	if err != nil {
-		return -1, err
-	}
-
-	if !r.server.user.HasFilePermission(qualifiedPath, "runcommands") {
-		return -1, fmt.Errorf("No permission to execute path: %s", qualifiedPath)
-	}
-
-	r.run = run.New(qualifiedPath, args)
-	pid, err := r.run.StartBackground(ctx, wg, ec, r.server.lines)
-	if err != nil {
-		return pid, err
-	}
-	return pid, nil
-}
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 5cf8041..3d1a53d 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -29,36 +29,34 @@ const (
 // the Bi-directional communication between SSH client and server.
 // This handler implements the handler of the SSH server.
 type ServerHandler struct {
-	done                *internal.Done
-	lines               chan line.Line
-	regex               string
-	aggregate           *server.Aggregate
-	aggregatedMessages  chan string
-	serverMessages      chan string
-	payload             []byte
-	hostname            string
-	user                *user.User
-	catLimiter          chan struct{}
-	tailLimiter         chan struct{}
-	globalServerWaitFor chan struct{}
-	ackCloseReceived    chan struct{}
-	activeCommands      int32
-	activeReaders       int32
+	done               *internal.Done
+	lines              chan line.Line
+	regex              string
+	aggregate          *server.Aggregate
+	aggregatedMessages chan string
+	serverMessages     chan string
+	payload            []byte
+	hostname           string
+	user               *user.User
+	catLimiter         chan struct{}
+	tailLimiter        chan struct{}
+	ackCloseReceived   chan struct{}
+	activeCommands     int32
+	activeReaders      int32
 }
 
 // NewServerHandler returns the server handler.
-func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}) *ServerHandler {
+func NewServerHandler(user *user.User, catLimiter, tailLimiter chan struct{}) *ServerHandler {
 	h := ServerHandler{
-		done:                internal.NewDone(),
-		lines:               make(chan line.Line, 100),
-		serverMessages:      make(chan string, 10),
-		aggregatedMessages:  make(chan string, 10),
-		ackCloseReceived:    make(chan struct{}),
-		catLimiter:          catLimiter,
-		tailLimiter:         tailLimiter,
-		globalServerWaitFor: globalServerWaitFor,
-		regex:               ".",
-		user:                user,
+		done:               internal.NewDone(),
+		lines:              make(chan line.Line, 100),
+		serverMessages:     make(chan string, 10),
+		aggregatedMessages: make(chan string, 10),
+		ackCloseReceived:   make(chan struct{}),
+		catLimiter:         catLimiter,
+		tailLimiter:        tailLimiter,
+		regex:              ".",
+		user:               user,
 	}
 
 	fqdn, err := os.Hostname()
diff --git a/internal/server/server.go b/internal/server/server.go
index 31fa85d..a20737e 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -7,7 +7,6 @@ import (
 	"io"
 	"net"
 	"strings"
-	"time"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
@@ -33,9 +32,6 @@ type Server struct {
 	sched *scheduler
 	// Mointor log files for pattern (if configured)
 	cont *continuous
-	// Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed.
-	// TODO: Remove this counter.
-	shutdownWaitFor chan struct{}
 }
 
 // New returns a new server.
@@ -46,7 +42,6 @@ func New() *Server {
 		sshServerConfig: &gossh.ServerConfig{},
 		catLimiter:      make(chan struct{}, config.Server.MaxConcurrentCats),
 		tailLimiter:     make(chan struct{}, config.Server.MaxConcurrentTails),
-		shutdownWaitFor: make(chan struct{}, 1000),
 		sched:           newScheduler(),
 		cont:            newContinuous(),
 	}
@@ -80,27 +75,12 @@ func (s *Server) Start(ctx context.Context) int {
 	go s.cont.start(ctx)
 	go s.listenerLoop(ctx, listener)
 
-	select {
-	case <-ctx.Done():
-		// Wait until all commands/jobs/children are no more!
-		s.wait()
-	}
+	<-ctx.Done()
 
 	// For future use.
 	return 0
 }
 
-func (s *Server) wait() {
-	for {
-		num := len(s.shutdownWaitFor)
-		logger.Debug("Waiting for stuff to finish", num)
-		if num <= 0 {
-			return
-		}
-		time.Sleep(time.Second)
-	}
-}
-
 func (s *Server) listenerLoop(ctx context.Context, listener net.Listener) {
 	logger.Debug("Starting listener loop")
 
@@ -180,7 +160,7 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch
 			case config.ControlUser:
 				handler = handlers.NewControlHandler(user)
 			default:
-				handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor)
+				handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter)
 			}
 
 			terminate := func() {