Introduce option for more granular decision making process to compres…

…s a response
remerge · May 6, 2022 · 502e570 · 502e570
1 parent 4407045
commit 502e570
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 10 deletions.
diff --git a/gzip_test.go b/gzip_test.go
@@ -11,6 +11,7 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/gin-gonic/gin"
@@ -104,6 +105,89 @@ func TestGzipPNG(t *testing.T) {
 	assert.Equal(t, w.Body.String(), "this is a PNG!")
 }
 
+func TestMatchSupportedRequests(t *testing.T) {
+	router := gin.New()
+	router.Use(
+		Gzip(DefaultCompression,
+			WithMatchSupportedRequestFn(func(req *http.Request) (bool, bool) {
+				xheader := req.Header.Get("X-Test-Header")
+				if xheader == "" {
+					return false, false
+				}
+
+				ok, supported := strings.HasPrefix(xheader, "+"), strings.HasSuffix(xheader, "compress me")
+				return ok, supported
+			}),
+			// For testing the precedence order
+			WithExcludedExtensions([]string{".php"}),
+			WithExcludedPaths([]string{"/api/"}),
+		))
+	router.GET("/index.html", func(c *gin.Context) {
+		c.String(200, "this is a HTML!")
+	})
+
+	t.Run("Is Compressed/matched header's value", func(t *testing.T) {
+		req, _ := http.NewRequestWithContext(context.Background(), "GET", "/index.html", nil)
+		req.Header.Add("Accept-Encoding", "gzip")
+		req.Header.Add("X-Test-Header", "+compress me")
+
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+
+		assert.Equal(t, w.Code, 200)
+		assert.Equal(t, w.Header().Get("Content-Encoding"), "gzip")
+		assert.Equal(t, w.Header().Get("Vary"), "Accept-Encoding")
+		assert.NotEqual(t, w.Header().Get("Content-Length"), "0")
+		assert.NotEqual(t, w.Body.Len(), 19)
+		assert.Equal(t, fmt.Sprint(w.Body.Len()), w.Header().Get("Content-Length"))
+	})
+
+	t.Run("Is Compressed/no header", func(t *testing.T) {
+		req, _ := http.NewRequestWithContext(context.Background(), "GET", "/index.html", nil)
+		req.Header.Add("Accept-Encoding", "gzip")
+
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+
+		assert.Equal(t, w.Code, 200)
+		assert.Equal(t, w.Header().Get("Content-Encoding"), "gzip")
+		assert.Equal(t, w.Header().Get("Vary"), "Accept-Encoding")
+		assert.NotEqual(t, w.Header().Get("Content-Length"), "0")
+		assert.NotEqual(t, w.Body.Len(), 19)
+		assert.Equal(t, fmt.Sprint(w.Body.Len()), w.Header().Get("Content-Length"))
+	})
+
+	t.Run("Is Not Compressed/no match", func(t *testing.T) {
+		req, _ := http.NewRequestWithContext(context.Background(), "GET", "/index.html", nil)
+		req.Header.Add("Accept-Encoding", "gzip")
+		req.Header.Add("X-Test-Header", "+skip me")
+
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Equal(t, "", w.Header().Get("Content-Encoding"))
+		assert.Equal(t, "", w.Header().Get("Vary"))
+		assert.Equal(t, "this is a HTML!", w.Body.String())
+		assert.Equal(t, "", w.Header().Get("Content-Length"))
+	})
+
+	t.Run("Is Not Compressed/Precedence over Exclusion rules", func(t *testing.T) {
+		req, _ := http.NewRequestWithContext(context.Background(), "GET", "/index.html", nil)
+		req.Header.Add("Accept-Encoding", "gzip")
+		req.Header.Add("X-Test-Header", "+compressme")
+
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Equal(t, "", w.Header().Get("Content-Encoding"))
+		assert.Equal(t, "", w.Header().Get("Vary"))
+		assert.Equal(t, "this is a HTML!", w.Body.String())
+		assert.Equal(t, "", w.Header().Get("Content-Length"))
+	})
+}
+
 func TestExcludedExtensions(t *testing.T) {
 	req, _ := http.NewRequestWithContext(context.Background(), "GET", "/index.html", nil)
 	req.Header.Add("Accept-Encoding", "gzip")

diff --git a/handler.go b/handler.go
@@ -67,17 +67,16 @@ func (g *gzipHandler) shouldCompress(req *http.Request) bool {
 		return false
 	}
 
-	extension := filepath.Ext(req.URL.Path)
-	if g.ExcludedExtensions.Contains(extension) {
-		return false
+	if g.MatchSupportedRequestFn != nil {
+		if ok, supported := g.MatchSupportedRequestFn(req); ok {
+			return supported
+		}
 	}
 
-	if g.ExcludedPaths.Contains(req.URL.Path) {
-		return false
-	}
-	if g.ExcludedPathesRegexs.Contains(req.URL.Path) {
-		return false
-	}
+	return !g.isPathExcluded(req.URL.Path)
+}
 
-	return true
+func (g *gzipHandler) isPathExcluded(path string) bool {
+	extension := filepath.Ext(path)
+	return g.ExcludedExtensions.Contains(extension) || g.ExcludedPaths.Contains(path) || g.ExcludedPathesRegexs.Contains(path)
 }
diff --git a/options.go b/options.go
@@ -23,6 +23,9 @@ type Options struct {
 	ExcludedPaths        ExcludedPaths
 	ExcludedPathesRegexs ExcludedPathesRegexs
 	DecompressFn         func(c *gin.Context)
+	// MatchSupportedRequestFn has a precedence over Excluded* options.
+	// `ok' value is to stop other checks and treat `supported` value as if the request should be compressed or not.
+	MatchSupportedRequestFn func(req *http.Request) (ok bool, supported bool)
 }
 
 type Option func(*Options)
@@ -51,6 +54,12 @@ func WithDecompressFn(decompressFn func(c *gin.Context)) Option {
 	}
 }
 
+func WithMatchSupportedRequestFn(matchSupportedRequestFn func(req *http.Request) (bool, bool)) Option {
+	return func(o *Options) {
+		o.MatchSupportedRequestFn = matchSupportedRequestFn
+	}
+}
+
 // Using map for better lookup performance
 type ExcludedExtensions map[string]bool