Browse Source

feat: support customized excluded extensions

Pei-Ming Wu 6 years ago
parent
commit
a5a15a8290
5 changed files with 93 additions and 14 deletions
  1. 25 0
      README.md
  2. 2 2
      gzip.go
  3. 20 0
      gzip_test.go
  4. 8 12
      handler.go
  5. 38 0
      options.go

+ 25 - 0
README.md

@@ -49,3 +49,28 @@ 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")
+}
+```

+ 2 - 2
gzip.go

@@ -13,8 +13,8 @@ const (
 	NoCompression      = gzip.NoCompression
 )
 
-func Gzip(level int) gin.HandlerFunc {
-	return newGzipHandler(level).Handle
+func Gzip(level int, options ...Option) gin.HandlerFunc {
+	return newGzipHandler(level, options...).Handle
 }
 
 type gzipWriter struct {

+ 20 - 0
gzip_test.go

@@ -110,6 +110,26 @@ 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 TestNoGzip(t *testing.T) {
 	req, _ := http.NewRequest("GET", "/", nil)
 

+ 8 - 12
handler.go

@@ -13,10 +13,11 @@ import (
 )
 
 type gzipHandler struct {
+	*Options
 	gzPool sync.Pool
 }
 
-func newGzipHandler(level int) *gzipHandler {
+func newGzipHandler(level int, options ...Option) *gzipHandler {
 	var gzPool sync.Pool
 	gzPool.New = func() interface{} {
 		gz, err := gzip.NewWriterLevel(ioutil.Discard, level)
@@ -26,7 +27,11 @@ func newGzipHandler(level int) *gzipHandler {
 		return gz
 	}
 	handler := &gzipHandler{
-		gzPool: gzPool,
+		Options: DefaultOptions,
+		gzPool:  gzPool,
+	}
+	for _, setter := range options {
+		setter(handler.Options)
 	}
 	return handler
 }
@@ -60,14 +65,5 @@ func (g *gzipHandler) shouldCompress(req *http.Request) bool {
 	}
 
 	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
-	}
+	return !g.ExcludedExtensions.Contains(extension)
 }

+ 38 - 0
options.go

@@ -0,0 +1,38 @@
+package gzip
+
+var (
+	DefaultExcludedExtentions = NewExcludedExtensions([]string{
+		".png", ".gif", ".jpeg", ".jpg",
+	})
+	DefaultOptions = &Options{
+		ExcludedExtensions: DefaultExcludedExtentions,
+	}
+)
+
+type Options struct {
+	ExcludedExtensions ExcludedExtensions
+}
+
+type Option func(*Options)
+
+func WithExcludedExtensions(args []string) Option {
+	return func(o *Options) {
+		o.ExcludedExtensions = NewExcludedExtensions(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
+}