392 lines
14 KiB
Go
392 lines
14 KiB
Go
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")
|
|
}
|
|
}
|