Merge pull request #22 from mkfsn/master

feat: customized excluded extensions and paths

README.md

@@ -49,3 +49,53 @@ func main() {
 }
 
 ```
+
+### Customized Excluded Extensions
+
+```go
+package main
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/gin-contrib/gzip"
+	"github.com/gin-gonic/gin"
+)
+
+func main() {
+	r := gin.Default()
+	r.Use(gzip.Gzip(gzip.DefaultCompression, gzip.WithExcludedExtensions([]string{".pdf", ".mp4"})))
+	r.GET("/ping", func(c *gin.Context) {
+		c.String(200, "pong "+fmt.Sprint(time.Now().Unix()))
+	})
+
+	// Listen and Server in 0.0.0.0:8080
+	r.Run(":8080")
+}
+```
+
+### Customized Excluded Paths
+
+```go
+package main
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/gin-contrib/gzip"
+	"github.com/gin-gonic/gin"
+)
+
+func main() {
+	r := gin.Default()
+	r.Use(gzip.Gzip(gzip.DefaultCompression, gzip.WithExcludedPaths([]string{"/api/"})))
+	r.GET("/ping", func(c *gin.Context) {
+		c.String(200, "pong "+fmt.Sprint(time.Now().Unix()))
+	})
+
+	// Listen and Server in 0.0.0.0:8080
+	r.Run(":8080")
+}
+```

gzip.go

@@ -2,12 +2,6 @@ package gzip
 
 import (
 	"compress/gzip"
-	"fmt"
-	"io/ioutil"
-	"net/http"
-	"path/filepath"
-	"strings"
-	"sync"
 
 	"github.com/gin-gonic/gin"
 )
@@ -19,34 +13,8 @@ const (
 	NoCompression      = gzip.NoCompression
 )
 
-func Gzip(level int) gin.HandlerFunc {
-	var gzPool sync.Pool
-	gzPool.New = func() interface{} {
-		gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
-		if err != nil {
-			panic(err)
-		}
-		return gz
-	}
-	return func(c *gin.Context) {
-		if !shouldCompress(c.Request) {
-			return
-		}
-
-		gz := gzPool.Get().(*gzip.Writer)
-		defer gzPool.Put(gz)
-		defer gz.Reset(ioutil.Discard)
-		gz.Reset(c.Writer)
-
-		c.Header("Content-Encoding", "gzip")
-		c.Header("Vary", "Accept-Encoding")
-		c.Writer = &gzipWriter{c.Writer, gz}
-		defer func() {
-			gz.Close()
-			c.Header("Content-Length", fmt.Sprint(c.Writer.Size()))
-		}()
-		c.Next()
-	}
+func Gzip(level int, options ...Option) gin.HandlerFunc {
+	return newGzipHandler(level, options...).Handle
 }
 
 type gzipWriter struct {
@@ -67,24 +35,3 @@ func (g *gzipWriter) WriteHeader(code int) {
 	g.Header().Del("Content-Length")
 	g.ResponseWriter.WriteHeader(code)
 }
-
-func shouldCompress(req *http.Request) bool {
-	if !strings.Contains(req.Header.Get("Accept-Encoding"), "gzip") ||
-		strings.Contains(req.Header.Get("Connection"), "Upgrade") ||
-		strings.Contains(req.Header.Get("Content-Type"), "text/event-stream") {
-
-		return false
-	}
-
-	extension := filepath.Ext(req.URL.Path)
-	if len(extension) < 4 { // fast path
-		return true
-	}
-
-	switch extension {
-	case ".png", ".gif", ".jpeg", ".jpg":
-		return false
-	default:
-		return true
-	}
-}

gzip_test.go

@@ -110,6 +110,46 @@ func TestGzipPNG(t *testing.T) {
 	assert.Equal(t, w.Body.String(), "this is a PNG!")
 }
 
+func TestExcludedExtensions(t *testing.T) {
+	req, _ := http.NewRequest("GET", "/index.html", nil)
+	req.Header.Add("Accept-Encoding", "gzip")
+
+	router := gin.New()
+	router.Use(Gzip(DefaultCompression, WithExcludedExtensions([]string{".html"})))
+	router.GET("/index.html", func(c *gin.Context) {
+		c.String(200, "this is a HTML!")
+	})
+
+	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 TestExcludedPaths(t *testing.T) {
+	req, _ := http.NewRequest("GET", "/api/books", nil)
+	req.Header.Add("Accept-Encoding", "gzip")
+
+	router := gin.New()
+	router.Use(Gzip(DefaultCompression, WithExcludedPaths([]string{"/api/"})))
+	router.GET("/api/books", func(c *gin.Context) {
+		c.String(200, "this is books!")
+	})
+
+	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 books!", w.Body.String())
+	assert.Equal(t, "", w.Header().Get("Content-Length"))
+}
+
 func TestNoGzip(t *testing.T) {
 	req, _ := http.NewRequest("GET", "/", nil)
 

handler.go

@@ -0,0 +1,73 @@
+package gzip
+
+import (
+	"compress/gzip"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/gin-gonic/gin"
+)
+
+type gzipHandler struct {
+	*Options
+	gzPool sync.Pool
+}
+
+func newGzipHandler(level int, options ...Option) *gzipHandler {
+	var gzPool sync.Pool
+	gzPool.New = func() interface{} {
+		gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
+		if err != nil {
+			panic(err)
+		}
+		return gz
+	}
+	handler := &gzipHandler{
+		Options: DefaultOptions,
+		gzPool:  gzPool,
+	}
+	for _, setter := range options {
+		setter(handler.Options)
+	}
+	return handler
+}
+
+func (g *gzipHandler) Handle(c *gin.Context) {
+	if !g.shouldCompress(c.Request) {
+		return
+	}
+
+	gz := g.gzPool.Get().(*gzip.Writer)
+	defer g.gzPool.Put(gz)
+	defer gz.Reset(ioutil.Discard)
+	gz.Reset(c.Writer)
+
+	c.Header("Content-Encoding", "gzip")
+	c.Header("Vary", "Accept-Encoding")
+	c.Writer = &gzipWriter{c.Writer, gz}
+	defer func() {
+		gz.Close()
+		c.Header("Content-Length", fmt.Sprint(c.Writer.Size()))
+	}()
+	c.Next()
+}
+
+func (g *gzipHandler) shouldCompress(req *http.Request) bool {
+	if !strings.Contains(req.Header.Get("Accept-Encoding"), "gzip") ||
+		strings.Contains(req.Header.Get("Connection"), "Upgrade") ||
+		strings.Contains(req.Header.Get("Content-Type"), "text/event-stream") {
+
+		return false
+	}
+
+	extension := filepath.Ext(req.URL.Path)
+	if g.ExcludedExtensions.Contains(extension) {
+		return false
+	}
+
+	return !g.ExcludedPaths.Contains(req.URL.Path)
+}

options.go

@@ -0,0 +1,64 @@
+package gzip
+
+import (
+	"strings"
+)
+
+var (
+	DefaultExcludedExtentions = NewExcludedExtensions([]string{
+		".png", ".gif", ".jpeg", ".jpg",
+	})
+	DefaultOptions = &Options{
+		ExcludedExtensions: DefaultExcludedExtentions,
+	}
+)
+
+type Options struct {
+	ExcludedExtensions ExcludedExtensions
+	ExcludedPaths      ExcludedPaths
+}
+
+type Option func(*Options)
+
+func WithExcludedExtensions(args []string) Option {
+	return func(o *Options) {
+		o.ExcludedExtensions = NewExcludedExtensions(args)
+	}
+}
+
+func WithExcludedPaths(args []string) Option {
+	return func(o *Options) {
+		o.ExcludedPaths = NewExcludedPaths(args)
+	}
+}
+
+// Using map for better lookup performance
+type ExcludedExtensions map[string]bool
+
+func NewExcludedExtensions(extensions []string) ExcludedExtensions {
+	res := make(ExcludedExtensions)
+	for _, e := range extensions {
+		res[e] = true
+	}
+	return res
+}
+
+func (e ExcludedExtensions) Contains(target string) bool {
+	_, ok := e[target]
+	return ok
+}
+
+type ExcludedPaths []string
+
+func NewExcludedPaths(paths []string) ExcludedPaths {
+	return ExcludedPaths(paths)
+}
+
+func (e ExcludedPaths) Contains(requestURI string) bool {
+	for _, path := range e {
+		if strings.HasPrefix(requestURI, path) {
+			return true
+		}
+	}
+	return false
+}