Skip to content

Commit

Permalink
refactor ssh agent
Browse files Browse the repository at this point in the history
  • Loading branch information
lonnywong committed Nov 18, 2023
1 parent 4fc7fe4 commit 3013f1b
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 166 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ GO_TEST := ${shell basename `which gotest 2>/dev/null` 2>/dev/null || echo go te

all: ${BIN_DIR}/${TSSH}

${BIN_DIR}/${TSSH}: $(wildcard ./cmd/tssh/*.go ./tssh/*.go)
${BIN_DIR}/${TSSH}: $(wildcard ./cmd/tssh/*.go ./tssh/*.go) go.mod go.sum
go build -o ${BIN_DIR}/ ./cmd/tssh

clean:
Expand Down
10 changes: 6 additions & 4 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module github.com/trzsz/trzsz-ssh
go 1.20

require (
github.com/Microsoft/go-winio v0.6.1
github.com/alessio/shellescape v1.4.2
github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
github.com/chzyer/readline v1.5.1
Expand All @@ -12,8 +13,7 @@ require (
github.com/skeema/knownhosts v1.2.1
github.com/stretchr/testify v1.8.4
github.com/trzsz/go-arg v1.5.2
github.com/trzsz/iterm2 v0.1.0
github.com/trzsz/npipe v0.1.0
github.com/trzsz/iterm2 v0.1.1
github.com/trzsz/promptui v0.10.3
github.com/trzsz/ssh_config v1.3.3
github.com/trzsz/trzsz-go v1.1.7-0.20231111144918-b45bed013817
Expand All @@ -27,18 +27,20 @@ require (
github.com/akavel/rsrc v0.10.2 // indirect
github.com/alexflint/go-scalar v1.2.0 // indirect
github.com/andybrewer/mack v0.0.0-20220307193339-22e922cc18af // indirect
github.com/creack/pty v1.1.20 // indirect
github.com/creack/pty v1.1.21 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dchest/jsmin v0.0.0-20220218165748-59f39799265f // indirect
github.com/gorilla/websocket v1.5.1 // indirect
github.com/josephspurrier/goversioninfo v1.4.0 // indirect
github.com/klauspost/compress v1.17.2 // indirect
github.com/klauspost/compress v1.17.3 // indirect
github.com/ncruces/zenity v0.10.10 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/randall77/makefat v0.0.0-20210315173500-7ddd0e42c844 // indirect
golang.org/x/image v0.14.0 // indirect
golang.org/x/mod v0.14.0 // indirect
golang.org/x/net v0.18.0 // indirect
golang.org/x/text v0.14.0 // indirect
golang.org/x/tools v0.15.0 // indirect
google.golang.org/protobuf v1.31.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
21 changes: 13 additions & 8 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow=
github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM=
github.com/UserExistsError/conpty v0.1.1 h1:cHDsU/XeoeDAQmVvCTV53SrXLG39YJ4++Pp3iAi1gXE=
github.com/UserExistsError/conpty v0.1.1/go.mod h1:PDglKIkX3O/2xVk0MV9a6bCWxRmPVfxqZoTG/5sSd9I=
github.com/akavel/rsrc v0.10.2 h1:Zxm8V5eI1hW4gGaYsJQUhxpjkENuG91ki8B4zCrvEsw=
Expand All @@ -16,8 +18,8 @@ github.com/chzyer/readline v1.5.1 h1:upd/6fQk4src78LMRzh5vItIt361/o4uq553V8B5sGI
github.com/chzyer/readline v1.5.1/go.mod h1:Eh+b79XXUwfKfcPLepksvw2tcLE/Ct21YObkaSkeBlk=
github.com/chzyer/test v1.0.0 h1:p3BQDXSxOhOG0P9z6/hGnII4LGiEPOYBhs8asl/fC04=
github.com/chzyer/test v1.0.0/go.mod h1:2JlltgoNkt4TW/z9V/IzDdFaMTM2JPIi26O1pF38GC8=
github.com/creack/pty v1.1.20 h1:VIPb/a2s17qNeQgDnkfZC35RScx+blkKF8GV68n80J4=
github.com/creack/pty v1.1.20/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
github.com/creack/pty v1.1.21 h1:1/QdRyBaHHJP61QkWMXlOIBfsgdDeeKfK8SYVUWJKf0=
github.com/creack/pty v1.1.21/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
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=
Expand All @@ -32,8 +34,8 @@ github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/
github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
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.17.2 h1:RlWWUY/Dr4fL8qk9YG7DTZ7PDgME2V4csBXA8L/ixi4=
github.com/klauspost/compress v1.17.2/go.mod h1:ntbaceVETuRiXiv4DpjP66DpAtAGkEQskQzEyD//IeE=
github.com/klauspost/compress v1.17.3 h1:qkRjuerhUU1EmXLYGkSH6EZL+vPSxIrYjLNAK4slzwA=
github.com/klauspost/compress v1.17.3/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
Expand All @@ -53,10 +55,8 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
github.com/trzsz/go-arg v1.5.2 h1:zGxCuTKvtC3jBf7HbvNk0HooUjv8uKAy2mY+bHVhRas=
github.com/trzsz/go-arg v1.5.2/go.mod h1:IC6Z/FiVH7uYvcbp1/gJhDYCFPS/GkL0APYakVvgY4I=
github.com/trzsz/iterm2 v0.1.0 h1:oMoxJPoeemSFvFhlz+2o9tL1HKmDDRztCdUsJGH64YI=
github.com/trzsz/iterm2 v0.1.0/go.mod h1:PMI+3JcT7J9D0T6e3mOWv8ICYdrrNZwuge3Tm7zDLws=
github.com/trzsz/npipe v0.1.0 h1:9/He1oey9scgy0SWNqMygT30ZeEbPApLTybaeoareiI=
github.com/trzsz/npipe v0.1.0/go.mod h1:SvsmyY6zmu0Btm8AW3Hpi9c3BB+Yn5z6K1MgLLXo1s0=
github.com/trzsz/iterm2 v0.1.1 h1:UZ+Su5xOuBNDXpRStfuMCsTWhajNnKEHChJ4tfd62Mc=
github.com/trzsz/iterm2 v0.1.1/go.mod h1:PMI+3JcT7J9D0T6e3mOWv8ICYdrrNZwuge3Tm7zDLws=
github.com/trzsz/promptui v0.10.3 h1:uhcLQsLZqMxEtGiYoeM2lR/Hd4pSxoYsd2eFctH8MCs=
github.com/trzsz/promptui v0.10.3/go.mod h1:GMZtu6ZTzU73CBFkzGtmB4wnTROIAbv4GFA74fV8V8g=
github.com/trzsz/ssh_config v1.3.3 h1:FodC5NBpnBSI2CJqU62d0Ya3yTUHWDpbWKugKX0LXZs=
Expand All @@ -68,8 +68,11 @@ golang.org/x/crypto v0.15.0 h1:frVn1TEaCEaZcn3Tmd7Y2b5KKPaZ+I32Q2OA3kYp5TA=
golang.org/x/crypto v0.15.0/go.mod h1:4ChreQoLWfG3xLDer1WdlH5NdlQ3+mwnQq1YTKY+72g=
golang.org/x/image v0.14.0 h1:tNgSxAFe3jC4uYqvZdTr84SZoM1KfwdC9SKIFrLjFn4=
golang.org/x/image v0.14.0/go.mod h1:HUYqC05R2ZcZ3ejNQsIHQDQiwWM4JBqmm6MKANTp4LE=
golang.org/x/mod v0.14.0 h1:dGoOF9QVLYng8IHTm7BAyWqCqSheQ5pYWGhzW00YJr0=
golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
golang.org/x/net v0.18.0 h1:mIYleuAkSbHh0tCv7RvjL3F6ZVbLjq4+R7zbOn3Kokg=
golang.org/x/net v0.18.0/go.mod h1:/czyP5RqHAH4odGYxBJ1qz0+CE5WZ+2j1YgoEo8F2jQ=
golang.org/x/sync v0.5.0 h1:60k92dhOjHxJkrqnwsfl8KuaHbn/5dl0lUPUklKo3qE=
golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
Expand All @@ -80,6 +83,8 @@ golang.org/x/term v0.14.0 h1:LGK9IlZ8T9jvdy6cTdfKUCltatMFOehAQo9SRC46UQ8=
golang.org/x/term v0.14.0/go.mod h1:TySc+nGkYR6qt8km8wUhuFRTVSMIX3XPR58y2lC8vww=
golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
golang.org/x/tools v0.15.0 h1:zdAyfUGbYmuVokhzVmghFl2ZJh5QhcfebBgmVPFYA+8=
golang.org/x/tools v0.15.0/go.mod h1:hpksKq4dtpQWS1uQ61JkdqWM3LscIS6Slf+VVkm+wQk=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
Expand Down
136 changes: 136 additions & 0 deletions tssh/agent.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
package tssh

/*
MIT License
Copyright (c) 2023 Lonny Wong <[email protected]>
Copyright (c) 2023 [Contributors](https://github.com/trzsz/trzsz-ssh/graphs/contributors)
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

import (
"fmt"
"io"
"net"
"os"
"strings"
"sync"

"golang.org/x/crypto/ssh"
"golang.org/x/crypto/ssh/agent"
)

var (
agentOnce sync.Once
agentClient agent.ExtendedAgent
)

func getAgentAddr(args *sshArgs) string {
if addr := getOptionConfig(args, "IdentityAgent"); addr != "" {
if strings.ToLower(addr) == "none" {
return ""
}
return addr
}
if addr := os.Getenv("SSH_AUTH_SOCK"); addr != "" {
return addr
}
if addr := defaultAgentAddr; addr != "" && isFileExist(addr) {
return addr
}
return ""
}

func getAgentClient(args *sshArgs) agent.ExtendedAgent {
agentOnce.Do(func() {
addr := resolveHomeDir(getAgentAddr(args))
if addr == "" {
debug("ssh agent address is not set")
return
}

conn, err := dialAgent(addr)
if err != nil {
debug("dial ssh agent [%s] failed: %v", addr, err)
return
}

agentClient = agent.NewClient(conn)
debug("new ssh agent client [%s] success", addr)

cleanupAfterLogined = append(cleanupAfterLogined, func() {
conn.Close()
agentClient = nil
})
})
return agentClient
}

const channelType = "[email protected]"

func forwardToRemote(client *ssh.Client, addr string) error {
channels := client.HandleChannelOpen(channelType)
if channels == nil {
return fmt.Errorf("agent: already have handler for %s", channelType)
}
conn, err := dialAgent(addr)
if err != nil {
return err
}
conn.Close()

go func() {
for ch := range channels {
channel, reqs, err := ch.Accept()
if err != nil {
continue
}
go ssh.DiscardRequests(reqs)
go forwardAgentRequest(channel, addr)
}
}()
return nil
}

func forwardAgentRequest(channel ssh.Channel, addr string) {
conn, err := dialAgent(addr)
if err != nil {
return
}

var wg sync.WaitGroup
wg.Add(2)
go func() {
_, _ = io.Copy(conn, channel)
if unixConn, ok := conn.(*net.UnixConn); ok {
_ = unixConn.CloseWrite()
}
wg.Done()
}()
go func() {
_, _ = io.Copy(channel, conn)
_ = channel.CloseWrite()
wg.Done()
}()

wg.Wait()
conn.Close()
channel.Close()
}
53 changes: 3 additions & 50 deletions tssh/agent_unix.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,58 +29,11 @@ package tssh

import (
"net"
"os"
"strings"
"sync"
"time"

"golang.org/x/crypto/ssh"
"golang.org/x/crypto/ssh/agent"
)

var (
agentOnce sync.Once
agentConn net.Conn
agentClient agent.ExtendedAgent
)

func getAgentAddr(args *sshArgs) string {
if addr := getOptionConfig(args, "IdentityAgent"); addr != "" {
if strings.ToLower(addr) == "none" {
return ""
}
return addr
}
return os.Getenv("SSH_AUTH_SOCK")
}

func getAgentClient(args *sshArgs) agent.ExtendedAgent {
agentOnce.Do(func() {
addr := resolveHomeDir(getAgentAddr(args))
if addr == "" {
debug("ssh agent unix socket is not set")
return
}

var err error
agentConn, err = net.DialTimeout("unix", addr, time.Second)
if err != nil {
debug("dial ssh agent unix socket [%s] failed: %v", addr, err)
return
}

agentClient = agent.NewClient(agentConn)
debug("new ssh agent client [%s] success", addr)

cleanupAfterLogined = append(cleanupAfterLogined, func() {
agentConn.Close()
agentConn = nil
agentClient = nil
})
})
return agentClient
}
const defaultAgentAddr = ""

func forwardToRemote(client *ssh.Client, addr string) error {
return agent.ForwardToRemote(client, addr)
func dialAgent(addr string) (net.Conn, error) {
return net.DialTimeout("unix", addr, time.Second)
}
Loading

0 comments on commit 3013f1b

Please sign in to comment.