From 8a8cead7f97ed5769c4978cb41bff75315dabb49 Mon Sep 17 00:00:00 2001
From: Lonny Wong <lonnywong@qq.com>
Date: Sat, 14 Oct 2023 09:00:59 +0800
Subject: [PATCH] improve ProxyCommand #46

---
 go.mod               |  1 +
 go.sum               |  2 ++
 tssh/login.go        | 15 +++++++++------
 tssh/term_unix.go    |  5 +++++
 tssh/term_windows.go |  4 ++++
 tssh/util_darwin.go  |  5 -----
 tssh/util_unix.go    |  5 -----
 tssh/util_windows.go |  9 ---------
 8 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/go.mod b/go.mod
index 973c7a3..023e170 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.20
 require (
 	github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
 	github.com/chzyer/readline v1.5.1
+	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/mattn/go-isatty v0.0.19
 	github.com/natefinch/npipe v0.0.0-20160621034901-c1b8fa8bdcce
 	github.com/skeema/knownhosts v1.2.1-0.20230916192230-09454b7d5683
diff --git a/go.sum b/go.sum
index 69cef94..fcb1513 100644
--- a/go.sum
+++ b/go.sum
@@ -22,6 +22,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dchest/jsmin v0.0.0-20220218165748-59f39799265f h1:OGqDDftRTwrvUoL6pOG7rYTmWsTCvyEWFsMjg+HcOaA=
 github.com/dchest/jsmin v0.0.0-20220218165748-59f39799265f/go.mod h1:Dv9D0NUlAsaQcGQZa5kc5mqR9ua72SmA8VXi4cd+cBw=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
+github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/josephspurrier/goversioninfo v1.4.0 h1:Puhl12NSHUSALHSuzYwPYQkqa2E1+7SrtAPJorKK0C8=
 github.com/josephspurrier/goversioninfo v1.4.0/go.mod h1:JWzv5rKQr+MmW+LvM412ToT/IkYDZjaclF2pKDss8IY=
 github.com/klauspost/compress v1.16.7 h1:2mk3MPGNzKyxErAw8YaohYh69+pa4sIQSC0fPGCFR9I=
diff --git a/tssh/login.go b/tssh/login.go
index 6b7627d..21e5af5 100644
--- a/tssh/login.go
+++ b/tssh/login.go
@@ -631,13 +631,16 @@ func execProxyCommand(param *loginParam) (net.Conn, string, error) {
 	command = strings.ReplaceAll(command, "%r", param.user)
 	debug("exec proxy command: %s", command)
 
-	var cmd *exec.Cmd
-	if !strings.ContainsAny(command, "'\"\\") {
-		tokens := strings.Fields(command)
-		cmd = exec.Command(tokens[0], tokens[1:]...)
-	} else {
-		cmd = createProxyCommand(command)
+	argv, err := splitCommandLine(command)
+	if err != nil || len(argv) == 0 {
+		return nil, command, fmt.Errorf("split proxy command failed: %v", err)
+	}
+	if enableDebugLogging {
+		for i, arg := range argv {
+			debug("proxy command argv[%d] = %s", i, arg)
+		}
 	}
+	cmd := exec.Command(argv[0], argv[1:]...)
 
 	cmdIn, err := cmd.StdinPipe()
 	if err != nil {
diff --git a/tssh/term_unix.go b/tssh/term_unix.go
index be00538..0534e8f 100644
--- a/tssh/term_unix.go
+++ b/tssh/term_unix.go
@@ -36,6 +36,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/google/shlex"
 	"golang.org/x/term"
 )
 
@@ -115,3 +116,7 @@ func isSshTmuxEnv() bool {
 
 	return isRemoteSshEnv(pid)
 }
+
+func splitCommandLine(command string) ([]string, error) {
+	return shlex.Split(command)
+}
diff --git a/tssh/term_windows.go b/tssh/term_windows.go
index 5b27136..88f9653 100644
--- a/tssh/term_windows.go
+++ b/tssh/term_windows.go
@@ -262,3 +262,7 @@ func getKeyboardInput() (*os.File, func(), error) {
 
 	return file, func() { _ = file.Close() }, nil
 }
+
+func splitCommandLine(command string) ([]string, error) {
+	return windows.DecomposeCommandLine(command)
+}
diff --git a/tssh/util_darwin.go b/tssh/util_darwin.go
index 57fd86d..411225b 100644
--- a/tssh/util_darwin.go
+++ b/tssh/util_darwin.go
@@ -28,7 +28,6 @@ package tssh
 import (
 	"bytes"
 	"os"
-	"os/exec"
 
 	"golang.org/x/sys/unix"
 )
@@ -67,7 +66,3 @@ func isNoGUI() bool {
 	}
 	return isDockerEnv() || isRemoteSshEnv(os.Getppid()) || isSshTmuxEnv()
 }
-
-func createProxyCommand(command string) *exec.Cmd {
-	return exec.Command("sh", "-c", command)
-}
diff --git a/tssh/util_unix.go b/tssh/util_unix.go
index 22099de..1ec2371 100644
--- a/tssh/util_unix.go
+++ b/tssh/util_unix.go
@@ -31,7 +31,6 @@ import (
 	"bytes"
 	"fmt"
 	"os"
-	"os/exec"
 	"strconv"
 )
 
@@ -90,7 +89,3 @@ func isNoGUI() bool {
 	}
 	return isDockerEnv() || isRemoteSshEnv(os.Getppid()) || isSshTmuxEnv()
 }
-
-func createProxyCommand(command string) *exec.Cmd {
-	return exec.Command("sh", "-c", command)
-}
diff --git a/tssh/util_windows.go b/tssh/util_windows.go
index 1f5a9a2..d939eea 100644
--- a/tssh/util_windows.go
+++ b/tssh/util_windows.go
@@ -26,11 +26,8 @@ SOFTWARE.
 package tssh
 
 import (
-	"fmt"
 	"os"
-	"os/exec"
 	"strings"
-	"syscall"
 	"unsafe"
 
 	"golang.org/x/sys/windows"
@@ -65,9 +62,3 @@ func isNoGUI() bool {
 	}
 	return false
 }
-
-func createProxyCommand(command string) *exec.Cmd {
-	cmd := exec.Command("cmd.exe")
-	cmd.SysProcAttr = &syscall.SysProcAttr{CmdLine: fmt.Sprintf("/c %s", command)}
-	return cmd
-}