From c7d6d4373fdfbef5d6f44df0a8ef410c198420ee Mon Sep 17 00:00:00 2001
From: Kai Ratzeburg <hello@kai-ratzeburg.de>
Date: Sun, 5 Nov 2023 17:01:01 +0100
Subject: [PATCH] proxy middleware: reuse echo request context (#2537)

---
 middleware/proxy.go      |  4 +++
 middleware/proxy_test.go | 60 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/middleware/proxy.go b/middleware/proxy.go
index e4f98d9e..16b00d64 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():
diff --git a/middleware/proxy_test.go b/middleware/proxy_test.go
index 415d68e7..1c93ba03 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"))
+}