From 9e24f89bf440afc24ebec2a4419382295fcf854c Mon Sep 17 00:00:00 2001
From: zensh <admin@zensh.com>
Date: Fri, 10 Jul 2020 16:58:11 +0800
Subject: [PATCH] canary middleware support config labels map

---
 pkg/config/dynamic/middlewares.go | 25 ++++++++++++++++---------
 pkg/middlewares/canary/canary.go  | 20 ++++++++++++++++++++
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/pkg/config/dynamic/middlewares.go b/pkg/config/dynamic/middlewares.go
index eadef1110b..c4e8be4142 100644
--- a/pkg/config/dynamic/middlewares.go
+++ b/pkg/config/dynamic/middlewares.go
@@ -508,13 +508,20 @@ func (c *ClientTLS) CreateTLSConfig() (*tls.Config, error) {
 
 // Canary middleware settings.
 type Canary struct {
-	Product              string         `json:"product,omitempty" toml:"product,omitempty" yaml:"product,omitempty" export:"true"`
-	Server               string         `json:"server,omitempty" toml:"server,omitempty" yaml:"server,omitempty" export:"true"`
-	UIDCookies           []string       `json:"uidCookies,omitempty" toml:"uidCookies,omitempty" yaml:"uidCookies,omitempty" export:"true"`
-	AddRequestID         bool           `json:"addRequestID,omitempty" toml:"addRequestID,omitempty" yaml:"addRequestID,omitempty" export:"true"`
-	CanaryResponseHeader bool           `json:"canaryResponseHeader,omitempty" toml:"canaryResponseHeader,omitempty" yaml:"canaryResponseHeader,omitempty" export:"true"`
-	MaxCacheSize         int            `json:"maxCacheSize,omitempty" toml:"maxCacheSize,omitempty" yaml:"maxCacheSize,omitempty" export:"true"`
-	CacheExpiration      types.Duration `json:"cacheExpiration,omitempty" toml:"cacheExpiration,omitempty" yaml:"cacheExpiration,omitempty" export:"true"`
-	CacheCleanDuration   types.Duration `json:"cacheCleanDuration,omitempty" toml:"cacheCleanDuration,omitempty" yaml:"cacheCleanDuration,omitempty" export:"true"`
-	Sticky               *Sticky        `json:"sticky,omitempty" toml:"sticky,omitempty" yaml:"sticky,omitempty" export:"true"`
+	Product              string         `json:"product,omitempty" toml:"product,omitempty" yaml:"product,omitempty"`
+	Server               string         `json:"server,omitempty" toml:"server,omitempty" yaml:"server,omitempty"`
+	UIDCookies           []string       `json:"uidCookies,omitempty" toml:"uidCookies,omitempty" yaml:"uidCookies,omitempty"`
+	AddRequestID         bool           `json:"addRequestID,omitempty" toml:"addRequestID,omitempty" yaml:"addRequestID,omitempty"`
+	CanaryResponseHeader bool           `json:"canaryResponseHeader,omitempty" toml:"canaryResponseHeader,omitempty" yaml:"canaryResponseHeader,omitempty"`
+	MaxCacheSize         int            `json:"maxCacheSize,omitempty" toml:"maxCacheSize,omitempty" yaml:"maxCacheSize,omitempty"`
+	CacheExpiration      types.Duration `json:"cacheExpiration,omitempty" toml:"cacheExpiration,omitempty" yaml:"cacheExpiration,omitempty"`
+	CacheCleanDuration   types.Duration `json:"cacheCleanDuration,omitempty" toml:"cacheCleanDuration,omitempty" yaml:"cacheCleanDuration,omitempty"`
+	Sticky               *Sticky        `json:"sticky,omitempty" toml:"sticky,omitempty" yaml:"sticky,omitempty"`
+	LabelsMap            *LabelsMap     `json:"labelsMap,omitempty" toml:"labelsMap,omitempty" yaml:"labelsMap,omitempty"`
+}
+
+// LabelsMap get canary labels from header with map
+type LabelsMap struct {
+	RequestHeaderName string            `json:"requestHeaderName,omitempty" toml:"requestHeaderName,omitempty" yaml:"requestHeaderName,omitempty"`
+	Labels            map[string]string `json:"labels,omitempty" toml:"labels,omitempty" yaml:"labels,omitempty"`
 }
diff --git a/pkg/middlewares/canary/canary.go b/pkg/middlewares/canary/canary.go
index f1f6e07708..a0338db701 100644
--- a/pkg/middlewares/canary/canary.go
+++ b/pkg/middlewares/canary/canary.go
@@ -47,6 +47,7 @@ type Canary struct {
 	loadLabels           bool
 	ls                   *LabelStore
 	sticky               *dynamic.Sticky
+	labelsMap            *dynamic.LabelsMap
 	next                 http.Handler
 }
 
@@ -71,6 +72,12 @@ func New(ctx context.Context, next http.Handler, cfg dynamic.Canary, name string
 		cfg.MaxCacheSize = defaultCacheSize
 	}
 
+	if cfg.LabelsMap != nil {
+		if cfg.LabelsMap.RequestHeaderName == "" || len(cfg.LabelsMap.Labels) == 0 {
+			cfg.LabelsMap = nil
+		}
+	}
+
 	c := &Canary{
 		name:                 name,
 		next:                 next,
@@ -80,6 +87,7 @@ func New(ctx context.Context, next http.Handler, cfg dynamic.Canary, name string
 		addRequestID:         cfg.AddRequestID,
 		canaryResponseHeader: cfg.CanaryResponseHeader,
 		sticky:               cfg.Sticky,
+		labelsMap:            cfg.LabelsMap,
 	}
 
 	if cfg.Sticky != nil {
@@ -147,15 +155,26 @@ func (c *Canary) processCanary(rw http.ResponseWriter, req *http.Request) {
 	} else {
 		// load user's labels and update to header when work as public gateway.
 		info.fromHeader(req.Header, false)
+
+		// try load labels from cookie when not exists in request X-Canary header.
 		if info.label == "" {
 			if cookie, _ := req.Cookie(headerXCanary); cookie != nil && cookie.Value != "" {
 				info.feed(strings.Split(cookie.Value, ","), false)
 			}
 		}
 
+		// try load labels from config with header when not exists.
+		if info.label == "" && c.labelsMap != nil {
+			key := req.Header.Get(c.labelsMap.RequestHeaderName)
+			if vals := c.labelsMap.Labels[key]; vals != "" {
+				info.feed(strings.Split(vals, ","), false)
+			}
+		}
+
 		info.product = c.product
 		info.uid = extractUserID(req, c.uidCookies)
 
+		// anonymous user
 		if info.uid == "" && c.sticky != nil {
 			addr := req.Header.Get("X-Real-Ip")
 			if addr == "" {
@@ -168,6 +187,7 @@ func (c *Canary) processCanary(rw http.ResponseWriter, req *http.Request) {
 			c.addSticky(info.uid, rw)
 		}
 
+		// try load labels from server
 		if info.label == "" && info.uid != "" {
 			labels := c.ls.MustLoadLabels(req.Context(), info.uid, req.Header.Get(headerXRequestID))
 			for _, l := range labels {