can list remote jobs and can also pass outer args to scripts

7 files changed, 273 insertions, 125 deletions

diff --git a/cmd/drun/main.go b/cmd/drun/main.go
index 1f0e0fe..3cc19e7 100644
--- a/cmd/drun/main.go
+++ b/cmd/drun/main.go
@@ -5,6 +5,7 @@ import (
 	"flag"
 	"io/ioutil"
 	"os"
+	"strings"
 
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
@@ -16,15 +17,15 @@ import (
 
 // The evil begins here.
 func main() {
+	var background string
 	var cfgFile string
 	var command string
 	var connectionsPerCPU int
 	var debugEnable bool
 	var discovery string
 	var displayVersion bool
+	var jobName string
 	var noColor bool
-	var background bool
-	var cancel bool
 	var serversStr string
 	var silentEnable bool
 	var sshPort int
@@ -38,14 +39,14 @@ func main() {
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
-	flag.BoolVar(&background, "background", false, "Command starts in background on the server")
-	flag.BoolVar(&cancel, "cancel", false, "Command will be cancelled if it runs in background")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.IntVar(&timeout, "timeout", 0, "Command execution timeout")
+	flag.StringVar(&background, "background", "", "Can be one of 'start', 'cancel', 'list' or empty")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&command, "command", "", "Command to run")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
+	flag.StringVar(&jobName, "name", "", "The job name (if run in background)")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
 
@@ -58,22 +59,26 @@ func main() {
 		version.PrintAndExit()
 	}
 
-	ctx := context.Background()
+	// TODO: Change other commands to use TODO contexts too.
+	ctx := context.TODO()
 	serverEnable := false
 
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
+	command, commandArgs := readCommand(command)
+
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		What:              readCommand(command),
+		What:              command,
+		Arguments:         commandArgs,
 		TrustAllHosts:     trustAllHosts,
 		Timeout:           timeout,
 	}
 
-	client, err := clients.NewRunClient(args, background, cancel)
+	client, err := clients.NewRunClient(args, background, jobName)
 	if err != nil {
 		panic(err)
 	}
@@ -83,15 +88,20 @@ func main() {
 	os.Exit(status)
 }
 
-func readCommand(command string) string {
-	if _, err := os.Stat(command); os.IsNotExist(err) {
-		return command
+func readCommand(command string) (string, []string) {
+	splitted := strings.Split(command, " ")
+
+	script := splitted[0]
+	if _, err := os.Stat(script); os.IsNotExist(err) {
+		var commandArgs []string
+		return command, commandArgs
 	}
+	commandArgs := splitted[1:]
 
-	bytes, err := ioutil.ReadFile(command)
+	bytes, err := ioutil.ReadFile(script)
 	if err != nil {
 		panic(err)
 	}
 
-	return string(bytes)
+	return string(bytes), commandArgs
 }
diff --git a/internal/clients/args.go b/internal/clients/args.go
index 4f3c86a..b4852d4 100644
--- a/internal/clients/args.go
+++ b/internal/clients/args.go
@@ -12,6 +12,7 @@ type Args struct {
 	ServersStr         string
 	UserName           string
 	What               string
+	Arguments          []string
 	Regex              string
 	TrustAllHosts      bool
 	Discovery          string
diff --git a/internal/clients/runclient.go b/internal/clients/runclient.go
index c2f6f62..543df15 100644
--- a/internal/clients/runclient.go
+++ b/internal/clients/runclient.go
@@ -1,32 +1,41 @@
 package clients
 
 import (
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
 	"fmt"
 	"runtime"
+	"strings"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/omode"
 )
 
 // RunClient is a client to run various commands on the server.
 type RunClient struct {
 	baseClient
-	background bool
-	cancel     bool
+	jobName    string
+	background string
 }
 
-// NewRunClient returns a new cat client.
-func NewRunClient(args Args, background, cancel bool) (*RunClient, error) {
+// NewRunClient returns a new run client to execute commands on the remote server.
+func NewRunClient(args Args, background, jobName string) (*RunClient, error) {
 	args.Mode = omode.RunClient
 
+	if jobName == "" {
+		jobName = hash(strings.Join(args.Arguments, " "))
+	}
+
 	c := RunClient{
 		baseClient: baseClient{
 			Args:       args,
 			throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()),
 			retry:      false,
 		},
+		jobName:    jobName,
 		background: background,
-		cancel:     cancel,
 	}
 
 	c.init(c)
@@ -39,20 +48,40 @@ func (c RunClient) makeHandler(server string) handlers.Handler {
 
 func (c RunClient) makeCommands() (commands []string) {
 	if c.Timeout > 0 {
-		commands = append(commands, fmt.Sprintf("timeout %d run%s %s", c.Timeout, c.flags(), c.What))
+		commands = append(commands, fmt.Sprintf("timeout %d run%s %s", c.Timeout, c.options(), c.What))
 		return
 	}
 
-	commands = append(commands, fmt.Sprintf("run%s %s", c.flags(), c.What))
+	commands = append(commands, fmt.Sprintf("run%s %s", c.options(), c.What))
+	logger.Debug(commands)
+
 	return
 }
 
-func (c RunClient) flags() string {
-	if c.background {
-		return ":background.start"
-	}
-	if c.cancel {
-		return ":background.cancel"
+func (c RunClient) options() string {
+	var sb strings.Builder
+
+	logger.Debug("options", fmt.Sprintf(":background=%s", c.background))
+	sb.WriteString(fmt.Sprintf(":background=%s", c.background))
+
+	logger.Debug("options", fmt.Sprintf(":jobName=%s", c.jobName))
+	sb.WriteString(fmt.Sprintf(":jobName=%s", c.jobName))
+
+	if len(c.Arguments) > 0 {
+		logger.Debug("options", fmt.Sprintf(":outerArgs=base64%%%s", strings.Join(c.Arguments, " ")))
+		sb.WriteString(fmt.Sprintf(":outerArgs=base64%%%s", encode64(strings.Join(c.Arguments, " "))))
 	}
-	return ""
+
+	return sb.String()
+}
+
+func encode64(str string) string {
+	return base64.StdEncoding.EncodeToString([]byte(str))
+}
+
+func hash(str string) string {
+	h := sha256.New()
+	h.Write([]byte(str))
+
+	return hex.EncodeToString(h.Sum(nil))
 }
diff --git a/internal/io/run/run.go b/internal/io/run/run.go
index 18e1eb9..186528d 100644
--- a/internal/io/run/run.go
+++ b/internal/io/run/run.go
@@ -31,11 +31,7 @@ func New(command string, args []string) Run {
 }
 
 // Start running the command.
-func (r Run) Start(ctx context.Context, lines chan<- line.Line) (pid int, ec int, err error) {
-	done := make(chan struct{})
-	defer close(done)
-
-	ec = 255
+func (r Run) StartBackground(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, lines chan<- line.Line) (pid int, err error) {
 	pid = -1
 
 	if len(r.args) > 0 {
@@ -45,48 +41,66 @@ func (r Run) Start(ctx context.Context, lines chan<- line.Line) (pid int, ec int
 		logger.Debug(r.command)
 		r.cmd = exec.CommandContext(ctx, r.command)
 	}
+
 	// Create a new process group, so that kill() will only kill this command + pgroup.
 	r.cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 
 	stdoutPipe, myErr := r.cmd.StdoutPipe()
 	if err != nil {
+		wg.Done()
 		err = myErr
 		return
 	}
 
 	stderrPipe, myErr := r.cmd.StderrPipe()
 	if myErr != nil {
+		wg.Done()
 		err = myErr
 		return
 	}
 
 	if myErr := r.cmd.Start(); err != nil {
+		wg.Done()
 		err = myErr
 		return
 	}
 
 	if r.cmd.Process != nil {
 		pid = r.cmd.Process.Pid
-		ec = 0
 	}
-	go r.killPgroup(ctx, done, pid)
 
-	var wg sync.WaitGroup
-	wg.Add(2)
+	commandExited := make(chan struct{})
 
-	go r.pipeToLines(done, &wg, pid, stdoutPipe, "STDOUT", lines)
-	go r.pipeToLines(done, &wg, pid, stderrPipe, "STDERR", lines)
+	var pipeWg sync.WaitGroup
+	pipeWg.Add(2)
 
-	if err = r.cmd.Wait(); err != nil {
-		if exitError, ok := err.(*exec.ExitError); ok {
-			ec = exitError.ExitCode()
+	go r.killPgroup(ctx, commandExited, pid)
+	go r.pipeToLines(commandExited, &pipeWg, pid, stdoutPipe, "STDOUT", lines)
+	go r.pipeToLines(commandExited, &pipeWg, pid, stderrPipe, "STDERR", lines)
+
+	go func() {
+		exitCode := 255
+		if waitErr := r.cmd.Wait(); waitErr != nil {
+			if exitError, ok := waitErr.(*exec.ExitError); ok {
+				exitCode = exitError.ExitCode()
+			}
 		}
-	}
+		ec <- exitCode
+
+		// Tell pipes we are done
+		close(commandExited)
+		// Wait for process group to be killed
+		<-r.pgroupKilled
+		// Wait for the pipes to flush the contents
+		pipeWg.Wait()
+		// Now the job is truly done
+		wg.Done()
+	}()
 
 	return
 }
 
-func (r Run) pipeToLines(done chan struct{}, wg *sync.WaitGroup, pid int, reader io.Reader, what string, lines chan<- line.Line) {
+func (r Run) pipeToLines(commandExited chan struct{}, wg *sync.WaitGroup, pid int, reader io.Reader, what string, lines chan<- line.Line) {
 	defer wg.Done()
 	bufReader := bufio.NewReader(reader)
 
@@ -102,7 +116,7 @@ func (r Run) pipeToLines(done chan struct{}, wg *sync.WaitGroup, pid int, reader
 			lineStr, err = bufReader.ReadString('\n')
 		}
 		select {
-		case <-done:
+		case <-commandExited:
 			return
 		default:
 		}
@@ -110,12 +124,7 @@ func (r Run) pipeToLines(done chan struct{}, wg *sync.WaitGroup, pid int, reader
 	}
 }
 
-// PgroupKilled identifies whether all subprocesses are killed or not.
-func (r Run) PgroupKilled() <-chan struct{} {
-	return r.pgroupKilled
-}
-
-func (r Run) killPgroup(ctx context.Context, done chan struct{}, pid int) {
+func (r Run) killPgroup(ctx context.Context, commandExited chan struct{}, pid int) {
 	if pid == -1 {
 		close(r.pgroupKilled)
 		return
@@ -125,7 +134,7 @@ func (r Run) killPgroup(ctx context.Context, done chan struct{}, pid int) {
 		// Kill process group when done
 		select {
 		case <-ctx.Done():
-		case <-done:
+		case <-commandExited:
 		}
 		syscall.Kill(-pgid, syscall.SIGKILL)
 		close(r.pgroupKilled)
diff --git a/internal/server/background/background.go b/internal/server/background/background.go
index 05a502f..51ef052 100644
--- a/internal/server/background/background.go
+++ b/internal/server/background/background.go
@@ -3,12 +3,16 @@ package background
 import (
 	"context"
 	"errors"
+	"fmt"
+	"strings"
 	"sync"
+
+	"github.com/mimecast/dtail/internal/io/logger"
 )
 
 type job struct {
 	cancel context.CancelFunc
-	done   <-chan struct{}
+	wg     *sync.WaitGroup
 }
 
 // Background specifies a job or command run in background on server side.
@@ -27,43 +31,81 @@ func New() Background {
 }
 
 // Add a background job.
-func (b Background) Add(name string, cancel context.CancelFunc, done <-chan struct{}) error {
+func (b Background) Add(userName, jobName string, cancel context.CancelFunc, wg *sync.WaitGroup) error {
+	key := b.key(userName, jobName)
+
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
 
-	if _, ok := b.jobs[name]; ok {
+	if _, ok := b.jobs[key]; ok {
 		return errors.New("job already exists")
 	}
 
-	b.jobs[name] = job{cancel, done}
+	b.jobs[key] = job{cancel, wg}
+
+	// Clean up background job database.
+	go func() {
+		wg.Wait()
+		b.cancel(key)
+	}()
+
 	return nil
 }
 
 // Cancel a background job.
-func (b Background) Cancel(name string) error {
-	job, ok := b.get(name)
+func (b Background) Cancel(userName, jobName string) error {
+	return b.cancel(b.key(userName, jobName))
+}
+
+func (b Background) cancel(key string) error {
+	job, ok := b.get(key)
 	if !ok {
 		return errors.New("no job to cancel")
 	}
 
 	job.cancel()
-	<-job.done
-	b.delete(name)
+	job.wg.Wait()
+	b.delete(key)
 
 	return nil
 }
 
-func (b Background) get(name string) (job, bool) {
+// ListJobsC returns a channel listing all jobs of the given user.
+func (b Background) ListJobsC(userName string) <-chan string {
+	ch := make(chan string)
+
+	go func() {
+		defer close(ch)
+
+		b.mutex.Lock()
+		defer b.mutex.Unlock()
+
+		for k, _ := range b.jobs {
+			logger.Debug("ListJobsC", k, userName)
+			if strings.HasPrefix(k, fmt.Sprintf("%s.", userName)) {
+				ch <- k
+			}
+		}
+	}()
+
+	return ch
+}
+
+func (b Background) get(key string) (job, bool) {
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
 
-	job, ok := b.jobs[name]
+	job, ok := b.jobs[key]
 	return job, ok
 }
 
-func (b Background) delete(name string) {
+func (b Background) delete(key string) {
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
 
-	delete(b.jobs, name)
+	delete(b.jobs, key)
+}
+
+func (Background) key(userName, jobName string) string {
+	return fmt.Sprintf("%s.%s", userName, jobName)
 }
diff --git a/internal/server/handlers/runcommand.go b/internal/server/handlers/runcommand.go
index 95db52f..8e5895b 100644
--- a/internal/server/handlers/runcommand.go
+++ b/internal/server/handlers/runcommand.go
@@ -2,11 +2,13 @@ package handlers
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
 	"os/exec"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mimecast/dtail/internal/config"
@@ -25,26 +27,39 @@ func newRunCommand(server *ServerHandler) runCommand {
 	}
 }
 
-func (r runCommand) Start(ctx context.Context, argc int, args []string) {
+func (r runCommand) StartBackground(ctx context.Context, wg *sync.WaitGroup, argc int, args, outerArgs []string) error {
 	if argc < 2 {
-		r.server.sendServerMessage(logger.Warn(r.server.user, commandParseWarning, args, argc))
-		return
+		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") {
-		r.startScript(ctx, command)
-		return
+		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
 	}
 
-	r.start(ctx, strings.TrimSpace(command))
+	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, script string) {
+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) {
-		logger.Error(r.server.user, err)
-		r.server.sendServerMessage(logger.Error(r.server.user, "Unable to execute command(s), check server logs"))
-		return
+		return -1, err
 	}
 
 	timestamp := time.Now().UnixNano()
@@ -55,45 +70,42 @@ func (r runCommand) startScript(ctx context.Context, script string) {
 
 	script = fmt.Sprintf("#!/bin/sh\n%s", script)
 	if err := ioutil.WriteFile(scriptPath, []byte(script), 0700); err != nil {
-		logger.Error(r.server.user, err)
-		r.server.sendServerMessage(logger.Error(r.server.user, "Unable to execute command(s), check server logs"))
-		return
+		return -1, err
 	}
 
-	r.start(ctx, scriptPath)
-	os.Remove(scriptPath)
+	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, command string) {
+func (r runCommand) start(ctx context.Context, wg *sync.WaitGroup, ec chan<- int, command string, outerArgs []string) (int, error) {
 	if len(command) == 0 {
-		return
+		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 {
-		logger.Error(r.server.user, err)
-		r.server.sendServerMessage(logger.Error(r.server.user, "Unable to execute command(s), check server logs"))
-		r.server.sendServerMessage(".run exitstatus 255")
-		return
+		return -1, err
 	}
 
 	if !r.server.user.HasFilePermission(qualifiedPath, "runcommands") {
-		logger.Error(r.server.user, "No permission to execute path", qualifiedPath)
-		r.server.sendServerMessage(logger.Error(r.server.user, "Unable to execute command(s), check server logs"))
-		r.server.sendServerMessage(".run exitstatus 255")
-		return
+		return -1, fmt.Errorf("No permission to execute path: %s", qualifiedPath)
 	}
 
 	r.run = run.New(qualifiedPath, args)
-	pid, ec, _ := r.run.Start(ctx, r.server.lines)
-
-	r.server.sendServerMessage(fmt.Sprintf(".run exitstatus %d", ec))
-	r.server.sendServerMessage(logger.Info(fmt.Sprintf("Process %d exited with status %d", pid, ec)))
-
-	logger.Debug(r.server.user, "Waiting for Pgroup to be killed")
-	<-r.run.PgroupKilled()
-	logger.Debug(r.server.user, "Pgroup killed")
+	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 01e4054..819cddd 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -2,9 +2,7 @@ package handlers
 
 import (
 	"context"
-	"crypto/sha256"
 	"encoding/base64"
-	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -251,7 +249,14 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 
 	splitted := strings.Split(args[0], ":")
 	command := splitted[0]
-	flags := splitted[1:]
+
+	// TODO: Refactor: Create an "options" clase, combine makeOptions and readOptions there.
+	options, err := readOptions(splitted[1:])
+	if err != nil {
+		h.sendServerMessage(logger.Error(h.user, err))
+		finished()
+		return
+	}
 
 	switch command {
 	case "grep", "cat":
@@ -287,10 +292,12 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 	case "run":
 		// TODO: Refactor this "run" case, move code to runcommand.go
 		command := newRunCommand(h)
-		jobName := fmt.Sprintf("%s%%%s", h.user.Name, hash(strings.Join(args[1:], " ")))
 
-		if contains(flags, "background.cancel") {
-			if err := h.background.Cancel(jobName); err != nil {
+		jobName, _ := options["jobName"]
+		logger.Debug(h.user, "run", options)
+
+		if val, ok := options["background"]; ok && val == "cancel" {
+			if err := h.background.Cancel(h.user.Name, jobName); err != nil {
 				h.sendServerMessage(logger.Error(h.user, err, jobName, args))
 			} else {
 				h.sendServerMessage(logger.Info(h.user, "job cancelled", jobName))
@@ -299,37 +306,64 @@ func (h *ServerHandler) handleUserCommand(ctx context.Context, argc int, args []
 			return
 		}
 
-		done := make(chan struct{})
+		if val, ok := options["background"]; ok && val == "list" {
+			h.sendServerMessage("Listing jobs")
+			count := 0
+			for jobName := range h.background.ListJobsC(h.user.Name) {
+				h.sendServerMessage(jobName)
+				count++
+			}
+			h.sendServerMessage(fmt.Sprintf("Found %d jobs", count))
+			finished()
+			return
+		}
+
+		str, _ := options["outerArgs"]
+		outerArgs := strings.Split(str, " ")
+
+		var background bool
+		if val, ok := options["background"]; ok && val == "start" {
+			background = true
+		}
+
+		var wg sync.WaitGroup
+		wg.Add(1)
 
-		if contains(flags, "background.start") {
+		if background {
 			commandCtx, cancel := context.WithCancel(h.serverCtx)
+
 			// TODO: For background jobs dont attempt to send data to dtail client as there might be no SSH connection
-			if err := h.background.Add(jobName, cancel, done); err != nil {
+			if err := h.background.Add(h.user.Name, jobName, cancel, &wg); err != nil {
 				h.sendServerMessage(logger.Error(h.user, err, jobName, args))
 				finished()
 				return
 			}
+			ctx = commandCtx
+		}
 
-			go func() { h.globalServerWaitFor <- struct{}{} }()
-			go func() {
-				command.Start(commandCtx, argc, args)
-				close(done)
-				<-h.globalServerWaitFor
-			}()
-
-			h.sendServerMessage(logger.Info(h.user, jobName, "job started in background"))
+		if err := command.StartBackground(ctx, &wg, argc, args, outerArgs); err != nil {
+			h.sendServerMessage(logger.Error(h.user, "Unable to execute command", argc, args, err))
 			finished()
 			return
 		}
 
+		// Make sure that server waits for all sub-processes to finish on shutdown
 		go func() { h.globalServerWaitFor <- struct{}{} }()
 		go func() {
-			command.Start(ctx, argc, args)
-			close(done)
+			wg.Wait()
 			<-h.globalServerWaitFor
-			finished()
 		}()
 
+		if background {
+			h.sendServerMessage(logger.Info(h.user, jobName, "job started in background"))
+			finished()
+			return
+		}
+
+		// Command run in foreground, wait for it to complete before finishing the connection.
+		wg.Wait()
+		finished()
+
 	case "ack", ".ack":
 		h.handleAckCommand(argc, args)
 		finished()
@@ -427,17 +461,28 @@ func (h *ServerHandler) decrementActiveCommands() int {
 	return h.activeCommands
 }
 
-func contains(haystack []string, needle string) bool {
-	for _, str := range haystack {
-		if str == needle {
-			return true
+func readOptions(opts []string) (map[string]string, error) {
+	options := make(map[string]string, len(opts))
+
+	for _, o := range opts {
+		kv := strings.SplitN(o, "=", 2)
+		if len(kv) != 2 {
+			return options, fmt.Errorf("Unable to parse options: %v", kv)
 		}
+		key := kv[0]
+		val := kv[1]
+
+		if strings.HasPrefix(val, "base64%") {
+			s := strings.SplitN(val, "%", 2)
+			decoded, err := base64.StdEncoding.DecodeString(s[1])
+			if err != nil {
+				return options, err
+			}
+			val = string(decoded)
+		}
+
+		options[key] = val
 	}
-	return false
-}
 
-func hash(str string) string {
-	h := sha256.New()
-	h.Write([]byte(str))
-	return hex.EncodeToString(h.Sum(nil))
+	return options, nil
 }