optimize session command

go.mod

@@ -4,10 +4,11 @@ go 1.21
 
 require (
 	github.com/UserExistsError/conpty v0.1.4
+	github.com/alessio/shellescape v1.4.2
 	github.com/creack/pty v1.1.21
 	github.com/quic-go/quic-go v0.45.1
 	github.com/trzsz/go-arg v1.5.4
-	github.com/xtaci/kcp-go/v5 v5.6.9
+	github.com/xtaci/kcp-go/v5 v5.6.11
 	github.com/xtaci/smux v1.5.25
 	golang.org/x/crypto v0.25.0
 	golang.org/x/sys v0.22.0
@@ -16,7 +17,7 @@ require (
 require (
 	github.com/alexflint/go-scalar v1.2.0 // indirect
 	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
-	github.com/google/pprof v0.0.0-20240722153945-304e4f0156b8 // indirect
+	github.com/google/pprof v0.0.0-20240727154555-813a5fbdbec8 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.8 // indirect
 	github.com/klauspost/reedsolomon v1.12.3 // indirect
 	github.com/onsi/ginkgo/v2 v2.19.1 // indirect

go.sum

@@ -2,6 +2,8 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/UserExistsError/conpty v0.1.4 h1:+3FhJhiqhyEJa+K5qaK3/w6w+sN3Nh9O9VbJyBS02to=
 github.com/UserExistsError/conpty v0.1.4/go.mod h1:PDglKIkX3O/2xVk0MV9a6bCWxRmPVfxqZoTG/5sSd9I=
+github.com/alessio/shellescape v1.4.2 h1:MHPfaU+ddJ0/bYWpgIeUnQUqKrlJ1S7BfEYPM4uEoM0=
+github.com/alessio/shellescape v1.4.2/go.mod h1:PZAiSCk0LJaZkiCSkPv8qIobYglO3FPpyFjDCtHLS30=
 github.com/alexflint/go-scalar v1.2.0 h1:WR7JPKkeNpnYIOfHRa7ivM21aWAdHD0gEWHCx+WQBRw=
 github.com/alexflint/go-scalar v1.2.0/go.mod h1:LoFvNMqS1CPrMVltza4LvnGKhaSpc3oyLEBUZVhhS2o=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
@@ -35,8 +37,8 @@ github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMyw
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/pprof v0.0.0-20240722153945-304e4f0156b8 h1:ssNFCCVmib/GQSzx3uCWyfMgOamLGWuGqlMS77Y1m3Y=
-github.com/google/pprof v0.0.0-20240722153945-304e4f0156b8/go.mod h1:K1liHPHnj73Fdn/EKuT8nrFqBihUSKXoLYU0BuatOYo=
+github.com/google/pprof v0.0.0-20240727154555-813a5fbdbec8 h1:FKHo8hFI3A+7w0aUQuYXQ+6EN5stWmeY/AZqtM8xk9k=
+github.com/google/pprof v0.0.0-20240727154555-813a5fbdbec8/go.mod h1:K1liHPHnj73Fdn/EKuT8nrFqBihUSKXoLYU0BuatOYo=
 github.com/klauspost/cpuid/v2 v2.2.8 h1:+StwCXwm9PdpiEkPyzBXIy+M9KUb4ODm0Zarf1kS5BM=
 github.com/klauspost/cpuid/v2 v2.2.8/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
 github.com/klauspost/reedsolomon v1.12.3 h1:tzUznbfc3OFwJaTebv/QdhnFf2Xvb7gZ24XaHLBPmdc=
@@ -64,8 +66,8 @@ github.com/tjfoc/gmsm v1.4.1 h1:aMe1GlZb+0bLjn+cKTPEvvn9oUEBlJitaZiiBwsbgho=
 github.com/tjfoc/gmsm v1.4.1/go.mod h1:j4INPkHWMrhJb38G+J6W4Tw0AbuN8Thu3PbdVYhVcTE=
 github.com/trzsz/go-arg v1.5.4 h1:8cuwV8F1UvlKRXJcdLG1cSZNGXkRped+U1rz17M7Kac=
 github.com/trzsz/go-arg v1.5.4/go.mod h1:IC6Z/FiVH7uYvcbp1/gJhDYCFPS/GkL0APYakVvgY4I=
-github.com/xtaci/kcp-go/v5 v5.6.9 h1:rSfwk6ofh0wrwtlOyzXFNUkXRu6vsXpTF1y8Kh/nof4=
-github.com/xtaci/kcp-go/v5 v5.6.9/go.mod h1:GSs9Z62r41kTb4CxKaRKr6ED+j7I0sUvj7Koql/RN6c=
+github.com/xtaci/kcp-go/v5 v5.6.11 h1:I2fPv8wSxV1hp21VVKUzABB8T/e3T177keaEzIicJXA=
+github.com/xtaci/kcp-go/v5 v5.6.11/go.mod h1:GSs9Z62r41kTb4CxKaRKr6ED+j7I0sUvj7Koql/RN6c=
 github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae h1:J0GxkO96kL4WF+AIT3M4mfUVinOCPgf2uUWYFUzN0sM=
 github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae/go.mod h1:gXtu8J62kEgmN++bm9BVICuT/e8yiLI2KFobd/TRFsE=
 github.com/xtaci/smux v1.5.25 h1:pNWlUQFcE0O4letC2+giIjxbGVQXWmjD5vAc/opm3A0=

tsshd/session.go

@@ -31,10 +31,13 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/alessio/shellescape"
 )
 
 type sessionContext struct {
@@ -286,7 +289,42 @@ func getSessionStartCmd(msg *StartMessage) (*exec.Cmd, error) {
 	}
 
 	if !msg.Shell {
-		cmd := exec.Command(msg.Name, msg.Args...)
+		name := msg.Name
+		args := msg.Args
+		wrap := false
+		if name == "cd" {
+			wrap = true
+		} else if _, err := exec.LookPath(name); err != nil {
+			wrap = true
+		} else {
+			for _, arg := range args {
+				if strings.HasPrefix(arg, "~/") {
+					wrap = true
+					break
+				}
+			}
+		}
+		if wrap {
+			re := regexp.MustCompile(`\s`)
+			var buf strings.Builder
+			buf.WriteString(name)
+			for _, arg := range args {
+				buf.WriteByte(' ')
+				if re.MatchString(arg) {
+					buf.WriteString(shellescape.Quote(arg))
+				} else {
+					buf.WriteString(arg)
+				}
+			}
+			if runtime.GOOS == "windows" {
+				name = "cmd"
+				args = []string{"/c", buf.String()}
+			} else {
+				name = "sh"
+				args = []string{"-c", buf.String()}
+			}
+		}
+		cmd := exec.Command(name, args...)
 		cmd.Env = envs
 		return cmd, nil
 	}