Browse Source

refactor: seperate handler and writer

Pei-Ming Wu 6 years ago
parent
commit
3fb40b6bb9
2 changed files with 74 additions and 54 deletions
  1. 1 54
      gzip.go
  2. 73 0
      handler.go

+ 1 - 54
gzip.go

@@ -2,12 +2,6 @@ package gzip
 
 
 import (
 import (
 	"compress/gzip"
 	"compress/gzip"
-	"fmt"
-	"io/ioutil"
-	"net/http"
-	"path/filepath"
-	"strings"
-	"sync"
 
 
 	"github.com/gin-gonic/gin"
 	"github.com/gin-gonic/gin"
 )
 )
@@ -20,33 +14,7 @@ const (
 )
 )
 
 
 func Gzip(level int) gin.HandlerFunc {
 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()
-	}
+	return newGzipHandler(level).Handle
 }
 }
 
 
 type gzipWriter struct {
 type gzipWriter struct {
@@ -67,24 +35,3 @@ func (g *gzipWriter) WriteHeader(code int) {
 	g.Header().Del("Content-Length")
 	g.Header().Del("Content-Length")
 	g.ResponseWriter.WriteHeader(code)
 	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
-	}
-}

+ 73 - 0
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 {
+	gzPool sync.Pool
+}
+
+func newGzipHandler(level int) *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{
+		gzPool: gzPool,
+	}
+	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 len(extension) < 4 { // fast path
+		return true
+	}
+
+	switch extension {
+	case ".png", ".gif", ".jpeg", ".jpg":
+		return false
+	default:
+		return true
+	}
+}