Use cobra to improve CLI

README.md

@@ -30,22 +30,26 @@ established successfully
 
 ```bash
 % ./tcpgoon --help
-Usage of ./tcpgoon:
+tcpgoon tests concurrent connections towards a server listening on a TCP port
+
+Usage:
+  tcpgoon [flags] <host> <port>
+
+Flags:
   -y, --assume-yes         Force execution without asking for confirmation
   -c, --connections int    Number of connections you want to open (default 100)
-  -d, --debug              Print debugging information to the standard error
-  -t, --dial-timeout int   Connection dialing timeout, in s (default 5)
-  -h, --host string        Host you want to open tcp connections against (Required)
+  -d, --dial-timeout int   Connection dialing timeout, in ms (default 5000)
+  -h, --help               help for tcpgoon
   -i, --interval int       Interval, in seconds, between stats updates (default 1)
-  -p, --port int           Port you want to open tcp connections against (Required)
   -s, --sleep int          Time you want to sleep between connections, in ms (default 10)
+  -v, --verbose            Print debugging information to the standard error
 ```
 
 ## Examples
 
 Successful execution (connections were opened as expected):
 ```bash
-% ./tcpgoon --host myhttpsamplehost.com --port 80 --connections 10 --sleep 999 -y 
+% ./tcpgoon myhttpsamplehost.com 80 --connections 10 --sleep 999 -y 
 Total: 10, Dialing: 0, Established: 0, Closed: 0, Error: 0, NotInitiated: 10
 Total: 10, Dialing: 1, Established: 1, Closed: 0, Error: 0, NotInitiated: 8
 Total: 10, Dialing: 1, Established: 2, Closed: 0, Error: 0, NotInitiated: 7
@@ -66,7 +70,7 @@ Response time stats for 10 established connections min/avg/max/dev = 17.929ms/19
 
 Partially succeeded execution (mix of successes and errors against the target):
 ```bash
-% ./tcpgoon --host myhttpsamplehost.com --port 8080 --connections 10 --sleep 999 -y
+% ./tcpgoon myhttpsamplehost.com 8080 --connections 10 --sleep 999 -y
 Total: 10, Dialing: 0, Established: 0, Closed: 0, Error: 0, NotInitiated: 10
 Total: 10, Dialing: 0, Established: 1, Closed: 0, Error: 0, NotInitiated: 9
 Total: 10, Dialing: 0, Established: 2, Closed: 0, Error: 0, NotInitiated: 8
@@ -93,7 +97,7 @@ Time to error stats for 8 failed connections min/avg/max/dev = 5.000819s/5.00249
 
 Unsuccessful execution (unable to open connections against the destination host:port):
 ```bash
-% ./tcpgoon --host myhttpsamplehost.com --port 81 --connections 10 --sleep 999 -y
+% ./tcpgoon myhttpsamplehost.com 81 --connections 10 --sleep 999 -y
 Total: 10, Dialing: 0, Established: 0, Closed: 0, Error: 0, NotInitiated: 10
 Total: 10, Dialing: 2, Established: 0, Closed: 0, Error: 0, NotInitiated: 8
 Total: 10, Dialing: 3, Established: 0, Closed: 0, Error: 0, NotInitiated: 7

cmd/tcpgoon.go

@@ -0,0 +1,100 @@
+package cmd
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+
+	"github.com/dachad/tcpgoon/cmdutil"
+	"github.com/dachad/tcpgoon/debugging"
+	"github.com/dachad/tcpgoon/mtcpclient"
+	"github.com/dachad/tcpgoon/tcpclient"
+	"github.com/spf13/cobra"
+)
+
+type tcpgoonParams struct {
+	hostPtr              string
+	portPtr              int
+	numberConnectionsPtr int
+	delayPtr             int
+	connDialTimeoutPtr   int
+	debugPtr             bool
+	reportingIntervalPtr int
+	assumeyesPtr         bool
+}
+
+var params tcpgoonParams
+
+var rootCmd = &cobra.Command{
+	Use:   "tcpgoon [flags] <host> <port>",
+	Short: "tcpgoon tests concurrent connections towards a server listening on a TCP port",
+	Long:  ``,
+	PreRun: func(cmd *cobra.Command, args []string) {
+		if err := validateRequiredArgs(&params, args); err != nil {
+			cmd.Println(cmd.UsageString())
+			os.Exit(1)
+		}
+		enableDebuggingIfFlagSet(params)
+		autorunValidation(params)
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		runroot(params)
+	},
+}
+
+func Execute() {
+	if err := rootCmd.Execute(); err != nil {
+		os.Exit(1)
+	}
+}
+
+func init() {
+	rootCmd.Flags().IntVarP(&params.numberConnectionsPtr, "connections", "c", 100, "Number of connections you want to open")
+	rootCmd.Flags().IntVarP(&params.delayPtr, "sleep", "s", 10, "Time you want to sleep between connections, in ms")
+	rootCmd.Flags().IntVarP(&params.connDialTimeoutPtr, "dial-timeout", "d", 5000, "Connection dialing timeout, in ms")
+	rootCmd.Flags().BoolVarP(&params.debugPtr, "verbose", "v", false, "Print debugging information to the standard error")
+	rootCmd.Flags().IntVarP(&params.reportingIntervalPtr, "interval", "i", 1, "Interval, in seconds, between stats updates")
+	rootCmd.Flags().BoolVarP(&params.assumeyesPtr, "assume-yes", "y", false, "Force execution without asking for confirmation")
+}
+
+func validateRequiredArgs(params *tcpgoonParams, args []string) error {
+	if len(args) != 2 {
+		return errors.New("Number of required parameters doesn't match")
+	}
+	params.hostPtr = args[0]
+	port, err := strconv.Atoi(args[1])
+	if err != nil && port <= 0 {
+		return errors.New("Port argument is not a valid integer")
+	}
+	params.portPtr = port
+
+	return nil
+}
+
+func enableDebuggingIfFlagSet(params tcpgoonParams) {
+	if params.debugPtr {
+		debugging.EnableDebug()
+	}
+}
+
+func autorunValidation(params tcpgoonParams) {
+	if !(params.assumeyesPtr || cmdutil.AskForUserConfirmation(params.hostPtr, params.portPtr, params.numberConnectionsPtr)) {
+		fmt.Fprintln(debugging.DebugOut, "Execution not approved by the user")
+		cmdutil.CloseAbruptly()
+	}
+}
+
+func runroot(params tcpgoonParams) {
+	tcpclient.DefaultDialTimeoutInMs = params.connDialTimeoutPtr
+
+	// TODO: we should decouple the caller from the mtcpclient package (too many structures being moved from
+	//  one side to the other.. everything in a single structure, or applying something like the builder pattern,
+	//  may help
+	connStatusCh, connStatusTracker := mtcpclient.StartBackgroundReporting(params.numberConnectionsPtr, params.reportingIntervalPtr)
+	closureCh := mtcpclient.StartBackgroundClosureTrigger(connStatusTracker)
+	mtcpclient.MultiTCPConnect(params.numberConnectionsPtr, params.delayPtr, params.hostPtr, params.portPtr, connStatusCh, closureCh)
+	fmt.Fprintln(debugging.DebugOut, "Tests execution completed")
+
+	cmdutil.CloseNicely(params.hostPtr, params.portPtr, connStatusTracker)
+}

cmdutil/userOutput.go

@@ -24,7 +24,7 @@ func printClosureReport(host string, port int, gc mtcpclient.GroupOfConnections)
 func AskForUserConfirmation(host string, port int, connections int) bool {
 	fmt.Println("****************************** WARNING ******************************")
 	fmt.Println("* You are going to run a TCP stress check with these arguments:")
-	fmt.Println("*	- Target: " + host)
+	fmt.Println("*	- Host: " + host)
 	fmt.Println("*	- TCP Port: " + strconv.Itoa(port))
 	fmt.Println("*	- # of concurrent connections: " + strconv.Itoa(connections))
 	fmt.Println("*********************************************************************")

main.go

@@ -0,0 +1,9 @@
+package main
+
+import (
+	"github.com/dachad/tcpgoon/cmd"
+)
+
+func main() {
+	cmd.Execute()
+}