Add a hook for errors occurring before the roundtrip.

rp.go

@@ -34,13 +34,21 @@
 	// RoundTrip performs the round trip of the request.
 	// It is necessary to implement the functions that http.Transport is responsible for (e.g. MaxIdleConnsPerHost).
 	RoundTrip(r *http.Request) (*http.Response, error)
+	RoundTripOnError(r *http.Request) (*http.Response, error)
+}
+
+type RoundTipperOnErrorer interface {
+	// RoundTripOnError performs the round trip of the request when the upstream returns an error.
+	// If this method is not implemented, the request will be sent to the default transport error.
+	RoundTripOnError(r *http.Request) (*http.Response, error)
 }
 
 type relayer struct {
 	Relayer
-	Rewrite        func(*httputil.ProxyRequest) error
-	GetCertificate func(*tls.ClientHelloInfo) (*tls.Certificate, error)
-	RoundTrip      func(*http.Request) (*http.Response, error)
+	Rewrite          func(*httputil.ProxyRequest) error
+	GetCertificate   func(*tls.ClientHelloInfo) (*tls.Certificate, error)
+	RoundTrip        func(*http.Request) (*http.Response, error)
+	RoundTripOnError func(*http.Request) (*http.Response, error)
 }
 
 func newRelayer(r Relayer) *relayer {
@@ -58,6 +66,9 @@
 	} else {
 		rr.RoundTrip = http.DefaultTransport.RoundTrip
 	}
+	if v, ok := r.(RoundTipperOnErrorer); ok {
+		rr.RoundTripOnError = v.RoundTripOnError
+	}
 	return rr
 }
 
@@ -144,6 +155,9 @@
 
 func (t *transport) RoundTrip(r *http.Request) (*http.Response, error) {
 	if v := r.Header.Get(errorKey); v != "" {
+		if t.rr.RoundTripOnError != nil {
+			return t.rr.RoundTripOnError(r)
+		}
 		// If errorKey is set, return error response.
 		body := v
 		resp := &http.Response{

rp_test.go

@@ -26,18 +26,26 @@
 }
 
 func TestHTTPRouting(t *testing.T) {
+	simpleRelayer := func(m map[string]string) rp.Relayer {
+		return testutil.NewSimpleRelayer(m)
+	}
+	roundTripErrorRelayer := func(m map[string]string) rp.Relayer {
+		return testutil.NewRoundTripOnErrorRelayer(m)
+	}
 	tests := []struct {
 		upstreams      []upstream
 		reqURL         string
 		want           string
 		wantStatusCode int
+		setupRelayer   func(m map[string]string) rp.Relayer
 	}{
-		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://a.example.com", "response from / [a.example.com]", http.StatusOK},
-		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://b.example.com/hello", "response from /hello [b.example.com]", http.StatusOK},
-		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://a.example.com/hello?foo=bar", "response from /hello?foo=bar [a.example.com]", http.StatusOK},
-		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://x.example.com/hello", "not found upstream: x.example.com", http.StatusBadGateway},
-		{[]upstream{{"a.example.com", "/A"}}, "http://a.example.com", "response from /A [a.example.com]", http.StatusOK},
-		{[]upstream{{"a.example.com", "/A"}}, "http://a.example.com/hello?foo=bar", "response from /A/hello?foo=bar [a.example.com]", http.StatusOK},
+		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://a.example.com", "response from / [a.example.com]", http.StatusOK, simpleRelayer},
+		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://b.example.com/hello", "response from /hello [b.example.com]", http.StatusOK, simpleRelayer},
+		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://a.example.com/hello?foo=bar", "response from /hello?foo=bar [a.example.com]", http.StatusOK, simpleRelayer},
+		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://x.example.com/hello", "not found upstream: x.example.com", http.StatusBadGateway, simpleRelayer},
+		{[]upstream{{"a.example.com", "/A"}}, "http://a.example.com", "response from /A [a.example.com]", http.StatusOK, simpleRelayer},
+		{[]upstream{{"a.example.com", "/A"}}, "http://a.example.com/hello?foo=bar", "response from /A/hello?foo=bar [a.example.com]", http.StatusOK, simpleRelayer},
+		{[]upstream{{"a.example.com", "/"}, {"b.example.com", "/"}}, "http://x.example.com/hello", "round trip on error: x.example.com", http.StatusBadGateway, roundTripErrorRelayer},
 	}
 	for _, tt := range tests {
 		t.Run(tt.reqURL, func(t *testing.T) {
@@ -47,7 +55,7 @@
 				us := testutil.NewUpstreamServer(t, u.hostname)
 				m[u.hostname] = us.URL + u.rootPath
 			}
-			r := testutil.NewSimpleRelayer(m)
+			r := tt.setupRelayer(m)
 			port, err := testutil.NewPort()
 			if err != nil {
 				t.Fatal(err)
@@ -62,10 +70,10 @@
 				}
 			}()
 			t.Cleanup(func() {
 				_ = proxy.Shutdown(context.Background())
 			})
 			t.Cleanup(func() {
 				_ = proxy.Shutdown(context.Background())
 			})
 			proxyURL, err := url.Parse(fmt.Sprintf("http://127.0.0.1:%d", port))
 			if err != nil {
@@ -142,7 +150,7 @@
 				}
 			}()
 			t.Cleanup(func() {
 				_ = proxy.Shutdown(context.Background())
 			})
 			proxyURL, err := url.Parse(fmt.Sprintf("https://127.0.0.1:%d", port))
 			if err != nil {
@@ -193,7 +201,7 @@
 				return
 			}
 			if tt.wantErr {
 				b, _ := httputil.DumpResponse(resp, true)
 				t.Errorf("want error, got resp:\n%s", string(b))
 				return
 			}

testutil/relayer.go

@@ -1,8 +1,10 @@
 package testutil
 
 import (
+	"bytes"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -80,3 +82,29 @@ func (r *SimpleRelayer) GetUpstream(req *http.Request) (*url.URL, error) {
 	}
 	return nil, fmt.Errorf("not found upstream: %v", host)
 }
+
+type RoundTripOnErrorRelayer struct {
+	SimpleRelayer
+}
+
+func NewRoundTripOnErrorRelayer(h map[string]string) *RoundTripOnErrorRelayer {
+	return &RoundTripOnErrorRelayer{
+		SimpleRelayer: SimpleRelayer{
+			h: h,
+		},
+	}
+}
+func (r *RoundTripOnErrorRelayer) RoundTripOnError(req *http.Request) (*http.Response, error) {
+	body := fmt.Sprintf("round trip on error: %v", req.Host)
+	return &http.Response{
+		Status:        http.StatusText(http.StatusBadGateway),
+		StatusCode:    http.StatusBadGateway,
+		Proto:         req.Proto,
+		ProtoMajor:    req.ProtoMajor,
+		ProtoMinor:    req.ProtoMinor,
+		Body:          io.NopCloser(bytes.NewBufferString(body)),
+		ContentLength: int64(len(body)),
+		Request:       req,
+		Header:        make(http.Header, 0),
+	}, nil
+}