test: make gzip decompression-bomb regression test actually bite

pkg/util/http_test.go

@@ -9,6 +9,7 @@ import (
 	"math/rand"
 	"net/http"
 	"net/http/httptest"
+	"runtime"
 	"strconv"
 	"testing"
 
@@ -223,24 +224,70 @@ func TestIsRequestBodyTooLargeRegression(t *testing.T) {
 }
 
 func TestParseProtoReader_GzipDecompressionBomb(t *testing.T) {
-	// Create a gzip payload where decompressed size far exceeds maxSize.
-	const maxSize = 4096                // 4 KB limit on decompressed output
-	uncompressed := make([]byte, 1<<20) // 1 MB of zeros
+	// Create a gzip payload that is small while compressed but decompresses to
+	// far more than maxSize, like a malicious "zip bomb".
+	const maxSize = 4096                 // 4 KB limit on decompressed output
+	uncompressed := make([]byte, 32<<20) // 32 MB of zeros, compresses to a few KB
+	require.Greater(t, len(uncompressed), maxSize)
 
 	var compressed bytes.Buffer
 	gzw := gzip.NewWriter(&compressed)
 	_, err := gzw.Write(uncompressed)
 	require.NoError(t, err)
 	require.NoError(t, gzw.Close())
 
-	// The compressed payload is small enough to pass the compressed-size limit,
-	// but decompresses to far more than maxSize.
-	require.Less(t, compressed.Len(), maxSize)
+	var m1, m2 runtime.MemStats
+	runtime.GC()
+	runtime.ReadMemStats(&m1)
 
 	var fromWire cortexpb.PreallocWriteRequest
 	err = util.ParseProtoReader(context.Background(), io.NopCloser(&compressed), 0, maxSize, &fromWire, util.Gzip)
-	// The decompressed output should be limited to maxSize+1 bytes, causing a
-	// proto unmarshal error (not an OOM). The key assertion is that we don't
-	// allocate 1 MB of memory.
-	assert.NotNil(t, err)
+
+	runtime.ReadMemStats(&m2)
+	allocated := m2.TotalAlloc - m1.TotalAlloc
+
+	// The request is rejected, not unmarshaled.
+	assert.Error(t, err)
+	// Decompression must stop early: with the cap in place only ~maxSize bytes are
+	// produced, so the whole call allocates well under 1 MB. Without the cap the
+	// reader decompresses megabytes before the size check rejects it, blowing past
+	// this bound, which is what makes this a real regression test for the bomb
+	// protection rather than a no-op.
+	assert.Less(t, allocated, uint64(1<<20), "decompression should be capped at ~maxSize, but allocated %d bytes", allocated)
+}
+
+func TestParseProtoReader_Gzip(t *testing.T) {
+	req := &cortexpb.PreallocWriteRequest{
+		WriteRequest: cortexpb.WriteRequest{
+			Timeseries: []cortexpb.PreallocTimeseries{
+				{TimeSeries: &cortexpb.TimeSeries{
+					Labels:     []cortexpb.LabelAdapter{{Name: "foo", Value: "bar"}},
+					Samples:    []cortexpb.Sample{{Value: 10, TimestampMs: 1}},
+					Exemplars:  []cortexpb.Exemplar{},
+					Histograms: []cortexpb.WrappedHistogram{},
+				}},
+			},
+		},
+	}
+	data, err := req.Marshal()
+	require.NoError(t, err)
+
+	var compressed bytes.Buffer
+	gzw := gzip.NewWriter(&compressed)
+	_, err = gzw.Write(data)
+	require.NoError(t, err)
+	require.NoError(t, gzw.Close())
+
+	// Under-cap payload decodes successfully. maxSize must clear both the
+	// compressed input and the decompressed body; gzip framing makes a tiny
+	// message larger compressed, so size off the compressed length.
+	maxSize := compressed.Len()
+	var fromWire cortexpb.PreallocWriteRequest
+	err = util.ParseProtoReader(context.Background(), io.NopCloser(bytes.NewReader(compressed.Bytes())), 0, maxSize, &fromWire, util.Gzip)
+	require.NoError(t, err)
+	require.Equal(t, req.Timeseries, fromWire.Timeseries)
+
+	// Capping the decompressed size below the message rejects it.
+	err = util.ParseProtoReader(context.Background(), io.NopCloser(bytes.NewReader(compressed.Bytes())), 0, len(data)-1, &fromWire, util.Gzip)
+	assert.Error(t, err)
 }