Cap payload allocations

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.

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,13 +404,17 @@ package main
 
 import (
 	"bytes"
+	"io"
 	"log"
 	"net"
 
 	"github.com/gobwas/ws"
 	"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)
 

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(`<error description here>`),
 				)),
 				ContentLength: 24,

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.
 */

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:

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)
 

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
+	}
+}