Content type sniffing for compress middleware

Signed-off-by: Vishal Rana <[email protected]>

echo.go

@@ -110,6 +110,7 @@ const (
 	ContentType        = "Content-Type"
 	Authorization      = "Authorization"
 	Upgrade            = "Upgrade"
+	Vary               = "Vary"
 
 	//-----------
 	// Protocols

middleware/auth.go

@@ -32,7 +32,6 @@ func BasicAuth(fn BasicValidateFunc) echo.HandlerFunc {
 		auth := c.Request().Header.Get(echo.Authorization)
 		l := len(Basic)
 		he := echo.NewHTTPError(http.StatusBadRequest)
-		println(auth)
 
 		if len(auth) > l+1 && auth[:l] == Basic {
 			b, err := base64.StdEncoding.DecodeString(auth[l+1:])

middleware/compress.go

@@ -1,16 +1,14 @@
 package middleware
 
 import (
-	"compress/gzip"
-	"strings"
-
-	"net/http"
-
 	"bufio"
+	"compress/gzip"
+	"io"
 	"net"
+	"net/http"
+	"strings"
 
 	"github.com/labstack/echo"
-	"io"
 )
 
 type (
@@ -21,6 +19,9 @@ type (
 )
 
 func (w gzipWriter) Write(b []byte) (int, error) {
+	if w.Header().Get(echo.ContentType) == "" {
+		w.Header().Set(echo.ContentType, http.DetectContentType(b))
+	}
 	return w.Writer.Write(b)
 }
 
@@ -43,6 +44,7 @@ func Gzip() echo.MiddlewareFunc {
 
 	return func(h echo.HandlerFunc) echo.HandlerFunc {
 		return func(c *echo.Context) error {
+			c.Response().Header().Add(echo.Vary, echo.AcceptEncoding)
 			if strings.Contains(c.Request().Header.Get(echo.AcceptEncoding), scheme) {
 				w := gzip.NewWriter(c.Response().Writer())
 				defer w.Close()

middleware/compress_test.go

@@ -16,22 +16,25 @@ func TestGzip(t *testing.T) {
 	rec := httptest.NewRecorder()
 	c := echo.NewContext(req, echo.NewResponse(rec), echo.New())
 	h := func(c *echo.Context) error {
-		return c.String(http.StatusOK, "test")
+		c.Response().Write([]byte("test")) // Content-Type sniffing
+		return nil
 	}
 
 	// Skip if no Accept-Encoding header
 	Gzip()(h)(c)
 	assert.Equal(t, http.StatusOK, rec.Code)
 	assert.Equal(t, "test", rec.Body.String())
 
-	// Gzip
 	req, _ = http.NewRequest(echo.GET, "/", nil)
 	req.Header.Set(echo.AcceptEncoding, "gzip")
 	rec = httptest.NewRecorder()
 	c = echo.NewContext(req, echo.NewResponse(rec), echo.New())
+
+	// Gzip
 	Gzip()(h)(c)
 	assert.Equal(t, http.StatusOK, rec.Code)
 	assert.Equal(t, "gzip", rec.Header().Get(echo.ContentEncoding))
+	assert.Contains(t, rec.Header().Get(echo.ContentType), echo.TextPlain)
 	r, err := gzip.NewReader(rec.Body)
 	defer r.Close()
 	if assert.NoError(t, err) {