From 4a13f31ffd4ff8648ccc4c2b2193abde22f7fd26 Mon Sep 17 00:00:00 2001 From: Emmanuel Odeke Date: Fri, 28 Oct 2016 08:11:39 -0700 Subject: 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 Run-TryBot: Brad Fitzpatrick TryBot-Result: Gobot Gobot --- src/net/http/httputil/reverseproxy_test.go | 41 ++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) (limited to 'src/net/http/httputil/reverseproxy_test.go') 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) { -- cgit v1.3-5-g9baa