From 6a79a2dc6ab1b0add7e614cc1e11518f5cac0de6 Mon Sep 17 00:00:00 2001 From: Sergey Kamardin Date: Mon, 26 Oct 2026 06:16:47 +0000 Subject: [PATCH 1/2] all: modernize io/ioutil use with io --- dialer_test.go | 5 ++--- doc.go | 2 +- example/autobahn/autobahn.go | 3 +-- http_test.go | 4 ++-- server_test.go | 5 ++--- write_test.go | 5 ++--- wsflate/cbuf_test.go | 3 +-- wsflate/writer_test.go | 3 +-- wsutil/cipher_test.go | 4 ++-- wsutil/dialer.go | 3 +-- wsutil/dialer_test.go | 5 ++--- wsutil/handler.go | 3 +-- wsutil/helper.go | 7 +++---- wsutil/reader.go | 7 +++---- wsutil/reader_test.go | 9 ++++----- wsutil/upgrader.go | 3 +-- wsutil/utf8_test.go | 3 +-- wsutil/wsutil.go | 4 ++-- 18 files changed, 32 insertions(+), 46 deletions(-) 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/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 bde89ec..a1d379a 100644 --- a/write_test.go +++ b/write_test.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "io" - "io/ioutil" "testing" ) @@ -127,7 +126,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) } } @@ -140,7 +139,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 842cba7..5b26759 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" @@ -207,7 +206,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..e780f72 100644 --- a/wsutil/helper.go +++ b/wsutil/helper.go @@ -3,7 +3,6 @@ package wsutil import ( "bytes" "io" - "io/ioutil" "github.com/gobwas/ws" ) @@ -30,7 +29,7 @@ func ReadMessage(r io.Reader, s ws.State, m []Message) ([]Message, error) { State: s, CheckUTF8: true, OnIntermediate: func(hdr ws.Header, src io.Reader) error { - bts, err := ioutil.ReadAll(src) + bts, err := io.ReadAll(src) if err != nil { return err } @@ -51,7 +50,7 @@ func ReadMessage(r io.Reader, s ws.State, m []Message) ([]Message, error) { // Thus we consistent here with io.Reader behavior. _, err = io.ReadFull(&rd, p) } else { - // Frame is fragmented, thus use ioutil.ReadAll behavior. + // Frame is fragmented, thus use io.ReadAll behavior. var buf bytes.Buffer _, err = buf.ReadFrom(&rd) p = buf.Bytes() @@ -272,7 +271,7 @@ func readData(rw io.ReadWriter, s ws.State, want ws.OpCode) ([]byte, ws.OpCode, continue } - bts, err := ioutil.ReadAll(&rd) + bts, err := io.ReadAll(&rd) return bts, hdr.OpCode, err } diff --git a/wsutil/reader.go b/wsutil/reader.go index f2710af..0f8cfb7 100644 --- a/wsutil/reader.go +++ b/wsutil/reader.go @@ -4,7 +4,6 @@ import ( "encoding/binary" "errors" "io" - "io/ioutil" "github.com/gobwas/ws" ) @@ -147,7 +146,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 +173,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 @@ -221,7 +220,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 } diff --git a/wsutil/reader_test.go b/wsutil/reader_test.go index 4204bf3..ab51c60 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) 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..f79e5a1 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 } From e5248a30e02f8320f848f29f301bdf4f83d5843e Mon Sep 17 00:00:00 2001 From: Sergey Kamardin Date: Mon, 26 Oct 2026 06:35:22 +0000 Subject: [PATCH 2/2] ws,wsutil: do not trust peer-declared payload length for allocation Payloads are now read through internal/saferio.ReadN(): at most 32 KiB are allocated before the bytes actually arrive, the buffer grows with them and ends with an exact capacity. ws.ReadFrame() and MustReadFrame() are deprecated: the package has no knob to bound a payload and none is in the RFC; it's now encouraged to use ReadHeader() and check Header.Length instead. wsutil.Reader gets MaxMessageSize, the total of all fragments (RFC 6455 section 10.4), which MaxFrameSize alone never bounded. The helpers are configured with goroutine-safe SetMaxFrameSize(), SetMaxMessageSize(), SetCheckUTF8() and SetPreallocLimit(), each with a getter, and DefaultReader() returns the Reader they use. README and the package doc recommend setting the limits. Refs #223. Closes #229. --- AGENTS.md | 19 ++ README.md | 53 +++- internal/saferio/saferio.go | 71 +++++ internal/saferio/saferio_test.go | 263 +++++++++++++++++ read.go | 35 ++- read_test.go | 105 ++++++- wsutil/helper.go | 187 ++++++++++--- wsutil/helper_test.go | 467 ++++++++++++++++++++++++++++++- wsutil/reader.go | 45 ++- wsutil/reader_test.go | 155 ++++++++++ wsutil/wsutil.go | 18 ++ 11 files changed, 1361 insertions(+), 57 deletions(-) create mode 100644 AGENTS.md create mode 100644 internal/saferio/saferio.go create mode 100644 internal/saferio/saferio_test.go 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/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/wsutil/helper.go b/wsutil/helper.go index e780f72..2d55b6a 100644 --- a/wsutil/helper.go +++ b/wsutil/helper.go @@ -3,8 +3,10 @@ package wsutil import ( "bytes" "io" + "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 @@ -15,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 @@ -22,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 := io.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 io.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). @@ -83,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) } @@ -246,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 @@ -271,7 +386,7 @@ func readData(rw io.ReadWriter, s ws.State, want ws.OpCode) ([]byte, ws.OpCode, continue } - bts, err := io.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 0f8cfb7..f87a702 100644 --- a/wsutil/reader.go +++ b/wsutil/reader.go @@ -16,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. // @@ -46,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. @@ -188,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. @@ -266,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 ab51c60..99f082a 100644 --- a/wsutil/reader_test.go +++ b/wsutil/reader_test.go @@ -350,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/wsutil.go b/wsutil/wsutil.go index f79e5a1..2365809 100644 --- a/wsutil/wsutil.go +++ b/wsutil/wsutil.go @@ -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