From 31257903a38cc27d1d10b784388938aa0b221afd Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sat, 3 Nov 2018 11:24:25 +0800
Subject: [PATCH] add timeout for TLS handshaking (#316)

---
 http2.go |  4 ++--
 tls.go   | 14 +++++++++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/http2.go b/http2.go
index c30fd96..6fc0392 100644
--- a/http2.go
+++ b/http2.go
@@ -114,7 +114,7 @@ func (tr *http2Transporter) Dial(addr string, options ...DialOption) (net.Conn,
 				if err != nil {
 					return nil, err
 				}
-				return wrapTLSClient(conn, cfg)
+				return wrapTLSClient(conn, cfg, opts.Timeout)
 			},
 		}
 		client = &http.Client{
@@ -182,7 +182,7 @@ func (tr *h2Transporter) Dial(addr string, options ...DialOption) (net.Conn, err
 				if tr.tlsConfig == nil {
 					return conn, nil
 				}
-				return wrapTLSClient(conn, cfg)
+				return wrapTLSClient(conn, cfg, opts.Timeout)
 			},
 		}
 		client = &http.Client{
diff --git a/tls.go b/tls.go
index 35791d6..9ad4594 100644
--- a/tls.go
+++ b/tls.go
@@ -30,7 +30,7 @@ func (tr *tlsTransporter) Handshake(conn net.Conn, options ...HandshakeOption) (
 	if opts.TLSConfig == nil {
 		opts.TLSConfig = &tls.Config{InsecureSkipVerify: true}
 	}
-	return wrapTLSClient(conn, opts.TLSConfig)
+	return wrapTLSClient(conn, opts.TLSConfig, opts.Timeout)
 }
 
 type mtlsTransporter struct {
@@ -113,7 +113,7 @@ func (tr *mtlsTransporter) initSession(addr string, conn net.Conn, opts *Handsha
 	if opts.TLSConfig == nil {
 		opts.TLSConfig = &tls.Config{InsecureSkipVerify: true}
 	}
-	conn, err := wrapTLSClient(conn, opts.TLSConfig)
+	conn, err := wrapTLSClient(conn, opts.TLSConfig, opts.Timeout)
 	if err != nil {
 		return nil, err
 	}
@@ -248,7 +248,7 @@ func (l *mtlsListener) Close() error {
 //
 // This code is taken from consul:
 // https://github.com/hashicorp/consul/blob/master/tlsutil/config.go
-func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
+func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config, timeout time.Duration) (net.Conn, error) {
 	var err error
 	var tlsConn *tls.Conn
 
@@ -264,6 +264,12 @@ func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
 		return tlsConn, nil
 	}
 
+	if timeout <= 0 {
+		timeout = 10 * time.Second // default timeout
+	}
+
+	tlsConn.SetDeadline(time.Now().Add(timeout))
+
 	// Otherwise perform handshake, but don't verify the domain
 	//
 	// The following is lightly-modified from the doFullHandshake
@@ -273,6 +279,8 @@ func wrapTLSClient(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
 		return nil, err
 	}
 
+	tlsConn.SetDeadline(time.Time{}) // clear timeout
+
 	opts := x509.VerifyOptions{
 		Roots:         tlsConfig.RootCAs,
 		CurrentTime:   time.Now(),