snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/websocket/compression_limit_test.go

@@ -0,0 +1,391 @@
+package websocket
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"math/rand"
+	"testing"
+)
+
+// writeCompressedFrame returns a buffer holding a single permessage-deflate
+// compressed binary message carrying payload. The frame is written as a client
+// (masked) so it can be consumed by a server-side reader.
+func writeCompressedFrame(t *testing.T, payload []byte, level int) *bytes.Buffer {
+	t.Helper()
+	var buf bytes.Buffer
+	wc := newTestConn(nil, &buf, false)
+	wc.newCompressionWriter = compressNoContextTakeover
+	if level != 0 {
+		if err := wc.SetCompressionLevel(level); err != nil {
+			t.Fatalf("SetCompressionLevel: %v", err)
+		}
+	}
+	w, err := wc.NextWriter(BinaryMessage)
+	if err != nil {
+		t.Fatalf("NextWriter: %v", err)
+	}
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	return &buf
+}
+
+// newCompressedReader builds a server-side reader for frames stored in src,
+// with permessage-deflate decompression enabled and the given decompressed read
+// limit. out captures any control frames (e.g. CloseMessageTooBig) the reader
+// writes back to the peer.
+func newCompressedReader(src *bytes.Buffer, decompressedLimit int64) (rc *Conn, out *bytes.Buffer) {
+	out = &bytes.Buffer{}
+	rc = newTestConn(src, out, true)
+	rc.newDecompressionReader = decompressNoContextTakeover
+	rc.SetDecompressedReadLimit(decompressedLimit)
+	return rc, out
+}
+
+func sentTooBig(out *bytes.Buffer) bool {
+	// FormatCloseMessage encodes the 2-byte close code; the server reader does
+	// not mask control frames, so the encoded code appears verbatim in out.
+	return bytes.Contains(out.Bytes(), FormatCloseMessage(CloseMessageTooBig, ""))
+}
+
+// TestDecompressedReadLimit_Bomb proves that a tiny compressed frame that
+// inflates to a huge payload is rejected with ErrReadLimit, that the close
+// frame is sent, and crucially that memory is bounded: io.ReadAll stops after
+// at most limit+1 bytes instead of materializing the whole bomb.
+func TestDecompressedReadLimit_Bomb(t *testing.T) {
+	const limit = 64 * 1024
+	const decompressedSize = 8 * 1024 * 1024 // 8MB inflates from a few KB.
+
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, decompressedSize), 9)
+	if src.Len() >= limit {
+		t.Fatalf("compressed frame %d bytes is not under the compressed read limit %d; test would not isolate the decompressed limit", src.Len(), limit)
+	}
+
+	rc, out := newCompressedReader(src, limit)
+
+	mt, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	if mt != BinaryMessage {
+		t.Fatalf("message type = %d, want %d", mt, BinaryMessage)
+	}
+
+	data, err := io.ReadAll(r)
+	if !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+	if int64(len(data)) > limit+1 {
+		t.Fatalf("read %d bytes, want at most %d (memory must stay bounded)", len(data), limit+1)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+	// The server-sent 1009 must be observable as an outgoing close code.
+	if code, incoming := rc.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestDecompressedReadLimit_Boundary proves the exact accept/reject boundary:
+// a message of exactly limit bytes is accepted, limit+1 is rejected.
+func TestDecompressedReadLimit_Boundary(t *testing.T) {
+	const limit = 4096
+
+	cases := []struct {
+		name       string
+		size       int
+		wantErr    bool
+		wantTooBig bool
+	}{
+		{"under", limit - 1, false, false},
+		{"exact", limit, false, false},
+		{"over_by_one", limit + 1, true, true},
+		{"over_by_many", limit * 4, true, true},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			payload := bytes.Repeat([]byte{'x'}, tc.size)
+			src := writeCompressedFrame(t, payload, 0)
+			rc, out := newCompressedReader(src, limit)
+
+			_, r, err := rc.NextReader()
+			if err != nil {
+				t.Fatalf("NextReader: %v", err)
+			}
+			data, err := io.ReadAll(r)
+
+			if tc.wantErr {
+				if !errors.Is(err, ErrReadLimit) {
+					t.Fatalf("error = %v, want ErrReadLimit", err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("unexpected error: %v", err)
+				}
+				if !bytes.Equal(data, payload) {
+					t.Fatalf("decompressed payload mismatch: got %d bytes, want %d", len(data), len(payload))
+				}
+			}
+			if got := sentTooBig(out); got != tc.wantTooBig {
+				t.Fatalf("sentTooBig = %v, want %v", got, tc.wantTooBig)
+			}
+		})
+	}
+}
+
+// TestDecompressedReadLimit_LegitCompressible proves the fix does not break the
+// legitimate case the report worried about: a message whose compressed size is
+// small but whose decompressed size is well above the compressed bytes is still
+// accepted, as long as it stays within the (higher) decompressed limit.
+func TestDecompressedReadLimit_LegitCompressible(t *testing.T) {
+	const compressedReadLimit = 8 * 1024    // mimics MessageSizeLimit
+	const decompressedReadLimit = 80 * 1024 // mimics MessageSizeLimit * multiplier
+	const payloadSize = 64 * 1024           // > compressed limit, < decompressed limit
+
+	payload := bytes.Repeat([]byte("centrifuge "), payloadSize/11+1)[:payloadSize]
+	src := writeCompressedFrame(t, payload, 9)
+	if src.Len() >= compressedReadLimit {
+		t.Fatalf("compressed size %d not under compressed read limit %d", src.Len(), compressedReadLimit)
+	}
+
+	rc, out := newCompressedReader(src, decompressedReadLimit)
+	rc.SetReadLimit(compressedReadLimit) // compressed-bytes limit also active
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("legit compressible message rejected: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("payload mismatch: got %d bytes, want %d", len(data), len(payload))
+	}
+	if sentTooBig(out) {
+		t.Fatalf("unexpected CloseMessageTooBig for a legit message")
+	}
+}
+
+// TestDecompressedReadLimit_ReproducesPreFixVulnerability documents the
+// vulnerability the fix addresses: with only the compressed-bytes limit
+// (SetReadLimit, the protection that existed before this change) and no
+// decompressed limit, a small compressed frame that satisfies the compressed
+// limit still inflates without bound. This is exactly the pre-fix behavior
+// (decompressedReadLimit == 0 is equivalent to the code before the fix).
+func TestDecompressedReadLimit_ReproducesPreFixVulnerability(t *testing.T) {
+	const compressedReadLimit = 64 * 1024    // the old message_size_limit default
+	const decompressedSize = 8 * 1024 * 1024 // 8MB from a few KB on the wire
+
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, decompressedSize), 9)
+	compressedLen := src.Len() // capture before reading drains the buffer
+	// The compressed frame is comfortably under the compressed read limit, so
+	// the pre-fix defense does not trigger.
+	if compressedLen >= compressedReadLimit {
+		t.Fatalf("compressed frame %d bytes not under compressed limit %d", compressedLen, compressedReadLimit)
+	}
+
+	rc, out := newCompressedReader(src, 0) // 0 == no decompressed limit (pre-fix)
+	rc.SetReadLimit(compressedReadLimit)   // only the old compressed-bytes defense
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	// Pre-fix: no error and the full 8MB is materialized despite the compressed
+	// read limit, demonstrating the unbounded amplification (the bug).
+	if err != nil {
+		t.Fatalf("pre-fix path unexpectedly errored: %v", err)
+	}
+	if len(data) != decompressedSize {
+		t.Fatalf("inflated to %d bytes, want full %d (unbounded amplification)", len(data), decompressedSize)
+	}
+	if sentTooBig(out) {
+		t.Fatalf("pre-fix path should not send CloseMessageTooBig")
+	}
+	t.Logf("reproduced: %d compressed wire bytes (<= %d limit) inflated to %d bytes unbounded",
+		compressedLen, compressedReadLimit, len(data))
+}
+
+// TestDecompressedReadLimit_Disabled proves that a zero limit means unlimited
+// (the fork primitive's contract): the bomb fully inflates without error.
+func TestDecompressedReadLimit_Disabled(t *testing.T) {
+	const decompressedSize = 1 * 1024 * 1024
+	payload := bytes.Repeat([]byte{'A'}, decompressedSize)
+	src := writeCompressedFrame(t, payload, 9)
+
+	rc, out := newCompressedReader(src, 0) // 0 == unlimited
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("unexpected error with limit disabled: %v", err)
+	}
+	if len(data) != decompressedSize {
+		t.Fatalf("read %d bytes, want %d", len(data), decompressedSize)
+	}
+	if sentTooBig(out) {
+		t.Fatalf("unexpected CloseMessageTooBig with limit disabled")
+	}
+}
+
+// TestDecompressedReadLimit_IgnoredWithoutCompression proves the decompressed
+// limit only applies to compressed frames: an uncompressed message larger than
+// the decompressed limit (but within the compressed read limit) passes through,
+// because the limit is gated on the RSV1/decompress path.
+func TestDecompressedReadLimit_IgnoredWithoutCompression(t *testing.T) {
+	const decompressedReadLimit = 1024
+	const payloadSize = 8 * 1024
+
+	var src bytes.Buffer
+	wc := newTestConn(nil, &src, false) // no compression writer -> plain frame
+	w, err := wc.NextWriter(BinaryMessage)
+	if err != nil {
+		t.Fatalf("NextWriter: %v", err)
+	}
+	payload := bytes.Repeat([]byte{'y'}, payloadSize)
+	if _, err := w.Write(payload); err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	rc, _ := newCompressedReader(&src, decompressedReadLimit)
+	// readLimit (compressed bytes) generous enough to admit the plain frame.
+	rc.SetReadLimit(payloadSize + 1024)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	data, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatalf("uncompressed message wrongly limited: %v", err)
+	}
+	if !bytes.Equal(data, payload) {
+		t.Fatalf("payload mismatch: got %d bytes, want %d", len(data), len(payload))
+	}
+}
+
+// TestDecompressedReadLimit_CompressedLimitStillApplies proves the two limits
+// are independent: with a decompressed limit set, the compressed-bytes limit
+// enforced in advanceFrame still trips for an oversized compressed frame.
+func TestDecompressedReadLimit_CompressedLimitStillApplies(t *testing.T) {
+	const compressedReadLimit = 1024
+	// Incompressible (deterministic pseudo-random) payload so compressed size
+	// stays large and exceeds the compressed read limit before decompression
+	// even matters.
+	payload := make([]byte, 8*1024)
+	rng := rand.New(rand.NewSource(1))
+	_, _ = rng.Read(payload)
+	src := writeCompressedFrame(t, payload, 0)
+	if src.Len() <= compressedReadLimit {
+		t.Fatalf("payload compressed too well (%d bytes) to exceed compressed limit", src.Len())
+	}
+
+	rc, _ := newCompressedReader(src, 10*1024*1024) // generous decompressed limit
+	rc.SetReadLimit(compressedReadLimit)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		// Limit may surface from NextReader itself (advanceFrame) ...
+		if !errors.Is(err, ErrReadLimit) {
+			t.Fatalf("NextReader error = %v, want ErrReadLimit", err)
+		}
+		return
+	}
+	// ... or while reading.
+	if _, err := io.ReadAll(r); !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+}
+
+// TestDecompressedReadLimit_ErrorIsPermanent proves that once the decompressed
+// limit trips, the error is sticky: subsequent NextReader calls keep returning
+// ErrReadLimit instead of advancing to the next frame. This mirrors the
+// compressed-bytes limit (enforced in advanceFrame, which sets c.readErr via
+// NextReader) and honors the gorilla contract that read errors are permanent.
+func TestDecompressedReadLimit_ErrorIsPermanent(t *testing.T) {
+	const limit = 4096
+
+	// Two back-to-back compressed frames in one stream. The first is a bomb that
+	// trips the decompressed limit; the second is a small legitimate message that
+	// must NOT be delivered once the connection is in its terminal error state.
+	src := writeCompressedFrame(t, bytes.Repeat([]byte{'A'}, 1*1024*1024), 9)
+	next := writeCompressedFrame(t, []byte("should never be read"), 9)
+	src.Write(next.Bytes())
+
+	rc, out := newCompressedReader(src, limit)
+
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+	if _, err := io.ReadAll(r); !errors.Is(err, ErrReadLimit) {
+		t.Fatalf("io.ReadAll error = %v, want ErrReadLimit", err)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+
+	// The error must now be permanent: NextReader keeps returning ErrReadLimit
+	// rather than reading the second frame.
+	for i := 0; i < 3; i++ {
+		mt, r, err := rc.NextReader()
+		if !errors.Is(err, ErrReadLimit) {
+			t.Fatalf("NextReader #%d after limit = (%d, %v), want ErrReadLimit", i, mt, err)
+		}
+		if r != nil {
+			t.Fatalf("NextReader #%d returned non-nil reader after permanent error", i)
+		}
+	}
+}
+
+// TestDecompressedReadLimit_StreamingReads proves the limit is enforced on the
+// incremental NextReader path (used by the bidirectional handler's streaming
+// command decoder), not only via io.ReadAll: reading small chunks still trips
+// the limit and keeps total consumed bytes bounded.
+func TestDecompressedReadLimit_StreamingReads(t *testing.T) {
+	const limit = 4096
+	payload := bytes.Repeat([]byte{'z'}, 1*1024*1024)
+	src := writeCompressedFrame(t, payload, 9)
+
+	rc, out := newCompressedReader(src, limit)
+	_, r, err := rc.NextReader()
+	if err != nil {
+		t.Fatalf("NextReader: %v", err)
+	}
+
+	buf := make([]byte, 256)
+	var total int
+	for {
+		n, err := r.Read(buf)
+		total += n
+		if err != nil {
+			if !errors.Is(err, ErrReadLimit) {
+				t.Fatalf("Read error = %v, want ErrReadLimit", err)
+			}
+			break
+		}
+		if total > limit+1 {
+			t.Fatalf("consumed %d bytes without hitting limit %d", total, limit)
+		}
+	}
+	if total > limit+1 {
+		t.Fatalf("consumed %d bytes, want at most %d", total, limit+1)
+	}
+	if !sentTooBig(out) {
+		t.Fatalf("expected CloseMessageTooBig control frame to be sent")
+	}
+}