Bug in gzip, fixed #95, #96.

Signed-off-by: Vishal Rana <vr@labstack.com>

echo.go

@@ -283,7 +283,7 @@ func (e *Echo) WebSocket(path string, h HandlerFunc) {
 				err = h(c)
 			},
 		}
-		wss.ServeHTTP(c.response.writer, c.request)
+		wss.ServeHTTP(c.response, c.request)
 		return err
 	})
 }
@@ -400,7 +400,6 @@ func (e *Echo) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	c.reset(r, w, e)
 	if h == nil {
-		c.response.status = http.StatusNotFound // Helpful to skip middleware
 		h = e.notFoundHandler
 	}
 

examples/website/server.go

@@ -64,6 +64,7 @@ func main() {
 	// Middleware
 	e.Use(mw.Logger())
 	e.Use(mw.Recover())
+	e.Use(mw.Gzip())
 
 	//------------------------
 	// Third-party middleware

examples/websocket/server.go

@@ -0,0 +1,18 @@
+package main
+
+import (
+	"github.com/labstack/echo"
+	"io"
+	mw "github.com/labstack/echo/middleware"
+)
+
+func main() {
+	e := echo.New()
+	e.Use(mw.Logger())
+	e.Use(mw.Gzip())
+	e.WebSocket("/ws", func(c *echo.Context) error {
+		io.Copy(c.Socket(), c.Socket())
+		return nil
+	})
+	e.Run(":1323")
+}

middleware/compress.go

@@ -2,12 +2,15 @@ package middleware
 
 import (
 	"compress/gzip"
-	"io"
 	"strings"
 
 	"net/http"
 
+	"bufio"
+	"net"
+
 	"github.com/labstack/echo"
+	"io"
 )
 
 type (
@@ -21,6 +24,18 @@ func (w gzipWriter) Write(b []byte) (int, error) {
 	return w.Writer.Write(b)
 }
 
+func (w gzipWriter) Flush() {
+	w.Writer.(*gzip.Writer).Flush()
+}
+
+func (w gzipWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return w.ResponseWriter.(http.Hijacker).Hijack()
+}
+
+func (w *gzipWriter) CloseNotify() <-chan bool {
+	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}
+
 // Gzip returns a middleware which compresses HTTP response using gzip compression
 // scheme.
 func Gzip() echo.MiddlewareFunc {
@@ -28,16 +43,17 @@ func Gzip() echo.MiddlewareFunc {
 
 	return func(h echo.HandlerFunc) echo.HandlerFunc {
 		return func(c *echo.Context) error {
-			if (c.Request().Header.Get(echo.Upgrade)) != echo.WebSocket && // Skip for WebSocket
+			if strings.Contains(c.Request().Header.Get(echo.AcceptEncoding), scheme) {
-				strings.Contains(c.Request().Header.Get(echo.AcceptEncoding), scheme) &&
-				c.Response().Status() != http.StatusNotFound { // Skip for "404 - Not Found"
 				w := gzip.NewWriter(c.Response().Writer())
 				defer w.Close()
 				gw := gzipWriter{Writer: w, ResponseWriter: c.Response().Writer()}
 				c.Response().Header().Set(echo.ContentEncoding, scheme)
 				c.Response().SetWriter(gw)
 			}
-			return h(c)
+			if err := h(c); err != nil {
+				c.Error(err)
+			}
+			return nil
 		}
 	}
 }

middleware/compress_test.go

@@ -24,14 +24,6 @@ func TestGzip(t *testing.T) {
 	assert.Equal(t, http.StatusOK, rec.Code)
 	assert.Equal(t, "test", rec.Body.String())
 
-	// Skip if WebSocket
-	rec = httptest.NewRecorder()
-	c = echo.NewContext(req, echo.NewResponse(rec), echo.New())
-	c.Request().Header.Set(echo.Upgrade, echo.WebSocket)
-	Gzip()(h)(c)
-	assert.Equal(t, http.StatusOK, rec.Code)
-	assert.NotEqual(t, "gzip", rec.Header().Get(echo.ContentEncoding))
-
 	// Gzip
 	req, _ = http.NewRequest(echo.GET, "/", nil)
 	req.Header.Set(echo.AcceptEncoding, "gzip")