From 657f8f1fd6d2f3ad0aa08505552cd54446d2ae43 Mon Sep 17 00:00:00 2001
From: Dingjun <fangdingjun@gmail.com>
Date: Sat, 16 Sep 2017 09:44:08 +0800
Subject: [PATCH] add timeout for dial

---
 client.go       | 6 ++++--
 obfssh/proxy.go | 6 +++---
 obfssh/ssh.go   | 4 +++-
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/client.go b/client.go
index c3f69fb..9be170b 100644
--- a/client.go
+++ b/client.go
@@ -98,7 +98,8 @@ func (cc *Client) Close() {
 		l.Close()
 	}
 	//Log(DEBUG, "close ssh connection")
-	//cc.sshConn.Close()
+	cc.sshConn.Close()
+	cc.conn.Close()
 }
 
 // RunCmd run a single command on server
@@ -304,7 +305,8 @@ func (cc *Client) keepAlive(interval time.Duration, maxCount int) {
 			if count >= maxCount {
 				cc.err = fmt.Errorf("keep alive detects connection hang up")
 				Log(ERROR, "keep alive hit max count, exit")
-				cc.sshConn.Close()
+				//cc.sshConn.Close()
+				//cc.conn.Close()
 				// send exit signal
 				close(cc.ch)
 				return
diff --git a/obfssh/proxy.go b/obfssh/proxy.go
index bfbdeb6..6c5a803 100644
--- a/obfssh/proxy.go
+++ b/obfssh/proxy.go
@@ -147,7 +147,7 @@ func httpProxyHandshake(c net.Conn, host string, port int) (net.Conn, error) {
 }
 
 func dialHTTPProxy(host string, port int, p proxy) (net.Conn, error) {
-	c, err := net.Dial("tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)))
+	c, err := dialer.Dial("tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)))
 	if err != nil {
 		return nil, err
 	}
@@ -171,7 +171,7 @@ func dialHTTPSProxy(host string, port int, p proxy) (net.Conn, error) {
 		InsecureSkipVerify: p.Insecure,
 	}
 
-	c, err := tls.Dial("tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)), tlsconfig)
+	c, err := tls.DialWithDialer(dialer, "tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)), tlsconfig)
 	if err != nil {
 		return nil, err
 	}
@@ -190,7 +190,7 @@ func dialHTTPSProxy(host string, port int, p proxy) (net.Conn, error) {
 }
 
 func dialSocks5Proxy(host string, port int, p proxy) (net.Conn, error) {
-	c, err := net.Dial("tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)))
+	c, err := dialer.Dial("tcp", net.JoinHostPort(p.Host, fmt.Sprintf("%d", p.Port)))
 	if err != nil {
 		return nil, err
 	}
diff --git a/obfssh/ssh.go b/obfssh/ssh.go
index 5da95cd..1ebaa00 100644
--- a/obfssh/ssh.go
+++ b/obfssh/ssh.go
@@ -16,6 +16,8 @@ import (
 	"time"
 )
 
+var dialer = &net.Dialer{Timeout: 10 * time.Second}
+
 func main() {
 	var configfile string
 	var cfg config
@@ -185,7 +187,7 @@ func main() {
 			err = fmt.Errorf("unsupported scheme: %s", cfg.Proxy.Scheme)
 		}
 	} else {
-		c, err = net.Dial("tcp", rhost)
+		c, err = dialer.Dial("tcp", rhost)
 	}
 
 	if err != nil {