Repository: labstack/echo
Pull request #2537 (2023-10-27 14:14:37)

proxy middleware: reuse echo request context
I have used the proxy middleware in one of my projects and need the context values in modifyResponse, which I had set before in my custom balancer.
Unfortunately, I had to realise that the context does not seem to be taken over and there is no option to get it.
With this change, the context from `http.Request` is reused.

If there is an alternative way to get the previous request context, please tell me.


x1h0: Hi @aldas, could you please check this PR?
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/labstack/echo/pull/2537?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`69a0de8`)](https://app.codecov.io/gh/labstack/echo/commit/69a0de84158fd7cad326599d145c2248bcc15a69?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) 92.88% compared to head [(`27a8c74`)](https://app.codecov.io/gh/labstack/echo/pull/2537?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack) 92.89%.

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


```diff
@@           Coverage Diff           @@
##           master    #2537   +/-   ##
=======================================
  Coverage   92.88%   92.89%           
=======================================
  Files          39       39           
  Lines        4655     4658    +3     
=======================================
+ Hits         4324     4327    +3     
  Misses        240      240           
  Partials       91       91           
```


| [Files](https://app.codecov.io/gh/labstack/echo/pull/2537?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://app.codecov.io/gh/labstack/echo/pull/2537?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack#diff-bWlkZGxld2FyZS9wcm94eS5nbw==) | `74.75% <100.00%> (+0.38%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/labstack/echo/pull/2537?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=labstack).   
:loudspeaker: Have feedback on the report? [Share it here](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 e4f98d9..16b00d6 100644
--- a/middleware/proxy.go
+++ b/middleware/proxy.go
@@ -359,6 +359,10 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 					c.Set("_error", nil)
 				}
 
+				// This is needed for ProxyConfig.ModifyResponse and/or ProxyConfig.Transport to be able to process the Request
+				// that Balancer may have replaced with c.SetRequest.
+				req = c.Request()
+
 				// Proxy
 				switch {
 				case c.IsWebSocket():


Test patch:
diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go
index 415d68e..1c93ba0 100644
--- a/middleware/proxy_test.go
+++ b/middleware/proxy_test.go
@@ -747,3 +747,63 @@ func TestProxyBalancerWithNoTargets(t *testing.T) {
 	rrb := NewRoundRobinBalancer([]*ProxyTarget{})
 	assert.Nil(t, rrb.Next(nil))
 }
+
+type testContextKey string
+
+type customBalancer struct {
+	target *ProxyTarget
+}
+
+func (b *customBalancer) AddTarget(target *ProxyTarget) bool {
+	return false
+}
+
+func (b *customBalancer) RemoveTarget(name string) bool {
+	return false
+}
+
+func (b *customBalancer) Next(c echo.Context) *ProxyTarget {
+	ctx := context.WithValue(c.Request().Context(), testContextKey("FROM_BALANCER"), "CUSTOM_BALANCER")
+	c.SetRequest(c.Request().WithContext(ctx))
+	return b.target
+}
+
+func TestModifyResponseUseContext(t *testing.T) {
+	server := httptest.NewServer(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("OK"))
+		}),
+	)
+	defer server.Close()
+
+	targetURL, _ := url.Parse(server.URL)
+	e := echo.New()
+	e.Use(ProxyWithConfig(
+		ProxyConfig{
+			Balancer: &customBalancer{
+				target: &ProxyTarget{
+					Name: "tst",
+					URL:  targetURL,
+				},
+			},
+			RetryCount: 1,
+			ModifyResponse: func(res *http.Response) error {
+				val := res.Request.Context().Value(testContextKey("FROM_BALANCER"))
+				if valStr, ok := val.(string); ok {
+					res.Header.Set("FROM_BALANCER", valStr)
+				}
+				return nil
+			},
+		},
+	))
+
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+
+	e.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "OK", rec.Body.String())
+	assert.Equal(t, "CUSTOM_BALANCER", rec.Header().Get("FROM_BALANCER"))
+}