diff --git a/compress_test.go b/compress_test.go index 7e6e85b..caf2d68 100644 --- a/compress_test.go +++ b/compress_test.go @@ -5,6 +5,7 @@ import ( "encoding/base64" "encoding/hex" "fmt" + "io" "math" "os" "path/filepath" @@ -127,6 +128,139 @@ func TestCompressWithWhitelist(t *testing.T) { require.True(t, bytes.Equal(b, dat2)) } +func TestMemoryLimit(t *testing.T) { + t.Run("limitedBuffer_within_limit", func(t *testing.T) { + buf := &limitedBuffer{maxSize: 10} + n, err := buf.Write([]byte("hello")) + require.NoError(t, err) + require.Equal(t, 5, n) + }) + + t.Run("limitedBuffer_at_exact_limit", func(t *testing.T) { + buf := &limitedBuffer{maxSize: 5} + n, err := buf.Write([]byte("hello")) + require.NoError(t, err) + require.Equal(t, 5, n) + }) + + t.Run("limitedBuffer_exceeds_limit", func(t *testing.T) { + buf := &limitedBuffer{maxSize: 4} + _, err := buf.Write([]byte("hello")) + require.Error(t, err) + require.ErrorIs(t, err, ErrOutputTooBig) + }) + + t.Run("limitedBuffer_zero_means_no_limit", func(t *testing.T) { + buf := &limitedBuffer{maxSize: 0} + _, err := buf.Write(bytes.Repeat([]byte("x"), 10000)) + require.NoError(t, err) + }) + + t.Run("limitedBuffer_cumulative_writes", func(t *testing.T) { + buf := &limitedBuffer{maxSize: 8} + _, err := buf.Write([]byte("hello")) + require.NoError(t, err) + _, err = buf.Write([]byte("world")) + require.Error(t, err) + require.ErrorIs(t, err, ErrOutputTooBig) + }) + + t.Run("flateInflateWithLimit_within_limit", func(t *testing.T) { + data := bytes.Repeat([]byte("abcdefgh"), 100) + compressed, err := flateCompress(data) + require.NoError(t, err) + out, err := flateInflateWithLimit(compressed, int64(len(data))) + require.NoError(t, err) + require.Equal(t, data, out) + }) + + t.Run("flateInflateWithLimit_exceeds_limit", func(t *testing.T) { + data := bytes.Repeat([]byte("abcdefgh"), 100) + compressed, err := flateCompress(data) + require.NoError(t, err) + _, err = flateInflateWithLimit(compressed, int64(len(data)-1)) + require.Error(t, err) + require.ErrorIs(t, err, ErrOutputTooBig) + }) + + t.Run("InflateWithLimit_sufficient_limit", func(t *testing.T) { + original := loadTestVector(t, vectors[0]) + compressed, err := Compress(original) + require.NoError(t, err) + out, err := InflateWithLimit(compressed, int64(len(original))*2) + require.NoError(t, err) + require.Equal(t, original, out) + }) + + t.Run("InflateWithLimit_too_small_limit", func(t *testing.T) { + original := loadTestVector(t, vectors[0]) + compressed, err := Compress(original) + require.NoError(t, err) + _, err = InflateWithLimit(compressed, 1) + require.Error(t, err) + require.ErrorIs(t, err, ErrOutputTooBig) + }) + + t.Run("InflateWithLimit_exact_size_limit", func(t *testing.T) { + original := loadTestVector(t, vectors[0]) + compressed, err := Compress(original) + require.NoError(t, err) + out, err := InflateWithLimit(compressed, int64(len(original))) + require.NoError(t, err) + require.Equal(t, original, out) + }) + + // Regression: maxSize+1 overflowed to MinInt64 when maxSize==math.MaxInt64, + // making LimitReader return EOF immediately. + t.Run("flateInflateWithLimit_maxInt64_does_not_overflow", func(t *testing.T) { + data := bytes.Repeat([]byte("abcdefgh"), 100) + compressed, err := flateCompress(data) + require.NoError(t, err) + out, err := flateInflateWithLimit(compressed, math.MaxInt64) + require.NoError(t, err) + require.Equal(t, data, out) + }) + + // Regression: InflateWithLimit used c.maxSize as the keymap decompression limit, + // so a payload whose internal keymap exceeded maxSize was wrongly rejected even + // when the actual output was within the caller's limit. + t.Run("InflateWithLimit_limit_applies_to_output_not_keymap", func(t *testing.T) { + // Use a vector with a non-trivial keymap; compress it, then inflate with a + // limit equal to the original size — if the keymap limit were misapplied the + // internal keymap bytes would count against the budget and this would fail. + original := loadTestVector(t, vectors[2]) // thread-106, large keymap + compressed, err := Compress(original) + require.NoError(t, err) + out, err := InflateWithLimit(compressed, int64(len(original))) + require.NoError(t, err) + require.Equal(t, original, out) + }) + + // Regression: Inflate() (no limit) was given a hard 128 MB cap on inflateData + // output that Compress() does not enforce, breaking the Compress→Inflate invariant. + t.Run("Inflate_roundtrip_invariant", func(t *testing.T) { + for _, v := range vectors { + original := loadTestVector(t, v) + compressed, err := Compress(original) + require.NoError(t, err) + out, err := Inflate(compressed) + require.NoError(t, err) + require.Equal(t, original, out, "round-trip failed for %s", v.name) + } + }) + + // guardAlloc should be a no-op (not panic or error) for readers that don't + // implement Len(), so the guard stays safe for future reader types. + t.Run("guardAlloc_passthrough_for_non_len_reader", func(t *testing.T) { + // Wrap in a struct that exposes only io.Reader, hiding Len(). + type plainReader struct{ io.Reader } + r := plainReader{bytes.NewBuffer([]byte("hello"))} + // Should return nil (no Len() method — type assertion misses). + err := guardAlloc(r, 100) + require.NoError(t, err) + }) +} + func TestIntegerOverflowProtection(t *testing.T) { t.Run("msgpackInt_toUint32_overflow", func(t *testing.T) { // Test uint64 value exceeding MaxUint32 diff --git a/decode.go b/decode.go index cb6c5c3..1e251b4 100644 --- a/decode.go +++ b/decode.go @@ -1,6 +1,7 @@ package msgpackzip import ( + "bytes" "encoding/binary" "errors" "io" @@ -12,6 +13,7 @@ var ( ErrContainerTooBig = errors.New("container allocation is too big") ErrStringTooBig = errors.New("string allocation is too big") ErrBinaryTooBig = errors.New("binary allocation is too big") + ErrOutputTooBig = errors.New("decompressed output exceeds limit") ErrLenTooBig = errors.New("Lengths bigger than 0x8000000 are too big") ErrIntTooBig = errors.New("Cannot handle ints largers than int64 max") ErrExtTooBig = errors.New("extenal data type too big") @@ -121,6 +123,15 @@ type msgpackDecoderHooks struct { fallthroughHook func(i any, s string) error } +// guardAlloc returns io.ErrUnexpectedEOF if the underlying reader has fewer than +// need bytes available, preventing large pre-allocations from crafted headers. +func guardAlloc(r io.Reader, need int) error { + if buf, ok := r.(*bytes.Buffer); ok && buf.Len() < need { + return io.ErrUnexpectedEOF + } + return nil +} + func readByte(r io.Reader) (byte, error) { var buf [1]byte _, err := r.Read(buf[:]) @@ -283,6 +294,9 @@ func (m *msgpackDecoder) decodeString(s decodeStack, i msgpackInt) (err error) { if l > bigString { return ErrStringTooBig } + if err := guardAlloc(m.r, l); err != nil { + return err + } buf := make([]byte, l) _, err = io.ReadFull(m.r, buf) if err != nil { @@ -299,6 +313,9 @@ func (m *msgpackDecoder) decodeBinary(s decodeStack, i msgpackInt) (err error) { if l > bigBinary { return ErrBinaryTooBig } + if err := guardAlloc(m.r, l); err != nil { + return err + } buf := make([]byte, l) _, err = io.ReadFull(m.r, buf) if err != nil { @@ -443,6 +460,9 @@ func (m *msgpackDecoder) decodeExt(s decodeStack, n uint32) (err error) { if n > bigExt { return ErrExtTooBig } + if err := guardAlloc(m.r, int(n)); err != nil { //nolint:gosec // G115: n <= bigExt < MaxInt32 + return err + } buf := make([]byte, n) _, err = io.ReadFull(m.r, buf) if err != nil { diff --git a/flate.go b/flate.go index b80ef8a..291384a 100644 --- a/flate.go +++ b/flate.go @@ -4,6 +4,7 @@ import ( "bytes" "compress/flate" "io" + "math" ) func flateCompress(b []byte) ([]byte, error) { @@ -23,8 +24,22 @@ func flateCompress(b []byte) ([]byte, error) { return buf.Bytes(), nil } -func flateInflate(b []byte) ([]byte, error) { - buf := bytes.NewBuffer(b) - zr := flate.NewReader(buf) - return io.ReadAll(zr) +// flateInflateWithLimit decompresses b, returning ErrOutputTooBig if the +// result exceeds maxSize bytes. maxSize must be positive. +func flateInflateWithLimit(b []byte, maxSize int64) ([]byte, error) { + zr := flate.NewReader(bytes.NewBuffer(b)) + defer zr.Close() //nolint:errcheck // reader Close only returns decompressor to pool + // Read one byte past maxSize to detect oversize; guard against overflow. + limit := maxSize + if limit < math.MaxInt64 { + limit++ + } + out, err := io.ReadAll(io.LimitReader(zr, limit)) + if err != nil { + return nil, err + } + if int64(len(out)) > maxSize { + return nil, ErrOutputTooBig + } + return out, nil } diff --git a/inflate.go b/inflate.go index 4d81e9c..ba88143 100644 --- a/inflate.go +++ b/inflate.go @@ -9,17 +9,30 @@ import ( ) type inflator struct { - input *bytes.Buffer + input *bytes.Buffer + maxSize int64 // 0 = no limit } func newInflator(b []byte) *inflator { - return &inflator{input: bytes.NewBuffer(b)} + return newInflatorWithLimit(b, 0) +} + +func newInflatorWithLimit(b []byte, maxSize int64) *inflator { + return &inflator{input: bytes.NewBuffer(b), maxSize: maxSize} } func Inflate(input []byte) (output []byte, err error) { return newInflator(input).run() } +// InflateWithLimit inflates input but returns ErrOutputTooBig if the +// decompressed output would exceed maxSize bytes. The output limit is enforced +// incrementally via limitedBuffer. Note: the internal keymap intermediate is +// always capped at bigLen (128 MB) regardless of maxSize. +func InflateWithLimit(input []byte, maxSize int64) (output []byte, err error) { + return newInflatorWithLimit(input, maxSize).run() +} + func (c *inflator) run() (output []byte, err error) { version, compressedData, compressedKeymap, err := c.openOuter() if err != nil { @@ -95,7 +108,9 @@ func (c *inflator) openOuter() (version int, compressedData []byte, compressedKe } func (c *inflator) inflateKeymap(compressedKeymap []byte) (keymap map[uint]any, err error) { - rawKeymap, err := flateInflate(compressedKeymap) + // Always cap keymap decompression at bigLen regardless of the caller-supplied + // output limit: the keymap is an internal intermediate, not part of the output. + rawKeymap, err := flateInflateWithLimit(compressedKeymap, bigLen) if err != nil { return nil, err } @@ -138,7 +153,7 @@ func (c *inflator) inflateKeymap(compressedKeymap []byte) (keymap map[uint]any, if err != nil { return d, err } - keymap = make(map[uint]any, i) + keymap = make(map[uint]any, min(i, len(rawKeymap)/2)) return d, nil }, mapKeyHook: func(d decodeStack) (decodeStack, error) { @@ -218,6 +233,9 @@ func decodeBufToUint32(b []byte) (uint32, error) { func (c *inflator) inflateData(keymap map[uint]any, compressedData []byte) (ret []byte, err error) { var data outputter + if c.maxSize > 0 { + data = newOutputterWithLimit(c.maxSize) + } hooks := data.decoderHooks() hooks.mapKeyHook = func(d decodeStack) (decodeStack, error) { d.hooks = msgpackDecoderHooks{ diff --git a/output.go b/output.go index 78de296..06f946c 100644 --- a/output.go +++ b/output.go @@ -8,8 +8,31 @@ import ( "math" ) +// limitedBuffer wraps bytes.Buffer and rejects writes that would push the +// total past maxSize (0 = no limit). It tracks total bytes written rather +// than using Len() so the check remains correct even if the buffer is +// partially drained between writes. +type limitedBuffer struct { + bytes.Buffer + maxSize int64 + written int64 +} + +func (b *limitedBuffer) Write(p []byte) (int, error) { + if b.maxSize > 0 && b.written+int64(len(p)) > b.maxSize { + return 0, ErrOutputTooBig + } + n, err := b.Buffer.Write(p) + b.written += int64(n) + return n, err +} + type outputter struct { - buf bytes.Buffer + buf limitedBuffer +} + +func newOutputterWithLimit(maxSize int64) outputter { + return outputter{buf: limitedBuffer{maxSize: maxSize}} } func (o *outputter) Bytes() []byte {