fix(util): enforce gzip decompression limits

Replace blind trust of gzip trailer size with io.LimitReader to
prevent memory exhaustion from forged payloads. Add a regression
test, and refactor to use test cases.

Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>

v1/server/server_test.go

@@ -2108,7 +2108,7 @@ func TestDataPostV1CompressedDecodingLimits(t *testing.T) {
 			payload:               mustGZIPPayload([]byte(`{"input": {"user": "alice"}}`)),
 			expRespHTTPStatus:     400,
 			forcePayloadSizeField: 134217728, // 128 MB
-			expErrorMsg:           "gzip payload too large",
+			expErrorMsg:           "gzip: invalid checksum",
 			gzipMaxLen:            1024,
 		},
 		{
@@ -2118,7 +2118,7 @@ func TestDataPostV1CompressedDecodingLimits(t *testing.T) {
 			payload:               mustGZIPPayload([]byte(`{"input": {"user": "alice"}}`)),
 			expRespHTTPStatus:     400,
 			forcePayloadSizeField: 134217728, // 128 MB
-			expErrorMsg:           "gzip payload too large",
+			expErrorMsg:           "gzip: invalid checksum",
 			gzipMaxLen:            1024,
 		},
 		{

v1/util/read_gzip_body.go

@@ -3,7 +3,6 @@ package util
 import (
 	"bytes"
 	"compress/gzip"
-	"encoding/binary"
 	"errors"
 	"io"
 	"net/http"
@@ -15,11 +14,10 @@ import (
 var gzipReaderPool = NewSyncPool[gzip.Reader]()
 
 // Note(philipc): Originally taken from server/server.go
-// The DecodingLimitHandler handles validating that the gzip payload is within the
-// allowed max size limit. Thus, in the event of a forged payload size trailer,
-// the worst that can happen is that we waste memory up to the allowed max gzip
-// payload size, but not an unbounded amount of memory, as was potentially
-// possible before.
+// The DecodingLimitHandler handles setting the max size limits in the context.
+// This function enforces those limits. For gzip payloads, we use a LimitReader
+// to ensure we don't decompress more than the allowed maximum, preventing
+// memory exhaustion from forged gzip trailers.
 func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
 	length := r.ContentLength
 	if maxLenConf, ok := decoding.GetServerDecodingMaxLen(r.Context()); ok {
@@ -34,15 +32,6 @@ func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
 	if strings.Contains(r.Header.Get("Content-Encoding"), "gzip") {
 		gzipMaxLength, _ := decoding.GetServerDecodingGzipMaxLen(r.Context())
 
-		// Note(philipc): The last 4 bytes of a well-formed gzip blob will
-		// always be a little-endian uint32, representing the decompressed
-		// content size, modulo 2^32. We validate that the size is safe,
-		// earlier in DecodingLimitHandler.
-		sizeDecompressed := int64(binary.LittleEndian.Uint32(content[len(content)-4:]))
-		if sizeDecompressed > gzipMaxLength {
-			return nil, errors.New("gzip payload too large")
-		}
-
 		gzReader := gzipReaderPool.Get()
 		defer func() {
 			gzReader.Close()
@@ -53,11 +42,16 @@ func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
 			return nil, err
 		}
 
-		decompressed := bytes.NewBuffer(make([]byte, 0, sizeDecompressed))
-		if _, err = io.CopyN(decompressed, gzReader, sizeDecompressed); err != nil {
+		decompressed := bytes.NewBuffer(make([]byte, 0, len(content)))
+		limitReader := io.LimitReader(gzReader, gzipMaxLength+1)
+		if _, err := decompressed.ReadFrom(limitReader); err != nil {
 			return nil, err
 		}
 
+		if int64(decompressed.Len()) > gzipMaxLength {
+			return nil, errors.New("gzip payload too large")
+		}
+
 		return decompressed.Bytes(), nil
 	}
 

v1/util/read_gzip_body_test.go

@@ -3,6 +3,7 @@ package util_test
 import (
 	"bytes"
 	"compress/gzip"
+	"encoding/binary"
 	"net/http/httptest"
 	"testing"
 
@@ -13,31 +14,72 @@ import (
 func TestReadMaybeCompressedBody(t *testing.T) {
 	t.Parallel()
 
-	exp := []byte(`{"input": {"foo": "bar"}}`)
-
-	bb := new(bytes.Buffer)
-	gz := gzip.NewWriter(bb)
-	if _, err := gz.Write(exp); err != nil {
-		t.Fatal(err)
-	}
-	if err := gz.Close(); err != nil {
-		t.Fatal(err)
-	}
-	compressed := bb.Bytes()
-
-	ctx := decoding.AddServerDecodingGzipMaxLen(decoding.AddServerDecodingMaxLen(t.Context(), 200), 200)
-	req := httptest.NewRequestWithContext(ctx, "POST", "/v1/data/test", bytes.NewReader(compressed))
-
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Content-Encoding", "gzip")
-
-	got, err := util.ReadMaybeCompressedBody(req)
-	if err != nil {
-		t.Fatal(err)
+	tests := []struct {
+		name          string
+		payload       []byte
+		forgedSize    uint32 // If > 0, overwrite trailer with this size
+		limit         int64
+		expectedError string
+	}{
+		{
+			name:          "valid payload",
+			payload:       []byte(`{"input": {"foo": "bar"}}`),
+			limit:         200,
+			expectedError: "",
+		},
+		{
+			name:          "forged small trailer, actual content larger",
+			payload:       bytes.Repeat([]byte("a"), 100),
+			forgedSize:    50,
+			limit:         200,
+			expectedError: "gzip: invalid checksum",
+		},
+		{
+			name:          "content exceeds limit",
+			payload:       bytes.Repeat([]byte("a"), 300),
+			limit:         200,
+			expectedError: "gzip payload too large",
+		},
 	}
 
-	if !bytes.Equal(got, exp) {
-		t.Fatalf("Expected %q, got: %q", exp, got)
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			bb := new(bytes.Buffer)
+			gz := gzip.NewWriter(bb)
+			if _, err := gz.Write(tc.payload); err != nil {
+				t.Fatal(err)
+			}
+			if err := gz.Close(); err != nil {
+				t.Fatal(err)
+			}
+			compressed := bb.Bytes()
+
+			if tc.forgedSize > 0 {
+				binary.LittleEndian.PutUint32(compressed[len(compressed)-4:], tc.forgedSize)
+			}
+
+			ctx := decoding.AddServerDecodingGzipMaxLen(decoding.AddServerDecodingMaxLen(t.Context(), tc.limit), tc.limit)
+			req := httptest.NewRequestWithContext(ctx, "POST", "/v1/data/test", bytes.NewReader(compressed))
+			req.Header.Set("Content-Type", "application/json")
+			req.Header.Set("Content-Encoding", "gzip")
+
+			got, err := util.ReadMaybeCompressedBody(req)
+			if tc.expectedError != "" {
+				if err == nil {
+					t.Fatal("Expected error, got nil")
+				}
+				if err.Error() != tc.expectedError {
+					t.Fatalf("Expected error %q, got: %v", tc.expectedError, err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("Expected no error, got: %v", err)
+				}
+				if !bytes.Equal(got, tc.payload) {
+					t.Fatalf("Expected %q, got: %q", tc.payload, got)
+				}
+			}
+		})
 	}
 }