From d3e33ec4dd0ded602689548138eab91fcf098852 Mon Sep 17 00:00:00 2001
From: Vishal Rana <vr@labstack.com>
Date: Sun, 27 Mar 2016 15:21:56 -0700
Subject: [PATCH] Closes #433

Signed-off-by: Vishal Rana <vr@labstack.com>
---
 engine/fasthttp/server.go   | 12 ++++++++++--
 engine/standard/response.go |  8 ++++++--
 engine/standard/server.go   | 20 ++++++++++----------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/engine/fasthttp/server.go b/engine/fasthttp/server.go
index f6eeb003..63a8f686 100644
--- a/engine/fasthttp/server.go
+++ b/engine/fasthttp/server.go
@@ -149,8 +149,12 @@ func (s *Server) ServeHTTP(c *fasthttp.RequestCtx) {
 // WrapHandler wraps `fasthttp.RequestHandler` into `echo.HandlerFunc`.
 func WrapHandler(h fasthttp.RequestHandler) echo.HandlerFunc {
 	return func(c echo.Context) error {
-		ctx := c.Request().(*Request).RequestCtx
+		rq := c.Request().(*Request)
+		rs := c.Response().(*Response)
+		ctx := rq.RequestCtx
 		h(ctx)
+		rs.status = ctx.Response.StatusCode()
+		rs.size = int64(ctx.Response.Header.ContentLength())
 		return nil
 	}
 }
@@ -159,8 +163,12 @@ func WrapHandler(h fasthttp.RequestHandler) echo.HandlerFunc {
 func WrapMiddleware(h fasthttp.RequestHandler) echo.MiddlewareFunc {
 	return func(next echo.Handler) echo.Handler {
 		return echo.HandlerFunc(func(c echo.Context) error {
-			ctx := c.Request().(*Request).RequestCtx
+			rq := c.Request().(*Request)
+			rs := c.Response().(*Response)
+			ctx := rq.RequestCtx
 			h(ctx)
+			rs.status = ctx.Response.StatusCode()
+			rs.size = int64(ctx.Response.Header.ContentLength())
 			return next.Handle(c)
 		})
 	}
diff --git a/engine/standard/response.go b/engine/standard/response.go
index 0a179136..f3d6d4c3 100644
--- a/engine/standard/response.go
+++ b/engine/standard/response.go
@@ -24,7 +24,7 @@ type (
 
 	responseAdapter struct {
 		http.ResponseWriter
-		writer io.Writer
+		Response *Response
 	}
 )
 
@@ -108,8 +108,12 @@ func (r *Response) reset(w http.ResponseWriter, h engine.Header) {
 	r.writer = w
 }
 
+func (r *responseAdapter) WriteHeader(code int) {
+	r.Response.WriteHeader(code)
+}
+
 func (r *responseAdapter) Write(b []byte) (n int, err error) {
-	return r.writer.Write(b)
+	return r.Response.Write(b)
 }
 
 func (r *responseAdapter) Flush() {
diff --git a/engine/standard/server.go b/engine/standard/server.go
index ac8eca01..fce541e6 100644
--- a/engine/standard/server.go
+++ b/engine/standard/server.go
@@ -138,12 +138,13 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 // WrapHandler wraps `http.Handler` into `echo.HandlerFunc`.
 func WrapHandler(h http.Handler) echo.HandlerFunc {
 	return func(c echo.Context) error {
+		rq := c.Request().(*Request)
+		rs := c.Response().(*Response)
 		w := &responseAdapter{
-			ResponseWriter: c.Response().(*Response).ResponseWriter,
-			writer:         c.Response(),
+			ResponseWriter: rs.ResponseWriter,
+			Response:       rs,
 		}
-		r := c.Request().(*Request).Request
-		h.ServeHTTP(w, r)
+		h.ServeHTTP(w, rq.Request)
 		return nil
 	}
 }
@@ -154,14 +155,13 @@ func WrapMiddleware(m func(http.Handler) http.Handler) echo.MiddlewareFunc {
 		return echo.HandlerFunc(func(c echo.Context) (err error) {
 			rq := c.Request().(*Request)
 			rs := c.Response().(*Response)
+			w := &responseAdapter{
+				ResponseWriter: rs.ResponseWriter,
+				Response:       rs,
+			}
 			m(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				rs.ResponseWriter = &responseAdapter{
-					ResponseWriter: rs.ResponseWriter,
-					writer:         c.Response(),
-				}
-				rq.Request = r
 				err = next.Handle(c)
-			})).ServeHTTP(rs.ResponseWriter, rq.Request)
+			})).ServeHTTP(w, rq.Request)
 			return
 		})
 	}