diff --git a/conn.go b/conn.go index 0569bd2c..6761da3f 100644 --- a/conn.go +++ b/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 diff --git a/conn_test.go b/conn_test.go index 22208520..a3ee5910 100644 --- a/conn_test.go +++ b/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) diff --git a/netconn.go b/netconn.go index a854faa4..be27b25a 100644 --- a/netconn.go +++ b/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, diff --git a/netconn_js_test.go b/netconn_js_test.go new file mode 100644 index 00000000..2daeab62 --- /dev/null +++ b/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) + }) + } +} diff --git a/read.go b/read.go index 06672e10..13a8d479 100644 --- a/read.go +++ b/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 { diff --git a/ws_js.go b/ws_js.go index 026b75fc..126f3ca3 100644 --- a/ws_js.go +++ b/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)