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>

diff --git a/src/net/http/httputil/reverseproxy.go b/src/net/http/httputil/reverseproxy.go
index f18dd886cc1cce9b7c3f75c1110d75a2b4c84993..78675057085c23bce6c5282af50ff67bf7aaf520 100644 (file)
--- a/src/net/http/httputil/reverseproxy.go
+++ b/src/net/http/httputil/reverseproxy.go
@@ -52,6 +52,11 @@ type ReverseProxy struct {
        // get byte slices for use by io.CopyBuffer when
        // copying HTTP response bodies.
        BufferPool BufferPool
+
+       // ModifyResponse is an optional function that
+       // modifies the Response from the backend.
+       // If it returns an error, the proxy returns a StatusBadGateway error.
+       ModifyResponse func(*http.Response) error
 }
 
 // A BufferPool is an interface for getting and returning temporary
@@ -216,6 +221,14 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
                res.Header.Del(h)
        }
 
+       if p.ModifyResponse != nil {
+               if err := p.ModifyResponse(res); err != nil {
+                       p.logf("http: proxy error: %v", err)
+                       rw.WriteHeader(http.StatusBadGateway)
+                       return
+               }
+       }
+
        copyHeader(rw.Header(), res.Header)
 
        // The "Trailer" header isn't included in the Transport's response,

diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index b3270a1a63fde5195623b7c6815f07c2f090f05b..20c4e16bcb8b30387e1e452d9b3d6ce195638ece 100644 (file)
--- 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) {