added ModifyResponse option to ProxyConfig (#1622)

Co-authored-by: Peter C <[email protected]>

Author: petoc

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
 	}
 

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
 }

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) {