Add --connect-timeout

Allow initial dial timeout duration to be configured and not just set as 'wait forever'

Author: NHAS

cmd/client/main.go

@@ -12,6 +12,7 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/NHAS/reverse_ssh/internal/client"
 	"github.com/NHAS/reverse_ssh/internal/client/keys"
@@ -66,6 +67,7 @@ func printHelp() {
 	fmt.Println("\t\t--log-level\tChange logging output levels, [INFO,WARNING,ERROR,FATAL,DISABLED]")
 	fmt.Println("\t\t--version-string\tSSH version string to use, i.e SSH-VERSION, defaults to internal.Version-runtime.GOOS_runtime.GOARCH")
 	fmt.Println("\t\t--private-key-path\tOptional path to unencrypted SSH key to use for connecting")
+	fmt.Println("\t\t--connect-timeout\tDuration to wait for initial connection seconds, default 180, set to 0 to wait indefinitely")
 
 	if runtime.GOOS == "windows" {
 		fmt.Println("\t\t--host-kerberos\tUse kerberos authentication on proxy server (if proxy server specified)")
@@ -173,6 +175,18 @@ func main() {
 		settings.SNI = userSpecifiedSNI
 	}
 
+	timeoutInt := 180
+	timeout, err := line.GetArgString("connect-timeout")
+	if err == nil {
+		timeoutInt, err = strconv.Atoi(timeout)
+		if err != nil {
+			log.Printf("could not parse --connect-timeout as number %v, setting default to 180", err)
+			timeoutInt = 180
+		}
+	}
+
+	settings.ConnectTimeout = time.Duration(timeoutInt) * time.Second
+
 	userSpecifiedNTLMCreds, err := line.GetArgString("ntlm-proxy-creds")
 	if err == nil {
 		if line.IsSet("host-kerberos") {

internal/client/client.go

@@ -310,6 +310,8 @@ type Settings struct {
 
 	VersionString string
 
+	ConnectTimeout time.Duration
+
 	ntlm *ntlmssp.Client
 }
 
@@ -363,7 +365,8 @@ func Run(settings *Settings) {
 	}
 
 	config := &ssh.ClientConfig{
-		User: fmt.Sprintf("%s.%s", username, hostname),
+		Timeout: settings.ConnectTimeout,
+		User:    fmt.Sprintf("%s.%s", username, hostname),
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeys(sshPriv),
 		},
@@ -396,8 +399,9 @@ func Run(settings *Settings) {
 		var conn net.Conn
 		if scheme != "stdio" {
 			log.Println("Connecting to", settings.Addr)
+
 			// First create raw TCP connection
-			conn, err = Connect(realAddr, settings.ProxyAddr, config.Timeout, settings.ProxyUseHostKerberos, settings.ntlm)
+			conn, err = Connect(realAddr, settings.ProxyAddr, settings.ConnectTimeout, settings.ProxyUseHostKerberos, settings.ntlm)
 			if err != nil {
 
 				if errMsg := err.Error(); strings.Contains(errMsg, "missing port in address") {
@@ -426,7 +430,7 @@ func Run(settings *Settings) {
 					continue
 				}
 
-				<-time.After(10 * time.Second)
+				time.Sleep(10 * time.Second)
 				continue
 			}
 
@@ -480,7 +484,7 @@ func Run(settings *Settings) {
 			case "http", "https":
 
 				conn, err = NewHTTPConn(scheme+"://"+realAddr, func() (net.Conn, error) {
-					return Connect(realAddr, settings.ProxyAddr, config.Timeout, settings.ProxyUseHostKerberos, settings.ntlm)
+					return Connect(realAddr, settings.ProxyAddr, settings.ConnectTimeout, settings.ProxyUseHostKerberos, settings.ntlm)
 				})
 
 				if err != nil {