proxy middleware: reuse echo request context

[x1h0]

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.

[TomKrouss]

I have run into a similar problem, any update (either on the change or an alternative) would be appreciated.

[x1h0]

Hi @aldas, could you please check this PR?

[aldas]

ok, I think I see your point.

package main

import (
	"context"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
)

type customBalancer struct {
	target *middleware.ProxyTarget
}

func (b *customBalancer) AddTarget(target *middleware.ProxyTarget) bool {
	panic("not implemented")
	return true
}

func (b *customBalancer) RemoveTarget(name string) bool {
	panic("not implemented")
	return false
}

func (b *customBalancer) Next(c echo.Context) *middleware.ProxyTarget {
	ctx := context.WithValue(c.Request().Context(), "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(middleware.ProxyWithConfig(
		middleware.ProxyConfig{
			Balancer: &customBalancer{
				target: &middleware.ProxyTarget{
					Name: "tst",
					URL:  targetURL,
				},
			},
			RetryCount: 1,
			ModifyResponse: func(res *http.Response) error {
				val := res.Request.Context().Value("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"))
}

Is it something like that?

[aldas]

maybe

				// 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()

1. This comment says why and where if line is useful.
2. do we need req = req.WithContext() at all we could just retake that pointer (pointing to new request instance) from echo.Context?

[aldas]

also - please add at least one test

[x1h0]

Hi @aldas,

thanks for your review. I have adjusted the PR according to your comments.

Is it something like that?

Yes, your test does exactly what I am trying to do.

This comment says why and where if line is useful.

I have adopted the comment.

do we need req = req.WithContext() at all we could just retake that pointer (pointing to new request instance) from echo.Context?

No, we don't need it. Your suggestion req = c.Request() is perfectly adequate.

also - please add at least one test

I have adopted your test and fixed the context key lint error. Is that okay?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (69a0de8) 92.88% compared to head (27a8c74) 92.89%.

Additional details and impacted files

@@           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	Coverage Δ
middleware/proxy.go	74.75% <100.00%> (+0.38%)	⬆️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[aldas]

LGTM

[aldas]

I will look if we have some other PRs to merge and maybe do a patch release containing this PR.

[aldas]

New version is out: https://github.com/labstack/echo/releases/tag/v4.11.3