fix: keep a read limit set before NetConn

conn.go

@@ -60,6 +60,11 @@ type Conn struct {
 	readControlBuf [maxControlPayload]byte
 	msgReader      *msgReader
 
+	readLimitMu sync.Mutex // Protects following. Never held during I/O.
+	// readLimitSet is true once SetReadLimit has been called, and NetConn
+	// then leaves the limit unchanged.
+	readLimitSet bool
+
 	// Write state.
 	msgWriter      *msgWriter
 	writeFrameMu   *mu

conn_test.go

@@ -285,7 +285,7 @@ func TestConn(t *testing.T) {
 		}
 	})
 
-	t.Run("netConn/readLimit", func(t *testing.T) {
+	t.Run("netConn/readLimit/Unset", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
 
 		n1 := websocket.NetConn(tt.ctx, c1, websocket.MessageBinary)
@@ -316,6 +316,83 @@ func TestConn(t *testing.T) {
 		assert.Equal(t, "read msg", s, string(b))
 	})
 
+	const (
+		readLimit        = 1024
+		overReadLimit    = 4 * readLimit
+		overDefaultLimit = 2 * 32768 // SetReadLimit documents the 32768 byte default.
+	)
+
+	// A default at or above overDefaultLimit fails this case instead of letting
+	// the cases that expect no error pass without a limit to disable.
+	t.Run("netConn/readLimit/DefaultWithoutNetConn", func(t *testing.T) {
+		tt, c1, c2 := newConnTest(t, nil, nil)
+
+		_ = c2.CloseRead(tt.ctx)
+		writeDone := xsync.Go(func() error {
+			return c2.Write(tt.ctx, websocket.MessageBinary, xrand.Bytes(overDefaultLimit))
+		})
+
+		_, _, err := c1.Read(tt.ctx)
+		assert.ErrorIs(t, websocket.ErrMessageTooBig, err)
+
+		_ = c2.CloseNow()
+		<-writeDone
+	})
+
+	for _, tc := range []struct {
+		name    string
+		netConn func(ctx context.Context, c *websocket.Conn) net.Conn
+		msgSize int
+		wantErr error
+	}{
+		{"DisabledBeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(-1)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, overDefaultLimit, nil},
+		{"BeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(readLimit)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, overReadLimit, websocket.ErrMessageTooBig},
+		{"AfterNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			nc := websocket.NetConn(ctx, c, websocket.MessageBinary)
+			c.SetReadLimit(readLimit)
+			return nc
+		}, overReadLimit, websocket.ErrMessageTooBig},
+		{"ConcurrentWithNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			set := xsync.Go(func() error {
+				c.SetReadLimit(readLimit)
+				return nil
+			})
+			nc := websocket.NetConn(ctx, c, websocket.MessageBinary)
+			<-set
+			return nc
+		}, overReadLimit, websocket.ErrMessageTooBig},
+	} {
+		t.Run("netConn/readLimit/"+tc.name, func(t *testing.T) {
+			tt, c1, c2 := newConnTest(t, nil, nil)
+
+			n1 := tc.netConn(tt.ctx, c1)
+			_ = c2.CloseRead(tt.ctx)
+
+			msg := xrand.Bytes(tc.msgSize)
+			writeDone := xsync.Go(func() error {
+				return c2.Write(tt.ctx, websocket.MessageBinary, msg)
+			})
+
+			got := make([]byte, len(msg))
+			_, err := io.ReadFull(n1, got)
+			if tc.wantErr != nil {
+				assert.ErrorIs(t, tc.wantErr, err)
+				_ = c2.CloseNow()
+				<-writeDone
+				return
+			}
+			assert.Success(t, err)
+			assert.Success(t, <-writeDone)
+			assert.Equal(t, "read msg", msg, got)
+		})
+	}
+
 	t.Run("netConn/pastDeadline", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
 

netconn.go

@@ -44,9 +44,10 @@ import (
 // A received StatusNormalClosure or StatusGoingAway close frame will be translated to
 // io.EOF when reading.
 //
-// Furthermore, the ReadLimit is set to -1 to disable it.
+// NetConn disables the read limit unless SetReadLimit has been called on c,
+// because each Write to a net.Conn becomes one message and can be any size.
 func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
-	c.SetReadLimit(-1)
+	c.disableReadLimitUnlessSet()
 
 	nc := &netConn{
 		c:       c,

netconn_js_test.go

@@ -0,0 +1,98 @@
+package websocket_test
+
+import (
+	"context"
+	"io"
+	"net"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/internal/test/assert"
+	"github.com/coder/websocket/internal/test/xrand"
+)
+
+func TestWasmNetConnReadLimit(t *testing.T) {
+	t.Parallel()
+
+	const (
+		readLimit        = 1024
+		overReadLimit    = 4 * readLimit
+		overDefaultLimit = 2 * 32768 // SetReadLimit documents the 32768 byte default.
+	)
+
+	// A default at or above overDefaultLimit fails this case instead of letting
+	// the cases that expect no error pass without a limit to disable.
+	t.Run("DefaultWithoutNetConn", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+		defer cancel()
+
+		c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{
+			Subprotocols: []string{"echo"},
+		})
+		assert.Success(t, err)
+		defer c.Close(websocket.StatusNormalClosure, "")
+
+		err = c.Write(ctx, websocket.MessageBinary, xrand.Bytes(overDefaultLimit))
+		assert.Success(t, err)
+
+		_, _, err = c.Read(ctx)
+		assert.ErrorIs(t, websocket.ErrMessageTooBig, err)
+	})
+
+	for _, tc := range []struct {
+		name    string
+		netConn func(ctx context.Context, c *websocket.Conn) net.Conn
+		msgSize int
+		wantErr error
+	}{
+		{"Unset", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, overDefaultLimit, nil},
+		{"DisabledBeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(-1)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, overDefaultLimit, nil},
+		{"BeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(readLimit)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, overReadLimit, websocket.ErrMessageTooBig},
+		{"AfterNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			nc := websocket.NetConn(ctx, c, websocket.MessageBinary)
+			c.SetReadLimit(readLimit)
+			return nc
+		}, overReadLimit, websocket.ErrMessageTooBig},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+			defer cancel()
+
+			c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{
+				Subprotocols: []string{"echo"},
+			})
+			assert.Success(t, err)
+			// The echo server fails TestWasm unless the close status is
+			// StatusNormalClosure.
+			defer c.Close(websocket.StatusNormalClosure, "")
+
+			nc := tc.netConn(ctx, c)
+			msg := []byte(xrand.String(tc.msgSize))
+			_, err = nc.Write(msg)
+			assert.Success(t, err)
+
+			got := make([]byte, len(msg))
+			_, err = io.ReadFull(nc, got)
+			if tc.wantErr != nil {
+				assert.ErrorIs(t, tc.wantErr, err)
+				return
+			}
+			assert.Success(t, err)
+			assert.Equal(t, "read msg", msg, got)
+		})
+	}
+}

read.go

@@ -95,6 +95,9 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 //
 // Set to -1 to disable.
 func (c *Conn) SetReadLimit(n int64) {
+	c.readLimitMu.Lock()
+	defer c.readLimitMu.Unlock()
+	c.readLimitSet = true
 	if n >= 0 {
 		// We read one more byte than the limit in case
 		// there is a fin frame that needs to be read.
@@ -104,6 +107,14 @@ func (c *Conn) SetReadLimit(n int64) {
 	c.msgReader.limitReader.limit.Store(n)
 }
 
+func (c *Conn) disableReadLimitUnlessSet() {
+	c.readLimitMu.Lock()
+	defer c.readLimitMu.Unlock()
+	if !c.readLimitSet {
+		c.msgReader.limitReader.limit.Store(-1)
+	}
+}
+
 const defaultReadLimit = 32768
 
 func newMsgReader(c *Conn) *msgReader {

ws_js.go

@@ -47,6 +47,11 @@ type Conn struct {
 	// read limit for a message in bytes.
 	msgReadLimit atomic.Int64
 
+	readLimitMu sync.Mutex // Protects following. Never held during I/O.
+	// readLimitSet is true once SetReadLimit has been called, and NetConn
+	// then leaves the limit unchanged.
+	readLimitSet bool
+
 	closeReadMu  sync.Mutex
 	closeReadCtx context.Context
 
@@ -414,9 +419,20 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 
 // SetReadLimit implements *Conn.SetReadLimit for wasm.
 func (c *Conn) SetReadLimit(n int64) {
+	c.readLimitMu.Lock()
+	defer c.readLimitMu.Unlock()
+	c.readLimitSet = true
 	c.msgReadLimit.Store(n)
 }
 
+func (c *Conn) disableReadLimitUnlessSet() {
+	c.readLimitMu.Lock()
+	defer c.readLimitMu.Unlock()
+	if !c.readLimitSet {
+		c.msgReadLimit.Store(-1)
+	}
+}
+
 func (c *Conn) setCloseErr(err error) {
 	c.closeErrOnce.Do(func() {
 		c.closeErr = fmt.Errorf("WebSocket closed: %w", err)