Move compression logic into its own middleware

Author: kevinmcconnell

internal/compression_handler.go

@@ -0,0 +1,37 @@
+package internal
+
+import (
+	"net/http"
+
+	"github.com/klauspost/compress/gzhttp"
+)
+
+func NewCompressionHandler(jitter int, disableOnAuth bool, next http.Handler) http.Handler {
+	var wrapper func(http.Handler) http.HandlerFunc
+	var err error
+
+	if jitter > 0 {
+		wrapper, err = gzhttp.NewWrapper(
+			gzhttp.MinSize(1024),
+			gzhttp.CompressionLevel(6),
+			gzhttp.RandomJitter(jitter, 0, false),
+		)
+	} else {
+		wrapper, err = gzhttp.NewWrapper(
+			gzhttp.MinSize(1024),
+			gzhttp.CompressionLevel(6),
+		)
+	}
+
+	if err != nil {
+		panic("failed to create gzip wrapper: " + err.Error())
+	}
+
+	handler := wrapper(next)
+
+	if disableOnAuth {
+		return NewCompressionGuardHandler(handler)
+	}
+
+	return handler
+}

internal/compression_handler_test.go

@@ -0,0 +1,88 @@
+package internal
+
+import (
+	"compress/gzip"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCompressionHandler(t *testing.T) {
+	largeBody := strings.Repeat("A", 2000)
+
+	upstream := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/plain")
+		_, err := w.Write([]byte(largeBody))
+		require.NoError(t, err)
+	})
+
+	t.Run("compresses responses", func(t *testing.T) {
+		handler := NewCompressionHandler(0, false, upstream)
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Accept-Encoding", "gzip")
+		rr := httptest.NewRecorder()
+
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, "gzip", rr.Header().Get("Content-Encoding"))
+
+		reader, err := gzip.NewReader(rr.Body)
+		require.NoError(t, err)
+		defer reader.Close()
+		body, err := io.ReadAll(reader)
+		require.NoError(t, err)
+		assert.Equal(t, largeBody, string(body))
+	})
+
+	t.Run("applies jitter when configured", func(t *testing.T) {
+		handler := NewCompressionHandler(32, false, upstream)
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Accept-Encoding", "gzip")
+		rr := httptest.NewRecorder()
+
+		handler.ServeHTTP(rr, req)
+
+		require.Equal(t, "gzip", rr.Header().Get("Content-Encoding"))
+
+		// Check for GZIP header with FCOMMENT flag (0x10)
+		bodyBytes := rr.Body.Bytes()
+		require.Greater(t, len(bodyBytes), 10)
+		hasComment := (bodyBytes[3] & 0x10) != 0
+		assert.True(t, hasComment, "Expected FCOMMENT flag due to jitter")
+	})
+
+	t.Run("wraps with guard when disableOnAuth is true", func(t *testing.T) {
+		handler := NewCompressionHandler(0, true, upstream)
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Accept-Encoding", "gzip")
+		req.Header.Set("Cookie", "session=secret")
+		rr := httptest.NewRecorder()
+
+		handler.ServeHTTP(rr, req)
+
+		// Should NOT be compressed due to Cookie header
+		assert.Empty(t, rr.Header().Get("Content-Encoding"))
+		assert.Equal(t, largeBody, rr.Body.String())
+	})
+
+	t.Run("compresses authenticated requests when disableOnAuth is false", func(t *testing.T) {
+		handler := NewCompressionHandler(0, false, upstream)
+
+		req := httptest.NewRequest("GET", "/", nil)
+		req.Header.Set("Accept-Encoding", "gzip")
+		req.Header.Set("Cookie", "session=secret")
+		rr := httptest.NewRecorder()
+
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, "gzip", rr.Header().Get("Content-Encoding"))
+	})
+}

internal/handler.go

@@ -4,8 +4,6 @@ import (
 	"log/slog"
 	"net/http"
 	"net/url"
-
-	"github.com/klauspost/compress/gzhttp"
 )
 
 type HandlerOptions struct {
@@ -29,35 +27,7 @@ func NewHandler(options HandlerOptions) http.Handler {
 	handler = NewRequestStartHandler(handler)
 
 	if options.gzipCompressionEnabled {
-		var wrapper func(http.Handler) http.HandlerFunc
-		var err error
-
-		if options.gzipCompressionJitter > 0 {
-			wrapper, err = gzhttp.NewWrapper(
-				gzhttp.MinSize(1024),
-				gzhttp.CompressionLevel(6),
-				gzhttp.RandomJitter(options.gzipCompressionJitter, 0, false),
-			)
-		} else {
-			wrapper, err = gzhttp.NewWrapper(
-				gzhttp.MinSize(1024),
-				gzhttp.CompressionLevel(6),
-			)
-		}
-
-		if err != nil {
-			// If we cannot create the wrapper with the requested configuration (including jitter),
-			// we must fail hard rather than silently downgrading security or performance.
-			panic("failed to create gzip wrapper: " + err.Error())
-		}
-
-		gzipHandler := wrapper(handler)
-
-		if options.gzipCompressionDisableOnAuth {
-			handler = NewCompressionGuardHandler(gzipHandler)
-		} else {
-			handler = gzipHandler
-		}
+		handler = NewCompressionHandler(options.gzipCompressionJitter, options.gzipCompressionDisableOnAuth, handler)
 	}
 
 	if options.maxRequestBody > 0 {