Update

vendor/github.com/coder/websocket/write.go

@@ -1,10 +1,10 @@
 //go:build !js
-// +build !js
 
 package websocket
 
 import (
 	"bufio"
+	"compress/flate"
 	"context"
 	"crypto/rand"
 	"encoding/binary"
@@ -14,8 +14,7 @@ import (
 	"net"
 	"time"
 
-	"compress/flate"
-
+	"github.com/coder/websocket/internal/bpool"
 	"github.com/coder/websocket/internal/errd"
 	"github.com/coder/websocket/internal/util"
 )
@@ -102,23 +101,17 @@ func (c *Conn) writer(ctx context.Context, typ MessageType) (io.WriteCloser, err
 }
 
 func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) (int, error) {
-	mw, err := c.writer(ctx, typ)
+	err := c.msgWriter.reset(ctx, typ)
 	if err != nil {
 		return 0, err
 	}
+	defer c.msgWriter.mu.unlock()
 
-	if !c.flate() {
-		defer c.msgWriter.mu.unlock()
+	if !c.flate() || len(p) < c.flateThreshold {
 		return c.writeFrame(ctx, true, false, c.msgWriter.opcode, p)
 	}
 
-	n, err := mw.Write(p)
-	if err != nil {
-		return n, err
-	}
-
-	err = mw.Close()
-	return n, err
+	return c.msgWriter.writeCompressedFrame(ctx, p)
 }
 
 func (mw *msgWriter) reset(ctx context.Context, typ MessageType) error {
@@ -144,6 +137,56 @@ func (mw *msgWriter) putFlateWriter() {
 	}
 }
 
+// writeCompressedFrame compresses and writes p as a single frame.
+func (mw *msgWriter) writeCompressedFrame(ctx context.Context, p []byte) (int, error) {
+	err := mw.writeMu.lock(mw.ctx)
+	if err != nil {
+		return 0, fmt.Errorf("failed to write: %w", err)
+	}
+	defer mw.writeMu.unlock()
+
+	if mw.closed {
+		return 0, errors.New("cannot use closed writer")
+	}
+
+	mw.ensureFlate()
+
+	buf := bpool.Get()
+	defer bpool.Put(buf)
+
+	// Buffer compressed output so we can write as
+	// a single frame instead of chunked frames.
+	origWriter := mw.trimWriter.w
+	mw.trimWriter.w = buf
+	defer func() {
+		mw.trimWriter.w = origWriter
+	}()
+
+	_, err = mw.flateWriter.Write(p)
+	if err != nil {
+		return 0, fmt.Errorf("failed to compress: %w", err)
+	}
+
+	err = mw.flateWriter.Flush()
+	if err != nil {
+		return 0, fmt.Errorf("failed to flush compression: %w", err)
+	}
+
+	mw.trimWriter.reset()
+
+	if !mw.flateContextTakeover() {
+		mw.putFlateWriter()
+	}
+
+	mw.closed = true
+
+	_, err = mw.c.writeFrame(ctx, true, true, mw.opcode, buf.Bytes())
+	if err != nil {
+		return 0, err
+	}
+	return len(p), nil
+}
+
 // Write writes the given bytes to the WebSocket connection.
 func (mw *msgWriter) Write(p []byte) (_ int, err error) {
 	err = mw.writeMu.lock(mw.ctx)
@@ -249,25 +292,36 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 	}
 	defer c.writeFrameMu.unlock()
 
-	select {
-	case <-c.closed:
-		return 0, net.ErrClosed
-	case c.writeTimeout <- ctx:
-	}
-
 	defer func() {
+		if c.isClosed() && opcode == opClose {
+			err = nil
+		}
 		if err != nil {
-			select {
-			case <-c.closed:
-				err = net.ErrClosed
-			case <-ctx.Done():
+			if ctx.Err() != nil {
 				err = ctx.Err()
-			default:
+			} else if c.isClosed() {
+				err = net.ErrClosed
 			}
 			err = fmt.Errorf("failed to write frame: %w", err)
 		}
 	}()
 
+	c.closeStateMu.Lock()
+	closeSentErr := c.closeSentErr
+	c.closeStateMu.Unlock()
+	if closeSentErr != nil {
+		return 0, net.ErrClosed
+	}
+
+	select {
+	case <-c.closed:
+		return 0, net.ErrClosed
+	default:
+	}
+	if c.setupWriteTimeout(ctx) {
+		defer c.clearWriteTimeout()
+	}
+
 	c.writeHeader.fin = fin
 	c.writeHeader.opcode = opcode
 	c.writeHeader.payloadLength = int64(len(p))
@@ -303,13 +357,16 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 		}
 	}
 
-	select {
-	case <-c.closed:
-		if opcode == opClose {
-			return n, nil
+	if opcode == opClose {
+		c.closeStateMu.Lock()
+		c.closeSentErr = fmt.Errorf("sent close frame: %w", net.ErrClosed)
+		closeReceived := c.closeReceivedErr != nil
+		c.closeStateMu.Unlock()
+
+		if closeReceived && !c.casClosing() {
+			c.writeFrameMu.unlock()
+			_ = c.close()
 		}
-		return n, net.ErrClosed
-	case c.writeTimeout <- context.Background():
 	}
 
 	return n, nil
@@ -335,10 +392,7 @@ func (c *Conn) writeFramePayload(p []byte) (n int, err error) {
 		// Start of next write in the buffer.
 		i := c.bw.Buffered()
 
-		j := len(p)
-		if j > c.bw.Available() {
-			j = c.bw.Available()
-		}
+		j := min(len(p), c.bw.Available())
 
 		_, err := c.bw.Write(p[:j])
 		if err != nil {