Introduce drun command, refactor code to use context package

6 files changed, 68 insertions, 53 deletions

diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go
index b02d369..1ec945d 100644
--- a/cmd/dcat/main.go
+++ b/cmd/dcat/main.go
@@ -1,12 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
+	"os"
 
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -27,7 +29,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 60
 	userName := user.Name()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
@@ -37,7 +38,6 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -54,9 +54,10 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
 	serverEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
+
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
@@ -67,14 +68,16 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              files,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 	}
 
 	client, err := clients.NewCatClient(args)
 	if err != nil {
 		panic(err)
 	}
-	client.Start()
+
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }
diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index d1a7d52..74a501f 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -1,12 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
+	"os"
 
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -28,7 +30,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 60
 	userName := user.Name()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
@@ -38,7 +39,6 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -56,9 +56,9 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
 	serverEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
@@ -69,9 +69,8 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              files,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 		Regex:             regex,
 	}
 
@@ -79,5 +78,8 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
-	client.Start()
+
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }
diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index 83dad50..f3f706a 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -1,13 +1,15 @@
 package main
 
 import (
+	"context"
 	"flag"
+	"os"
 
-	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -29,7 +31,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 900
 	userName := user.Name()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
@@ -39,7 +40,6 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -57,10 +57,10 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
 	serverEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
 
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
 	}
@@ -70,9 +70,8 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              files,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 		Mode:              omode.MapClient,
 	}
 
@@ -81,5 +80,7 @@ func main() {
 		panic(err)
 	}
 
-	client.Start()
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }
diff --git a/cmd/dexec/main.go b/cmd/drun/main.go
index 7a7ab1f..b1936d4 100644
--- a/cmd/dexec/main.go
+++ b/cmd/drun/main.go
@@ -1,12 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
+	"os"
 
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -15,11 +17,11 @@ import (
 // The evil begins here.
 func main() {
 	var cfgFile string
+	var command string
 	var connectionsPerCPU int
 	var debugEnable bool
 	var discovery string
 	var displayVersion bool
-	var command string
 	var noColor bool
 	var pprofEnable bool
 	var serversStr string
@@ -27,7 +29,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 60
 	userName := user.Name()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
@@ -37,11 +38,10 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
-	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&command, "command", "", "Command to run")
+	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
 
@@ -54,10 +54,10 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
 	serverEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
 
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
 	}
@@ -67,14 +67,16 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              command,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 	}
 
-	client, err := clients.NewExecClient(args)
+	client, err := clients.NewRunClient(args)
 	if err != nil {
 		panic(err)
 	}
-	client.Start()
+
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }
diff --git a/cmd/dserver/main.go b/cmd/dserver/main.go
index 489910b..aa209a8 100644
--- a/cmd/dserver/main.go
+++ b/cmd/dserver/main.go
@@ -1,13 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"os"
 	"time"
 
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/server"
 	"github.com/mimecast/dtail/internal/user"
@@ -24,7 +25,7 @@ func main() {
 	var shutdownAfter int
 	var sshPort int
 
-	userName := user.Name()
+	user.NoRootCheck()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
@@ -43,19 +44,23 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
+
 	serverEnable := true
 	silentEnable := false
 	nothingEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, nothingEnable)
-	defer logger.Stop()
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, nothingEnable)
 
 	if shutdownAfter > 0 {
 		go func() {
 			defer os.Exit(1)
 
 			logger.Info("Enabling auto shutdown timer", shutdownAfter)
-			time.Sleep(time.Duration(shutdownAfter) * time.Second)
-			logger.Info("Auto shutdown timer reached, shutting down now")
+			select {
+			case <-time.After(time.Duration(shutdownAfter) * time.Second):
+				logger.Info("Auto shutdown timer reached, shutting down now")
+			case <-ctx.Done():
+			}
 		}()
 	}
 
@@ -63,7 +68,8 @@ func main() {
 		pprof.Start()
 	}
 
-	logger.Info("Launching server", version.String(), userName)
 	sshServer := server.New()
-	sshServer.Start()
+	status := sshServer.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index 1bf77c7..76070ff 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -1,13 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"os"
 
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
@@ -32,7 +33,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 5
 	userName := user.Name()
 
 	flag.BoolVar(&checkHealth, "checkHealth", false, "Only check for server health")
@@ -43,7 +43,6 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -62,17 +61,18 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
+
 	if checkHealth {
 		healthClient, _ := clients.NewHealthClient(omode.HealthClient)
-		os.Exit(healthClient.Start())
+		os.Exit(healthClient.Start(ctx))
 	}
 
 	serverEnable := false
 	if checkHealth {
 		silentEnable = true
 	}
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
@@ -83,9 +83,8 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              files,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 		Regex:             regex,
 		Mode:              omode.TailClient,
 	}
@@ -104,5 +103,7 @@ func main() {
 		}
 	}
 
-	client.Start()
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }