From ff09c45fb6463c7c687f60baf7f3c20bd58f61f6 Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Mon, 6 Dec 2021 23:53:10 +0800
Subject: [PATCH] ssh: init conn in Dial

---
 go.mod           |  2 +-
 go.sum           |  4 ++--
 proxy/ssh/ssh.go | 24 ++++++++++++------------
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/go.mod b/go.mod
index c83541a..93e1f6b 100644
--- a/go.mod
+++ b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/nadoo/ipset v0.3.0
 	github.com/xtaci/kcp-go/v5 v5.6.1
 	golang.org/x/crypto v0.0.0-20211202192323-5770296d904e
-	golang.org/x/sys v0.0.0-20211204120058-94396e421777
+	golang.org/x/sys v0.0.0-20211205182925-97ca703d548d
 )
 
 require (
diff --git a/go.sum b/go.sum
index b4aaef3..c679173 100644
--- a/go.sum
+++ b/go.sum
@@ -157,8 +157,8 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210525143221-35b2ab0089ea/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211204120058-94396e421777 h1:QAkhGVjOxMa+n4mlsAWeAU+BMZmimQAaNiMu+iUi94E=
-golang.org/x/sys v0.0.0-20211204120058-94396e421777/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20211205182925-97ca703d548d h1:FjkYO/PPp4Wi0EAUOVLxePm7qVW4r4ctbWpURyuOD0E=
+golang.org/x/sys v0.0.0-20211205182925-97ca703d548d/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
diff --git a/proxy/ssh/ssh.go b/proxy/ssh/ssh.go
index 0bec94d..591a0a4 100644
--- a/proxy/ssh/ssh.go
+++ b/proxy/ssh/ssh.go
@@ -19,11 +19,9 @@ type SSH struct {
 	proxy  proxy.Proxy
 	addr   string
 
-	mu      sync.Mutex
-	sshCfg  *ssh.ClientConfig
-	sshConn ssh.Conn
-	sshChan <-chan ssh.NewChannel
-	sshReq  <-chan *ssh.Request
+	mu        sync.Mutex
+	sshCfg    *ssh.ClientConfig
+	sshClient *ssh.Client
 }
 
 func init() {
@@ -75,7 +73,7 @@ func NewSSH(s string, d proxy.Dialer, p proxy.Proxy) (*SSH, error) {
 		t.addr = net.JoinHostPort(t.addr, "22")
 	}
 
-	return t, t.initConn()
+	return t, nil
 }
 
 // NewSSHDialer returns a ssh proxy dialer.
@@ -98,12 +96,12 @@ func (s *SSH) initConn() error {
 		return err
 	}
 
-	s.sshConn, s.sshChan, s.sshReq, err = ssh.NewClientConn(c, s.addr, s.sshCfg)
+	sshConn, sshChan, sshReq, err := ssh.NewClientConn(c, s.addr, s.sshCfg)
 	if err != nil {
 		log.F("[ssh]: initial connection to %s error: %s", s.addr, err)
 		return err
 	}
-
+	s.sshClient = ssh.NewClient(sshConn, sshChan, sshReq)
 	return nil
 }
 
@@ -112,14 +110,16 @@ func (s *SSH) Dial(network, addr string) (net.Conn, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	if c, err := ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr); err == nil {
-		return c, nil
+	if s.sshClient != nil {
+		if c, err := s.sshClient.Dial(network, addr); err == nil {
+			return c, nil
+		}
+		s.sshClient.Conn.Close()
 	}
-	s.sshConn.Close()
 	if err := s.initConn(); err != nil {
 		return nil, err
 	}
-	return ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr)
+	return s.sshClient.Dial(network, addr)
 }
 
 // DialUDP connects to the given address via the proxy.