From ae721bd2b2e55a105f3fb9974d4785fb5a897c98 Mon Sep 17 00:00:00 2001
From: Vishal Rana <vr@labstack.com>
Date: Tue, 9 Jun 2015 15:24:01 -0700
Subject: [PATCH] Bug in gzip, fixed #95, #96.

Signed-off-by: Vishal Rana <vr@labstack.com>
---
 echo.go                      |  3 +--
 examples/website/server.go   |  1 +
 examples/websocket/server.go | 18 ++++++++++++++++++
 middleware/compress.go       | 26 +++++++++++++++++++++-----
 middleware/compress_test.go  |  8 --------
 5 files changed, 41 insertions(+), 15 deletions(-)
 create mode 100644 examples/websocket/server.go

diff --git a/echo.go b/echo.go
index 50a8fcec..cc6ac6bd 100644
--- a/echo.go
+++ b/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
 	}
 
diff --git a/examples/website/server.go b/examples/website/server.go
index 40d95d42..38886b92 100644
--- a/examples/website/server.go
+++ b/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
diff --git a/examples/websocket/server.go b/examples/websocket/server.go
new file mode 100644
index 00000000..3daaaf48
--- /dev/null
+++ b/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")
+}
diff --git a/middleware/compress.go b/middleware/compress.go
index 1c6b8064..aea51f19 100644
--- a/middleware/compress.go
+++ b/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
-				strings.Contains(c.Request().Header.Get(echo.AcceptEncoding), scheme) &&
-				c.Response().Status() != http.StatusNotFound { // Skip for "404 - Not Found"
+			if strings.Contains(c.Request().Header.Get(echo.AcceptEncoding), scheme) {
 				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
 		}
 	}
 }
diff --git a/middleware/compress_test.go b/middleware/compress_test.go
index e6de08b4..febc05d8 100644
--- a/middleware/compress_test.go
+++ b/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")