feat: introduce custom compression decision function in handler (#96)

* feat: introduce custom compression decision function in handler

- Add detailed comments explaining the `Handle` function in `handler.go`
- Modify the condition in `Handle` to include a custom compression decision function
- Consolidate exclusion checks in `shouldCompress` into a single condition
- Add a new field `customShouldCompressFn` to the `config` struct in `options.go`
- Introduce `WithCustomShouldCompressFn` option for setting a custom compression decision function

Signed-off-by: appleboy <[email protected]>

* test: improve test coverage and optimize API performance

- Add a test for custom compression function in gzip middleware

Signed-off-by: appleboy <[email protected]>

---------

Signed-off-by: appleboy <[email protected]>

Author: appleboy

gzip_test.go

@@ -319,3 +319,28 @@ func TestGzipWithDecompressOnly(t *testing.T) {
 	assert.Equal(t, w.Header().Get("Content-Encoding"), "")
 	assert.Equal(t, w.Body.String(), testResponse)
 }
+
+func TestCustomShouldCompressFn(t *testing.T) {
+	req, _ := http.NewRequestWithContext(context.Background(), "GET", "/", nil)
+	req.Header.Add("Accept-Encoding", "gzip")
+
+	router := gin.New()
+	router.Use(Gzip(
+		DefaultCompression,
+		WithCustomShouldCompressFn(func(_ *gin.Context) bool {
+			return false
+		}),
+	))
+	router.GET("/", func(c *gin.Context) {
+		c.Header("Content-Length", strconv.Itoa(len(testResponse)))
+		c.String(200, testResponse)
+	})
+
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, req)
+
+	assert.Equal(t, 200, w.Code)
+	assert.Equal(t, "", w.Header().Get("Content-Encoding"))
+	assert.Equal(t, "19", w.Header().Get("Content-Length"))
+	assert.Equal(t, testResponse, w.Body.String())
+}

handler.go

@@ -42,12 +42,21 @@ func newGzipHandler(level int, opts ...Option) *gzipHandler {
 	return handler
 }
 
+// Handle is a middleware function for handling gzip compression in HTTP requests and responses.
+// It first checks if the request has a "Content-Encoding" header set to "gzip" and if a decompression
+// function is provided, it will call the decompression function. If the handler is set to decompress only,
+// or if the custom compression decision function indicates not to compress, it will return early.
+// Otherwise, it retrieves a gzip.Writer from the pool, sets the necessary response headers for gzip encoding,
+// and wraps the response writer with a gzipWriter. After the request is processed, it ensures the gzip.Writer
+// is properly closed and the "Content-Length" header is set based on the response size.
 func (g *gzipHandler) Handle(c *gin.Context) {
 	if fn := g.decompressFn; fn != nil && c.Request.Header.Get("Content-Encoding") == "gzip" {
 		fn(c)
 	}
 
-	if g.decompressOnly || !g.shouldCompress(c.Request) {
+	if g.decompressOnly ||
+		(g.customShouldCompressFn != nil && !g.customShouldCompressFn(c)) ||
+		(g.customShouldCompressFn == nil && !g.shouldCompress(c.Request)) {
 		return
 	}
 
@@ -76,15 +85,11 @@ func (g *gzipHandler) shouldCompress(req *http.Request) bool {
 		return false
 	}
 
+	// Check if the request path is excluded from compression
 	extension := filepath.Ext(req.URL.Path)
-	if g.excludedExtensions.Contains(extension) {
-		return false
-	}
-
-	if g.excludedPaths.Contains(req.URL.Path) {
-		return false
-	}
-	if g.excludedPathesRegexs.Contains(req.URL.Path) {
+	if g.excludedExtensions.Contains(extension) ||
+		g.excludedPaths.Contains(req.URL.Path) ||
+		g.excludedPathesRegexs.Contains(req.URL.Path) {
 		return false
 	}
 

options.go

@@ -34,11 +34,12 @@ func (o optionFunc) apply(c *config) {
 }
 
 type config struct {
-	excludedExtensions   ExcludedExtensions
-	excludedPaths        ExcludedPaths
-	excludedPathesRegexs ExcludedPathesRegexs
-	decompressFn         func(c *gin.Context)
-	decompressOnly       bool
+	excludedExtensions     ExcludedExtensions
+	excludedPaths          ExcludedPaths
+	excludedPathesRegexs   ExcludedPathesRegexs
+	decompressFn           func(c *gin.Context)
+	decompressOnly         bool
+	customShouldCompressFn func(c *gin.Context) bool
 }
 
 // WithExcludedExtensions returns an Option that sets the ExcludedExtensions field of the Options struct.
@@ -87,6 +88,28 @@ func WithDecompressOnly() Option {
 	})
 }
 
+// WithCustomShouldCompressFn returns an Option that sets the CustomShouldCompressFn field of the Options struct.
+// Parameters:
+//   - fn: func(c *gin.Context) bool - A function to determine if a request should be compressed.
+//     The function should return true if the request should be compressed, false otherwise.
+//     If the function returns false, the middleware will not compress the response.
+//     If the function is nil, the middleware will use the default logic to determine
+//     if the response should be compressed.
+//
+// Returns:
+//   - Option - An option that sets the CustomShouldCompressFn field of the Options struct.
+//
+// Example:
+//
+//	router.Use(gzip.Gzip(gzip.DefaultCompression, gzip.WithCustomShouldCompressFn(func(c *gin.Context) bool {
+//		return c.Request.URL.Path != "/no-compress"
+//	})))
+func WithCustomShouldCompressFn(fn func(c *gin.Context) bool) Option {
+	return optionFunc(func(o *config) {
+		o.customShouldCompressFn = fn
+	})
+}
+
 // Using map for better lookup performance
 type ExcludedExtensions map[string]struct{}