mask_asm.go

+//go:build amd64 || arm64
+
+package websocket
+
+func mask(b []byte, key uint32) uint32 {
+	// TODO: Will enable in v1.9.0.
+	return maskGo(b, key)
+	/*
+		if len(b) > 0 {
+			return maskAsm(&b[0], len(b), key)
+		}
+		return key
+	*/
+}
+
+// @nhooyr: I am not confident that the amd64 or the arm64 implementations of this
+// function are perfect. There are almost certainly missing optimizations or
+// opportunities for simplification. I'm confident there are no bugs though.
+// For example, the arm64 implementation doesn't align memory like the amd64.
+// Or the amd64 implementation could use AVX512 instead of just AVX2.
+// The AVX2 code I had to disable anyway as it wasn't performing as expected.
+// See https://github.com/nhooyr/websocket/pull/326#issuecomment-1771138049
+//
+//go:noescape
+//lint:ignore U1000 disabled till v1.9.0
+func maskAsm(b *byte, len int, key uint32) uint32

mask_asm_test.go

+//go:build amd64 || arm64
+
+package websocket
+
+import "testing"
+
+func TestMaskASM(t *testing.T) {
+	t.Parallel()
+
+	testMask(t, "maskASM", mask)
+}

mask_go.go

+//go:build !amd64 && !arm64 && !js
+
+package websocket
+
+func mask(b []byte, key uint32) uint32 {
+	return maskGo(b, key)
+}

mask_test.go

+package websocket
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/binary"
+	"math/big"
+	"math/bits"
+	"testing"
+
+	"github.com/coder/websocket/internal/test/assert"
+)
+
+func basicMask(b []byte, key uint32) uint32 {
+	for i := range b {
+		b[i] ^= byte(key)
+		key = bits.RotateLeft32(key, -8)
+	}
+	return key
+}
+
+func basicMask2(b []byte, key uint32) uint32 {
+	keyb := binary.LittleEndian.AppendUint32(nil, key)
+	pos := 0
+	for i := range b {
+		b[i] ^= keyb[pos&3]
+		pos++
+	}
+	return bits.RotateLeft32(key, (pos&3)*-8)
+}
+
+func TestMask(t *testing.T) {
+	t.Parallel()
+
+	testMask(t, "basicMask", basicMask)
+	testMask(t, "maskGo", maskGo)
+	testMask(t, "basicMask2", basicMask2)
+}
+
+func testMask(t *testing.T, name string, fn func(b []byte, key uint32) uint32) {
+	t.Run(name, func(t *testing.T) {
+		t.Parallel()
+		for i := 0; i < 9999; i++ {
+			keyb := make([]byte, 4)
+			_, err := rand.Read(keyb)
+			assert.Success(t, err)
+			key := binary.LittleEndian.Uint32(keyb)
+
+			n, err := rand.Int(rand.Reader, big.NewInt(1<<16))
+			assert.Success(t, err)
+
+			b := make([]byte, 1+n.Int64())
+			_, err = rand.Read(b)
+			assert.Success(t, err)
+
+			b2 := make([]byte, len(b))
+			copy(b2, b)
+			b3 := make([]byte, len(b))
+			copy(b3, b)
+
+			key2 := basicMask(b2, key)
+			key3 := fn(b3, key)
+
+			if key2 != key3 {
+				t.Errorf("expected key %X but got %X", key2, key3)
+			}
+			if !bytes.Equal(b2, b3) {
+				t.Error("bad bytes")
+				return
+			}
+		}
+	})
+}

netconn.go

@@ -68,7 +68,7 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
 		defer nc.writeMu.unlock()
 
 		// Prevents future writes from writing until the deadline is reset.
-		atomic.StoreInt64(&nc.writeExpired, 1)
+		nc.writeExpired.Store(1)
 	})
 	if !nc.writeTimer.Stop() {
 		<-nc.writeTimer.C
@@ -84,7 +84,7 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
 		defer nc.readMu.unlock()
 
 		// Prevents future reads from reading until the deadline is reset.
-		atomic.StoreInt64(&nc.readExpired, 1)
+		nc.readExpired.Store(1)
 	})
 	if !nc.readTimer.Stop() {
 		<-nc.readTimer.C
@@ -99,13 +99,13 @@ type netConn struct {
 
 	writeTimer   *time.Timer
 	writeMu      *mu
-	writeExpired int64
+	writeExpired atomic.Int64
 	writeCtx     context.Context
 	writeCancel  context.CancelFunc
 
 	readTimer   *time.Timer
 	readMu      *mu
-	readExpired int64
+	readExpired atomic.Int64
 	readCtx     context.Context
 	readCancel  context.CancelFunc
 	readEOFed   bool
@@ -126,7 +126,7 @@ func (nc *netConn) Write(p []byte) (int, error) {
 	nc.writeMu.forceLock()
 	defer nc.writeMu.unlock()
 
-	if atomic.LoadInt64(&nc.writeExpired) == 1 {
+	if nc.writeExpired.Load() == 1 {
 		return 0, fmt.Errorf("failed to write: %w", context.DeadlineExceeded)
 	}
 
@@ -154,7 +154,7 @@ func (nc *netConn) Read(p []byte) (int, error) {
 }
 
 func (nc *netConn) read(p []byte) (int, error) {
-	if atomic.LoadInt64(&nc.readExpired) == 1 {
+	if nc.readExpired.Load() == 1 {
 		return 0, fmt.Errorf("failed to read: %w", context.DeadlineExceeded)
 	}
 
@@ -206,21 +206,29 @@ func (nc *netConn) SetDeadline(t time.Time) error {
 }
 
 func (nc *netConn) SetWriteDeadline(t time.Time) error {
-	atomic.StoreInt64(&nc.writeExpired, 0)
+	nc.writeExpired.Store(0)
 	if t.IsZero() {
 		nc.writeTimer.Stop()
 	} else {
-		nc.writeTimer.Reset(time.Until(t))
+		dur := time.Until(t)
+		if dur <= 0 {
+			dur = 1
+		}
+		nc.writeTimer.Reset(dur)
 	}
 	return nil
 }
 
 func (nc *netConn) SetReadDeadline(t time.Time) error {
-	atomic.StoreInt64(&nc.readExpired, 0)
+	nc.readExpired.Store(0)
 	if t.IsZero() {
 		nc.readTimer.Stop()
 	} else {
-		nc.readTimer.Reset(time.Until(t))
+		dur := time.Until(t)
+		if dur <= 0 {
+			dur = 1
+		}
+		nc.readTimer.Reset(dur)
 	}
 	return nil
 }

read.go

@@ -9,12 +9,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"strings"
+	"sync/atomic"
 	"time"
 
-	"nhooyr.io/websocket/internal/errd"
-	"nhooyr.io/websocket/internal/util"
-	"nhooyr.io/websocket/internal/xsync"
+	"github.com/coder/websocket/internal/errd"
+	"github.com/coder/websocket/internal/util"
 )
 
 // Reader reads from the connection until there is a WebSocket
@@ -59,12 +60,28 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
 // Call CloseRead when you do not expect to read any more messages.
 // Since it actively reads from the connection, it will ensure that ping, pong and close
 // frames are responded to. This means c.Ping and c.Close will still work as expected.
+//
+// This function is idempotent.
 func (c *Conn) CloseRead(ctx context.Context) context.Context {
+	c.closeReadMu.Lock()
+	ctx2 := c.closeReadCtx
+	if ctx2 != nil {
+		c.closeReadMu.Unlock()
+		return ctx2
+	}
 	ctx, cancel := context.WithCancel(ctx)
+	c.closeReadCtx = ctx
+	c.closeReadDone = make(chan struct{})
+	c.closeReadMu.Unlock()
+
 	go func() {
+		defer close(c.closeReadDone)
 		defer cancel()
-		c.Reader(ctx)
-		c.Close(StatusPolicyViolation, "unexpected data message")
+		defer c.close()
+		_, _, err := c.Reader(ctx)
+		if err == nil {
+			c.Close(StatusPolicyViolation, "unexpected data message")
+		}
 	}()
 	return ctx
 }
@@ -200,60 +217,68 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) {
 	}
 }
 
-func (c *Conn) readFrameHeader(ctx context.Context) (header, error) {
+// prepareRead sets the readTimeout context and returns a done function
+// to be called after the read is done. It also returns an error if the
+// connection is closed. The reference to the error is used to assign
+// an error depending on if the connection closed or the context timed
+// out during use. Typically the referenced error is a named return
+// variable of the function calling this method.
+func (c *Conn) prepareRead(ctx context.Context, err *error) (func(), error) {
 	select {
 	case <-c.closed:
-		return header{}, c.closeErr
+		return nil, net.ErrClosed
 	case c.readTimeout <- ctx:
 	}
 
-	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
-	if err != nil {
+	done := func() {
 		select {
 		case <-c.closed:
-			return header{}, c.closeErr
-		case <-ctx.Done():
-			return header{}, ctx.Err()
-		default:
-			c.close(err)
-			return header{}, err
+			if *err != nil {
+				*err = net.ErrClosed
+			}
+		case c.readTimeout <- context.Background():
+		}
+		if *err != nil && ctx.Err() != nil {
+			*err = ctx.Err()
 		}
 	}
 
-	select {
-	case <-c.closed:
-		return header{}, c.closeErr
-	case c.readTimeout <- context.Background():
+	c.closeStateMu.Lock()
+	closeReceivedErr := c.closeReceivedErr
+	c.closeStateMu.Unlock()
+	if closeReceivedErr != nil {
+		defer done()
+		return nil, closeReceivedErr
 	}
 
-	return h, nil
+	return done, nil
 }
 
-func (c *Conn) readFramePayload(ctx context.Context, p []byte) (int, error) {
-	select {
-	case <-c.closed:
-		return 0, c.closeErr
-	case c.readTimeout <- ctx:
+func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {
+	readDone, err := c.prepareRead(ctx, &err)
+	if err != nil {
+		return header{}, err
 	}
+	defer readDone()
 
-	n, err := io.ReadFull(c.br, p)
+	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
 	if err != nil {
-		select {
-		case <-c.closed:
-			return n, c.closeErr
-		case <-ctx.Done():
-			return n, ctx.Err()
-		default:
-			err = fmt.Errorf("failed to read frame payload: %w", err)
-			c.close(err)
-			return n, err
-		}
+		return header{}, err
 	}
 
-	select {
-	case <-c.closed:
-		return n, c.closeErr
-	case c.readTimeout <- context.Background():
+	return h, nil
+}
+
+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 {
+		return n, fmt.Errorf("failed to read frame payload: %w", err)
 	}
 
 	return n, err
@@ -282,13 +307,21 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 	}
 
 	if h.masked {
-		mask(h.maskKey, b)
+		mask(b, h.maskKey)
 	}
 
 	switch h.opcode {
 	case opPing:
+		if c.onPingReceived != nil {
+			if !c.onPingReceived(ctx, b) {
+				return nil
+			}
+		}
 		return c.writeControl(ctx, opPong, b)
 	case opPong:
+		if c.onPongReceived != nil {
+			c.onPongReceived(ctx, b)
+		}
 		c.activePingsMu.Lock()
 		pong, ok := c.activePings[string(b)]
 		c.activePingsMu.Unlock()
@@ -301,9 +334,7 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 		return nil
 	}
 
-	defer func() {
-		c.readCloseFrameErr = err
-	}()
+	// opClose
 
 	ce, err := parseClosePayload(b)
 	if err != nil {
@@ -313,9 +344,22 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
 	}
 
 	err = fmt.Errorf("received close frame: %w", ce)
-	c.setCloseErr(err)
-	c.writeClose(ce.Code, ce.Reason)
-	c.close(err)
+	c.closeStateMu.Lock()
+	c.closeReceivedErr = err
+	closeSent := c.closeSentErr != nil
+	c.closeStateMu.Unlock()
+
+	// Only unlock readMu if this connection is being closed becaue
+	// c.close will try to acquire the readMu lock. We unlock for
+	// writeClose as well because it may also call c.close.
+	if !closeSent {
+		c.readMu.unlock()
+		_ = c.writeClose(ce.Code, ce.Reason)
+	}
+	if !c.casClosing() {
+		c.readMu.unlock()
+		_ = c.close()
+	}
 	return err
 }
 
@@ -329,9 +373,7 @@ func (c *Conn) reader(ctx context.Context) (_ MessageType, _ io.Reader, err erro
 	defer c.readMu.unlock()
 
 	if !c.msgReader.fin {
-		err = errors.New("previous message not read to completion")
-		c.close(fmt.Errorf("failed to get reader: %w", err))
-		return 0, nil, err
+		return 0, nil, errors.New("previous message not read to completion")
 	}
 
 	h, err := c.readLoop(ctx)
@@ -404,10 +446,9 @@ func (mr *msgReader) Read(p []byte) (n int, err error) {
 		return n, io.EOF
 	}
 	if err != nil {
-		err = fmt.Errorf("failed to read: %w", err)
-		mr.c.close(err)
+		return n, fmt.Errorf("failed to read: %w", err)
 	}
-	return n, err
+	return n, nil
 }
 
 func (mr *msgReader) read(p []byte) (int, error) {
@@ -446,7 +487,7 @@ func (mr *msgReader) read(p []byte) (int, error) {
 		mr.payloadLength -= int64(n)
 
 		if !mr.c.client {
-			mr.maskKey = mask(mr.maskKey, p)
+			mr.maskKey = mask(p, mr.maskKey)
 		}
 
 		return n, nil
@@ -456,7 +497,7 @@ func (mr *msgReader) read(p []byte) (int, error) {
 type limitReader struct {
 	c     *Conn
 	r     io.Reader
-	limit xsync.Int64
+	limit atomic.Int64
 	n     int64
 }
 

write.go

@@ -5,18 +5,18 @@ package websocket
 
 import (
 	"bufio"
+	"compress/flate"
 	"context"
 	"crypto/rand"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"time"
 
-	"compress/flate"
-
-	"nhooyr.io/websocket/internal/errd"
-	"nhooyr.io/websocket/internal/util"
+	"github.com/coder/websocket/internal/errd"
+	"github.com/coder/websocket/internal/util"
 )
 
 // Writer returns a writer bounded by the context that will write
@@ -38,7 +38,7 @@ func (c *Conn) Writer(ctx context.Context, typ MessageType) (io.WriteCloser, err
 //
 // See the Writer method if you want to stream a message.
 //
-// If compression is disabled or the threshold is not met, then it
+// If compression is disabled or the compression threshold is not met, then it
 // will write the message in a single frame.
 func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
 	_, err := c.write(ctx, typ, p)
@@ -49,30 +49,11 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
 }
 
 type msgWriter struct {
-	mw     *msgWriterState
-	closed bool
-}
-
-func (mw *msgWriter) Write(p []byte) (int, error) {
-	if mw.closed {
-		return 0, errors.New("cannot use closed writer")
-	}
-	return mw.mw.Write(p)
-}
-
-func (mw *msgWriter) Close() error {
-	if mw.closed {
-		return errors.New("cannot use closed writer")
-	}
-	mw.closed = true
-	return mw.mw.Close()
-}
-
-type msgWriterState struct {
 	c *Conn
 
 	mu      *mu
 	writeMu *mu
+	closed  bool
 
 	ctx    context.Context
 	opcode opcode
@@ -82,8 +63,8 @@ type msgWriterState struct {
 	flateWriter *flate.Writer
 }
 
-func newMsgWriterState(c *Conn) *msgWriterState {
-	mw := &msgWriterState{
+func newMsgWriter(c *Conn) *msgWriter {
+	mw := &msgWriter{
 		c:       c,
 		mu:      newMu(c),
 		writeMu: newMu(c),
@@ -91,7 +72,7 @@ func newMsgWriterState(c *Conn) *msgWriterState {
 	return mw
 }
 
-func (mw *msgWriterState) ensureFlate() {
+func (mw *msgWriter) ensureFlate() {
 	if mw.trimWriter == nil {
 		mw.trimWriter = &trimLastFourBytesWriter{
 			w: util.WriterFunc(mw.write),
@@ -104,7 +85,7 @@ func (mw *msgWriterState) ensureFlate() {
 	mw.flate = true
 }
 
-func (mw *msgWriterState) flateContextTakeover() bool {
+func (mw *msgWriter) flateContextTakeover() bool {
 	if mw.c.client {
 		return !mw.c.copts.clientNoContextTakeover
 	}
@@ -112,14 +93,11 @@ func (mw *msgWriterState) flateContextTakeover() bool {
 }
 
 func (c *Conn) writer(ctx context.Context, typ MessageType) (io.WriteCloser, error) {
-	err := c.msgWriterState.reset(ctx, typ)
+	err := c.msgWriter.reset(ctx, typ)
 	if err != nil {
 		return nil, err
 	}
-	return &msgWriter{
-		mw:     c.msgWriterState,
-		closed: false,
-	}, nil
+	return c.msgWriter, nil
 }
 
 func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) (int, error) {
@@ -129,8 +107,8 @@ func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) (int, error
 	}
 
 	if !c.flate() {
-		defer c.msgWriterState.mu.unlock()
-		return c.writeFrame(ctx, true, false, c.msgWriterState.opcode, p)
+		defer c.msgWriter.mu.unlock()
+		return c.writeFrame(ctx, true, false, c.msgWriter.opcode, p)
 	}
 
 	n, err := mw.Write(p)
@@ -142,7 +120,7 @@ func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) (int, error
 	return n, err
 }
 
-func (mw *msgWriterState) reset(ctx context.Context, typ MessageType) error {
+func (mw *msgWriter) reset(ctx context.Context, typ MessageType) error {
 	err := mw.mu.lock(ctx)
 	if err != nil {
 		return err
@@ -151,13 +129,14 @@ func (mw *msgWriterState) reset(ctx context.Context, typ MessageType) error {
 	mw.ctx = ctx
 	mw.opcode = opcode(typ)
 	mw.flate = false
+	mw.closed = false
 
 	mw.trimWriter.reset()
 
 	return nil
 }
 
-func (mw *msgWriterState) putFlateWriter() {
+func (mw *msgWriter) putFlateWriter() {
 	if mw.flateWriter != nil {
 		putFlateWriter(mw.flateWriter)
 		mw.flateWriter = nil
@@ -165,17 +144,20 @@ func (mw *msgWriterState) putFlateWriter() {
 }
 
 // Write writes the given bytes to the WebSocket connection.
-func (mw *msgWriterState) Write(p []byte) (_ int, err error) {
+func (mw *msgWriter) Write(p []byte) (_ int, err 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")
+	}
+
 	defer func() {
 		if err != nil {
 			err = fmt.Errorf("failed to write: %w", err)
-			mw.c.close(err)
 		}
 	}()
 
@@ -194,7 +176,7 @@ func (mw *msgWriterState) Write(p []byte) (_ int, err error) {
 	return mw.write(p)
 }
 
-func (mw *msgWriterState) write(p []byte) (int, error) {
+func (mw *msgWriter) write(p []byte) (int, error) {
 	n, err := mw.c.writeFrame(mw.ctx, false, mw.flate, mw.opcode, p)
 	if err != nil {
 		return n, fmt.Errorf("failed to write data frame: %w", err)
@@ -204,7 +186,7 @@ func (mw *msgWriterState) write(p []byte) (int, error) {
 }
 
 // Close flushes the frame to the connection.
-func (mw *msgWriterState) Close() (err error) {
+func (mw *msgWriter) Close() (err error) {
 	defer errd.Wrap(&err, "failed to close writer")
 
 	err = mw.writeMu.lock(mw.ctx)
@@ -213,6 +195,11 @@ func (mw *msgWriterState) Close() (err error) {
 	}
 	defer mw.writeMu.unlock()
 
+	if mw.closed {
+		return errors.New("writer already closed")
+	}
+	mw.closed = true
+
 	if mw.flate {
 		err = mw.flateWriter.Flush()
 		if err != nil {
@@ -232,7 +219,7 @@ func (mw *msgWriterState) Close() (err error) {
 	return nil
 }
 
-func (mw *msgWriterState) close() {
+func (mw *msgWriter) close() {
 	if mw.c.client {
 		mw.c.writeFrameMu.forceLock()
 		putBufioWriter(mw.c.bw)
@@ -253,48 +240,44 @@ func (c *Conn) writeControl(ctx context.Context, opcode opcode, p []byte) error
 	return nil
 }
 
-// frame handles all writes to the connection.
+// 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)
 	if err != nil {
 		return 0, err
 	}
+	defer c.writeFrameMu.unlock()
 
-	// If the state says a close has already been written, we wait until
-	// the connection is closed and return that error.
-	//
-	// However, if the frame being written is a close, that means its the close from
-	// the state being set so we let it go through.
-	c.closeMu.Lock()
-	wroteClose := c.wroteClose
-	c.closeMu.Unlock()
-	if wroteClose && opcode != opClose {
-		c.writeFrameMu.unlock()
-		select {
-		case <-ctx.Done():
-			return 0, ctx.Err()
-		case <-c.closed:
-			return 0, c.closeErr
+	defer func() {
+		if c.isClosed() && opcode == opClose {
+			err = nil
+		}
+		if err != nil {
+			if ctx.Err() != nil {
+				err = ctx.Err()
+			} 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
 	}
-	defer c.writeFrameMu.unlock()
 
 	select {
 	case <-c.closed:
-		return 0, c.closeErr
+		return 0, net.ErrClosed
 	case c.writeTimeout <- ctx:
 	}
-
 	defer func() {
-		if err != nil {
-			select {
-			case <-c.closed:
-				err = c.closeErr
-			case <-ctx.Done():
-				err = ctx.Err()
-			}
-			c.close(err)
-			err = fmt.Errorf("failed to write frame: %w", err)
+		select {
+		case <-c.closed:
+		case c.writeTimeout <- context.Background():
 		}
 	}()
 
@@ -333,10 +316,16 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 		}
 	}
 
-	select {
-	case <-c.closed:
-		return n, c.closeErr
-	case c.writeTimeout <- context.Background():
+	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, nil
@@ -372,7 +361,7 @@ func (c *Conn) writeFramePayload(p []byte) (n int, err error) {
 			return n, err
 		}
 
-		maskKey = mask(maskKey, c.writeBuf[i:c.bw.Buffered()])
+		maskKey = mask(c.writeBuf[i:c.bw.Buffered()], maskKey)
 
 		p = p[j:]
 		n += j
@@ -399,7 +388,5 @@ func extractBufioWriterBuf(bw *bufio.Writer, w io.Writer) []byte {
 }
 
 func (c *Conn) writeError(code StatusCode, err error) {
-	c.setCloseErr(err)
 	c.writeClose(code, err.Error())
-	c.close(nil)
 }

ws_js.go

-package websocket // import "nhooyr.io/websocket"
+package websocket // import "github.com/coder/websocket"
 
 import (
 	"bytes"
@@ -6,16 +6,17 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"reflect"
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall/js"
 
-	"nhooyr.io/websocket/internal/bpool"
-	"nhooyr.io/websocket/internal/wsjs"
-	"nhooyr.io/websocket/internal/xsync"
+	"github.com/coder/websocket/internal/bpool"
+	"github.com/coder/websocket/internal/wsjs"
 )
 
 // opcode represents a WebSocket opcode.
@@ -40,14 +41,16 @@ const (
 
 // Conn provides a wrapper around the browser WebSocket API.
 type Conn struct {
-	noCopy
-	ws wsjs.WebSocket
+	noCopy noCopy
+	ws     wsjs.WebSocket
 
 	// read limit for a message in bytes.
-	msgReadLimit xsync.Int64
+	msgReadLimit atomic.Int64
+
+	closeReadMu  sync.Mutex
+	closeReadCtx context.Context
 
 	closingMu     sync.Mutex
-	isReadClosed  xsync.Int64
 	closeOnce     sync.Once
 	closed        chan struct{}
 	closeErrOnce  sync.Once
@@ -55,6 +58,7 @@ type Conn struct {
 	closeWasClean bool
 
 	releaseOnClose   func()
+	releaseOnError   func()
 	releaseOnMessage func()
 
 	readSignal chan struct{}
@@ -92,9 +96,15 @@ func (c *Conn) init() {
 		c.close(err, e.WasClean)
 
 		c.releaseOnClose()
+		c.releaseOnError()
 		c.releaseOnMessage()
 	})
 
+	c.releaseOnError = c.ws.OnError(func(v js.Value) {
+		c.setCloseErr(errors.New(v.Get("message").String()))
+		c.closeWithInternal()
+	})
+
 	c.releaseOnMessage = c.ws.OnMessage(func(e wsjs.MessageEvent) {
 		c.readBufMu.Lock()
 		defer c.readBufMu.Unlock()
@@ -121,7 +131,10 @@ func (c *Conn) closeWithInternal() {
 // Read attempts to read a message from the connection.
 // The maximum time spent waiting is bounded by the context.
 func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
-	if c.isReadClosed.Load() == 1 {
+	c.closeReadMu.Lock()
+	closedRead := c.closeReadCtx != nil
+	c.closeReadMu.Unlock()
+	if closedRead {
 		return 0, nil, errors.New("WebSocket connection read closed")
 	}
 
@@ -129,7 +142,8 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
 	if err != nil {
 		return 0, nil, fmt.Errorf("failed to read: %w", err)
 	}
-	if int64(len(p)) > c.msgReadLimit.Load() {
+	readLimit := c.msgReadLimit.Load()
+	if readLimit >= 0 && int64(len(p)) > readLimit {
 		err := fmt.Errorf("read limited at %v bytes", c.msgReadLimit.Load())
 		c.Close(StatusMessageTooBig, err.Error())
 		return 0, nil, err
@@ -144,7 +158,7 @@ func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) {
 		return 0, nil, ctx.Err()
 	case <-c.readSignal:
 	case <-c.closed:
-		return 0, nil, c.closeErr
+		return 0, nil, net.ErrClosed
 	}
 
 	c.readBufMu.Lock()
@@ -198,7 +212,7 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
 
 func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) error {
 	if c.isClosed() {
-		return c.closeErr
+		return net.ErrClosed
 	}
 	switch typ {
 	case MessageBinary:
@@ -222,19 +236,28 @@ func (c *Conn) Close(code StatusCode, reason string) error {
 	return nil
 }
 
+// CloseNow closes the WebSocket connection without attempting a close handshake.
+// Use when you do not want the overhead of the close handshake.
+//
+// note: No different from Close(StatusGoingAway, "") in WASM as there is no way to close
+// a WebSocket without the close handshake.
+func (c *Conn) CloseNow() error {
+	return c.Close(StatusGoingAway, "")
+}
+
 func (c *Conn) exportedClose(code StatusCode, reason string) error {
 	c.closingMu.Lock()
 	defer c.closingMu.Unlock()
 
+	if c.isClosed() {
+		return net.ErrClosed
+	}
+
 	ce := fmt.Errorf("sent close: %w", CloseError{
 		Code:   code,
 		Reason: reason,
 	})
 
-	if c.isClosed() {
-		return fmt.Errorf("tried to close with %q but connection already closed: %w", ce, c.closeErr)
-	}
-
 	c.setCloseErr(ce)
 	err := c.ws.Close(int(code), reason)
 	if err != nil {
@@ -305,7 +328,7 @@ func dial(ctx context.Context, url string, opts *DialOptions) (*Conn, *http.Resp
 			StatusCode: http.StatusSwitchingProtocols,
 		}, nil
 	case <-c.closed:
-		return nil, nil, c.closeErr
+		return nil, nil, net.ErrClosed
 	}
 }
 
@@ -368,13 +391,23 @@ func (w *writer) Close() error {
 
 // CloseRead implements *Conn.CloseRead for wasm.
 func (c *Conn) CloseRead(ctx context.Context) context.Context {
-	c.isReadClosed.Store(1)
-
+	c.closeReadMu.Lock()
+	ctx2 := c.closeReadCtx
+	if ctx2 != nil {
+		c.closeReadMu.Unlock()
+		return ctx2
+	}
 	ctx, cancel := context.WithCancel(ctx)
+	c.closeReadCtx = ctx
+	c.closeReadMu.Unlock()
+
 	go func() {
 		defer cancel()
-		c.read(ctx)
-		c.Close(StatusPolicyViolation, "unexpected data message")
+		defer c.CloseNow()
+		_, _, err := c.read(ctx)
+		if err != nil {
+			c.Close(StatusPolicyViolation, "unexpected data message")
+		}
 	}()
 	return ctx
 }

ws_js_test.go

@@ -7,9 +7,9 @@ import (
 	"testing"
 	"time"
 
-	"nhooyr.io/websocket"
-	"nhooyr.io/websocket/internal/test/assert"
-	"nhooyr.io/websocket/internal/test/wstest"
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/internal/test/assert"
+	"github.com/coder/websocket/internal/test/wstest"
 )
 
 func TestWasm(t *testing.T) {

wsjson/wsjson.go

 // Package wsjson provides helpers for reading and writing JSON messages.
-package wsjson // import "nhooyr.io/websocket/wsjson"
+package wsjson // import "github.com/coder/websocket/wsjson"
 
 import (
 	"context"
 	"encoding/json"
 	"fmt"
 
-	"nhooyr.io/websocket"
-	"nhooyr.io/websocket/internal/bpool"
-	"nhooyr.io/websocket/internal/errd"
-	"nhooyr.io/websocket/internal/util"
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/internal/bpool"
+	"github.com/coder/websocket/internal/errd"
+	"github.com/coder/websocket/internal/util"
 )
 
 // Read reads a JSON message from c into v.

wsjson/wsjson_test.go

+package wsjson_test
+
+import (
+	"encoding/json"
+	"io"
+	"strconv"
+	"testing"
+
+	"github.com/coder/websocket/internal/test/xrand"
+)
+
+func BenchmarkJSON(b *testing.B) {
+	sizes := []int{
+		8,
+		16,
+		32,
+		128,
+		256,
+		512,
+		1024,
+		2048,
+		4096,
+		8192,
+		16384,
+	}
+
+	b.Run("json.Encoder", func(b *testing.B) {
+		for _, size := range sizes {
+			b.Run(strconv.Itoa(size), func(b *testing.B) {
+				msg := xrand.String(size)
+				b.SetBytes(int64(size))
+				b.ReportAllocs()
+				b.ResetTimer()
+				for i := 0; i < b.N; i++ {
+					json.NewEncoder(io.Discard).Encode(msg)
+				}
+			})
+		}
+	})
+	b.Run("json.Marshal", func(b *testing.B) {
+		for _, size := range sizes {
+			b.Run(strconv.Itoa(size), func(b *testing.B) {
+				msg := xrand.String(size)
+				b.SetBytes(int64(size))
+				b.ReportAllocs()
+				b.ResetTimer()
+				for i := 0; i < b.N; i++ {
+					json.Marshal(msg)
+				}
+			})
+		}
+	})
+}