Skip to content

Commit

Permalink
netconn: Disable read limit on WebSocket
Browse files Browse the repository at this point in the history
Closes #245
  • Loading branch information
nhooyr committed Jan 9, 2021
1 parent 085d46c commit ea87744
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 4 deletions.
4 changes: 4 additions & 0 deletions netconn.go
Expand Up @@ -38,7 +38,11 @@ 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.
func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
c.SetReadLimit(-1)

nc := &netConn{
c: c,
msgType: msgType,
Expand Down
21 changes: 17 additions & 4 deletions read.go
Expand Up @@ -74,10 +74,16 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
// By default, the connection has a message read limit of 32768 bytes.
//
// When the limit is hit, the connection will be closed with StatusMessageTooBig.
//
// Set to -1 to disable.
func (c *Conn) SetReadLimit(n int64) {
// We add read one more byte than the limit in case
// there is a fin frame that needs to be read.
c.msgReader.limitReader.limit.Store(n + 1)
if n >= 0 {
// We read one more byte than the limit in case
// there is a fin frame that needs to be read.
n++
}

c.msgReader.limitReader.limit.Store(n)
}

const defaultReadLimit = 32768
Expand Down Expand Up @@ -455,7 +461,11 @@ func (lr *limitReader) reset(r io.Reader) {
}

func (lr *limitReader) Read(p []byte) (int, error) {
if lr.n <= 0 {
if lr.n < 0 {
return lr.r.Read(p)
}

if lr.n == 0 {
err := fmt.Errorf("read limited at %v bytes", lr.limit.Load())
lr.c.writeError(StatusMessageTooBig, err)
return 0, err
Expand All @@ -466,6 +476,9 @@ func (lr *limitReader) Read(p []byte) (int, error) {
}
n, err := lr.r.Read(p)
lr.n -= int64(n)
if lr.n < 0 {
lr.n = 0
}
return n, err
}

Expand Down

0 comments on commit ea87744

Please sign in to comment.