net/http/httputil: add ModifyResponse to reverseProxy

Adds ModifyResponse, an optional func to ReverseProxy
that modifies a response in the backend, right before
the headers of the response are written to the internal
response writer.
If ModifyResponse returns an error, the proxy returns
a StatusBadGateway error.

Fixes #14237.

Change-Id: I8e03139e34dea0084512ccbd8cc49e941bf9fb5d
Reviewed-on: https://go-review.googlesource.com/32356
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

1 files changed, 41 insertions, 0 deletions

diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index b3270a1a63..20c4e16bcb 100644
--- a/src/net/http/httputil/reverseproxy_test.go
+++ b/src/net/http/httputil/reverseproxy_test.go
@@ -583,6 +583,47 @@ func TestReverseProxy_NilBody(t *testing.T) {
 	}
 }
 
+// Issue 14237. Test ModifyResponse and that an error from it
+// causes the proxy to return StatusBadGateway, or StatusOK otherwise.
+func TestReverseProxyModifyResponse(t *testing.T) {
+	backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Add("X-Hit-Mod", fmt.Sprintf("%v", r.URL.Path == "/mod"))
+	}))
+	defer backendServer.Close()
+
+	rpURL, _ := url.Parse(backendServer.URL)
+	rproxy := NewSingleHostReverseProxy(rpURL)
+	rproxy.ErrorLog = log.New(ioutil.Discard, "", 0) // quiet for tests
+	rproxy.ModifyResponse = func(resp *http.Response) error {
+		if resp.Header.Get("X-Hit-Mod") != "true" {
+			return fmt.Errorf("tried to by-pass proxy")
+		}
+		return nil
+	}
+
+	frontendProxy := httptest.NewServer(rproxy)
+	defer frontendProxy.Close()
+
+	tests := []struct {
+		url      string
+		wantCode int
+	}{
+		{frontendProxy.URL + "/mod", http.StatusOK},
+		{frontendProxy.URL + "/schedule", http.StatusBadGateway},
+	}
+
+	for i, tt := range tests {
+		resp, err := http.Get(tt.url)
+		if err != nil {
+			t.Fatalf("failed to reach proxy: %v", err)
+		}
+		if g, e := resp.StatusCode, tt.wantCode; g != e {
+			t.Errorf("#%d: got res.StatusCode %d; expected %d", i, g, e)
+		}
+		resp.Body.Close()
+	}
+}
+
 // Issue 16659: log errors from short read
 func TestReverseProxy_CopyBuffer(t *testing.T) {
 	backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {