diff --git a/conn.go b/conn.go index 0569bd2c..1df92623 100644 --- a/conn.go +++ b/conn.go @@ -59,6 +59,7 @@ type Conn struct { readHeaderBuf [8]byte readControlBuf [maxControlPayload]byte msgReader *msgReader + readLimitSet atomic.Bool // Write state. msgWriter *msgWriter diff --git a/conn_test.go b/conn_test.go index 22208520..10405a70 100644 --- a/conn_test.go +++ b/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) diff --git a/netconn.go b/netconn.go index a854faa4..1dc1e79f 100644 --- a/netconn.go +++ b/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, diff --git a/netconn_js_test.go b/netconn_js_test.go new file mode 100644 index 00000000..046162ca --- /dev/null +++ b/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) + }) + } +} diff --git a/read.go b/read.go index 06672e10..70ae54e7 100644 --- a/read.go +++ b/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. diff --git a/ws_js.go b/ws_js.go index 026b75fc..880530b0 100644 --- a/ws_js.go +++ b/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) }