support canary response header; support custom labels server URL.

teambition · Mar 16, 2020 · deb9649 · deb9649
1 parent 4d69b83
commit deb9649
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 27 deletions.
diff --git a/pkg/config/dynamic/middlewares.go b/pkg/config/dynamic/middlewares.go
@@ -515,11 +515,12 @@ 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"`
-	Cookie             string         `json:"cookie,omitempty" toml:"cookie,omitempty" yaml:"cookie,omitempty" export:"true"`
-	AddRequestID       bool           `json:"addRequestID,omitempty" toml:"addRequestID,omitempty" yaml:"addRequestID,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"`
+	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"`
+	Cookie               string         `json:"cookie,omitempty" toml:"cookie,omitempty" yaml:"cookie,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"`
 }
diff --git a/pkg/middlewares/canary/canary.go b/pkg/middlewares/canary/canary.go
@@ -35,12 +35,13 @@ var validLabelReg = regexp.MustCompile(`^[a-z][0-9a-z-]{1,62}$`)
 
 // Canary ...
 type Canary struct {
-	name         string
-	product      string
-	cookie       string
-	addRequestID bool
-	ls           *LabelStore
-	next         http.Handler
+	name                 string
+	product              string
+	cookie               string
+	addRequestID         bool
+	canaryResponseHeader bool
+	ls                   *LabelStore
+	next                 http.Handler
 }
 
 // New returns a Canary instance.
@@ -69,8 +70,8 @@ func New(ctx context.Context, next http.Handler, cfg dynamic.Canary, name string
 	}
 
 	ls := NewLabelStore(logger, cfg, expiration, cacheCleanDuration)
-	return &Canary{name: name, product: cfg.Product,
-		cookie: cfg.Cookie, addRequestID: cfg.AddRequestID, ls: ls, next: next}, nil
+	return &Canary{name: name, product: cfg.Product, cookie: cfg.Cookie,
+		addRequestID: cfg.AddRequestID, canaryResponseHeader: cfg.CanaryResponseHeader, ls: ls, next: next}, nil
 }
 
 // GetTracingInformation implements Tracable interface
@@ -80,7 +81,7 @@ func (c *Canary) GetTracingInformation() (string, ext.SpanKindEnum) {
 
 func (c *Canary) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	c.processRequestID(req)
-	c.processCanary(req)
+	c.processCanary(rw, req)
 	c.next.ServeHTTP(rw, req)
 }
 
@@ -94,7 +95,7 @@ func (c *Canary) processRequestID(req *http.Request) {
 	}
 }
 
-func (c *Canary) processCanary(req *http.Request) {
+func (c *Canary) processCanary(rw http.ResponseWriter, req *http.Request) {
 	info := &canaryHeader{}
 	info.fromHeader(req.Header, false)
 
@@ -121,6 +122,9 @@ func (c *Canary) processCanary(req *http.Request) {
 		}
 	}
 	info.intoHeader(req.Header)
+	if c.canaryResponseHeader {
+		info.intoHeader(rw.Header())
+	}
 }
 
 type userInfo struct {

diff --git a/pkg/middlewares/canary/canary_test.go b/pkg/middlewares/canary/canary_test.go
@@ -175,55 +175,78 @@ func TestCanary(t *testing.T) {
 		a.Nil(err)
 
 		req := httptest.NewRequest("GET", "http://example.com/foo", nil)
-		c.processCanary(req)
+		rw := httptest.NewRecorder()
+		c.processCanary(rw, req)
 		ch := &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("", ch.label)
 
 		req = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rw = httptest.NewRecorder()
 		req.Header.Set(headerXCanary, "stable")
-		c.processCanary(req)
+		c.processCanary(rw, req)
 		ch = &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("stable", ch.label)
 		a.Equal("Urbs", ch.product)
 
 		req = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rw = httptest.NewRecorder()
 		req.Header.Set(headerXCanary, "label=beta")
-		c.processCanary(req)
+		c.processCanary(rw, req)
 		ch = &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("beta", ch.label)
 		a.Equal("Urbs", ch.product)
 
 		req = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rw = httptest.NewRecorder()
 		req.AddCookie(&http.Cookie{Name: headerXCanary, Value: "beta"})
-		c.processCanary(req)
+		c.processCanary(rw, req)
 		ch = &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("beta", ch.label)
 		a.Equal("Urbs", ch.product)
+		ch = &canaryHeader{}
+		ch.fromHeader(rw.Header(), true)
+		a.Equal("", ch.label)
+		a.Equal("", ch.product)
+
+		c.canaryResponseHeader = true
 
 		req = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rw = httptest.NewRecorder()
 		req.Header.Set(headerXCanary, "label=beta")
 		req.Header.Add(headerXCanary, "client=iOS")
 		req.AddCookie(&http.Cookie{Name: headerXCanary, Value: "stable"})
-		c.processCanary(req)
+		c.processCanary(rw, req)
 		ch = &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("beta", ch.label)
 		a.Equal("Urbs", ch.product)
 		a.Equal("iOS", ch.client)
+		ch = &canaryHeader{}
+		ch.fromHeader(rw.Header(), true)
+		a.Equal("beta", ch.label)
+		a.Equal("Urbs", ch.product)
+		a.Equal("iOS", ch.client)
 
 		req = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		rw = httptest.NewRecorder()
 		req.Header.Set("Authorization", fmt.Sprintf("OAuth %s", testToken))
 		req.Header.Set(headerXCanary, "client=iOS")
-		c.processCanary(req)
+		c.processCanary(rw, req)
 		ch = &canaryHeader{}
 		ch.fromHeader(req.Header, true)
 		a.Equal("someuid", ch.label)
 		a.Equal("Urbs", ch.product)
 		a.Equal("iOS", ch.client)
 		a.Equal("someuid", ch.uid)
+		ch = &canaryHeader{}
+		ch.fromHeader(rw.Header(), true)
+		a.Equal("someuid", ch.label)
+		a.Equal("Urbs", ch.product)
+		a.Equal("iOS", ch.client)
+		a.Equal("someuid", ch.uid)
 	})
 }
diff --git a/pkg/middlewares/canary/label.go b/pkg/middlewares/canary/label.go
@@ -3,6 +3,7 @@ package canary
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -76,10 +77,13 @@ func NewLabelStore(logger log.Logger, cfg dynamic.Canary, expiration, cacheClean
 
 	product := cfg.Product
 	apiURL := cfg.Server
-	if apiURL[len(apiURL)-1] == '/' {
-		apiURL = apiURL[:len(apiURL)-1]
+	// apiURL ex. https://labelServerHost/api/labels?uid=%s&product=%s
+	if !strings.Contains(apiURL, "%s") { // append default API path.
+		if apiURL[len(apiURL)-1] == '/' {
+			apiURL = apiURL[:len(apiURL)-1]
+		}
+		apiURL += "/users/%s/labels:cache?product=%s"
 	}
-	apiURL += "/users/%s/labels:cache?product=%s"
 
 	ls.mustFetchLabels = func(ctx context.Context, uid, requestID string) ([]Label, int64) {
 		url := fmt.Sprintf(apiURL, uid, product)

diff --git a/pkg/middlewares/canary/request.go b/pkg/middlewares/canary/request.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/containous/traefik/v2/pkg/log"
 	"github.com/containous/traefik/v2/pkg/tracing"
+	"github.com/containous/traefik/v2/pkg/version"
 	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
 )
@@ -22,7 +23,7 @@ func init() {
 	if hostname == "" {
 		hostname = "unknown"
 	}
-	userAgent = fmt.Sprintf("golang/%v hostname/%s Traefik Canary Middleware", runtime.Version(), hostname)
+	userAgent = fmt.Sprintf("Go/%v Hostname/%s Traefik/%s (Canary Middleware)", runtime.Version(), hostname, version.Version)
 }
 
 var userAgent string