read: Fix CloseRead to have its own done channel

Context can be cancelled by parent. Doesn't indicate the CloseRead goroutine
has exited.

Author: nhooyr

close.go

@@ -239,11 +239,11 @@ func (c *Conn) waitGoroutines() error {
 	}
 
 	c.closeReadMu.Lock()
-	ctx := c.closeReadCtx
+	closeRead := c.closeReadCtx != nil
 	c.closeReadMu.Unlock()
-	if ctx != nil {
+	if closeRead {
 		select {
-		case <-ctx.Done():
+		case <-c.closeReadDone:
 		case <-t.C:
 			return errors.New("failed to wait for close read goroutine to exit")
 		}

conn.go

@@ -57,10 +57,10 @@ type Conn struct {
 	timeoutLoopDone chan struct{}
 
 	// Read state.
-	readMu            *mu
-	readHeaderBuf     [8]byte
-	readControlBuf    [maxControlPayload]byte
-	msgReader         *msgReader
+	readMu         *mu
+	readHeaderBuf  [8]byte
+	readControlBuf [maxControlPayload]byte
+	msgReader      *msgReader
 
 	// Write state.
 	msgWriter      *msgWriter
@@ -69,8 +69,9 @@ type Conn struct {
 	writeHeaderBuf [8]byte
 	writeHeader    header
 
-	closeReadMu  sync.Mutex
-	closeReadCtx context.Context
+	closeReadMu   sync.Mutex
+	closeReadCtx  context.Context
+	closeReadDone chan struct{}
 
 	closed  chan struct{}
 	closeMu sync.Mutex

mask_asm.go

@@ -3,7 +3,7 @@
 package websocket
 
 func mask(b []byte, key uint32) uint32 {
-    // TODO: Will enable in v1.9.0.
+	// TODO: Will enable in v1.9.0.
 	return maskGo(b, key)
 	if len(b) > 0 {
 		return maskAsm(&b[0], len(b), key)

read.go

@@ -71,9 +71,11 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 	}
 	ctx, cancel := context.WithCancel(ctx)
 	c.closeReadCtx = ctx
+	c.closeReadDone = make(chan struct{})
 	c.closeReadMu.Unlock()
 
 	go func() {
+		defer close(c.closeReadDone)
 		defer cancel()
 		defer c.close()
 		_, _, err := c.Reader(ctx)