extend proxy middleware. closes #1202 (#1203)

* extend proxy middleware. closes #1202

middleware/proxy.go

@@ -38,6 +38,10 @@ type (
 		// "/users/*/orders/*": "/user/$1/order/$2",
 		Rewrite map[string]string
 
+		// Context key to store selected ProxyTarget into context.
+		// Optional. Default value "target".
+		ContextKey string
+
 		rewriteRegex map[*regexp.Regexp]string
 	}
 
@@ -45,13 +49,14 @@ type (
 	ProxyTarget struct {
 		Name string
 		URL  *url.URL
+		Meta echo.Map
 	}
 
 	// ProxyBalancer defines an interface to implement a load balancing technique.
 	ProxyBalancer interface {
 		AddTarget(*ProxyTarget) bool
 		RemoveTarget(string) bool
-		Next() *ProxyTarget
+		Next(echo.Context) *ProxyTarget
 	}
 
 	commonBalancer struct {
@@ -76,6 +81,7 @@ var (
 	// DefaultProxyConfig is the default Proxy middleware config.
 	DefaultProxyConfig = ProxyConfig{
 		Skipper:    DefaultSkipper,
+		ContextKey: "target",
 	}
 )
 
@@ -164,7 +170,7 @@ func (b *commonBalancer) RemoveTarget(name string) bool {
 }
 
 // Next randomly returns an upstream target.
-func (b *randomBalancer) Next() *ProxyTarget {
+func (b *randomBalancer) Next(c echo.Context) *ProxyTarget {
 	if b.random == nil {
 		b.random = rand.New(rand.NewSource(int64(time.Now().Nanosecond())))
 	}
@@ -174,7 +180,7 @@ func (b *randomBalancer) Next() *ProxyTarget {
 }
 
 // Next returns an upstream target using round-robin technique.
-func (b *roundRobinBalancer) Next() *ProxyTarget {
+func (b *roundRobinBalancer) Next(c echo.Context) *ProxyTarget {
 	b.i = b.i % uint32(len(b.targets))
 	t := b.targets[b.i]
 	atomic.AddUint32(&b.i, 1)
@@ -216,7 +222,8 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 
 			req := c.Request()
 			res := c.Response()
-			tgt := config.Balancer.Next()
+			tgt := config.Balancer.Next(c)
+			c.Set(config.ContextKey, tgt)
 
 			// Rewrite
 			for k, v := range config.rewriteRegex {

middleware/proxy_test.go

@@ -124,4 +124,19 @@ func TestProxy(t *testing.T) {
 	req.URL.Path = "/users/jack/orders/1"
 	e.ServeHTTP(rec, req)
 	assert.Equal(t, "/user/jack/order/1", req.URL.Path)
+
+	// ProxyTarget is set in context
+	contextObserver := func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) (err error) {
+			next(c)
+			assert.Contains(t, targets, c.Get("target"), "target is not set in context")
+			return nil
+		}
+	}
+	rrb1 := NewRoundRobinBalancer(targets)
+	e = echo.New()
+	e.Use(contextObserver)
+	e.Use(Proxy(rrb1))
+	rec = newCloseNotifyRecorder()
+	e.ServeHTTP(rec, req)
 }