read.go: Avoid handling ping after close frame has been sent

bhallionOhbibi · nhooyr · commit a856e5eca11e · 2023-10-18T21:57:30.000-07:00
Closes #298
diff --git a/conn_test.go b/conn_test.go
@@ -553,3 +553,26 @@ func assertClose(tb testing.TB, c *websocket.Conn) {
 	err := c.Close(websocket.StatusNormalClosure, "")
 	assert.Success(tb, err)
 }
+
+func TestConcurrentClosePing(t *testing.T) {
+	t.Parallel()
+	for i := 0; i < 999; i++ {
+		func() {
+			c1, c2 := wstest.Pipe(nil, nil)
+			defer c1.CloseNow()
+			defer c2.CloseNow()
+			c1.CloseRead(context.Background())
+			c2.CloseRead(context.Background())
+			go func() {
+				for range time.Tick(time.Millisecond) {
+					if err := c1.Ping(context.Background()); err != nil {
+						return
+					}
+				}
+			}()
+
+			time.Sleep(30 * time.Millisecond)
+			assert.Success(t, c1.Close(websocket.StatusNormalClosure, ""))
+		}()
+	}
+}
diff --git a/read.go b/read.go
@@ -1,4 +1,3 @@
-//go:build !js
 // +build !js
 
 package websocket
@@ -287,6 +286,13 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 
 	switch h.opcode {
 	case opPing:
+		c.closeMu.Lock()
+		wroteClose := c.wroteClose
+		c.closeMu.Unlock()
+		if wroteClose {
+			// Cannot respond to ping with a pong because we already sent a close frame.
+			return nil
+		}
 		return c.writeControl(ctx, opPong, b)
 	case opPong:
 		c.activePingsMu.Lock()
