Repository: labstack/echo
Pull request #2316 (2022-10-25 18:46:11)

[suggestion] Add helper interface for ProxyBalancer interface
Idea is discussed in this [issue](https://github.com/labstack/echo/issues/2313)

codecov[bot]: # [Codecov](https://codecov.io/gh/labstack/echo/pull/2316?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) Report
Base: **92.35**% // Head: **92.30**% // Decreases project coverage by **`-0.05%`** :warning:
> Coverage data is based on head [(`d70346c`)](https://codecov.io/gh/labstack/echo/pull/2316?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) compared to base [(`8f2bf82`)](https://codecov.io/gh/labstack/echo/commit/8f2bf82982a6005447b68bbbae739c6fcced5140?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack).
> Patch coverage: 72.72% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #2316      +/-   ##
==========================================
- Coverage   92.35%   92.30%   -0.06%     
==========================================
  Files          37       37              
  Lines        4436     4446      +10     
==========================================
+ Hits         4097     4104       +7     
- Misses        247      249       +2     
- Partials       92       93       +1     
```


| [Impacted Files](https://codecov.io/gh/labstack/echo/pull/2316?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) | Coverage Δ | |
|---|---|---|
| [middleware/proxy.go](https://codecov.io/gh/labstack/echo/pull/2316/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack#diff-bWlkZGxld2FyZS9wcm94eS5nbw==) | `66.66% <72.72%> (+0.23%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/labstack/echo/pull/2316?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack).


Patch:
diff --git a/middleware/proxy.go b/middleware/proxy.go
index 6cfd673..d2cd2aa 100644
--- a/middleware/proxy.go
+++ b/middleware/proxy.go
@@ -72,6 +72,11 @@ type (
 		Next(echo.Context) *ProxyTarget
 	}
 
+	// TargetProvider defines an interface that gives the opportunity for balancer to return custom errors when selecting target.
+	TargetProvider interface {
+		NextTarget(echo.Context) (*ProxyTarget, error)
+	}
+
 	commonBalancer struct {
 		targets []*ProxyTarget
 		mutex   sync.RWMutex
@@ -223,6 +228,7 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 		}
 	}
 
+	provider, isTargetProvider := config.Balancer.(TargetProvider)
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) (err error) {
 			if config.Skipper(c) {
@@ -231,7 +237,16 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 
 			req := c.Request()
 			res := c.Response()
-			tgt := config.Balancer.Next(c)
+
+			var tgt *ProxyTarget
+			if isTargetProvider {
+				tgt, err = provider.NextTarget(c)
+				if err != nil {
+					return err
+				}
+			} else {
+				tgt = config.Balancer.Next(c)
+			}
 			c.Set(config.ContextKey, tgt)
 
 			if err := rewriteURL(config.RegexRewrite, req); err != nil {


Test patch:
diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go
index 7939fc5..0ded50a 100644
--- a/middleware/proxy_test.go
+++ b/middleware/proxy_test.go
@@ -18,7 +18,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-//Assert expected with url.EscapedPath method to obtain the path.
+// Assert expected with url.EscapedPath method to obtain the path.
 func TestProxy(t *testing.T) {
 	// Setup
 	t1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -31,7 +31,6 @@ func TestProxy(t *testing.T) {
 	}))
 	defer t2.Close()
 	url2, _ := url.Parse(t2.URL)
-
 	targets := []*ProxyTarget{
 		{
 			Name: "target 1",
@@ -122,6 +121,55 @@ func TestProxy(t *testing.T) {
 	e.ServeHTTP(rec, req)
 }
 
+type testProvider struct {
+	*commonBalancer
+	target *ProxyTarget
+	err    error
+}
+
+func (p *testProvider) Next(c echo.Context) *ProxyTarget {
+	return &ProxyTarget{}
+}
+
+func (p *testProvider) NextTarget(c echo.Context) (*ProxyTarget, error) {
+	return p.target, p.err
+}
+
+func TestTargetProvider(t *testing.T) {
+	t1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(w, "target 1")
+	}))
+	defer t1.Close()
+	url1, _ := url.Parse(t1.URL)
+
+	e := echo.New()
+	tp := &testProvider{commonBalancer: new(commonBalancer)}
+	tp.target = &ProxyTarget{Name: "target 1", URL: url1}
+	e.Use(Proxy(tp))
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	e.ServeHTTP(rec, req)
+	body := rec.Body.String()
+	assert.Equal(t, "target 1", body)
+}
+
+func TestFailNextTarget(t *testing.T) {
+	url1, err := url.Parse("http://dummy:8080")
+	assert.Nil(t, err)
+
+	e := echo.New()
+	tp := &testProvider{commonBalancer: new(commonBalancer)}
+	tp.target = &ProxyTarget{Name: "target 1", URL: url1}
+	tp.err = echo.NewHTTPError(http.StatusInternalServerError, "method could not select target")
+
+	e.Use(Proxy(tp))
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	e.ServeHTTP(rec, req)
+	body := rec.Body.String()
+	assert.Equal(t, "{\"message\":\"method could not select target\"}\n", body)
+}
+
 func TestProxyRealIPHeader(t *testing.T) {
 	// Setup
 	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))