fix: rewrite close handshake flow when initiated from other side

Author: mafredri

close.go

@@ -100,7 +100,7 @@ func CloseStatus(err error) StatusCode {
 func (c *Conn) Close(code StatusCode, reason string) (err error) {
 	defer errd.Wrap(&err, "failed to close WebSocket")
 
-	if !c.casClosing() {
+	if c.casClosing() {
 		err = c.waitGoroutines()
 		if err != nil {
 			return err
@@ -133,7 +133,7 @@ func (c *Conn) Close(code StatusCode, reason string) (err error) {
 func (c *Conn) CloseNow() (err error) {
 	defer errd.Wrap(&err, "failed to immediately close WebSocket")
 
-	if !c.casClosing() {
+	if c.casClosing() {
 		err = c.waitGoroutines()
 		if err != nil {
 			return err
@@ -206,10 +206,6 @@ func (c *Conn) waitCloseHandshake() error {
 	}
 	defer c.readMu.unlock()
 
-	if c.readCloseErr != nil {
-		return c.readCloseErr
-	}
-
 	for i := int64(0); i < c.msgReader.payloadLength; i++ {
 		_, err := c.br.ReadByte()
 		if err != nil {
@@ -333,13 +329,7 @@ func (ce CloseError) bytesErr() ([]byte, error) {
 }
 
 func (c *Conn) casClosing() bool {
-	c.closeMu.Lock()
-	defer c.closeMu.Unlock()
-	if !c.closing {
-		c.closing = true
-		return true
-	}
-	return false
+	return c.closing.Swap(true)
 }
 
 func (c *Conn) isClosed() bool {

conn.go

@@ -61,24 +61,27 @@ type Conn struct {
 	readHeaderBuf  [8]byte
 	readControlBuf [maxControlPayload]byte
 	msgReader      *msgReader
-	readCloseErr   error
 
 	// Write state.
 	msgWriter      *msgWriter
 	writeFrameMu   *mu
 	writeBuf       []byte
 	writeHeaderBuf [8]byte
 	writeHeader    header
-	closeSent      bool
+
+	// Close handshake state.
+	closeStateMu     sync.RWMutex
+	closeReceivedErr error
+	closeSentErr     error
 
 	// CloseRead state.
 	closeReadMu   sync.Mutex
 	closeReadCtx  context.Context
 	closeReadDone chan struct{}
 
-	closed  chan struct{}
+	closing atomic.Bool
 	closeMu sync.Mutex
-	closing bool
+	closed  chan struct{}
 
 	pingCounter   atomic.Int64
 	activePingsMu sync.Mutex

conn_test.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -625,3 +626,149 @@ func TestConcurrentClosePing(t *testing.T) {
 		}()
 	}
 }
+
+func TestConnClosePropagation(t *testing.T) {
+	t.Parallel()
+
+	want := []byte("hello")
+	keepWriting := func(c *websocket.Conn) <-chan error {
+		return xsync.Go(func() error {
+			for {
+				err := c.Write(context.Background(), websocket.MessageText, want)
+				if err != nil {
+					return err
+				}
+			}
+		})
+	}
+	keepReading := func(c *websocket.Conn) <-chan error {
+		return xsync.Go(func() error {
+			for {
+				_, got, err := c.Read(context.Background())
+				if err != nil {
+					return err
+				}
+				if !bytes.Equal(want, got) {
+					return fmt.Errorf("unexpected message: want %q, got %q", want, got)
+				}
+			}
+		})
+	}
+	checkReadErr := func(t *testing.T, err error) {
+		var ce websocket.CloseError
+		if errors.As(err, &ce) {
+			assert.Equal(t, "", websocket.StatusNormalClosure, ce.Code)
+		} else {
+			assert.ErrorIs(t, net.ErrClosed, err)
+		}
+	}
+	checkConnErrs := func(t *testing.T, conn ...*websocket.Conn) {
+		for _, c := range conn {
+			// Check write error.
+			err := c.Write(context.Background(), websocket.MessageText, want)
+			assert.ErrorIs(t, net.ErrClosed, err)
+
+			// Check read error (output depends on when read is called in relation to connection closure).
+			_, _, err = c.Read(context.Background())
+			checkReadErr(t, err)
+		}
+	}
+
+	t.Run("CloseOtherSideDuringWrite", func(t *testing.T) {
+		tt, this, other := newConnTest(t, nil, nil)
+
+		_ = this.CloseRead(tt.ctx)
+		thisWriteErr := keepWriting(this)
+
+		_, got, err := other.Read(tt.ctx)
+		assert.Success(t, err)
+		assert.Equal(t, "msg", want, got)
+
+		err = other.Close(websocket.StatusNormalClosure, "")
+		assert.Success(t, err)
+
+		select {
+		case err := <-thisWriteErr:
+			assert.ErrorIs(t, net.ErrClosed, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		checkConnErrs(t, this, other)
+	})
+	t.Run("CloseThisSideDuringWrite", func(t *testing.T) {
+		tt, this, other := newConnTest(t, nil, nil)
+
+		_ = this.CloseRead(tt.ctx)
+		thisWriteErr := keepWriting(this)
+		otherReadErr := keepReading(other)
+
+		err := this.Close(websocket.StatusNormalClosure, "")
+		assert.Success(t, err)
+
+		select {
+		case err := <-thisWriteErr:
+			assert.ErrorIs(t, net.ErrClosed, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		select {
+		case err := <-otherReadErr:
+			checkReadErr(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		checkConnErrs(t, this, other)
+	})
+	t.Run("CloseOtherSideDuringRead", func(t *testing.T) {
+		tt, this, other := newConnTest(t, nil, nil)
+
+		_ = other.CloseRead(tt.ctx)
+		errs := keepReading(this)
+
+		err := other.Write(tt.ctx, websocket.MessageText, want)
+		assert.Success(t, err)
+
+		err = other.Close(websocket.StatusNormalClosure, "")
+		assert.Success(t, err)
+
+		select {
+		case err := <-errs:
+			checkReadErr(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		checkConnErrs(t, this, other)
+	})
+	t.Run("CloseThisSideDuringRead", func(t *testing.T) {
+		tt, this, other := newConnTest(t, nil, nil)
+
+		thisReadErr := keepReading(this)
+		otherReadErr := keepReading(other)
+
+		err := other.Write(tt.ctx, websocket.MessageText, want)
+		assert.Success(t, err)
+
+		err = this.Close(websocket.StatusNormalClosure, "")
+		assert.Success(t, err)
+
+		select {
+		case err := <-thisReadErr:
+			checkReadErr(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		select {
+		case err := <-otherReadErr:
+			checkReadErr(t, err)
+		case <-tt.ctx.Done():
+			t.Fatal(tt.ctx.Err())
+		}
+
+		checkConnErrs(t, this, other)
+	})
+}

read.go

@@ -181,15 +181,6 @@ func (c *Conn) readRSV1Illegal(h header) bool {
 }
 
 func (c *Conn) readLoop(ctx context.Context) (header, error) {
-	if c.readCloseErr != nil {
-		select {
-		case <-c.closed:
-			return header{}, net.ErrClosed
-		default:
-		}
-		return header{}, c.readCloseErr
-	}
-
 	for {
 		h, err := c.readFrameHeader(ctx)
 		if err != nil {
@@ -226,57 +217,59 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) {
 	}
 }
 
-func (c *Conn) readFrameHeader(ctx context.Context) (header, error) {
+func (c *Conn) prepareRead(ctx context.Context, err *error) (func(), error) {
 	select {
 	case <-c.closed:
-		return header{}, net.ErrClosed
+		return nil, net.ErrClosed
 	case c.readTimeout <- ctx:
 	}
 
-	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
-	if err != nil {
+	c.closeStateMu.Lock()
+	closeReceivedErr := c.closeReceivedErr
+	c.closeStateMu.Unlock()
+	if closeReceivedErr != nil {
+		return nil, closeReceivedErr
+	}
+
+	return func() {
 		select {
 		case <-c.closed:
-			return header{}, net.ErrClosed
-		case <-ctx.Done():
-			return header{}, ctx.Err()
-		default:
-			return header{}, err
+			if *err != nil {
+				*err = net.ErrClosed
+			}
+		case c.writeTimeout <- context.Background():
 		}
+		if *err != nil && ctx.Err() != nil {
+			*err = ctx.Err()
+		}
+	}, nil
+}
+
+func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {
+	readDone, err := c.prepareRead(ctx, &err)
+	if err != nil {
+		return header{}, err
 	}
+	defer readDone()
 
-	select {
-	case <-c.closed:
-		return header{}, net.ErrClosed
-	case c.readTimeout <- context.Background():
+	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
+	if err != nil {
+		return header{}, err
 	}
 
 	return h, nil
 }
 
-func (c *Conn) readFramePayload(ctx context.Context, p []byte) (int, error) {
-	select {
-	case <-c.closed:
-		return 0, net.ErrClosed
-	case c.readTimeout <- ctx:
+func (c *Conn) readFramePayload(ctx context.Context, p []byte) (_ int, err error) {
+	readDone, err := c.prepareRead(ctx, &err)
+	if err != nil {
+		return 0, err
 	}
+	defer readDone()
 
 	n, err := io.ReadFull(c.br, p)
 	if err != nil {
-		select {
-		case <-c.closed:
-			return n, net.ErrClosed
-		case <-ctx.Done():
-			return n, ctx.Err()
-		default:
-			return n, fmt.Errorf("failed to read frame payload: %w", err)
-		}
-	}
-
-	select {
-	case <-c.closed:
-		return n, net.ErrClosed
-	case c.readTimeout <- context.Background():
+		return n, fmt.Errorf("failed to read frame payload: %w", err)
 	}
 
 	return n, err
@@ -333,18 +326,20 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 		return err
 	}
 
-	if c.readCloseErr == nil {
-		c.readCloseErr = ce
-	}
-
 	err = fmt.Errorf("received close frame: %w", ce)
-	if err2 := c.writeClose(ce.Code, ce.Reason); errors.Is(err2, errCloseSent) {
-		// The close handshake has already been initiated, connection
-		// close should be handled elsewhere.
-		return err
+	c.closeStateMu.Lock()
+	c.closeReceivedErr = err
+	closeSent := c.closeSentErr != nil
+	c.closeStateMu.Unlock()
+
+	if !closeSent {
+		c.readMu.unlock()
+		_ = c.writeClose(ce.Code, ce.Reason)
+	}
+	if !c.casClosing() {
+		c.readMu.unlock()
+		_ = c.close()
 	}
-	c.readMu.unlock()
-	c.close()
 	return err
 }