bug fixes

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

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 {

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
 }

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{

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 {