Configure Content-Length check to avoid buffering completely

Author: takanuva15

README.md

@@ -85,6 +85,12 @@ curl -i --compressed 'http://localhost:8080/ping?size=2047'
 curl -i --compressed 'http://localhost:8080/ping?size=2048'
 ```
 
+Notes:
+- If a "Content-Length" header is set, that will be used to determine whether to compress based on the given min length.
+- If no "Content-Length" header is set, a buffer is used to temporarily store writes until the min length is met or the request completes.
+  - Setting a high min length will result in more buffering (2048 bytes is a recommended default for most cases)
+  - The handler performs optimizations to avoid unnecessary operations, such as testing if `len(data)` exceeds min length before writing to the buffer, and reusing buffers between requests.
+
 ### Customized Excluded Extensions
 
 ```go

gzip.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"net"
 	"net/http"
+	"strconv"
 
 	"github.com/gin-gonic/gin"
 )
@@ -35,12 +36,23 @@ type gzipWriter struct {
 }
 
 func (g *gzipWriter) WriteString(s string) (int, error) {
-	g.Header().Del("Content-Length")
 	return g.Write([]byte(s))
 }
 
+// Write writes the given data to the appropriate underlying writer.
+// Note that this method can be called multiple times within a single request.
 func (g *gzipWriter) Write(data []byte) (int, error) {
-	g.Header().Del("Content-Length")
+	// If a Content-Length header is set, use that to decide whether to compress the response.
+	if g.Header().Get("Content-Length") != "" {
+		contentLen, _ := strconv.Atoi(g.Header().Get("Content-Length"))
+		if contentLen >= g.minLength {
+			g.wasMinLengthMetForCompression = true
+		} else {
+			return g.ResponseWriter.Write(data)
+		}
+		g.Header().Del("Content-Length")
+	}
+
 	// Check if the response body is large enough to be compressed. If so, skip this condition and proceed with the
 	// normal write process. If not, store the data in the buffer in case more data is written later.
 	// (At the end, if the response body is still too small, the caller should check wasMinLengthMetForCompression and
@@ -64,12 +76,6 @@ func (g *gzipWriter) Flush() {
 	g.ResponseWriter.Flush()
 }
 
-// Fix: https://github.com/mholt/caddy/issues/38
-func (g *gzipWriter) WriteHeader(code int) {
-	g.Header().Del("Content-Length")
-	g.ResponseWriter.WriteHeader(code)
-}
-
 // Ensure gzipWriter implements the http.Hijacker interface.
 // This will cause a compile-time error if gzipWriter does not implement all methods of the http.Hijacker interface.
 var _ http.Hijacker = (*gzipWriter)(nil)

gzip_test.go

@@ -385,7 +385,6 @@ func TestMinLengthShortResponse(t *testing.T) {
 	router := gin.New()
 	router.Use(Gzip(DefaultCompression, WithMinLength(2048)))
 	router.GET("/", func(c *gin.Context) {
-		c.Header("Content-Length", strconv.Itoa(len(testResponse)))
 		c.String(200, testResponse)
 	})
 
@@ -405,7 +404,6 @@ func TestMinLengthLongResponse(t *testing.T) {
 	router := gin.New()
 	router.Use(Gzip(DefaultCompression, WithMinLength(2048)))
 	router.GET("/", func(c *gin.Context) {
-		c.Header("Content-Length", "2048")
 		c.String(200, strings.Repeat("a", 2048))
 	})
 
@@ -439,6 +437,28 @@ func TestMinLengthMultiWriteResponse(t *testing.T) {
 	assert.Less(t, w.Body.Len(), 2048)
 }
 
+// Note this test intentionally triggers gzipping even when the actual response doesn't meet min length. This is because
+// we use the Content-Length header as the primary determinant of compression to avoid the cost of buffering.
+func TestMinLengthUsesContentLengthHeaderInsteadOfBuffering(t *testing.T) {
+	req, _ := http.NewRequestWithContext(context.Background(), "GET", "/", nil)
+	req.Header.Add(headerAcceptEncoding, "gzip")
+
+	router := gin.New()
+	router.Use(Gzip(DefaultCompression, WithMinLength(2048)))
+	router.GET("/", func(c *gin.Context) {
+		c.Header("Content-Length", "2048")
+		c.String(200, testResponse)
+	})
+
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	assert.Equal(t, 200, w.Code)
+	assert.Equal(t, "gzip", w.Header().Get(headerContentEncoding))
+	assert.NotEmpty(t, w.Header().Get("Content-Length"))
+	assert.NotEqual(t, "19", w.Header().Get("Content-Length"))
+}
+
 type hijackableResponse struct {
 	Hijacked bool
 	header   http.Header