From f90b62c3ce74143fb49486e32777ba418d1f05a5 Mon Sep 17 00:00:00 2001
From: dottyjones <38705838+dottyjones@users.noreply.github.com>
Date: Mon, 28 May 2018 07:28:04 -0700
Subject: [PATCH 1/2] Add test for handshake deadline

---
 client_server_test.go | 57 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/client_server_test.go b/client_server_test.go
index dcaab3d1..20b648be 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -344,6 +344,63 @@ func TestDialTimeout(t *testing.T) {
 	}
 }
 
+// netConnDeadlineObserver fails test if read or write called without deadline.
+type netConnDeadlineObserver struct {
+	t           *testing.T
+	c           net.Conn
+	read, write bool
+}
+
+func (c *netConnDeadlineObserver) SetDeadline(t time.Time) error {
+	c.write = !t.Equal(time.Time{})
+	c.read = c.write
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) SetReadDeadline(t time.Time) error {
+	c.read = !t.Equal(time.Time{})
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) SetWriteDeadline(t time.Time) error {
+	c.write = !t.Equal(time.Time{})
+	return c.c.SetDeadline(t)
+}
+
+func (c *netConnDeadlineObserver) Write(p []byte) (int, error) {
+	if !c.write {
+		c.t.Fatalf("write with no deadline")
+	}
+	return c.c.Write(p)
+}
+
+func (c *netConnDeadlineObserver) Read(p []byte) (int, error) {
+	if !c.read {
+		c.t.Fatalf("read with no deadline")
+	}
+	return c.c.Read(p)
+}
+
+func (c *netConnDeadlineObserver) Close() error         { return c.c.Close() }
+func (c *netConnDeadlineObserver) LocalAddr() net.Addr  { return c.c.LocalAddr() }
+func (c *netConnDeadlineObserver) RemoteAddr() net.Addr { return c.c.RemoteAddr() }
+
+func TestHandshakeTimeout(t *testing.T) {
+	s := newServer(t)
+	defer s.Close()
+
+	d := cstDialer
+	d.NetDial = func(n, a string) (net.Conn, error) {
+		c, err := net.Dial(n, a)
+		return &netConnDeadlineObserver{c: c, t: t}, err
+	}
+	ws, _, err := d.Dial(s.URL, nil)
+	if err != nil {
+		t.Fatal("Dial:", err)
+	}
+	ws.Close()
+}
+
 func TestDialBadScheme(t *testing.T) {
 	s := newServer(t)
 	defer s.Close()

From badcf8718599c68a4868b0c16322c0069046d997 Mon Sep 17 00:00:00 2001
From: dottyjones <38705838+dottyjones@users.noreply.github.com>
Date: Mon, 28 May 2018 12:07:56 -0700
Subject: [PATCH 2/2] Improve names in handshake deadline test

---
 client_server_test.go | 46 ++++++++++++++++++++++---------------------
 1 file changed, 24 insertions(+), 22 deletions(-)

diff --git a/client_server_test.go b/client_server_test.go
index 20b648be..cdbc9d4b 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -344,46 +344,48 @@ func TestDialTimeout(t *testing.T) {
 	}
 }
 
-// netConnDeadlineObserver fails test if read or write called without deadline.
-type netConnDeadlineObserver struct {
-	t           *testing.T
-	c           net.Conn
-	read, write bool
-}
-
-func (c *netConnDeadlineObserver) SetDeadline(t time.Time) error {
-	c.write = !t.Equal(time.Time{})
-	c.read = c.write
+// requireDeadlineNetConn fails the current test when Read or Write are called
+// with no deadline.
+type requireDeadlineNetConn struct {
+	t                  *testing.T
+	c                  net.Conn
+	readDeadlineIsSet  bool
+	writeDeadlineIsSet bool
+}
+
+func (c *requireDeadlineNetConn) SetDeadline(t time.Time) error {
+	c.writeDeadlineIsSet = !t.Equal(time.Time{})
+	c.readDeadlineIsSet = c.writeDeadlineIsSet
 	return c.c.SetDeadline(t)
 }
 
-func (c *netConnDeadlineObserver) SetReadDeadline(t time.Time) error {
-	c.read = !t.Equal(time.Time{})
+func (c *requireDeadlineNetConn) SetReadDeadline(t time.Time) error {
+	c.readDeadlineIsSet = !t.Equal(time.Time{})
 	return c.c.SetDeadline(t)
 }
 
-func (c *netConnDeadlineObserver) SetWriteDeadline(t time.Time) error {
-	c.write = !t.Equal(time.Time{})
+func (c *requireDeadlineNetConn) SetWriteDeadline(t time.Time) error {
+	c.writeDeadlineIsSet = !t.Equal(time.Time{})
 	return c.c.SetDeadline(t)
 }
 
-func (c *netConnDeadlineObserver) Write(p []byte) (int, error) {
-	if !c.write {
+func (c *requireDeadlineNetConn) Write(p []byte) (int, error) {
+	if !c.writeDeadlineIsSet {
 		c.t.Fatalf("write with no deadline")
 	}
 	return c.c.Write(p)
 }
 
-func (c *netConnDeadlineObserver) Read(p []byte) (int, error) {
-	if !c.read {
+func (c *requireDeadlineNetConn) Read(p []byte) (int, error) {
+	if !c.readDeadlineIsSet {
 		c.t.Fatalf("read with no deadline")
 	}
 	return c.c.Read(p)
 }
 
-func (c *netConnDeadlineObserver) Close() error         { return c.c.Close() }
-func (c *netConnDeadlineObserver) LocalAddr() net.Addr  { return c.c.LocalAddr() }
-func (c *netConnDeadlineObserver) RemoteAddr() net.Addr { return c.c.RemoteAddr() }
+func (c *requireDeadlineNetConn) Close() error         { return c.c.Close() }
+func (c *requireDeadlineNetConn) LocalAddr() net.Addr  { return c.c.LocalAddr() }
+func (c *requireDeadlineNetConn) RemoteAddr() net.Addr { return c.c.RemoteAddr() }
 
 func TestHandshakeTimeout(t *testing.T) {
 	s := newServer(t)
@@ -392,7 +394,7 @@ func TestHandshakeTimeout(t *testing.T) {
 	d := cstDialer
 	d.NetDial = func(n, a string) (net.Conn, error) {
 		c, err := net.Dial(n, a)
-		return &netConnDeadlineObserver{c: c, t: t}, err
+		return &requireDeadlineNetConn{c: c, t: t}, err
 	}
 	ws, _, err := d.Dial(s.URL, nil)
 	if err != nil {