diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..cb716fb --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,19 @@ +# Rules for coding agents + +- No limits or defaults that are not in the WebSocket RFCs. If it is not in the + spec, it is not in the code; expose a knob instead. +- Never allocate on a peer-declared length. Grow as bytes actually arrive (see + `internal/saferio`). +- Package-level settings must be safe to change concurrently (atomics or a + lock), or not exist. +- Prefer the standard library. Replace it with hand-written code only with + measurements, and say why in a comment. +- Helpers shared by `ws` and `wsutil` live under `internal/`. +- Style: no multi-line `if err := f(...); err != nil`; assign first, then + check. Single-line init statements are fine. +- Style: `t.Fatalf`/`t.Errorf`/`fmt.Errorf`/etc. calls are multi-line: the + format string on its own line, all arguments on the next line. +- Style: table-driven test cases are struct literals with one field per line, + never a whole case on one line. +- Before reporting: `gofmt -l`, `go vet`, `go test -race` on the touched + packages. diff --git a/README.md b/README.md index 0bd0f6b..6aa10ac 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,10 @@ import ( ) func main() { + // Bound the messages read by the wsutil helpers; there is no limit + // otherwise (see Limits below). + wsutil.SetMaxMessageSize(1 << 20) + http.ListenAndServe(":8080", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, _, _, err := ws.UpgradeHTTP(r, w) if err != nil { @@ -226,6 +230,37 @@ func main() { } ``` +# Limits + +The protocol does not limit the size of a frame or a message, and neither does +this library by default: it is up to you. A peer can declare a 16 GiB frame or +send a message in a million fragments, and the `wsutil` helpers read a whole +message into memory. It is highly recommended to set the limits that fit your +application: + +```go +// Reject a message, in all its fragments, bigger than 1 MiB. +wsutil.SetMaxMessageSize(1 << 20) +// Reject a single frame bigger than 64 KiB. +wsutil.SetMaxFrameSize(64 << 10) +``` + +These apply to `wsutil.ReadMessage()`, `wsutil.ReadData()` and their shortcuts +such as `wsutil.ReadClientData()`. `wsutil.DefaultReader()` returns the +`wsutil.Reader` they use. The same limits are the `MaxMessageSize` and +`MaxFrameSize` fields of a `wsutil.Reader` you build yourself. + +The declared length of a frame is never trusted for allocation: at most +`wsutil.PreallocLimit()` bytes (32 KiB by default, see +`wsutil.SetPreallocLimit()`) are allocated before the payload actually arrives, +so a peer that lies about the length can not make you allocate what it has not +sent. + +The lower level `ws` package has no such knobs: read the header with +`ws.ReadHeader()`, check `Header.Length` against your own limit, and only then +read the payload (see the compression example below). `ws.ReadFrame()` is +deprecated for exactly that reason. + # Zero-copy upgrade Zero-copy upgrade helps to avoid unnecessary allocations and copying while @@ -369,6 +404,7 @@ package main import ( "bytes" + "io" "log" "net" @@ -376,6 +412,9 @@ import ( "github.com/gobwas/ws/wsflate" ) +// maxPayloadLength bounds the frames we accept; ws does not do it for us. +const maxPayloadLength = 1 << 20 + func main() { ln, err := net.Listen("tcp", "localhost:8080") if err != nil { @@ -414,11 +453,23 @@ func main() { go func() { defer conn.Close() for { - frame, err := ws.ReadFrame(conn) + header, err := ws.ReadHeader(conn) if err != nil { // Handle error. return } + if header.Length > maxPayloadLength { + // Handle too large frame: ws does not limit it for you. + return + } + frame := ws.Frame{ + Header: header, + Payload: make([]byte, header.Length), + } + if _, err := io.ReadFull(conn, frame.Payload); err != nil { + // Handle error. + return + } frame = ws.UnmaskFrameInPlace(frame) diff --git a/dialer_test.go b/dialer_test.go index d622726..b9a7380 100644 --- a/dialer_test.go +++ b/dialer_test.go @@ -7,7 +7,6 @@ import ( "encoding/base64" "fmt" "io" - "io/ioutil" "math/rand" "net" "net/http" @@ -290,7 +289,7 @@ func TestDialerHandshake(t *testing.T) { headerConnection: []string{"Upgrade"}, headerUpgrade: []string{"websocket"}, }, - Body: ioutil.NopCloser(bytes.NewReader([]byte(`hello, gopher!`))), + Body: io.NopCloser(bytes.NewReader([]byte(`hello, gopher!`))), ContentLength: 14, }, accept: acceptValid, @@ -327,7 +326,7 @@ func TestDialerHandshake(t *testing.T) { ProtoMajor: 1, ProtoMinor: 1, Header: make(http.Header), - Body: ioutil.NopCloser(bytes.NewReader( + Body: io.NopCloser(bytes.NewReader( []byte(``), )), ContentLength: 24, diff --git a/doc.go b/doc.go index 0118ce2..b4e52ba 100644 --- a/doc.go +++ b/doc.go @@ -74,7 +74,7 @@ Or: // handle err } - io.CopyN(ioutil.Discard, conn, header.Length) + io.CopyN(io.Discard, conn, header.Length) For more info see the documentation. */ diff --git a/example/autobahn/autobahn.go b/example/autobahn/autobahn.go index 182f8b2..2769864 100644 --- a/example/autobahn/autobahn.go +++ b/example/autobahn/autobahn.go @@ -6,7 +6,6 @@ import ( "flag" "fmt" "io" - "io/ioutil" "log" "net" "net/http" @@ -326,7 +325,7 @@ func wsHandler(w http.ResponseWriter, r *http.Request) { continue case ws.OpPong: - io.CopyN(ioutil.Discard, conn, header.Length) + io.CopyN(io.Discard, conn, header.Length) continue case ws.OpClose: diff --git a/http_test.go b/http_test.go index 7e2ff1e..f037830 100644 --- a/http_test.go +++ b/http_test.go @@ -2,7 +2,7 @@ package ws import ( "bufio" - "io/ioutil" + "io" "net/textproto" "net/url" "testing" @@ -109,7 +109,7 @@ func BenchmarkHttpWriteUpgradeRequest(b *testing.B) { host: "test-host", }, } { - bw := bufio.NewWriter(ioutil.Discard) + bw := bufio.NewWriter(io.Discard) nonce := make([]byte, nonceSize) initNonce(nonce) diff --git a/internal/saferio/saferio.go b/internal/saferio/saferio.go new file mode 100644 index 0000000..04f81b1 --- /dev/null +++ b/internal/saferio/saferio.go @@ -0,0 +1,71 @@ +// Package saferio provides I/O functions that avoid allocating large amounts +// of memory unnecessarily. This is intended for reading data from a network +// peer where the size is part of the input data but the peer may be lying +// about it, as a WebSocket frame header does with the payload length. +// +// It mirrors the internal/saferio package of the standard library, which is +// used by encoding/gob, debug/elf and the like for the same reason. +package saferio + +import ( + "io" +) + +// ReadN reads n bytes from r, but avoids allocating all n bytes before they +// arrive if n is larger than lim. This avoids crashing the program by +// allocating all n bytes in cases where n is incorrect or malicious. +// +// The buffer starts at most lim bytes long and doubles as bytes actually +// arrive, never exceeding n, so a peer can not make the caller allocate more +// than it has actually sent (plus the growth slack), and the returned slice +// has exactly n bytes of capacity. +// +// It panics if lim is not positive, since trusting nothing is not a sane +// configuration and is always a bug at the call site. +// +// The error is io.EOF only if no bytes were read. If an io.EOF happens after +// reading some but not all the bytes, ReadN returns io.ErrUnexpectedEOF. +// +// Unlike the standard library version, which allocates up to 10 MiB without +// concern and grows by append() above that, ReadN takes the limit as an +// argument and keeps the final capacity exact, which is what the callers +// measured for. +// +// It panics if the given limit is <=0. +// +// NOTE: io.ReadAll(io.LimitReader(r, n)) is not used here because it ignores +// n: it starts from a 512-byte buffer and grows it by the append() policy, +// which costs several times more allocations even for small payloads, and it +// treats a truncated payload as a clean EOF. +func ReadN(r io.Reader, lim, n int64) ([]byte, error) { + if lim <= 0 { + panic("saferio: ReadN() limit must be greater than zero") + } + c := n + if c > lim { + c = lim + } + p := make([]byte, 0, c) + for { + // NOTE: cap(p) never exceeds n, so reading up to the capacity never + // reads past the payload. + m, err := io.ReadFull(r, p[len(p):cap(p)]) + p = p[:len(p)+m] + if err != nil { + if err == io.EOF && len(p) > 0 { + err = io.ErrUnexpectedEOF + } + return p, err + } + if int64(len(p)) == n { + return p, nil + } + c = 2 * int64(cap(p)) + if c > n { + c = n + } + q := make([]byte, len(p), c) + copy(q, p) + p = q + } +} diff --git a/internal/saferio/saferio_test.go b/internal/saferio/saferio_test.go new file mode 100644 index 0000000..4ed2104 --- /dev/null +++ b/internal/saferio/saferio_test.go @@ -0,0 +1,263 @@ +package saferio + +import ( + "bytes" + "errors" + "io" + "testing" + "testing/iotest" +) + +func TestReadN(t *testing.T) { + const lim = 8 + for _, test := range []struct { + name string + lim int64 + n int64 + avail int64 // Bytes available in the reader. + chunk bool // Reader returns one byte per Read(). + err error + }{ + { + name: "zero n", + lim: lim, + n: 0, + avail: 0, + }, + { + name: "zero n with data", + lim: lim, + n: 0, + avail: 5, + }, + { + name: "one byte", + lim: lim, + n: 1, + avail: 1, + }, + { + name: "below lim", + lim: lim, + n: lim - 1, + avail: lim - 1, + }, + { + name: "exactly lim", + lim: lim, + n: lim, + avail: lim, + }, + { + name: "above lim", + lim: lim, + n: lim + 1, + avail: lim + 1, + }, + { + name: "twice lim", + lim: lim, + n: 2 * lim, + avail: 2 * lim, + }, + { + name: "twice lim plus one", + lim: lim, + n: 2*lim + 1, + avail: 2*lim + 1, + }, + { + name: "many lim", + lim: lim, + n: 5*lim + 3, + avail: 5*lim + 3, + }, + { + name: "extra data below lim", + lim: lim, + n: 3, + avail: 4, + }, + { + name: "extra data above lim", + lim: lim, + n: lim + 1, + avail: lim + 2, + }, + { + name: "chunked below lim", + lim: lim, + n: lim - 1, + avail: lim - 1, + chunk: true, + }, + { + name: "chunked above lim", + lim: lim, + n: 5*lim + 3, + avail: 5*lim + 3, + chunk: true, + }, + + { + name: "no data", + lim: lim, + n: 1, + avail: 0, + err: io.EOF, + }, + { + name: "no data above lim", + lim: lim, + n: lim + 1, + avail: 0, + err: io.EOF, + }, + { + name: "short by one below lim", + lim: lim, + n: lim - 1, + avail: lim - 2, + err: io.ErrUnexpectedEOF, + }, + { + name: "short by one at lim", + lim: lim, + n: lim, + avail: lim - 1, + err: io.ErrUnexpectedEOF, + }, + { + name: "short by one above lim", + lim: lim, + n: lim + 1, + avail: lim, + err: io.ErrUnexpectedEOF, + }, + { + name: "short at growth", + lim: lim, + n: 5*lim + 3, + avail: 2 * lim, + err: io.ErrUnexpectedEOF, + }, + { + name: "short chunked", + lim: lim, + n: 5*lim + 3, + avail: 3*lim + 1, + chunk: true, + err: io.ErrUnexpectedEOF, + }, + } { + t.Run(test.name, func(t *testing.T) { + data := make([]byte, test.avail) + for i := range data { + data[i] = byte(i + 1) + } + src := bytes.NewReader(data) + r := io.Reader(src) + if test.chunk { + r = iotest.OneByteReader(r) + } + + p, err := ReadN(r, test.lim, test.n) + if err != test.err { + t.Fatalf( + "unexpected error: %v; want %v", + err, test.err, + ) + } + want := test.n + if test.avail < want { + want = test.avail + } + if !bytes.Equal(p, data[:want]) { + t.Fatalf( + "unexpected bytes: %v; want %v", + p, data[:want], + ) + } + if err == nil && int64(cap(p)) != test.n { + t.Fatalf( + "capacity is %d; want exactly %d", + cap(p), test.n, + ) + } + if rest := int64(src.Len()); rest != test.avail-want { + t.Fatalf( + "%d bytes left unread; want %d", + rest, test.avail-want, + ) + } + }) + } +} + +func TestReadNAllocation(t *testing.T) { + // The buffer must not be allocated on trust: with no data the allocation + // is bounded by lim however big n is. + for _, n := range []int64{9, 1 << 20, 1 << 40, 1<<63 - 1} { + p, err := ReadN(bytes.NewReader(nil), 8, n) + if err != io.EOF { + t.Fatalf( + "n=%d: unexpected error: %v; want %v", + n, err, io.EOF, + ) + } + if cap(p) > 8 { + t.Fatalf( + "n=%d: allocated %d bytes; want at most 8", + n, cap(p), + ) + } + } + // Below lim the payload is allocated at once and exactly. + data := make([]byte, 8) + src := bytes.NewReader(data) + allocs := testing.AllocsPerRun(100, func() { + src.Reset(data) + ReadN(src, 8, 8) + }) + if allocs != 1 { + t.Fatalf( + "ReadN() made %.0f allocations; want 1", + allocs, + ) + } +} + +func TestReadNError(t *testing.T) { + // Reader errors other than EOF are returned as is, with what was read. + errTest := errors.New("test error") + r := io.MultiReader(bytes.NewReader([]byte{1, 2, 3}), iotest.ErrReader(errTest)) + p, err := ReadN(r, 2, 10) + if err != errTest { + t.Fatalf( + "unexpected error: %v; want %v", + err, errTest, + ) + } + if !bytes.Equal(p, []byte{1, 2, 3}) { + t.Fatalf( + "unexpected bytes: %v", + p, + ) + } +} + +func TestReadNPanicsOnLimit(t *testing.T) { + for _, lim := range []int64{0, -1} { + func() { + defer func() { + if recover() == nil { + t.Fatalf( + "ReadN(lim=%d) did not panic", + lim, + ) + } + }() + ReadN(bytes.NewReader(make([]byte, 5)), lim, 5) + }() + } +} diff --git a/read.go b/read.go index 437c088..138842c 100644 --- a/read.go +++ b/read.go @@ -4,6 +4,8 @@ import ( "encoding/binary" "fmt" "io" + + "github.com/gobwas/ws/internal/saferio" ) // Errors used by frame reader. @@ -20,7 +22,6 @@ func ReadHeader(r io.Reader) (h Header, err error) { // after first hop that reads first 2 constant bytes, we could reuse 2 bytes. // So 14 - 2 = 12. bts := make([]byte, 2, MaxHeaderSize-2) - return ReadHeaderBuffer(r, bts) } @@ -78,12 +79,12 @@ func ReadHeaderBuffer(r io.Reader, bts []byte) (h Header, err error) { return h, err } - switch { - case length == 126: + switch length { + case 126: h.Length = int64(binary.BigEndian.Uint16(bts[:2])) bts = bts[2:] - case length == 127: + case 127: if bts[0]&0x80 != 0 { err = ErrHeaderLengthMSB return h, err @@ -99,28 +100,38 @@ func ReadHeaderBuffer(r io.Reader, bts []byte) (h Header, err error) { return h, nil } +// payloadPreallocLimit is the maximum peer-declared payload length that +// ReadFrame() trusts for allocation before the payload is read; a payload +// declared bigger than that is read into a buffer that grows as bytes actually +// arrive. The value is the same as io.Copy() chunk size. +const payloadPreallocLimit = 32 * 1024 + // ReadFrame reads a frame from r. -// It is not designed for high optimized use case cause it makes allocation -// for frame.Header.Length size inside to read frame payload into. // // Note that ReadFrame does not unmask payload. +// +// Deprecated: the payload length is declared by the peer and ReadFrame reads +// it whole into memory, so a peer can make the process run out of memory by +// actually sending that much data. (The declared length is not trusted for +// allocation though: at most 32 KiB are allocated before the data actually +// arrives.) The only way to bound it is to wrap r with io.LimitReader(), which +// fails in the middle of the payload with io.ErrUnexpectedEOF; it is simpler +// and more precise to use ReadHeader(), check Header.Length against your own +// limit and read the payload yourself. func ReadFrame(r io.Reader) (f Frame, err error) { f.Header, err = ReadHeader(r) if err != nil { return f, err } - if f.Header.Length > 0 { - // int(f.Header.Length) is safe here cause we have - // checked it for overflow above in ReadHeader. - f.Payload = make([]byte, int(f.Header.Length)) - _, err = io.ReadFull(r, f.Payload) + f.Payload, err = saferio.ReadN(r, payloadPreallocLimit, f.Header.Length) } - return f, err } // MustReadFrame is like ReadFrame but panics if frame can not be read. +// +// Deprecated: see ReadFrame. func MustReadFrame(r io.Reader) Frame { f, err := ReadFrame(r) if err != nil { diff --git a/read_test.go b/read_test.go index 7e6bb59..120ee5b 100644 --- a/read_test.go +++ b/read_test.go @@ -22,16 +22,24 @@ func TestReadHeader(t *testing.T) { r := bytes.NewReader(test.Data) h, err := ReadHeader(r) if test.Err && err == nil { - t.Errorf("expected error, got nil") + t.Errorf( + "expected error, got nil", + ) } if !test.Err && err != nil { - t.Errorf("unexpected error: %s", err) + t.Errorf( + "unexpected error: %s", + err, + ) } if test.Err { return } if !reflect.DeepEqual(h, test.Header) { - t.Errorf("ReadHeader()\nread:\n\t%#v\nwant:\n\t%#v", h, test.Header) + t.Errorf( + "ReadHeader()\nread:\n\t%#v\nwant:\n\t%#v", + h, test.Header, + ) } }) } @@ -79,3 +87,94 @@ func BenchmarkReadHeader(b *testing.B) { }) } } + +func TestReadFrameDeclaredLength(t *testing.T) { + // Peer-declared payload length must not be trusted for allocation: with + // no payload bytes available ReadFrame() must fail without allocating + // more than payloadPreallocLimit. + for _, length := range []int64{ + payloadPreallocLimit + 1, + 1 << 34, // 16 GiB + 1<<63 - 1, + } { + t.Run(fmt.Sprintf("len=%d", length), func(t *testing.T) { + var buf bytes.Buffer + err := WriteHeader(&buf, Header{ + Fin: true, + OpCode: OpBinary, + Length: length, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + f, err := ReadFrame(&buf) + if err != io.EOF { + t.Fatalf( + "unexpected ReadFrame() error: %v; want %v", + err, io.EOF, + ) + } + if n := int64(cap(f.Payload)); n > payloadPreallocLimit { + t.Fatalf( + "ReadFrame() allocated %d bytes; want at most %d", + n, payloadPreallocLimit, + ) + } + }) + } +} + +func TestReadFramePayloadGrowth(t *testing.T) { + lim := int(payloadPreallocLimit) + for _, size := range []int{ + 0, + 1, + lim - 1, + lim, + lim + 1, + 3*lim + 7, + } { + t.Run(fmt.Sprintf("size=%d", size), func(t *testing.T) { + data := make([]byte, size) + for i := range data { + data[i] = byte(i) + } + raw := MustCompileFrame(NewBinaryFrame(data)) + + f, err := ReadFrame(bytes.NewReader(raw)) + if err != nil { + t.Fatalf( + "unexpected ReadFrame() error: %v", + err, + ) + } + if !bytes.Equal(f.Payload, data) { + t.Fatalf( + "unexpected payload: %d bytes; want %d", + len(f.Payload), len(data), + ) + } + if size > 0 && cap(f.Payload) != size { + t.Fatalf( + "payload capacity is %d; want exactly %d", + cap(f.Payload), size, + ) + } + + // Truncated payload must be reported as unexpected EOF, as before. + // (Cutting the only byte of a 1-byte payload is a clean EOF.) + if size > 1 { + _, err = ReadFrame(bytes.NewReader(raw[:len(raw)-1])) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "unexpected truncated ReadFrame() error: %v; want %v", + err, io.ErrUnexpectedEOF, + ) + } + } + }) + } +} diff --git a/server_test.go b/server_test.go index 0e44d7c..dc9ba52 100644 --- a/server_test.go +++ b/server_test.go @@ -5,7 +5,6 @@ import ( "bytes" "fmt" "io" - "io/ioutil" "math/rand" "net" "net/http" @@ -534,7 +533,7 @@ func BenchmarkUpgrader(b *testing.B) { b.Run(bench.label, func(b *testing.B) { conn := make([]io.ReadWriter, b.N) for i := 0; i < b.N; i++ { - conn[i] = benchReadWriter{bytes.NewReader(reqBytes), ioutil.Discard} + conn[i] = benchReadWriter{bytes.NewReader(reqBytes), io.Discard} } i := new(int64) @@ -806,7 +805,7 @@ func mustMakeErrResponse(code int, err error, headers http.Header) *http.Respons ProtoMinor: 1, ContentLength: int64(len(body)), } - res.Body = ioutil.NopCloser( + res.Body = io.NopCloser( strings.NewReader(body), ) for k, v := range headers { diff --git a/write_test.go b/write_test.go index e24f3c0..fd34892 100644 --- a/write_test.go +++ b/write_test.go @@ -3,7 +3,7 @@ package ws import ( "bytes" "fmt" - "io/ioutil" + "io" "testing" ) @@ -35,7 +35,7 @@ func BenchmarkWriteHeader(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - if err := WriteHeader(ioutil.Discard, bench.header); err != nil { + if err := WriteHeader(io.Discard, bench.header); err != nil { b.Fatal(err) } } @@ -48,7 +48,7 @@ func BenchmarkWriteHeader(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - if err := WriteHeaderBuffer(ioutil.Discard, bench.header, bts); err != nil { + if err := WriteHeaderBuffer(io.Discard, bench.header, bts); err != nil { b.Fatal(err) } } diff --git a/wsflate/cbuf_test.go b/wsflate/cbuf_test.go index 2b0c219..39b3aa1 100644 --- a/wsflate/cbuf_test.go +++ b/wsflate/cbuf_test.go @@ -4,7 +4,6 @@ import ( "bytes" "fmt" "io" - "io/ioutil" "testing" ) @@ -148,7 +147,7 @@ func BenchmarkCBuf(b *testing.B) { } { b.Run(test.name, func(b *testing.B) { w := &cbuf{ - dst: ioutil.Discard, + dst: io.Discard, } for i := 0; i < b.N; i++ { w.Write(test.chunk) diff --git a/wsflate/writer_test.go b/wsflate/writer_test.go index f088b83..cbc5f40 100644 --- a/wsflate/writer_test.go +++ b/wsflate/writer_test.go @@ -5,7 +5,6 @@ import ( "compress/flate" "fmt" "io" - "io/ioutil" "net" "net/url" "testing" @@ -35,7 +34,7 @@ func TestWriter(t *testing.T) { r := NewReader(&buf, func(r io.Reader) Decompressor { return flate.NewReader(r) }) - act, err := ioutil.ReadAll(r) + act, err := io.ReadAll(r) if err != nil { t.Fatalf("unexpected Reader error: %v", err) } diff --git a/wsutil/cipher_test.go b/wsutil/cipher_test.go index f1460d1..c892d72 100644 --- a/wsutil/cipher_test.go +++ b/wsutil/cipher_test.go @@ -3,7 +3,7 @@ package wsutil import ( "bytes" "fmt" - "io/ioutil" + "io" "reflect" "testing" @@ -36,7 +36,7 @@ func TestCipherReader(t *testing.T) { src := &chopReader{bytes.NewReader(masked), test.chop} rd := NewCipherReader(src, mask) - bts, err := ioutil.ReadAll(rd) + bts, err := io.ReadAll(rd) if err != nil { t.Fatalf("unexpected error: %s", err) } diff --git a/wsutil/dialer.go b/wsutil/dialer.go index 4f8788f..d8d1ab3 100644 --- a/wsutil/dialer.go +++ b/wsutil/dialer.go @@ -5,7 +5,6 @@ import ( "bytes" "context" "io" - "io/ioutil" "net" "net/http" @@ -134,7 +133,7 @@ func (r *prefetchResponseReader) Read(p []byte) (int, error) { io.TeeReader(r.source, r.buffer), ), nil) if err == nil { - *r.contentLength, _ = io.Copy(ioutil.Discard, resp.Body) + *r.contentLength, _ = io.Copy(io.Discard, resp.Body) resp.Body.Close() } bts := r.buffer.Bytes() diff --git a/wsutil/dialer_test.go b/wsutil/dialer_test.go index 885f05b..c20ace4 100644 --- a/wsutil/dialer_test.go +++ b/wsutil/dialer_test.go @@ -5,7 +5,6 @@ import ( "bytes" "context" "io" - "io/ioutil" "net" "net/http" "testing" @@ -75,7 +74,7 @@ func TestDebugDialer(t *testing.T) { StatusCode: http.StatusOK, ProtoMajor: 1, ProtoMinor: 1, - Body: ioutil.NopCloser(bytes.NewReader( + Body: io.NopCloser(bytes.NewReader( bytes.Repeat([]byte("x"), 5000), )), ContentLength: 5000, @@ -150,7 +149,7 @@ func TestDebugDialer(t *testing.T) { t.Errorf("returned connection is non raw") } if br != nil { - body, err := ioutil.ReadAll(br) + body, err := io.ReadAll(br) if err != nil { t.Fatal(err) } diff --git a/wsutil/handler.go b/wsutil/handler.go index 44fd360..9b4208b 100644 --- a/wsutil/handler.go +++ b/wsutil/handler.go @@ -3,7 +3,6 @@ package wsutil import ( "errors" "io" - "io/ioutil" "strconv" "github.com/gobwas/pool/pbytes" @@ -123,7 +122,7 @@ func (c ControlHandler) HandlePong(h ws.Header) error { // A Pong frame MAY be sent unsolicited. This serves as a // unidirectional heartbeat. A response to an unsolicited Pong frame // is not expected. - _, err := io.CopyBuffer(ioutil.Discard, c.Src, buf) + _, err := io.CopyBuffer(io.Discard, c.Src, buf) return err } diff --git a/wsutil/helper.go b/wsutil/helper.go index 231760b..2d55b6a 100644 --- a/wsutil/helper.go +++ b/wsutil/helper.go @@ -3,9 +3,10 @@ package wsutil import ( "bytes" "io" - "io/ioutil" + "sync/atomic" "github.com/gobwas/ws" + "github.com/gobwas/ws/internal/saferio" ) // Message represents a message from peer, that could be presented in one or @@ -16,6 +17,103 @@ type Message struct { Payload []byte } +// defaultPreallocLimit is the default for SetPreallocLimit(): the same as +// io.Copy() chunk size. +const defaultPreallocLimit = 32 * 1024 + +// Settings of the Reader used by ReadMessage(), ReadData() and their +// shortcuts. They are accessed atomically, so that the setters below are safe +// to call concurrently with the helpers. +var ( + helperMaxFrameSize int64 + helperMaxMessageSize int64 + helperCheckUTF8 int32 = 1 + helperPreallocLimit int64 = defaultPreallocLimit +) + +// SetMaxFrameSize sets the Reader.MaxFrameSize used by ReadMessage(), +// ReadData() and their shortcuts. Zero, the default, means no limit. +// +// It is safe to call concurrently with the helpers. +func SetMaxFrameSize(n int64) { + atomic.StoreInt64(&helperMaxFrameSize, n) +} + +// MaxFrameSize returns the value set with SetMaxFrameSize(). +func MaxFrameSize() int64 { + return atomic.LoadInt64(&helperMaxFrameSize) +} + +// SetMaxMessageSize sets the Reader.MaxMessageSize used by ReadMessage(), +// ReadData() and their shortcuts: the total length of all fragments of a +// message. Zero, the default, means no limit. +// +// It is safe to call concurrently with the helpers. +func SetMaxMessageSize(n int64) { + atomic.StoreInt64(&helperMaxMessageSize, n) +} + +// MaxMessageSize returns the value set with SetMaxMessageSize(). +func MaxMessageSize() int64 { + return atomic.LoadInt64(&helperMaxMessageSize) +} + +// SetCheckUTF8 sets the [Reader.CheckUTF8] used by [ReadMessage], [ReadData] +// and their shortcuts. It is true by default. +// +// It is safe to call concurrently with the helpers. +func SetCheckUTF8(v bool) { + var x int32 + if v { + x = 1 + } + atomic.StoreInt32(&helperCheckUTF8, x) +} + +// CheckUTF8 returns the value set with SetCheckUTF8(). +func CheckUTF8() bool { + return atomic.LoadInt32(&helperCheckUTF8) != 0 +} + +// SetPreallocLimit sets the maximum peer-declared payload length that +// ReadMessage(), ReadData() and their shortcuts trust for allocation before +// the payload is read; a payload declared bigger than that is read into a +// buffer that grows as bytes actually arrive. The default is 32 KiB, the +// io.Copy() chunk size. Limits set with SetMaxFrameSize() and +// SetMaxMessageSize() reject frames; they do not raise this one. +// +// It panics if n is not positive: trusting nothing is not a sane +// configuration, and the mistake is better reported here than in the middle of +// a read. +// +// It is safe to call concurrently with the helpers. +func SetPreallocLimit(n int64) { + if n <= 0 { + panic("wsutil: SetPreallocLimit() limit must be greater than zero") + } + atomic.StoreInt64(&helperPreallocLimit, n) +} + +// PreallocLimit returns the value set with SetPreallocLimit(). +func PreallocLimit() int64 { + return atomic.LoadInt64(&helperPreallocLimit) +} + +// DefaultReader returns a Reader for r and s configured with the package-level +// settings, that is, the Reader that ReadMessage(), ReadData() and their +// shortcuts read with. Unlike NewReader(), it applies MaxFrameSize() and +// MaxMessageSize() and CheckUTF8(). The OnIntermediate field is left for the +// caller to set. +func DefaultReader(r io.Reader, s ws.State) *Reader { + return &Reader{ + Source: r, + State: s, + CheckUTF8: CheckUTF8(), + MaxFrameSize: MaxFrameSize(), + MaxMessageSize: MaxMessageSize(), + } +} + // ReadMessage is a helper function that reads next message from r. It appends // received message(s) to the third argument and returns the result of it and // an error if some failure happened. That is, it probably could receive more @@ -23,45 +121,60 @@ type Message struct { // want to send some control frame between fragments. Then returned slice will // contain those control frames at first, and then result of gluing fragments. // -// TODO(gobwas): add DefaultReader with buffer size options. +// Note that there is no limit on the message size by default, since the +// message is read whole into memory; to set one, use SetMaxMessageSize(). +// Allocation does not trust the peer-declared size beyond the limit set with +// SetPreallocLimit() and allocates as soon as the actual data is read. +// DefaultReader() returns the Reader these settings apply to. func ReadMessage(r io.Reader, s ws.State, m []Message) ([]Message, error) { - rd := Reader{ - Source: r, - State: s, - CheckUTF8: true, - OnIntermediate: func(hdr ws.Header, src io.Reader) error { - bts, err := ioutil.ReadAll(src) - if err != nil { - return err - } - m = append(m, Message{hdr.OpCode, bts}) - return nil - }, + rd := DefaultReader(r, s) + rd.OnIntermediate = func(hdr ws.Header, src io.Reader) error { + bts, err := io.ReadAll(src) + if err != nil { + return err + } + m = append(m, Message{hdr.OpCode, bts}) + return nil } h, err := rd.NextFrame() if err != nil { return m, err } - var p []byte - if h.Fin { - // No more frames will be read. Use fixed sized buffer to read payload. - p = make([]byte, h.Length) - // It is not possible to receive io.EOF here because Reader does not - // return EOF if frame payload was successfully fetched. - // Thus we consistent here with io.Reader behavior. - _, err = io.ReadFull(&rd, p) - } else { - // Frame is fragmented, thus use ioutil.ReadAll behavior. - var buf bytes.Buffer - _, err = buf.ReadFrom(&rd) - p = buf.Bytes() - } + p, err := readMessagePayload(rd, h) if err != nil { return m, err } return append(m, Message{h.OpCode, p}), nil } +// readMessagePayload reads from rd the payload of a message whose first frame +// header is hdr. +func readMessagePayload(rd *Reader, hdr ws.Header) ([]byte, error) { + if hdr.Fin { + // No more frames will be read, so the message size is known. Note that + // it is declared by the peer and can not be trusted for allocation + // beyond the prealloc limit, which saferio.ReadN() takes care of. + // + // MaxFrameSize and MaxMessageSize do not raise that trust: a bound + // above the limit still allocates as bytes arrive, and a bound below + // it has already rejected in NextFrame() any frame the limit could + // matter for. + return saferio.ReadN(rd, PreallocLimit(), hdr.Length) + } + // Message is fragmented, so its total size is unknown: let the buffer + // grow as bytes actually arrive, like io.ReadAll() does. The first + // fragment's declared length is a hint for the start size, trusted up to + // the same limit as a whole frame would be. + n := hdr.Length + if lim := PreallocLimit(); n > lim { + n = lim + } + var buf bytes.Buffer + buf.Grow(int(n)) + _, err := buf.ReadFrom(rd) + return buf.Bytes(), err +} + // ReadClientMessage reads next message from r, considering that caller // represents server side. // It is a shortcut for ReadMessage(r, ws.StateServerSide, m). @@ -84,6 +197,12 @@ func ReadServerMessage(r io.Reader, m []Message) ([]Message, error) { // // Note this may handle and write control frames into the writer part of a // given io.ReadWriter. +// +// Note that there is no limit on the message size by default, since the +// message is read whole into memory; to set one, use SetMaxMessageSize(). +// Allocation does not trust the peer-declared size beyond the limit set with +// SetPreallocLimit() and allocates as soon as the actual data is read. +// DefaultReader() returns the Reader these settings apply to. func ReadData(rw io.ReadWriter, s ws.State) ([]byte, ws.OpCode, error) { return readData(rw, s, ws.OpText|ws.OpBinary) } @@ -247,20 +366,15 @@ func ControlFrameHandler(w io.Writer, state ws.State) FrameHandlerFunc { func readData(rw io.ReadWriter, s ws.State, want ws.OpCode) ([]byte, ws.OpCode, error) { controlHandler := ControlFrameHandler(rw, s) - rd := Reader{ - Source: rw, - State: s, - CheckUTF8: true, - SkipHeaderCheck: false, - OnIntermediate: controlHandler, - } + rd := DefaultReader(rw, s) + rd.OnIntermediate = controlHandler for { hdr, err := rd.NextFrame() if err != nil { return nil, 0, err } if hdr.OpCode.IsControl() { - if err := controlHandler(hdr, &rd); err != nil { + if err := controlHandler(hdr, rd); err != nil { return nil, 0, err } continue @@ -272,7 +386,7 @@ func readData(rw io.ReadWriter, s ws.State, want ws.OpCode) ([]byte, ws.OpCode, continue } - bts, err := ioutil.ReadAll(&rd) + bts, err := readMessagePayload(rd, hdr) return bts, hdr.OpCode, err } diff --git a/wsutil/helper_test.go b/wsutil/helper_test.go index 901e673..f2a0ab3 100644 --- a/wsutil/helper_test.go +++ b/wsutil/helper_test.go @@ -57,10 +57,16 @@ func TestReadMessageEOF(t *testing.T) { t.Run("", func(t *testing.T) { ms, err := ReadMessage(test.source(), 0, nil) if err != test.err { - t.Errorf("unexpected error: %v; want %v", err, test.err) + t.Errorf( + "unexpected error: %v; want %v", + err, test.err, + ) } if n := len(ms); n != len(test.messages) { - t.Fatalf("unexpected number of read messages: %d; want %d", n, 0) + t.Fatalf( + "unexpected number of read messages: %d; want %d", + n, 0, + ) } for i, exp := range test.messages { act := ms[i] @@ -80,3 +86,460 @@ func TestReadMessageEOF(t *testing.T) { }) } } + +func TestReadMessageDeclaredLength(t *testing.T) { + // Peer-declared payload length must not be trusted for allocation: with + // no payload bytes available ReadMessage() must fail without allocating + // more than defaultPreallocLimit. + for _, length := range []int64{ + defaultPreallocLimit + 1, + 1 << 34, // 16 GiB + 1<<63 - 1, + } { + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: true, + OpCode: ws.OpBinary, + Length: length, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + ms, err := ReadMessage(&buf, 0, nil) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "len=%d: unexpected ReadMessage() error: %v; want %v", + length, err, io.ErrUnexpectedEOF, + ) + } + if len(ms) != 0 { + t.Fatalf( + "len=%d: ReadMessage() returned %d messages; want 0", + length, len(ms), + ) + } + } +} + +func TestReadMessagePayloadGrowth(t *testing.T) { + lim := int(defaultPreallocLimit) + for _, size := range []int{ + 0, + 1, + lim, + 3*lim + 7, + } { + payload := make([]byte, size) + for i := range payload { + payload[i] = byte(i) + } + raw := ws.MustCompileFrame(ws.NewBinaryFrame(payload)) + + ms, err := ReadMessage(bytes.NewReader(raw), 0, nil) + if err != nil { + t.Fatalf( + "size=%d: unexpected ReadMessage() error: %v", + size, err, + ) + } + if len(ms) != 1 { + t.Fatalf( + "size=%d: ReadMessage() returned %d messages; want 1", + size, len(ms), + ) + } + if !bytes.Equal(ms[0].Payload, payload) { + t.Fatalf( + "size=%d: unexpected payload: %d bytes; want %d", + size, len(ms[0].Payload), len(payload), + ) + } + if size > 0 && cap(ms[0].Payload) != size { + t.Fatalf( + "size=%d: payload capacity is %d; want exactly %d", + size, cap(ms[0].Payload), size, + ) + } + } +} + +func TestReadDataDeclaredLength(t *testing.T) { + // Same as TestReadMessageDeclaredLength, but for the ReadData() family, + // which reads the payload through the same path. + for _, length := range []int64{ + defaultPreallocLimit + 1, + 1 << 34, // 16 GiB + 1<<63 - 1, + } { + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: true, + OpCode: ws.OpBinary, + Length: length, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + p, _, err := ReadData(&buf, ws.StateClientSide) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "len=%d: unexpected ReadData() error: %v; want %v", + length, err, io.ErrUnexpectedEOF, + ) + } + if n := int64(cap(p)); n > defaultPreallocLimit { + t.Fatalf( + "len=%d: ReadData() allocated %d bytes; want at most %d", + length, n, defaultPreallocLimit, + ) + } + } +} + +func TestSetMaxFrameSize(t *testing.T) { + SetMaxFrameSize(16) + defer SetMaxFrameSize(0) + + raw := ws.MustCompileFrame(ws.NewBinaryFrame(make([]byte, 17))) + if _, err := ReadMessage(bytes.NewReader(raw), 0, nil); err != ErrFrameTooLarge { + t.Fatalf( + "unexpected ReadMessage() error: %v; want %v", + err, ErrFrameTooLarge, + ) + } + var buf bytes.Buffer + buf.Write(raw) + if _, _, err := ReadData(&buf, ws.StateClientSide); err != ErrFrameTooLarge { + t.Fatalf( + "unexpected ReadData() error: %v; want %v", + err, ErrFrameTooLarge, + ) + } + + raw = ws.MustCompileFrame(ws.NewBinaryFrame(make([]byte, 16))) + if _, err := ReadMessage(bytes.NewReader(raw), 0, nil); err != nil { + t.Fatalf( + "unexpected ReadMessage() error: %v", + err, + ) + } +} + +func TestSetMaxFrameSizeAllocation(t *testing.T) { + // A bound above the prealloc limit does not raise what the declared + // length is trusted for: with no payload bytes available the allocation + // is still bounded by the prealloc limit. + size := 3*int64(defaultPreallocLimit) + 7 + SetMaxFrameSize(size) + defer SetMaxFrameSize(0) + + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: true, + OpCode: ws.OpBinary, + Length: size, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + p, _, err := ReadData(&buf, ws.StateClientSide) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "unexpected ReadData() error: %v; want %v", + err, io.ErrUnexpectedEOF, + ) + } + if n := cap(p); n > defaultPreallocLimit { + t.Fatalf( + "ReadData() allocated %d bytes; want at most %d", + n, defaultPreallocLimit, + ) + } + + // A complete payload is still read whole and exactly. + raw := ws.MustCompileFrame(ws.NewBinaryFrame(make([]byte, size))) + ms, err := ReadMessage(bytes.NewReader(raw), 0, nil) + if err != nil { + t.Fatalf( + "unexpected ReadMessage() error: %v", + err, + ) + } + if n := int64(cap(ms[0].Payload)); n != size { + t.Fatalf( + "payload capacity is %d; want exactly %d", + n, size, + ) + } +} + +func TestSetPreallocLimit(t *testing.T) { + SetPreallocLimit(8) + defer SetPreallocLimit(defaultPreallocLimit) + + // With no payload bytes available the allocation is bounded by the limit. + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: true, + OpCode: ws.OpBinary, + Length: 1 << 40, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + p, _, err := ReadData(&buf, ws.StateClientSide) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "unexpected ReadData() error: %v; want %v", + err, io.ErrUnexpectedEOF, + ) + } + if n := cap(p); n > 8 { + t.Fatalf( + "ReadData() allocated %d bytes; want at most 8", + n, + ) + } + + // A complete payload is still read whole and exactly. + data := make([]byte, 1024) + raw := ws.MustCompileFrame(ws.NewBinaryFrame(data)) + ms, err := ReadMessage(bytes.NewReader(raw), 0, nil) + if err != nil { + t.Fatalf( + "unexpected ReadMessage() error: %v", + err, + ) + } + if n := cap(ms[0].Payload); n != len(data) { + t.Fatalf( + "payload capacity is %d; want exactly %d", + n, len(data), + ) + } +} + +func TestSetPreallocLimitPanics(t *testing.T) { + for _, n := range []int64{0, -1} { + func() { + defer func() { + if recover() == nil { + t.Fatalf( + "SetPreallocLimit(%d) did not panic", + n, + ) + } + }() + SetPreallocLimit(n) + }() + } +} + +func TestSettingsGetters(t *testing.T) { + defer SetMaxFrameSize(MaxFrameSize()) + defer SetMaxMessageSize(MaxMessageSize()) + defer SetCheckUTF8(CheckUTF8()) + defer SetPreallocLimit(PreallocLimit()) + + if v := MaxFrameSize(); v != 0 { + t.Fatalf( + "MaxFrameSize() default: %d; want 0", + v, + ) + } + if v := MaxMessageSize(); v != 0 { + t.Fatalf( + "MaxMessageSize() default: %d; want 0", + v, + ) + } + if v := CheckUTF8(); !v { + t.Fatalf( + "CheckUTF8() default: %t; want true", + v, + ) + } + if v := PreallocLimit(); v != defaultPreallocLimit { + t.Fatalf( + "PreallocLimit() default: %d; want %d", + v, defaultPreallocLimit, + ) + } + + SetMaxFrameSize(42) + SetMaxMessageSize(84) + SetCheckUTF8(false) + SetPreallocLimit(7) + + if v := MaxFrameSize(); v != 42 { + t.Fatalf( + "MaxFrameSize(): %d; want 42", + v, + ) + } + if v := MaxMessageSize(); v != 84 { + t.Fatalf( + "MaxMessageSize(): %d; want 84", + v, + ) + } + if v := CheckUTF8(); v { + t.Fatalf( + "CheckUTF8(): %t; want false", + v, + ) + } + if v := PreallocLimit(); v != 7 { + t.Fatalf( + "PreallocLimit(): %d; want 7", + v, + ) + } + + src := bytes.NewReader(nil) + rd := DefaultReader(src, ws.StateServerSide) + if rd.Source != src || rd.State != ws.StateServerSide { + t.Fatalf( + "DefaultReader(): Source=%v State=%v; want the given ones", + rd.Source, rd.State, + ) + } + if rd.MaxFrameSize != 42 || rd.MaxMessageSize != 84 || rd.CheckUTF8 { + t.Fatalf( + "DefaultReader(): MaxFrameSize=%d MaxMessageSize=%d CheckUTF8=%t; want 42 84 false", + rd.MaxFrameSize, rd.MaxMessageSize, rd.CheckUTF8, + ) + } + if rd.OnIntermediate != nil { + t.Fatalf( + "DefaultReader(): OnIntermediate is set; want nil", + ) + } +} + +func TestSetMaxMessageSize(t *testing.T) { + SetMaxMessageSize(16) + defer SetMaxMessageSize(0) + + // frames writes fragments of a message with the given payload sizes. + frames := func(sizes ...int) []byte { + var buf bytes.Buffer + for i, n := range sizes { + op := ws.OpContinuation + if i == 0 { + op = ws.OpBinary + } + f := ws.NewFrame(op, i == len(sizes)-1, make([]byte, n)) + err := ws.WriteFrame(&buf, f) + if err != nil { + t.Fatal(err) + } + } + return buf.Bytes() + } + + _, err := ReadMessage(bytes.NewReader(frames(10, 10)), 0, nil) + if err != ErrMessageTooLarge { + t.Fatalf( + "ReadMessage() error: %v; want %v", + err, ErrMessageTooLarge, + ) + } + buf := bytes.NewBuffer(frames(10, 10)) + _, _, err = ReadData(buf, ws.StateClientSide) + if err != ErrMessageTooLarge { + t.Fatalf( + "ReadData() error: %v; want %v", + err, ErrMessageTooLarge, + ) + } + ms, err := ReadMessage(bytes.NewReader(frames(8, 8)), 0, nil) + if err != nil { + t.Fatalf( + "ReadMessage() error: %v; want nil", + err, + ) + } + if n := len(ms[0].Payload); n != 16 { + t.Fatalf( + "ReadMessage() payload is %d bytes; want 16", + n, + ) + } +} + +func TestSetMaxMessageSizeAllocation(t *testing.T) { + // Same as for MaxFrameSize: the bound rejects, it does not raise trust. + size := 3*int64(defaultPreallocLimit) + 7 + SetMaxMessageSize(size) + defer SetMaxMessageSize(0) + + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: true, + OpCode: ws.OpBinary, + Length: size, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + p, _, err := ReadData(&buf, ws.StateClientSide) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "unexpected ReadData() error: %v; want %v", + err, io.ErrUnexpectedEOF, + ) + } + if n := cap(p); n > defaultPreallocLimit { + t.Fatalf( + "ReadData() allocated %d bytes; want at most %d", + n, defaultPreallocLimit, + ) + } +} + +func TestReadMessageFragmentedDeclaredLength(t *testing.T) { + // The first fragment's declared length is only a hint for the buffer + // start size: with no payload bytes available the allocation is bounded + // by the prealloc limit. + var buf bytes.Buffer + err := ws.WriteHeader(&buf, ws.Header{ + Fin: false, + OpCode: ws.OpBinary, + Length: 1 << 40, + }) + if err != nil { + t.Fatalf( + "unexpected WriteHeader() error: %v", + err, + ) + } + p, _, err := ReadData(&buf, ws.StateClientSide) + if err != io.ErrUnexpectedEOF { + t.Fatalf( + "unexpected ReadData() error: %v; want %v", + err, io.ErrUnexpectedEOF, + ) + } + if n := cap(p); n > defaultPreallocLimit { + t.Fatalf( + "ReadData() allocated %d bytes; want at most %d", + n, defaultPreallocLimit, + ) + } +} diff --git a/wsutil/reader.go b/wsutil/reader.go index f2710af..f87a702 100644 --- a/wsutil/reader.go +++ b/wsutil/reader.go @@ -4,7 +4,6 @@ import ( "encoding/binary" "errors" "io" - "io/ioutil" "github.com/gobwas/ws" ) @@ -17,6 +16,10 @@ var ErrNoFrameAdvance = errors.New("no frame advance") // MaxFrameSize was being read. var ErrFrameTooLarge = errors.New("frame too large") +// ErrMessageTooLarge indicates that a message of total length higher than +// MaxMessageSize, that is, of all its fragments together, was being read. +var ErrMessageTooLarge = errors.New("message too large") + // FrameHandlerFunc handles parsed frame header and its body represented by // io.Reader. // @@ -47,17 +50,38 @@ type Reader struct { // header RSV segment. Extensions []RecvExtension - // MaxFrameSize controls the maximum frame size in bytes - // that can be read. A message exceeding that size will return - // a ErrFrameTooLarge to the application. + // MaxFrameSize controls the maximum frame size in bytes that can be read. + // A message exceeding that size will return a ErrFrameTooLarge to the + // application. // // Not setting this field means there is no limit. MaxFrameSize int64 + // MaxMessageSize controls the maximum message size in bytes that can be + // read, that is, the total length of all fragments of a message. A message + // exceeding that size will return a ErrMessageTooLarge to the application + // from the NextFrame() call for the first fragment crossing the limit. + // + // This is the limit RFC 6455 section 10.4 talks about. Note that + // MaxFrameSize alone does not bound a fragmented message. + // + // Not setting this field means there is no limit. + MaxMessageSize int64 + + // OnContinuation is called by NextFrame() for every continuation frame of + // a fragmented message, with its header and body. The body is the same + // reader the fragment is then read from, so bytes consumed by the callback + // are not returned by Read(). OnContinuation FrameHandlerFunc + + // OnIntermediate is called by NextFrame() for a control frame received in + // the middle of a fragmented message, with its header and body; the rest + // of the body is discarded after the callback returns. If it is nil, such + // frames are discarded silently, and only their header is returned. OnIntermediate FrameHandlerFunc opCode ws.OpCode // Used to store message op code on fragmentation. + msgLen int64 // Declared length of the current message fragments so far. frame io.Reader // Used to as frame reader. raw io.LimitedReader // Used to discard frames without cipher. utf8 UTF8Reader // Used to check UTF8 sequences if CheckUTF8 is true. @@ -147,7 +171,7 @@ func (r *Reader) Read(p []byte) (n int, err error) { // It discards all frames of fragmented message. func (r *Reader) Discard() (err error) { for { - _, err = io.Copy(ioutil.Discard, &r.raw) + _, err = io.Copy(io.Discard, &r.raw) if err != nil { break } @@ -174,7 +198,7 @@ func (r *Reader) NextFrame() (hdr ws.Header, err error) { // unexpected. // // NOTE: This is necessary to prevent callers such that - // ioutil.ReadAll to receive some amount of bytes without an error. + // io.ReadAll to receive some amount of bytes without an error. // ReadAll() ignores an io.EOF error, thus caller may think that // whole message fetched, but actually only part of it. err = io.ErrUnexpectedEOF @@ -189,6 +213,19 @@ func (r *Reader) NextFrame() (hdr ws.Header, err error) { if n := r.MaxFrameSize; n > 0 && hdr.Length > n { return hdr, ErrFrameTooLarge } + if !r.fragmented() { + // First frame of a message. + r.msgLen = 0 + } + if !r.fragmented() || !hdr.OpCode.IsControl() { + // Control frames in the middle of a fragmented message do not belong + // to it, everything else does. The subtraction can not overflow and + // the sum can not exceed the limit. + if n := r.MaxMessageSize; n > 0 && hdr.Length > n-r.msgLen { + return hdr, ErrMessageTooLarge + } + r.msgLen += hdr.Length + } // Save raw reader to use it on discarding frame without ciphering and // other streaming checks. @@ -221,7 +258,7 @@ func (r *Reader) NextFrame() (hdr ws.Header, err error) { } if err == nil { // Ensure that src is empty. - _, err = io.Copy(ioutil.Discard, &r.raw) + _, err = io.Copy(io.Discard, &r.raw) } return hdr, err } @@ -267,6 +304,7 @@ func (r *Reader) reset() { r.frame = nil r.utf8 = UTF8Reader{} r.opCode = 0 + r.msgLen = 0 } // readHeader reads a frame header from in. diff --git a/wsutil/reader_test.go b/wsutil/reader_test.go index 4204bf3..99f082a 100644 --- a/wsutil/reader_test.go +++ b/wsutil/reader_test.go @@ -3,7 +3,6 @@ package wsutil import ( "bytes" "io" - "io/ioutil" "testing" "unicode/utf8" @@ -28,7 +27,7 @@ func TestReadFromWithIntermediateControl(t *testing.T) { r := Reader{ Source: &buf, OnIntermediate: func(h ws.Header, r io.Reader) error { - bts, err := ioutil.ReadAll(r) + bts, err := io.ReadAll(r) if err != nil { t.Fatal(err) } @@ -53,7 +52,7 @@ func TestReadFromWithIntermediateControl(t *testing.T) { t.Fatalf("unexpected NextFrame() header: %+v; want %+v", act, exp) } - act, err := ioutil.ReadAll(&r) + act, err := io.ReadAll(&r) if err != nil { t.Fatal(err) } @@ -215,7 +214,7 @@ func TestReaderUTF8(t *testing.T) { if _, err := r.NextFrame(); err != nil { t.Fatal(err) } - bts, err := ioutil.ReadAll(&r) + bts, err := io.ReadAll(&r) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -321,7 +320,7 @@ func TestNextReader(t *testing.T) { var bts []byte _, reader, err := NextReader(conn, 0) if err == nil { - bts, err = ioutil.ReadAll(reader) + bts, err = io.ReadAll(reader) } if err != test.err { t.Fatalf("unexpected error; got %v; want %v", err, test.err) @@ -351,3 +350,158 @@ func (c chopReader) Read(p []byte) (n int, err error) { } return c.src.Read(p[:sz]) } + +func TestMaxMessageSize(t *testing.T) { + // frames writes the given frames into a buffer the Reader reads from. + frames := func(fs ...ws.Frame) *bytes.Buffer { + var buf bytes.Buffer + for _, f := range fs { + err := ws.WriteFrame(&buf, f) + if err != nil { + t.Fatal(err) + } + } + return &buf + } + // fragment returns a frame carrying n bytes of a fragmented message. + fragment := func(op ws.OpCode, fin bool, n int) ws.Frame { + return ws.NewFrame(op, fin, make([]byte, n)) + } + for _, test := range []struct { + name string + limit int64 + frames []ws.Frame + // Payload sizes of the messages read before the error, if any. + messages []int + err error + }{ + { + name: "single frame at limit", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, true, 10), + }, + messages: []int{10}, + }, + { + name: "single frame above limit", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, true, 11), + }, + err: ErrMessageTooLarge, + }, + { + name: "fragments at limit", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 4), + fragment(ws.OpContinuation, false, 4), + fragment(ws.OpContinuation, true, 2), + }, + messages: []int{10}, + }, + { + name: "fragments above limit", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 4), + fragment(ws.OpContinuation, false, 4), + fragment(ws.OpContinuation, true, 3), + }, + err: ErrMessageTooLarge, + }, + { + name: "many small fragments above limit", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, false, 1), + fragment(ws.OpContinuation, true, 1), + }, + err: ErrMessageTooLarge, + }, + { + name: "intermediate control does not count", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 5), + ws.NewPingFrame(make([]byte, 5)), + fragment(ws.OpContinuation, true, 5), + }, + messages: []int{10}, + }, + { + name: "counter resets between messages", + limit: 10, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 5), + fragment(ws.OpContinuation, true, 5), + fragment(ws.OpBinary, false, 5), + fragment(ws.OpContinuation, true, 5), + }, + messages: []int{10, 10}, + }, + { + name: "no limit", + limit: 0, + frames: []ws.Frame{ + fragment(ws.OpBinary, false, 5000), + fragment(ws.OpContinuation, true, 5000), + }, + messages: []int{10000}, + }, + } { + t.Run(test.name, func(t *testing.T) { + r := Reader{ + Source: frames(test.frames...), + MaxMessageSize: test.limit, + } + var messages []int + var err error + for err == nil { + _, err = r.NextFrame() + if err != nil { + break + } + var p []byte + p, err = io.ReadAll(&r) + if err != nil { + break + } + messages = append(messages, len(p)) + } + if err == io.EOF { + err = nil + } + if err != test.err { + t.Fatalf( + "unexpected error: %v; want %v", + err, test.err, + ) + } + if len(messages) != len(test.messages) { + t.Fatalf( + "read %d messages: %v; want %d: %v", + len(messages), messages, len(test.messages), test.messages, + ) + } + for i, n := range messages { + if n != test.messages[i] { + t.Fatalf( + "message #%d has %d bytes; want %d", + i, n, test.messages[i], + ) + } + } + }) + } +} diff --git a/wsutil/upgrader.go b/wsutil/upgrader.go index 2ed351e..9d0e559 100644 --- a/wsutil/upgrader.go +++ b/wsutil/upgrader.go @@ -4,7 +4,6 @@ import ( "bufio" "bytes" "io" - "io/ioutil" "net/http" "github.com/gobwas/ws" @@ -40,7 +39,7 @@ func (d *DebugUpgrader) Upgrade(conn io.ReadWriter) (hs ws.Handshake, err error) )) if err == nil { // Fulfill the buffer with the response body. - io.Copy(ioutil.Discard, req.Body) + io.Copy(io.Discard, req.Body) req.Body.Close() } onRequest(buf.Bytes()) diff --git a/wsutil/utf8_test.go b/wsutil/utf8_test.go index 09caa40..9c86596 100644 --- a/wsutil/utf8_test.go +++ b/wsutil/utf8_test.go @@ -5,7 +5,6 @@ import ( "encoding/hex" "fmt" "io" - "io/ioutil" "testing" "unicode/utf8" ) @@ -199,7 +198,7 @@ func BenchmarkUTF8Reader(b *testing.B) { sz: bench.chop, } r := NewUTF8Reader(cr) - _, err := ioutil.ReadAll(r) + _, err := io.ReadAll(r) if !bench.err && err != nil { b.Fatal(err) } diff --git a/wsutil/wsutil.go b/wsutil/wsutil.go index 86211f3..2365809 100644 --- a/wsutil/wsutil.go +++ b/wsutil/wsutil.go @@ -14,7 +14,7 @@ Overview: r = wsutil.NewCipherReader(r, header.Mask) r = wsutil.NewUTF8Reader(r) - payload, err := ioutil.ReadAll(r) + payload, err := io.ReadAll(r) if err != nil { // handle err } @@ -26,7 +26,7 @@ You could get the same behavior using just `wsutil.Reader`: CheckUTF8: true, } - payload, err := ioutil.ReadAll(r) + payload, err := io.ReadAll(r) if err != nil { // handle err } @@ -38,6 +38,24 @@ Or even simplest: // handle err } +Limits: + +The protocol has no limit on the size of a frame or a message, and neither has +this package by default, while ReadMessage(), ReadData() and their shortcuts +such as ReadClientText() read a whole message into memory. It is highly +recommended to set the limits that fit your application, once, before reading: + + // Reject a message, in all its fragments, bigger than 1 MiB. + wsutil.SetMaxMessageSize(1 << 20) + // Reject a single frame bigger than 64 KiB. + wsutil.SetMaxFrameSize(64 << 10) + +The same limits are the MaxMessageSize and MaxFrameSize fields of a Reader you +build yourself; DefaultReader() returns the Reader the helpers read with. The +declared length of a frame is never trusted for allocation beyond +PreallocLimit(), 32 KiB by default, so a peer that lies about the length can +not make you allocate what it has not sent. + Package is also exports tools for buffered writing: // Create buffered writer, that will buffer output bytes and send them as