From d3b03e423176af2c3585730a10fbc4e539be2f3a Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Mon, 24 Dec 2018 13:14:53 +0800
Subject: [PATCH] add default timeout

---
 obfs.go |  8 ++++++++
 ssh.go  | 21 ++++++++++++++++++++-
 tls.go  | 27 +++++++++++++--------------
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/obfs.go b/obfs.go
index 2c7da09..b0cea1a 100644
--- a/obfs.go
+++ b/obfs.go
@@ -320,6 +320,14 @@ func (tr *obfs4Transporter) Handshake(conn net.Conn, options ...HandshakeOption)
 	for _, option := range options {
 		option(opts)
 	}
+
+	timeout := opts.Timeout
+	if timeout <= 0 {
+		timeout = HandshakeTimeout
+	}
+	conn.SetDeadline(time.Now().Add(timeout))
+	defer conn.SetDeadline(time.Time{})
+
 	return obfs4ClientConn(opts.Addr, conn)
 }
 
diff --git a/ssh.go b/ssh.go
index c2dc5a5..6361804 100644
--- a/ssh.go
+++ b/ssh.go
@@ -40,10 +40,24 @@ func SSHDirectForwardConnector() Connector {
 }
 
 func (c *sshDirectForwardConnector) Connect(conn net.Conn, raddr string, options ...ConnectOption) (net.Conn, error) {
+	opts := &ConnectOptions{}
+	for _, option := range options {
+		option(opts)
+	}
+
 	cc, ok := conn.(*sshNopConn) // TODO: this is an ugly type assertion, need to find a better solution.
 	if !ok {
 		return nil, errors.New("ssh: wrong connection type")
 	}
+
+	timeout := opts.Timeout
+	if timeout <= 0 {
+		timeout = ConnectTimeout
+	}
+
+	cc.session.conn.SetDeadline(time.Now().Add(timeout))
+	defer cc.session.conn.SetDeadline(time.Time{})
+
 	conn, err := cc.session.client.Dial("tcp", raddr)
 	if err != nil {
 		log.Logf("[ssh-tcp] %s -> %s : %s", cc.session.addr, raddr, err)
@@ -177,6 +191,9 @@ func (tr *sshForwardTransporter) Handshake(conn net.Conn, options ...HandshakeOp
 	tr.sessionMutex.Lock()
 	defer tr.sessionMutex.Unlock()
 
+	conn.SetDeadline(time.Now().Add(timeout))
+	defer conn.SetDeadline(time.Time{})
+
 	session, ok := tr.sessions[opts.Addr]
 	if !ok || session.client == nil {
 		sshConn, chans, reqs, err := ssh.NewClientConn(conn, opts.Addr, &config)
@@ -269,7 +286,6 @@ func (tr *sshTunnelTransporter) Handshake(conn net.Conn, options ...HandshakeOpt
 	}
 
 	config := ssh.ClientConfig{
-		Timeout:         timeout,
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
 	// TODO: support pubkey auth.
@@ -284,6 +300,9 @@ func (tr *sshTunnelTransporter) Handshake(conn net.Conn, options ...HandshakeOpt
 	tr.sessionMutex.Lock()
 	defer tr.sessionMutex.Unlock()
 
+	conn.SetDeadline(time.Now().Add(timeout))
+	defer conn.SetDeadline(time.Time{})
+
 	session, ok := tr.sessions[opts.Addr]
 	if !ok || session.client == nil {
 		sshConn, chans, reqs, err := ssh.NewClientConn(conn, opts.Addr, &config)
diff --git a/tls.go b/tls.go
index f8e45cc..df5a2bf 100644
--- a/tls.go
+++ b/tls.go
@@ -271,23 +271,14 @@ func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config, timeout time.Duration)
 	var err error
 	var tlsConn *tls.Conn
 
-	tlsConn = tls.Client(conn, tlsConfig)
-
-	// If crypto/tls is doing verification, there's no need to do our own.
-	if tlsConfig.InsecureSkipVerify == false {
-		return tlsConn, nil
-	}
-
-	// Similarly if we use host's CA, we can do full handshake
-	if tlsConfig.RootCAs == nil {
-		return tlsConn, nil
-	}
-
 	if timeout <= 0 {
 		timeout = HandshakeTimeout // default timeout
 	}
 
-	tlsConn.SetDeadline(time.Now().Add(timeout))
+	conn.SetDeadline(time.Now().Add(timeout))
+	defer conn.SetDeadline(time.Time{})
+
+	tlsConn = tls.Client(conn, tlsConfig)
 
 	// Otherwise perform handshake, but don't verify the domain
 	//
@@ -298,7 +289,15 @@ func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config, timeout time.Duration)
 		return nil, err
 	}
 
-	tlsConn.SetDeadline(time.Time{}) // clear timeout
+	// If crypto/tls is doing verification, there's no need to do our own.
+	if tlsConfig.InsecureSkipVerify == false {
+		return tlsConn, nil
+	}
+
+	// Similarly if we use host's CA, we can do full handshake
+	if tlsConfig.RootCAs == nil {
+		return tlsConn, nil
+	}
 
 	opts := x509.VerifyOptions{
 		Roots:         tlsConfig.RootCAs,