Issue 950 (#955)

* Added implementation #950

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

* Added test for body-dump middleware

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

middleware/body_dump.go

@@ -0,0 +1,112 @@
+package middleware
+
+import (
+	"bufio"
+	"bytes"
+	"io/ioutil"
+	"net"
+	"net/http"
+
+	"io"
+
+	"github.com/labstack/echo"
+)
+
+type (
+	// BodyDumpConfig defines the config for BodyDump middleware.
+	BodyDumpConfig struct {
+		// Skipper defines a function to skip middleware.
+		Skipper Skipper
+
+		// Handler receives request and response payload.
+		// Required.
+		Handler BodyDumpHandler
+	}
+
+	// BodyDumpHandler receives the request and response payload.
+	BodyDumpHandler func(echo.Context, []byte, []byte)
+
+	bodyDumpResponseWriter struct {
+		io.Writer
+		http.ResponseWriter
+	}
+)
+
+var (
+	// DefaultBodyDumpConfig is the default Gzip middleware config.
+	DefaultBodyDumpConfig = BodyDumpConfig{
+		Skipper: DefaultSkipper,
+	}
+)
+
+// BodyDump returns a BodyDump middleware.
+//
+// BodyLimit middleware captures the request and response payload and calls the
+// registered handler.
+func BodyDump(handler BodyDumpHandler) echo.MiddlewareFunc {
+	c := DefaultBodyDumpConfig
+	c.Handler = handler
+	return BodyDumpWithConfig(c)
+}
+
+// BodyDumpWithConfig returns a BodyDump middleware with config.
+// See: `BodyConfig()`.
+func BodyDumpWithConfig(config BodyDumpConfig) echo.MiddlewareFunc {
+	// Defaults
+	if config.Handler == nil {
+		panic("body-dump middleware requires a handler function")
+	}
+	if config.Skipper == nil {
+		config.Skipper = DefaultBodyDumpConfig.Skipper
+	}
+
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) (err error) {
+			if config.Skipper(c) {
+				return next(c)
+			}
+
+			// Request
+			reqBody := []byte{}
+			if c.Request().Body != nil { // Read
+				reqBody, _ = ioutil.ReadAll(c.Request().Body)
+			}
+			c.Request().Body = ioutil.NopCloser(bytes.NewBuffer(reqBody)) // Reset
+
+			// Response
+			resBody := new(bytes.Buffer)
+			mw := io.MultiWriter(c.Response().Writer, resBody)
+			writer := &bodyDumpResponseWriter{Writer: mw, ResponseWriter: c.Response().Writer}
+			c.Response().Writer = writer
+
+			if err = next(c); err != nil {
+				c.Error(err)
+			}
+
+			// Callback
+			config.Handler(c, reqBody, resBody.Bytes())
+
+			return
+		}
+	}
+}
+
+func (w *bodyDumpResponseWriter) WriteHeader(code int) {
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *bodyDumpResponseWriter) Write(b []byte) (int, error) {
+	return w.Writer.Write(b)
+}
+
+func (w *bodyDumpResponseWriter) Flush() {
+	w.ResponseWriter.(http.Flusher).Flush()
+}
+
+func (w *bodyDumpResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return w.ResponseWriter.(http.Hijacker).Hijack()
+}
+
+func (w *bodyDumpResponseWriter) CloseNotify() <-chan bool {
+	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}

middleware/body_dump_test.go

@@ -0,0 +1,40 @@
+package middleware
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/labstack/echo"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBodyDump(t *testing.T) {
+	e := echo.New()
+	hw := "Hello, World!"
+	req := httptest.NewRequest(echo.POST, "/", strings.NewReader(hw))
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec)
+	h := func(c echo.Context) error {
+		body, err := ioutil.ReadAll(c.Request().Body)
+		if err != nil {
+			return err
+		}
+		return c.String(http.StatusOK, string(body))
+	}
+
+	requestBody := ""
+	responseBody := ""
+	mw := BodyDump(func(c echo.Context, reqBody, resBody []byte) {
+		requestBody = string(reqBody)
+		responseBody = string(resBody)
+	})
+	if assert.NoError(t, mw(h)(c)) {
+		assert.Equal(t, requestBody, hw)
+		assert.Equal(t, responseBody, hw)
+		assert.Equal(t, http.StatusOK, rec.Code)
+		assert.Equal(t, hw, rec.Body.String())
+	}
+}