fix timeout

Author: mafredri

read.go

@@ -224,25 +224,28 @@ func (c *Conn) prepareRead(ctx context.Context, err *error) (func(), error) {
 	case c.readTimeout <- ctx:
 	}
 
-	c.closeStateMu.Lock()
-	closeReceivedErr := c.closeReceivedErr
-	c.closeStateMu.Unlock()
-	if closeReceivedErr != nil {
-		return nil, closeReceivedErr
-	}
-
-	return func() {
+	done := func() {
 		select {
 		case <-c.closed:
 			if *err != nil {
 				*err = net.ErrClosed
 			}
-		case c.writeTimeout <- context.Background():
+		case c.readTimeout <- context.Background():
 		}
 		if *err != nil && ctx.Err() != nil {
 			*err = ctx.Err()
 		}
-	}, nil
+	}
+
+	c.closeStateMu.Lock()
+	closeReceivedErr := c.closeReceivedErr
+	c.closeStateMu.Unlock()
+	if closeReceivedErr != nil {
+		defer done()
+		return nil, closeReceivedErr
+	}
+
+	return done, nil
 }
 
 func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {

write.go

@@ -5,6 +5,7 @@ package websocket
 
 import (
 	"bufio"
+	"compress/flate"
 	"context"
 	"crypto/rand"
 	"encoding/binary"
@@ -14,8 +15,6 @@ import (
 	"net"
 	"time"
 
-	"compress/flate"
-
 	"github.com/coder/websocket/internal/errd"
 	"github.com/coder/websocket/internal/util"
 )
@@ -250,6 +249,9 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 	defer c.writeFrameMu.unlock()
 
 	defer func() {
+		if errors.Is(err, net.ErrClosed) && opcode == opClose {
+			err = nil
+		}
 		if err != nil {
 			if ctx.Err() != nil {
 				err = ctx.Err()
@@ -275,7 +277,6 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 	defer func() {
 		select {
 		case <-c.closed:
-			err = net.ErrClosed
 		case c.writeTimeout <- context.Background():
 		}
 	}()
@@ -315,15 +316,6 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 		}
 	}
 
-	select {
-	case <-c.closed:
-		if opcode == opClose {
-			return n, nil
-		}
-		return n, net.ErrClosed
-	case c.writeTimeout <- context.Background():
-	}
-
 	if opcode == opClose {
 		c.closeStateMu.Lock()
 		c.closeSentErr = fmt.Errorf("sent close frame: %w", net.ErrClosed)