fix: keep a read limit set before NetConn

conn.go

@@ -59,6 +59,7 @@ type Conn struct {
 	readHeaderBuf  [8]byte
 	readControlBuf [maxControlPayload]byte
 	msgReader      *msgReader
+	readLimitSet   atomic.Bool
 
 	// Write state.
 	msgWriter      *msgWriter

conn_test.go

@@ -316,6 +316,60 @@ func TestConn(t *testing.T) {
 		assert.Equal(t, "read msg", s, string(b))
 	})
 
+	for _, tc := range []struct {
+		name     string
+		setLimit func(ctx context.Context, c *websocket.Conn) net.Conn
+	}{
+		{"BeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(1024)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}},
+		{"AfterNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			nc := websocket.NetConn(ctx, c, websocket.MessageBinary)
+			c.SetReadLimit(1024)
+			return nc
+		}},
+	} {
+		t.Run("netConn/readLimitKept/"+tc.name, func(t *testing.T) {
+			tt, c1, c2 := newConnTest(t, nil, nil)
+
+			n1 := tc.setLimit(tt.ctx, c1)
+			_ = c2.CloseRead(tt.ctx)
+
+			writeDone := xsync.Go(func() error {
+				return c2.Write(tt.ctx, websocket.MessageBinary, make([]byte, 4096))
+			})
+
+			_, err := io.ReadFull(n1, make([]byte, 4096))
+			assert.ErrorIs(t, websocket.ErrMessageTooBig, err)
+
+			_ = c2.CloseNow()
+			<-writeDone
+		})
+	}
+
+	t.Run("netConn/readLimitDisabledBeforeNetConn", func(t *testing.T) {
+		tt, c1, c2 := newConnTest(t, nil, nil)
+
+		c1.SetReadLimit(-1)
+		n1 := websocket.NetConn(tt.ctx, c1, websocket.MessageBinary)
+		n2 := websocket.NetConn(tt.ctx, c2, websocket.MessageBinary)
+
+		s := strings.Repeat("papa", 1<<20)
+		errs := xsync.Go(func() error {
+			_, err := n2.Write([]byte(s))
+			if err != nil {
+				return err
+			}
+			return n2.Close()
+		})
+
+		b, err := io.ReadAll(n1)
+		assert.Success(t, err)
+		assert.Success(t, <-errs)
+		assert.Equal(t, "read msg", s, string(b))
+	})
+
 	t.Run("netConn/pastDeadline", func(t *testing.T) {
 		tt, c1, c2 := newConnTest(t, nil, nil)
 

netconn.go

@@ -44,9 +44,12 @@ 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.
+// If SetReadLimit has not been called on c, NetConn disables the read limit.
+// A limit set with SetReadLimit before or after NetConn is kept.
 func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
-	c.SetReadLimit(-1)
+	if !c.readLimitSet.Load() {
+		c.setReadLimit(-1)
+	}
 
 	nc := &netConn{
 		c:       c,

netconn_js_test.go

@@ -0,0 +1,70 @@
+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()
+
+	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)
+		}, 65536, nil},
+		{"DisabledBeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(-1)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, 65536, nil},
+		{"BeforeNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			c.SetReadLimit(1024)
+			return websocket.NetConn(ctx, c, websocket.MessageBinary)
+		}, 4096, websocket.ErrMessageTooBig},
+		{"AfterNetConn", func(ctx context.Context, c *websocket.Conn) net.Conn {
+			nc := websocket.NetConn(ctx, c, websocket.MessageBinary)
+			c.SetReadLimit(1024)
+			return nc
+		}, 4096, 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 expects 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,11 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 //
 // Set to -1 to disable.
 func (c *Conn) SetReadLimit(n int64) {
+	c.readLimitSet.Store(true)
+	c.setReadLimit(n)
+}
+
+func (c *Conn) setReadLimit(n int64) {
 	if n >= 0 {
 		// We read one more byte than the limit in case
 		// there is a fin frame that needs to be read.

ws_js.go

@@ -46,6 +46,7 @@ type Conn struct {
 
 	// read limit for a message in bytes.
 	msgReadLimit atomic.Int64
+	readLimitSet atomic.Bool
 
 	closeReadMu  sync.Mutex
 	closeReadCtx context.Context
@@ -414,6 +415,11 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 
 // SetReadLimit implements *Conn.SetReadLimit for wasm.
 func (c *Conn) SetReadLimit(n int64) {
+	c.readLimitSet.Store(true)
+	c.setReadLimit(n)
+}
+
+func (c *Conn) setReadLimit(n int64) {
 	c.msgReadLimit.Store(n)
 }