Wrap handler and middleware functions

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

echo.go

@@ -551,14 +551,22 @@ func (e *HTTPError) Error() string {
 	return e.Message
 }
 
-// WrapMiddleware wrap `echo.HandlerFunc` into `echo.MiddlewareFunc`.
-func WrapMiddleware(h HandlerFunc) MiddlewareFunc {
+// WrapHandler wraps `http.Handler` into `echo.HandlerFunc`.
+func WrapHandler(h http.Handler) HandlerFunc {
+	return func(c Context) error {
+		h.ServeHTTP(c.Response(), c.Request())
+		return nil
+	}
+}
+
+// WrapMiddleware wraps `func(http.Handler) http.Handler` into `echo.MiddlewareFunc`
+func WrapMiddleware(m func(http.Handler) http.Handler) MiddlewareFunc {
 	return func(next HandlerFunc) HandlerFunc {
-		return func(c Context) error {
-			if err := h(c); err != nil {
-				return err
-			}
-			return next(c)
+		return func(c Context) (err error) {
+			m(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				err = next(c)
+			})).ServeHTTP(c.Response(), c.Request())
+			return
 		}
 	}
 }

echo_test.go

@@ -85,26 +85,34 @@ func TestEchoMiddleware(t *testing.T) {
 	e := New()
 	buf := new(bytes.Buffer)
 
-	e.Pre(WrapMiddleware(func(c Context) error {
-		assert.Empty(t, c.Path())
-		buf.WriteString("-1")
-		return nil
-	}))
+	e.Pre(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			assert.Empty(t, c.Path())
+			buf.WriteString("-1")
+			return next(c)
+		}
+	})
 
-	e.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("1")
-		return nil
-	}))
+	e.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("1")
+			return next(c)
+		}
+	})
 
-	e.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("2")
-		return nil
-	}))
+	e.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("2")
+			return next(c)
+		}
+	})
 
-	e.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("3")
-		return nil
-	}))
+	e.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("3")
+			return next(c)
+		}
+	})
 
 	// Route
 	e.GET("/", func(c Context) error {
@@ -119,9 +127,11 @@ func TestEchoMiddleware(t *testing.T) {
 
 func TestEchoMiddlewareError(t *testing.T) {
 	e := New()
-	e.Use(WrapMiddleware(func(c Context) error {
-		return errors.New("error")
-	}))
+	e.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return errors.New("error")
+		}
+	})
 	e.GET("/", NotFoundHandler)
 	c, _ := request(GET, "/", e)
 	assert.Equal(t, http.StatusInternalServerError, c)
@@ -140,6 +150,43 @@ func TestEchoHandler(t *testing.T) {
 	assert.Equal(t, "OK", b)
 }
 
+func TestEchoWrapHandler(t *testing.T) {
+	e := New()
+	req, _ := http.NewRequest(GET, "", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+	h := WrapHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("test"))
+	}))
+	if assert.NoError(t, h(c)) {
+		assert.Equal(t, http.StatusOK, rec.Code)
+		assert.Equal(t, "test", rec.Body.String())
+	}
+}
+
+func TestEchoWrapMiddleware(t *testing.T) {
+	e := New()
+	req, _ := http.NewRequest(GET, "", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+	buf := new(bytes.Buffer)
+	mw := WrapMiddleware(func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			buf.Write([]byte("mw"))
+			h.ServeHTTP(w, r)
+		})
+	})
+	h := mw(func(c Context) error {
+		return c.String(http.StatusOK, "OK")
+	})
+	if assert.NoError(t, h(c)) {
+		assert.Equal(t, "mw", buf.String())
+		assert.Equal(t, http.StatusOK, rec.Code)
+		assert.Equal(t, "OK", rec.Body.String())
+	}
+}
+
 func TestEchoConnect(t *testing.T) {
 	e := New()
 	testMethod(t, CONNECT, "/", e)
@@ -248,10 +295,10 @@ func TestEchoRoutes(t *testing.T) {
 func TestEchoGroup(t *testing.T) {
 	e := New()
 	buf := new(bytes.Buffer)
-	e.Use(MiddlewareFunc(func(h HandlerFunc) HandlerFunc {
+	e.Use(MiddlewareFunc(func(next HandlerFunc) HandlerFunc {
 		return func(c Context) error {
 			buf.WriteString("0")
-			return h(c)
+			return next(c)
 		}
 	}))
 	h := func(c Context) error {
@@ -266,23 +313,29 @@ func TestEchoGroup(t *testing.T) {
 
 	// Group
 	g1 := e.Group("/group1")
-	g1.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("1")
-		return h(c)
-	}))
+	g1.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("1")
+			return next(c)
+		}
+	})
 	g1.Get("", h)
 
 	// Nested groups with middleware
 	g2 := e.Group("/group2")
-	g2.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("2")
-		return nil
-	}))
+	g2.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("2")
+			return next(c)
+		}
+	})
 	g3 := g2.Group("/group3")
-	g3.Use(WrapMiddleware(func(c Context) error {
-		buf.WriteString("3")
-		return nil
-	}))
+	g3.Use(func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			buf.WriteString("3")
+			return next(c)
+		}
+	})
 	g3.Get("", h)
 
 	request(GET, "/users", e)

group_test.go

@@ -1,8 +1,9 @@
 package echo
 
 import (
-	"github.com/stretchr/testify/assert"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 // TODO: Fix me
@@ -38,11 +39,31 @@ func TestGroupRouteMiddleware(t *testing.T) {
 	e := New()
 	g := e.Group("/group")
 	h := func(Context) error { return nil }
-	m1 := WrapMiddleware(func(c Context) error { return nil })
-	m2 := WrapMiddleware(func(c Context) error { return nil })
-	m3 := WrapMiddleware(func(c Context) error { return nil })
-	m4 := WrapMiddleware(func(c Context) error { return c.NoContent(404) })
-	m5 := WrapMiddleware(func(c Context) error { return c.NoContent(405) })
+	m1 := func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return next(c)
+		}
+	}
+	m2 := func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return next(c)
+		}
+	}
+	m3 := func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return next(c)
+		}
+	}
+	m4 := func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return c.NoContent(404)
+		}
+	}
+	m5 := func(next HandlerFunc) HandlerFunc {
+		return func(c Context) error {
+			return c.NoContent(405)
+		}
+	}
 	g.Use(m1, m2, m3)
 	g.GET("/404", h, m4)
 	g.GET("/405", h, m5)