http/http/httputil: add ReverseProxy.ErrorHandler

This permits specifying an ErrorHandler to customize the RoundTrip
error handling if the backend fails to return a response.

Fixes #22700
Fixes #21255

Change-Id: I8879f0956e2472a07f584660afa10105ef23bf11
Reviewed-on: https://go-review.googlesource.com/77410
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

1 files changed, 87 insertions, 0 deletions

diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index 2a12e753b5..2f75b4e34e 100644
--- a/src/net/http/httputil/reverseproxy_test.go
+++ b/src/net/http/httputil/reverseproxy_test.go
@@ -637,6 +637,93 @@ func TestReverseProxyModifyResponse(t *testing.T) {
 	}
 }
 
+type failingRoundTripper struct{}
+
+func (failingRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	return nil, errors.New("some error")
+}
+
+type staticResponseRoundTripper struct{ res *http.Response }
+
+func (rt staticResponseRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
+	return rt.res, nil
+}
+
+func TestReverseProxyErrorHandler(t *testing.T) {
+	tests := []struct {
+		name           string
+		wantCode       int
+		errorHandler   func(http.ResponseWriter, *http.Request, error)
+		transport      http.RoundTripper // defaults to failingRoundTripper
+		modifyResponse func(*http.Response) error
+	}{
+		{
+			name:     "default",
+			wantCode: http.StatusBadGateway,
+		},
+		{
+			name:         "errorhandler",
+			wantCode:     http.StatusTeapot,
+			errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
+		},
+		{
+			name: "modifyresponse_noerr",
+			transport: staticResponseRoundTripper{
+				&http.Response{StatusCode: 345, Body: http.NoBody},
+			},
+			modifyResponse: func(res *http.Response) error {
+				res.StatusCode++
+				return nil
+			},
+			errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
+			wantCode:     346,
+		},
+		{
+			name: "modifyresponse_err",
+			transport: staticResponseRoundTripper{
+				&http.Response{StatusCode: 345, Body: http.NoBody},
+			},
+			modifyResponse: func(res *http.Response) error {
+				res.StatusCode++
+				return errors.New("some error to trigger errorHandler")
+			},
+			errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
+			wantCode:     http.StatusTeapot,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			target := &url.URL{
+				Scheme: "http",
+				Host:   "dummy.tld",
+				Path:   "/",
+			}
+			rproxy := NewSingleHostReverseProxy(target)
+			rproxy.Transport = tt.transport
+			rproxy.ModifyResponse = tt.modifyResponse
+			if rproxy.Transport == nil {
+				rproxy.Transport = failingRoundTripper{}
+			}
+			rproxy.ErrorLog = log.New(ioutil.Discard, "", 0) // quiet for tests
+			if tt.errorHandler != nil {
+				rproxy.ErrorHandler = tt.errorHandler
+			}
+			frontendProxy := httptest.NewServer(rproxy)
+			defer frontendProxy.Close()
+
+			resp, err := http.Get(frontendProxy.URL + "/test")
+			if err != nil {
+				t.Fatalf("failed to reach proxy: %v", err)
+			}
+			if g, e := resp.StatusCode, tt.wantCode; g != e {
+				t.Errorf("got res.StatusCode %d; expected %d", 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) {