fix closenow race

Author: alixander

accept_test.go

@@ -6,10 +6,12 @@ package websocket
 import (
 	"bufio"
 	"errors"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"testing"
 
 	"nhooyr.io/websocket/internal/test/assert"
@@ -142,6 +144,43 @@ func TestAccept(t *testing.T) {
 		_, err := Accept(w, r, nil)
 		assert.Contains(t, err, `failed to hijack connection`)
 	})
+	t.Run("closeRace", func(t *testing.T) {
+		t.Parallel()
+
+		server, _ := net.Pipe()
+
+		pr, pw := io.Pipe()
+		rw := bufio.NewReadWriter(bufio.NewReader(pr), bufio.NewWriter(pw))
+		newResponseWriter := func() http.ResponseWriter {
+			return mockHijacker{
+				ResponseWriter: httptest.NewRecorder(),
+				hijack: func() (net.Conn, *bufio.ReadWriter, error) {
+					return server, rw, nil
+				},
+			}
+		}
+		w := newResponseWriter()
+
+		r := httptest.NewRequest("GET", "/", nil)
+		r.Header.Set("Connection", "Upgrade")
+		r.Header.Set("Upgrade", "websocket")
+		r.Header.Set("Sec-WebSocket-Version", "13")
+		r.Header.Set("Sec-WebSocket-Key", xrand.Base64(16))
+
+		c, err := Accept(w, r, nil)
+		wg := &sync.WaitGroup{}
+		wg.Add(2)
+		go func() {
+			c.Close(StatusInternalError, "the sky is falling")
+			wg.Done()
+		}()
+		go func() {
+			c.CloseNow()
+			wg.Done()
+		}()
+		wg.Wait()
+		assert.Success(t, err)
+	})
 }
 
 func Test_verifyClientHandshake(t *testing.T) {

close.go

@@ -114,6 +114,8 @@ func (c *Conn) CloseNow() (err error) {
 	}
 
 	c.close(nil)
+	c.closeMu.Lock()
+	defer c.closeMu.Unlock()
 	return c.closeErr
 }