diff --git a/internal/util/wirecodec/wirecodec.go b/internal/util/wirecodec/wirecodec.go index 0b3a130a1..7a69bfd1f 100644 --- a/internal/util/wirecodec/wirecodec.go +++ b/internal/util/wirecodec/wirecodec.go @@ -7,6 +7,7 @@ import ( "crypto/sha256" "encoding/hex" "errors" + "runtime" "github.com/klauspost/compress/zstd" ) @@ -35,8 +36,13 @@ const ( // EncodeAll/DecodeAll on these shared instances are safe for concurrent use. var ( + decoderWorkers = runtime.GOMAXPROCS(0) zstdEncoder, _ = zstd.NewWriter(nil) - zstdDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecodeBytes)) + zstdDecoder, _ = zstd.NewReader(nil, zstd.WithDecoderMaxMemory(maxDecodeBytes), + zstd.WithDecodeAllCapLimit(true), zstd.WithDecoderConcurrency(decoderWorkers)) + // decodeSlots keeps callers waiting for a decoder worker from each holding an output buffer. + decodeSlots = make(chan struct{}, decoderWorkers) + errDecodedTooLarge = errors.New("wirecodec: decompressed body exceeds limit") ) // Compress zstd-compresses b. @@ -47,12 +53,30 @@ func Compress(b []byte) []byte { // Decompress zstd-decompresses src, rejecting output larger than maxOut (and any // input that would blow the in-memory bomb ceiling). func Decompress(src []byte, maxOut int) ([]byte, error) { - out, err := zstdDecoder.DecodeAll(src, nil) - if err != nil { + if maxOut <= 0 || maxOut > maxDecodeBytes { + maxOut = maxDecodeBytes + } + var header zstd.Header + if err := header.Decode(src); err != nil { return nil, err } - if maxOut > 0 && len(out) > maxOut { - return nil, errors.New("wirecodec: decompressed body exceeds limit") + // The output buffer is sized once, from the first frame (Compress writes a single + // frame): letting it grow per frame copied all output so far for every extra frame. + size := maxOut + if header.HasFCS { + if header.FrameContentSize > uint64(maxOut) { + return nil, errDecodedTooLarge + } + size = int(header.FrameContentSize) + } + decodeSlots <- struct{}{} + defer func() { <-decodeSlots }() + out, err := zstdDecoder.DecodeAll(src, make([]byte, 0, size)) + if errors.Is(err, zstd.ErrDecoderSizeExceeded) { + return nil, errDecodedTooLarge + } + if err != nil { + return nil, err } return out, nil } diff --git a/internal/util/wirecodec/wirecodec_test.go b/internal/util/wirecodec/wirecodec_test.go index 8bd9137df..08ad22b30 100644 --- a/internal/util/wirecodec/wirecodec_test.go +++ b/internal/util/wirecodec/wirecodec_test.go @@ -2,6 +2,7 @@ package wirecodec import ( "bytes" + "runtime" "strings" "testing" ) @@ -49,3 +50,17 @@ func TestSha256HexStableAndSensitive(t *testing.T) { t.Fatalf("expected 64 hex chars, got %d", len(a)) } } + +// Every extra frame made the decoder copy all output so far into a larger buffer, so a +// few KiB of tiny frames cost gigabytes of allocation for one authenticated request. +func TestDecompressOfManyFramesAllocatesTheOutputOnce(t *testing.T) { + frame := Compress(bytes.Repeat([]byte("A"), 4096)) + src := bytes.Repeat(frame, 2000) + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + _, _ = Decompress(src, 8<<20) + runtime.ReadMemStats(&after) + if allocated := after.TotalAlloc - before.TotalAlloc; allocated > 64<<20 { + t.Fatalf("decoding %d bytes of frames allocated %d MiB", len(src), allocated>>20) + } +}