fix: improve gzip request decompression middleware (#2077)

Author: endigma

router/core/graph_server.go

@@ -237,12 +237,13 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC
 		)
 	})))
 
-	// Request traffic shaping related middlewares
-	httpRouter.Use(rmiddleware.RequestSize(int64(s.routerTrafficConfig.MaxRequestBodyBytes)))
 	if s.routerTrafficConfig.DecompressionEnabled {
 		httpRouter.Use(rmiddleware.HandleCompression(s.logger))
 	}
 
+	// Request traffic shaping related middlewares, happens after decompression to prevent unbounded decompression attacks
+	httpRouter.Use(rmiddleware.RequestSize(int64(s.routerTrafficConfig.MaxRequestBodyBytes)))
+
 	httpRouter.Use(middleware.RequestID)
 	httpRouter.Use(middleware.RealIP)
 	if s.corsOptions.Enabled {

router/internal/middleware/compression.go

@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"compress/gzip"
+
 	"net/http"
 	"strings"
 
@@ -32,14 +33,23 @@ func HandleCompression(logger *zap.Logger) func(http.Handler) http.Handler {
 					return
 				}
 
+				originalBody := r.Body
+
 				defer func() {
 					if err := gzr.Close(); err != nil {
 						logger.Error("failed to close gzip reader", zap.Error(err))
 					}
+
+					if err := originalBody.Close(); err != nil {
+						logger.Error("failed to close original body", zap.Error(err))
+					}
 				}()
 
 				r.Body = gzr
 
+				// Content-Length is no longer valid after decompression
+				r.Header.Del("Content-Length")
+				r.ContentLength = -1
 			case "":
 			default:
 				http.Error(w, "unsupported media type", http.StatusUnsupportedMediaType)

router/internal/middleware/compression_test.go

@@ -1,13 +1,15 @@
 package middleware
 
 import (
-	"compress/gzip"
+	"bytes"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"compress/gzip"
+
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 )
@@ -45,8 +47,7 @@ func TestHandleCompression(t *testing.T) {
 					w.WriteHeader(http.StatusOK)
 				})
 
-				req, err := http.NewRequest(tc.method, "/", strings.NewReader("test"))
-				require.NoError(t, err)
+				req := httptest.NewRequest(tc.method, "/", strings.NewReader("test"))
 
 				HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
 
@@ -64,8 +65,7 @@ func TestHandleCompression(t *testing.T) {
 			w.WriteHeader(http.StatusOK)
 		})
 
-		req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
-		require.NoError(t, err)
+		req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
 		req.Header.Set("Content-Encoding", "gzip, deflate")
 
 		HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
@@ -92,16 +92,14 @@ func TestHandleCompression(t *testing.T) {
 
 		// create gzip compressed request
 
-		var sb strings.Builder
-		w := gzip.NewWriter(&sb)
+		var buf bytes.Buffer
+		w := gzip.NewWriter(&buf)
 
 		_, err := w.Write([]byte("test"))
 		require.NoError(t, err)
-
 		require.NoError(t, w.Close())
 
-		req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader(sb.String()))
-		require.NoError(t, err)
+		req := httptest.NewRequest(http.MethodPost, "/", &buf)
 		req.Header.Set("Content-Encoding", "gzip")
 
 		HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
@@ -118,8 +116,7 @@ func TestHandleCompression(t *testing.T) {
 			w.WriteHeader(http.StatusOK)
 		})
 
-		req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
-		require.NoError(t, err)
+		req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
 		req.Header.Set("Content-Encoding", "gzip")
 
 		HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
@@ -136,8 +133,7 @@ func TestHandleCompression(t *testing.T) {
 			w.WriteHeader(http.StatusOK)
 		})
 
-		req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
-		require.NoError(t, err)
+		req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
 		req.Header.Set("Content-Encoding", "deflate")
 
 		HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
@@ -155,8 +151,7 @@ func TestHandleCompression(t *testing.T) {
 			w.WriteHeader(http.StatusOK)
 		})
 
-		req, err := http.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
-		require.NoError(t, err)
+		req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("test"))
 
 		HandleCompression(zap.NewNop())(next).ServeHTTP(recorder, req)
 

router/internal/middleware/request_size_test.go

@@ -0,0 +1,77 @@
+package middleware
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+func TestRequestSizeAndCompression(t *testing.T) {
+	// Define maximum allowed request size
+	const maxRequestSize = 1024 // Example: 1 KB
+
+	// Chain the compression and request_size middlewares
+	r := chi.NewMux()
+	r.Use(HandleCompression(zap.NewNop()))
+	r.Use(RequestSize(maxRequestSize))
+
+	t.Run("request size limiter should not allow requests exceeding allowed maximum request size", func(t *testing.T) {
+		// Recorder to capture the response
+		recorder := httptest.NewRecorder()
+
+		// Write a large payload that exceeds the maxRequestSize after decompression
+		largePayload := strings.Repeat("A", maxRequestSize*10) // 10x the max size
+
+		// Create the request with the gzip bomb payload
+		req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(largePayload))
+
+		r.HandleFunc("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			_, err := io.Copy(io.Discard, r.Body)
+			require.ErrorContains(t, err, "http: request body too large")
+			w.WriteHeader(http.StatusRequestEntityTooLarge)
+		}))
+
+		r.ServeHTTP(recorder, req)
+
+		// Assert that the response status is 413 Payload Too Large
+		require.Equal(t, http.StatusRequestEntityTooLarge, recorder.Code)
+	})
+
+	t.Run("request size limiter should reject gzip bomb exceeding allowed maximum request size", func(t *testing.T) {
+		// Recorder to capture the response
+		recorder := httptest.NewRecorder()
+
+		// Create a gzip bomb payload
+		var buf bytes.Buffer
+		w := gzip.NewWriter(&buf)
+
+		// Write a large payload that exceeds the maxRequestSize after decompression
+		largePayload := strings.Repeat("A", maxRequestSize*10) // 10x the max size
+		_, err := w.Write([]byte(largePayload))
+		require.NoError(t, err)
+		require.NoError(t, w.Close())
+
+		// Create the request with the gzip bomb payload
+		req := httptest.NewRequest(http.MethodPost, "/", &buf)
+		req.Header.Set("Content-Encoding", "gzip")
+
+		r.HandleFunc("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			_, err := io.Copy(io.Discard, r.Body)
+			require.ErrorContains(t, err, "http: request body too large")
+			w.WriteHeader(http.StatusRequestEntityTooLarge)
+		}))
+
+		r.ServeHTTP(recorder, req)
+
+		// Assert that the response status is 413 Payload Too Large
+		require.Equal(t, http.StatusRequestEntityTooLarge, recorder.Code)
+	})
+}