Skip to content

Commit 7d7c644

Browse files
refactor: add ErrMessageTooBig sentinel error for limited reads (#535)
--------- Co-authored-by: Mathias Fredriksson <[email protected]>
1 parent c7846ea commit 7d7c644

File tree

4 files changed

+35
-7
lines changed

4 files changed

+35
-7
lines changed

conn_test.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,25 @@ func TestConn(t *testing.T) {
421421
err = c1.Close(websocket.StatusNormalClosure, "")
422422
assert.Success(t, err)
423423
})
424+
425+
t.Run("ReadLimitExceededReturnsErrMessageTooBig", func(t *testing.T) {
426+
tt, c1, c2 := newConnTest(t, nil, nil)
427+
428+
c1.SetReadLimit(1024)
429+
_ = c2.CloseRead(tt.ctx)
430+
431+
writeDone := xsync.Go(func() error {
432+
payload := strings.Repeat("x", 4096)
433+
return c2.Write(tt.ctx, websocket.MessageText, []byte(payload))
434+
})
435+
436+
_, _, err := c1.Read(tt.ctx)
437+
assert.ErrorIs(t, websocket.ErrMessageTooBig, err)
438+
assert.Contains(t, err, "read limited at 1025 bytes")
439+
440+
_ = c2.CloseNow()
441+
<-writeDone
442+
})
424443
}
425444

426445
func TestWasm(t *testing.T) {

errors.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package websocket
2+
3+
import (
4+
"errors"
5+
)
6+
7+
// ErrMessageTooBig is returned when a message exceeds the read limit.
8+
var ErrMessageTooBig = errors.New("websocket: message too big")

read.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,8 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
9090
//
9191
// By default, the connection has a message read limit of 32768 bytes.
9292
//
93-
// When the limit is hit, the connection will be closed with StatusMessageTooBig.
93+
// When the limit is hit, reads return an error wrapping ErrMessageTooBig and
94+
// the connection is closed with StatusMessageTooBig.
9495
//
9596
// Set to -1 to disable.
9697
func (c *Conn) SetReadLimit(n int64) {
@@ -522,9 +523,9 @@ func (lr *limitReader) Read(p []byte) (int, error) {
522523
}
523524

524525
if lr.n == 0 {
525-
err := fmt.Errorf("read limited at %v bytes", lr.limit.Load())
526-
lr.c.writeError(StatusMessageTooBig, err)
527-
return 0, err
526+
reason := fmt.Errorf("read limited at %d bytes", lr.limit.Load())
527+
lr.c.writeError(StatusMessageTooBig, reason)
528+
return 0, fmt.Errorf("%w: %v", ErrMessageTooBig, reason)
528529
}
529530

530531
if int64(len(p)) > lr.n {

ws_js.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,9 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
144144
}
145145
readLimit := c.msgReadLimit.Load()
146146
if readLimit >= 0 && int64(len(p)) > readLimit {
147-
err := fmt.Errorf("read limited at %v bytes", c.msgReadLimit.Load())
148-
c.Close(StatusMessageTooBig, err.Error())
149-
return 0, nil, err
147+
reason := fmt.Errorf("read limited at %d bytes", c.msgReadLimit.Load())
148+
c.Close(StatusMessageTooBig, reason.Error())
149+
return 0, nil, fmt.Errorf("%w: %v", ErrMessageTooBig, reason)
150150
}
151151
return typ, p, nil
152152
}

0 commit comments

Comments
 (0)