From 87352bcaddfbf00f16d376a920ab5b8b409f95a9 Mon Sep 17 00:00:00 2001
From: Martin Greenwald <greenwald@outlook.com>
Date: Wed, 14 Feb 2024 19:06:54 -0800
Subject: [PATCH] Do not timeout when WriteControl deadline is zero

A zero value for the Conn.WriteControl deadline specifies no timeout,
but the feature was implemented as a very long timeout (1000 hours).
This PR updates the code to use no timeout when the deadline is zero.

See the discussion in #895 for more details.
---
 conn.go      | 25 +++++++++++++------------
 conn_test.go | 16 ++++++++++++++++
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/conn.go b/conn.go
index ac64831..5bc9e80 100644
--- a/conn.go
+++ b/conn.go
@@ -435,23 +435,24 @@ func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) er
 		maskBytes(key, 0, buf[6:])
 	}
 
-	d := 1000 * time.Hour
-	if !deadline.IsZero() {
-		d = deadline.Sub(time.Now())
+	if deadline.IsZero() {
+		// No timeout for zero time.
+		<-c.mu
+	} else {
+		d := time.Until(deadline)
 		if d < 0 {
 			return errWriteTimeout
 		}
-	}
-
-	select {
-	case <-c.mu:
-	default:
-		timer := time.NewTimer(d)
 		select {
 		case <-c.mu:
-			timer.Stop()
-		case <-timer.C:
-			return errWriteTimeout
+		default:
+			timer := time.NewTimer(d)
+			select {
+			case <-c.mu:
+				timer.Stop()
+			case <-timer.C:
+				return errWriteTimeout
+			}
 		}
 	}
 
diff --git a/conn_test.go b/conn_test.go
index fae9f85..19776c6 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -148,6 +148,22 @@ func TestFraming(t *testing.T) {
 	}
 }
 
+func TestWriteControlDeadline(t *testing.T) {
+	t.Parallel()
+	message := []byte("hello")
+	var connBuf bytes.Buffer
+	c := newTestConn(nil, &connBuf, true)
+	if err := c.WriteControl(PongMessage, message, time.Time{}); err != nil {
+		t.Errorf("WriteControl(..., zero deadline) = %v, want nil", err)
+	}
+	if err := c.WriteControl(PongMessage, message, time.Now().Add(time.Second)); err != nil {
+		t.Errorf("WriteControl(..., future deadline) = %v, want nil", err)
+	}
+	if err := c.WriteControl(PongMessage, message, time.Now().Add(-time.Second)); err == nil {
+		t.Errorf("WriteControl(..., past deadline) = nil, want timeout error")
+	}
+}
+
 func TestConcurrencyWriteControl(t *testing.T) {
 	const message = "this is a ping/pong messsage"
 	loop := 10