From e3d9535a080bc393bf507b3478673c23e25260c7 Mon Sep 17 00:00:00 2001 From: Lonny Wong Date: Sun, 28 Jul 2024 20:37:55 +0800 Subject: [PATCH] optimize session command --- go.mod | 5 +++-- go.sum | 10 ++++++---- tsshd/session.go | 40 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 48 insertions(+), 7 deletions(-) diff --git a/go.mod b/go.mod index 45716e4..9d48e32 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 2f904ab..9c7a1af 100644 --- a/go.sum +++ b/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= diff --git a/tsshd/session.go b/tsshd/session.go index ad94e21..6417c13 100644 --- a/tsshd/session.go +++ b/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 }