From 34a053b875df0de33c411e8c9281d0fbb7f2fd08 Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Tue, 20 Apr 2021 20:55:40 +0800
Subject: [PATCH] ssh: close unused connection (#251)

---
 go.mod           |  4 ++--
 go.sum           |  8 ++++----
 proxy/ssh/ssh.go | 18 +++++++++++-------
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/go.mod b/go.mod
index eee4434..313b23c 100644
--- a/go.mod
+++ b/go.mod
@@ -17,8 +17,8 @@ require (
 	github.com/tjfoc/gmsm v1.4.0 // indirect
 	github.com/xtaci/kcp-go/v5 v5.6.1
 	golang.org/x/crypto v0.0.0-20210415154028-4f45737414dc
-	golang.org/x/net v0.0.0-20210415231046-e915ea6b2b7d // indirect
-	golang.org/x/sys v0.0.0-20210419170143-37df388d1f33 // indirect
+	golang.org/x/net v0.0.0-20210420072503-d25e30425868 // indirect
+	golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe // indirect
 )
 
 // Replace dependency modules with local developing copy
diff --git a/go.sum b/go.sum
index 7f4db6e..0bb9ef8 100644
--- a/go.sum
+++ b/go.sum
@@ -128,8 +128,8 @@ golang.org/x/net v0.0.0-20200707034311-ab3426394381/go.mod h1:/O7V0waA8r7cgGh81R
 golang.org/x/net v0.0.0-20201010224723-4f7140c49acb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20210415231046-e915ea6b2b7d h1:BgJvlyh+UqCUaPlscHJ+PN8GcpfrFdr7NHjd1JL0+Gs=
-golang.org/x/net v0.0.0-20210415231046-e915ea6b2b7d/go.mod h1:9tjilg8BloeKEkVJvy7fQ90B1CfIiPueXVOjqfkSzI8=
+golang.org/x/net v0.0.0-20210420072503-d25e30425868 h1:mHVdVrNGft0Bv5N0WIf3/ujpDOQOe6KxvwlIikPbMr0=
+golang.org/x/net v0.0.0-20210420072503-d25e30425868/go.mod h1:9tjilg8BloeKEkVJvy7fQ90B1CfIiPueXVOjqfkSzI8=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -154,8 +154,8 @@ golang.org/x/sys v0.0.0-20201101102859-da207088b7d1/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201112073958-5cba982894dd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210419170143-37df388d1f33 h1:zah5VTTvBlVRELjcDwGLLaWRHZJQsBtplweVYCii0KM=
-golang.org/x/sys v0.0.0-20210419170143-37df388d1f33/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe h1:WdX7u8s3yOigWAhHEaDl8r9G+4XwFQEQFtBMYyN+kXQ=
+golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 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 59ea7d3..e08b147 100644
--- a/proxy/ssh/ssh.go
+++ b/proxy/ssh/ssh.go
@@ -4,6 +4,7 @@ import (
 	"net"
 	"net/url"
 	"os"
+	"sync"
 	"time"
 
 	"golang.org/x/crypto/ssh"
@@ -18,6 +19,7 @@ type SSH struct {
 	proxy  proxy.Proxy
 	addr   string
 
+	mu      sync.Mutex
 	sshCfg  *ssh.ClientConfig
 	sshConn ssh.Conn
 	sshChan <-chan ssh.NewChannel
@@ -90,6 +92,9 @@ func (s *SSH) Addr() string {
 }
 
 func (s *SSH) initConn() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
 	c, err := s.dialer.Dial("tcp", s.addr)
 	if err != nil {
 		log.F("[ssh]: dial to %s error: %s", s.addr, err)
@@ -106,14 +111,13 @@ func (s *SSH) initConn() error {
 }
 
 // Dial connects to the address addr on the network net via the proxy.
-func (s *SSH) Dial(network, addr string) (c net.Conn, err error) {
-	c, err = ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr)
-	if err != nil {
-		log.F("[ssh]: create conn to %s via %s error: %s, try again", addr, s.addr, err)
-		s.initConn()
-		c, err = ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr)
+func (s *SSH) Dial(network, addr string) (net.Conn, error) {
+	if c, err := ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr); err == nil {
+		return c, nil
 	}
-	return c, err
+	s.sshConn.Close()
+	s.initConn()
+	return ssh.NewClient(s.sshConn, s.sshChan, s.sshReq).Dial(network, addr)
 }
 
 // DialUDP connects to the given address via the proxy.