mirror of
https://github.com/golang/go
synced 2024-11-20 10:34:42 -07:00
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>
This commit is contained in:
parent
032d150bd4
commit
4a13f31ffd
@ -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,
|
||||
|
@ -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) {
|
||||
|
Loading…
Reference in New Issue
Block a user