From c50a3aa2adc85031b7d560f18889222d75c4fee7 Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Tue, 25 Oct 2022 21:41:46 +0300 Subject: [PATCH 1/6] [suggestion] Add helper interface for ProxyBalancer interface --- middleware/proxy.go | 17 ++++++++++++++++- middleware/proxy_test.go | 26 ++++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/middleware/proxy.go b/middleware/proxy.go index 6cfd6731e..9d857edaf 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 { + GetTarget(echo.Context) (*ProxyTarget, error) + } + commonBalancer struct { targets []*ProxyTarget mutex sync.RWMutex @@ -231,7 +236,17 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc { req := c.Request() res := c.Response() - tgt := config.Balancer.Next(c) + + var tgt *ProxyTarget + provider, ok := config.Balancer.(TargetProvider) + if ok { + tgt, err = provider.GetTarget(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 { diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index 7939fc5c2..cd03c1e5c 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -18,7 +18,19 @@ import ( "github.com/stretchr/testify/assert" ) -//Assert expected with url.EscapedPath method to obtain the path. +type testProvider struct { + *commonBalancer +} + +func (p *testProvider) Next(c echo.Context) *ProxyTarget { + return p.commonBalancer.targets[0] +} + +func (p *testProvider) GetTarget(c echo.Context) (*ProxyTarget, error) { + return p.commonBalancer.targets[1], nil +} + +// 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 +43,7 @@ func TestProxy(t *testing.T) { })) defer t2.Close() url2, _ := url.Parse(t2.URL) - + fmt.Println(url1, url2) targets := []*ProxyTarget{ { Name: "target 1", @@ -89,6 +101,16 @@ func TestProxy(t *testing.T) { body = rec.Body.String() assert.Equal(t, "target 2", body) + // Provider + e = echo.New() + tp := &testProvider{commonBalancer: new(commonBalancer)} + tp.targets = targets + e.Use(Proxy(tp)) + rec = httptest.NewRecorder() + e.ServeHTTP(rec, req) + body = rec.Body.String() + assert.Equal(t, "target 2", body) + // ModifyResponse e = echo.New() e.Use(ProxyWithConfig(ProxyConfig{ From eceed58164a0fd7e9a6ec34f7ea5e8c4447fbc3f Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Tue, 25 Oct 2022 21:48:06 +0300 Subject: [PATCH 2/6] Update proxy_test.go --- middleware/proxy_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index cd03c1e5c..686dc258d 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -43,7 +43,6 @@ func TestProxy(t *testing.T) { })) defer t2.Close() url2, _ := url.Parse(t2.URL) - fmt.Println(url1, url2) targets := []*ProxyTarget{ { Name: "target 1", From d70346c0b39011e531ecdce54cf52402ed9ac5f7 Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Tue, 25 Oct 2022 22:50:32 +0300 Subject: [PATCH 3/6] addressed code review comments --- middleware/proxy.go | 8 +++--- middleware/proxy_test.go | 58 +++++++++++++++++++++++++--------------- 2 files changed, 40 insertions(+), 26 deletions(-) diff --git a/middleware/proxy.go b/middleware/proxy.go index 9d857edaf..d2cd2aa6d 100644 --- a/middleware/proxy.go +++ b/middleware/proxy.go @@ -74,7 +74,7 @@ type ( // TargetProvider defines an interface that gives the opportunity for balancer to return custom errors when selecting target. TargetProvider interface { - GetTarget(echo.Context) (*ProxyTarget, error) + NextTarget(echo.Context) (*ProxyTarget, error) } commonBalancer struct { @@ -228,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) { @@ -238,9 +239,8 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc { res := c.Response() var tgt *ProxyTarget - provider, ok := config.Balancer.(TargetProvider) - if ok { - tgt, err = provider.GetTarget(c) + if isTargetProvider { + tgt, err = provider.NextTarget(c) if err != nil { return err } diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index 686dc258d..20552dcd4 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -18,18 +18,6 @@ import ( "github.com/stretchr/testify/assert" ) -type testProvider struct { - *commonBalancer -} - -func (p *testProvider) Next(c echo.Context) *ProxyTarget { - return p.commonBalancer.targets[0] -} - -func (p *testProvider) GetTarget(c echo.Context) (*ProxyTarget, error) { - return p.commonBalancer.targets[1], nil -} - // Assert expected with url.EscapedPath method to obtain the path. func TestProxy(t *testing.T) { // Setup @@ -100,16 +88,6 @@ func TestProxy(t *testing.T) { body = rec.Body.String() assert.Equal(t, "target 2", body) - // Provider - e = echo.New() - tp := &testProvider{commonBalancer: new(commonBalancer)} - tp.targets = targets - e.Use(Proxy(tp)) - rec = httptest.NewRecorder() - e.ServeHTTP(rec, req) - body = rec.Body.String() - assert.Equal(t, "target 2", body) - // ModifyResponse e = echo.New() e.Use(ProxyWithConfig(ProxyConfig{ @@ -143,6 +121,42 @@ func TestProxy(t *testing.T) { e.ServeHTTP(rec, req) } +type testProvider struct { + *commonBalancer +} + +func (p *testProvider) Next(c echo.Context) *ProxyTarget { + return &ProxyTarget{} +} + +func (p *testProvider) NextTarget(c echo.Context) (*ProxyTarget, error) { + return p.commonBalancer.targets[0], nil +} + +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) + targets := []*ProxyTarget{ + { + Name: "target 1", + URL: url1, + }, + } + + e := echo.New() + tp := &testProvider{commonBalancer: new(commonBalancer)} + tp.targets = targets + 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 TestProxyRealIPHeader(t *testing.T) { // Setup upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) From bac6d80d5ac6a631b297749f05972e16bae697df Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Thu, 27 Oct 2022 21:48:42 +0300 Subject: [PATCH 4/6] address pr comments --- middleware/proxy.go | 3 ++- middleware/proxy_test.go | 30 ++++++++++++++++++++++-------- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/middleware/proxy.go b/middleware/proxy.go index d2cd2aa6d..b94e1320e 100644 --- a/middleware/proxy.go +++ b/middleware/proxy.go @@ -242,7 +242,8 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc { if isTargetProvider { tgt, err = provider.NextTarget(c) if err != nil { - return err + c.String(http.StatusInternalServerError, err.Error()) + return } } else { tgt = config.Balancer.Next(c) diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index 20552dcd4..41433ea7a 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -123,6 +123,8 @@ func TestProxy(t *testing.T) { type testProvider struct { *commonBalancer + target *ProxyTarget + err error } func (p *testProvider) Next(c echo.Context) *ProxyTarget { @@ -130,7 +132,7 @@ func (p *testProvider) Next(c echo.Context) *ProxyTarget { } func (p *testProvider) NextTarget(c echo.Context) (*ProxyTarget, error) { - return p.commonBalancer.targets[0], nil + return p.target, p.err } func TestTargetProvider(t *testing.T) { @@ -139,16 +141,10 @@ func TestTargetProvider(t *testing.T) { })) defer t1.Close() url1, _ := url.Parse(t1.URL) - targets := []*ProxyTarget{ - { - Name: "target 1", - URL: url1, - }, - } e := echo.New() tp := &testProvider{commonBalancer: new(commonBalancer)} - tp.targets = targets + tp.target = &ProxyTarget{Name: "target 1", URL: url1} e.Use(Proxy(tp)) rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/", nil) @@ -157,6 +153,24 @@ func TestTargetProvider(t *testing.T) { 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() + msg := "method could not select target" + tp := &testProvider{commonBalancer: new(commonBalancer)} + tp.target = &ProxyTarget{Name: "target 1", URL: url1} + tp.err = fmt.Errorf("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, msg, body) +} + func TestProxyRealIPHeader(t *testing.T) { // Setup upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) From aad17b624d6eb0cd3ede048abc8d1cc241f0c356 Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Thu, 27 Oct 2022 21:50:20 +0300 Subject: [PATCH 5/6] clean up --- middleware/proxy_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index 41433ea7a..e35d98ad4 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -161,7 +161,7 @@ func TestFailNextTarget(t *testing.T) { msg := "method could not select target" tp := &testProvider{commonBalancer: new(commonBalancer)} tp.target = &ProxyTarget{Name: "target 1", URL: url1} - tp.err = fmt.Errorf("method could not select target") + tp.err = fmt.Errorf(msg) e.Use(Proxy(tp)) rec := httptest.NewRecorder() From 2c28513d5c8da3ac48750f5a6b9a61479dc7865d Mon Sep 17 00:00:00 2001 From: Hristo Hristov Date: Thu, 27 Oct 2022 22:46:39 +0300 Subject: [PATCH 6/6] return error --- middleware/proxy.go | 3 +-- middleware/proxy_test.go | 5 ++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/middleware/proxy.go b/middleware/proxy.go index b94e1320e..d2cd2aa6d 100644 --- a/middleware/proxy.go +++ b/middleware/proxy.go @@ -242,8 +242,7 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc { if isTargetProvider { tgt, err = provider.NextTarget(c) if err != nil { - c.String(http.StatusInternalServerError, err.Error()) - return + return err } } else { tgt = config.Balancer.Next(c) diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go index e35d98ad4..0ded50a1f 100644 --- a/middleware/proxy_test.go +++ b/middleware/proxy_test.go @@ -158,17 +158,16 @@ func TestFailNextTarget(t *testing.T) { assert.Nil(t, err) e := echo.New() - msg := "method could not select target" tp := &testProvider{commonBalancer: new(commonBalancer)} tp.target = &ProxyTarget{Name: "target 1", URL: url1} - tp.err = fmt.Errorf(msg) + 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, msg, body) + assert.Equal(t, "{\"message\":\"method could not select target\"}\n", body) } func TestProxyRealIPHeader(t *testing.T) {