fix: echo read error after close received

coder · mafredri · Dec 4, 2024 · Sep 5, 2024 · Sep 6, 2024 · Sep 6, 2024
commit c3613fc66341cffa20fc9eacb3017bec47a50f59
diff --git a/close.go b/close.go
@@ -206,6 +206,10 @@ 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 {

diff --git a/conn.go b/conn.go
@@ -61,6 +61,7 @@ type Conn struct {
 	readHeaderBuf  [8]byte
 	readControlBuf [maxControlPayload]byte
 	msgReader      *msgReader
+	readCloseErr   error
 
 	// Write state.
 	msgWriter      *msgWriter
@@ -70,6 +71,7 @@ type Conn struct {
 	writeHeader    header
 	closeSent      bool
 
+	// CloseRead state.
 	closeReadMu   sync.Mutex
 	closeReadCtx  context.Context
 	closeReadDone chan struct{}

diff --git a/read.go b/read.go
@@ -181,6 +181,15 @@ 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 {
@@ -324,8 +333,16 @@ 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)
-	c.writeClose(ce.Code, ce.Reason)
+	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.readMu.unlock()
 	c.close()
 	return err

diff --git a/write.go b/write.go
@@ -241,6 +241,8 @@ func (c *Conn) writeControl(ctx context.Context, opcode opcode, p []byte) error
 	return nil
 }
 
+var errCloseSent = errors.New("close sent")
+
 // writeFrame handles all writes to the connection.
 func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opcode, p []byte) (_ int, err error) {
 	err = c.writeFrameMu.lock(ctx)
@@ -255,7 +257,7 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
    			return 0, net.ErrClosed
 		default:
 		}
-		return 0, errors.New("close sent")
+		return 0, errCloseSent
 	}
 
 	select {