Fix bug that Content-Encoding won't be written to final response header for gzip and deflate

bodyWriter.go

@@ -20,16 +20,18 @@ import (
 type respWriter struct {
 	gin.ResponseWriter
 	threshold    int
+	encoding     string
 	algo         algorithm
 	buf          *bytes.Buffer
 	bytesWritten int
 	compressor   io.WriteCloser
 }
 
-func newResponseWriter(c *gin.Context, swapSize int, algo algorithm) *respWriter {
+func newResponseWriter(c *gin.Context, swapSize int, encoding string, algo algorithm) *respWriter {
 	return &respWriter{
 		c.Writer,
 		swapSize,
+		encoding,
 		algo,
 		bytes.NewBuffer(nil),
 		0,
@@ -45,6 +47,8 @@ func (rw *respWriter) Write(b []byte) (int, error) {
 	rw.Header().Del("Content-Length")
 
 	if !rw.Swapped() && rw.buf.Len()+len(b) >= rw.threshold {
+		rw.ResponseWriter.Header().Set("Content-Encoding", rw.encoding)
+		rw.ResponseWriter.Header().Set("Vary", "Accept-Encoding")
 		rw.compressor = rw.algo.getWriter(rw.ResponseWriter)
 		if copied, err := io.Copy(rw.compressor, rw.buf); err != nil {
 			return int(copied), err

compress_test.go

@@ -12,7 +12,7 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/aurowora/compress"
+	compress "github.com/lf4096/gin-compress"
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
 )

decompress_test.go

@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"github.com/andybalholm/brotli"
-	"github.com/aurowora/compress"
+	compress "github.com/lf4096/gin-compress"
 	"github.com/klauspost/compress/gzip"
 	"github.com/klauspost/compress/zlib"
 	"github.com/klauspost/compress/zstd"

go.mod

@@ -1,4 +1,4 @@
-module github.com/aurowora/compress
+module github.com/lf4096/gin-compress
 
 go 1.16
 

middleware.go

@@ -44,15 +44,10 @@ func (cm *compressMiddleware) Handler(c *gin.Context) {
 		return
 	}
 
-	rw := newResponseWriter(c, cm.cfg.minCompressBytes, algorithms[algo])
+	rw := newResponseWriter(c, cm.cfg.minCompressBytes, algo, algorithms[algo])
 	c.Writer = rw
 	c.Next()
 
-	if rw.Swapped() {
-		rw.ResponseWriter.Header().Set("Content-Encoding", algo)
-		rw.ResponseWriter.Header().Set("Vary", "Accept-Encoding")
-	}
-
 	_ = rw.Close()
 }