diff --git a/pkg/util/http_test.go b/pkg/util/http_test.go index d2a5e128d4..f715ae7bf6 100644 --- a/pkg/util/http_test.go +++ b/pkg/util/http_test.go @@ -9,6 +9,7 @@ import ( "math/rand" "net/http" "net/http/httptest" + "runtime" "strconv" "testing" @@ -223,9 +224,11 @@ 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) @@ -233,14 +236,58 @@ func TestParseProtoReader_GzipDecompressionBomb(t *testing.T) { 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) }