From bcef8431c98087addcb2f0ab484ff295abe41a74 Mon Sep 17 00:00:00 2001
From: Gary Burd <gary@beagledreams.com>
Date: Sat, 1 Jan 2022 08:43:22 -0800
Subject: [PATCH] Use context.Context in TLS handshake (#751)

Continued work on #730.
---
 .circleci/config.yml |  2 +-
 client.go            | 23 ++++++-----------------
 tls_handshake.go     | 21 +++++++++++++++++++++
 tls_handshake_116.go | 21 +++++++++++++++++++++
 trace.go             | 20 --------------------
 trace_17.go          | 13 -------------
 6 files changed, 49 insertions(+), 51 deletions(-)
 create mode 100644 tls_handshake.go
 create mode 100644 tls_handshake_116.go
 delete mode 100644 trace.go
 delete mode 100644 trace_17.go

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 554a446d..a0eb0ed2 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -67,4 +67,4 @@ workflows:
       - test:
           matrix:
             parameters:
-              version: ["latest", "1.15", "1.14", "1.13", "1.12", "1.11"]
+              version: ["latest", "1.17", "1.16", "1.15", "1.14", "1.13", "1.12", "1.11"]
diff --git a/client.go b/client.go
index 196a659c..a24c3ce2 100644
--- a/client.go
+++ b/client.go
@@ -314,11 +314,12 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
 		tlsConn := tls.Client(netConn, cfg)
 		netConn = tlsConn
 
-		var err error
-		if trace != nil {
-			err = doHandshakeWithTrace(trace, tlsConn, cfg)
-		} else {
-			err = doHandshake(tlsConn, cfg)
+		if trace != nil && trace.TLSHandshakeStart != nil {
+			trace.TLSHandshakeStart()
+		}
+		err := doHandshake(ctx, tlsConn, cfg)
+		if trace != nil && trace.TLSHandshakeDone != nil {
+			trace.TLSHandshakeDone(tlsConn.ConnectionState(), err)
 		}
 
 		if err != nil {
@@ -383,15 +384,3 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
 	netConn = nil // to avoid close in defer.
 	return conn, resp, nil
 }
-
-func doHandshake(tlsConn *tls.Conn, cfg *tls.Config) error {
-	if err := tlsConn.Handshake(); err != nil {
-		return err
-	}
-	if !cfg.InsecureSkipVerify {
-		if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil {
-			return err
-		}
-	}
-	return nil
-}
diff --git a/tls_handshake.go b/tls_handshake.go
new file mode 100644
index 00000000..a62b68cc
--- /dev/null
+++ b/tls_handshake.go
@@ -0,0 +1,21 @@
+//go:build go1.17
+// +build go1.17
+
+package websocket
+
+import (
+	"context"
+	"crypto/tls"
+)
+
+func doHandshake(ctx context.Context, tlsConn *tls.Conn, cfg *tls.Config) error {
+	if err := tlsConn.HandshakeContext(ctx); err != nil {
+		return err
+	}
+	if !cfg.InsecureSkipVerify {
+		if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/tls_handshake_116.go b/tls_handshake_116.go
new file mode 100644
index 00000000..e1b2b44f
--- /dev/null
+++ b/tls_handshake_116.go
@@ -0,0 +1,21 @@
+//go:build !go1.17
+// +build !go1.17
+
+package websocket
+
+import (
+	"context"
+	"crypto/tls"
+)
+
+func doHandshake(ctx context.Context, tlsConn *tls.Conn, cfg *tls.Config) error {
+	if err := tlsConn.Handshake(); err != nil {
+		return err
+	}
+	if !cfg.InsecureSkipVerify {
+		if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/trace.go b/trace.go
deleted file mode 100644
index 246a5d33..00000000
--- a/trace.go
+++ /dev/null
@@ -1,20 +0,0 @@
-//go:build go1.8
-// +build go1.8
-
-package websocket
-
-import (
-	"crypto/tls"
-	"net/http/httptrace"
-)
-
-func doHandshakeWithTrace(trace *httptrace.ClientTrace, tlsConn *tls.Conn, cfg *tls.Config) error {
-	if trace.TLSHandshakeStart != nil {
-		trace.TLSHandshakeStart()
-	}
-	err := doHandshake(tlsConn, cfg)
-	if trace.TLSHandshakeDone != nil {
-		trace.TLSHandshakeDone(tlsConn.ConnectionState(), err)
-	}
-	return err
-}
diff --git a/trace_17.go b/trace_17.go
deleted file mode 100644
index f4be940a..00000000
--- a/trace_17.go
+++ /dev/null
@@ -1,13 +0,0 @@
-//go:build !go1.8
-// +build !go1.8
-
-package websocket
-
-import (
-	"crypto/tls"
-	"net/http/httptrace"
-)
-
-func doHandshakeWithTrace(trace *httptrace.ClientTrace, tlsConn *tls.Conn, cfg *tls.Config) error {
-	return doHandshake(tlsConn, cfg)
-}