From 6463bcb190302063a2375c3268c78de5861e8d8c Mon Sep 17 00:00:00 2001
From: Peter C <63091190+petoc@users.noreply.github.com>
Date: Fri, 28 Aug 2020 02:51:27 +0200
Subject: [PATCH] added ModifyResponse option to ProxyConfig (#1622)

Co-authored-by: Peter C <petoc@users.noreply.github.com>
---
 middleware/proxy.go      |  3 +++
 middleware/proxy_1_11.go |  1 +
 middleware/proxy_test.go | 19 ++++++++++++++++++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/middleware/proxy.go b/middleware/proxy.go
index 1956e91e..a9b91f6c 100644
--- a/middleware/proxy.go
+++ b/middleware/proxy.go
@@ -45,6 +45,9 @@ type (
 		// Examples: If custom TLS certificates are required.
 		Transport http.RoundTripper
 
+		// ModifyResponse defines function to modify response from ProxyTarget.
+		ModifyResponse func(*http.Response) error
+
 		rewriteRegex map[*regexp.Regexp]string
 	}
 
diff --git a/middleware/proxy_1_11.go b/middleware/proxy_1_11.go
index 12b7568b..a4392781 100644
--- a/middleware/proxy_1_11.go
+++ b/middleware/proxy_1_11.go
@@ -20,5 +20,6 @@ func proxyHTTP(tgt *ProxyTarget, c echo.Context, config ProxyConfig) http.Handle
 		c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("remote %s unreachable, could not forward: %v", desc, err)))
 	}
 	proxy.Transport = config.Transport
+	proxy.ModifyResponse = config.ModifyResponse
 	return proxy
 }
diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go
index 5ef11bc8..b19bf4f2 100644
--- a/middleware/proxy_test.go
+++ b/middleware/proxy_test.go
@@ -1,7 +1,9 @@
 package middleware
 
 import (
+	"bytes"
 	"fmt"
+	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -104,11 +106,26 @@ func TestProxy(t *testing.T) {
 	e.ServeHTTP(rec, req)
 	assert.Equal(t, "/user/jack/order/1", req.URL.Path)
 	assert.Equal(t, http.StatusOK, rec.Code)
-  req.URL.Path = "/users/jill/orders/T%2FcO4lW%2Ft%2FVp%2F"
+	req.URL.Path = "/users/jill/orders/T%2FcO4lW%2Ft%2FVp%2F"
 	e.ServeHTTP(rec, req)
 	assert.Equal(t, "/user/jill/order/T%2FcO4lW%2Ft%2FVp%2F", req.URL.Path)
 	assert.Equal(t, http.StatusOK, rec.Code)
 
+	// ModifyResponse
+	e = echo.New()
+	e.Use(ProxyWithConfig(ProxyConfig{
+		Balancer: rrb,
+		ModifyResponse: func(res *http.Response) error {
+			res.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("modified")))
+			res.Header.Set("X-Modified", "1")
+			return nil
+		},
+	}))
+	rec = httptest.NewRecorder()
+	e.ServeHTTP(rec, req)
+	assert.Equal(t, "modified", rec.Body.String())
+	assert.Equal(t, "1", rec.Header().Get("X-Modified"))
+
 	// ProxyTarget is set in context
 	contextObserver := func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) (err error) {