Skip to content

Commit

Permalink
support anonymous user with sticky cookie
Browse files Browse the repository at this point in the history
  • Loading branch information
zensh committed Jul 8, 2020
1 parent 73cec76 commit d8afb49
Show file tree
Hide file tree
Showing 5 changed files with 118 additions and 8 deletions.
2 changes: 1 addition & 1 deletion pkg/config/dynamic/http_config.go
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ type WRRService struct {
// +k8s:deepcopy-gen=true

// LabeledRoundRobin defines a labeled load-balancer of services, which select service by label.
// Label will be extract from request header or cookie, with key `X-Canary-Label`.
// Label will be extract from request header or cookie, with key `X-Canary`.
// services should be named as `{defaultService}-{label}`. Ex. "myservice-stable", "myservice-beta", "myservice-dev"
type LabeledRoundRobin struct {
ServiceName string `json:"serviceName,omitempty" toml:"serviceName,omitempty" yaml:"serviceName,omitempty"`
Expand Down
1 change: 1 addition & 0 deletions pkg/config/dynamic/middlewares.go
Original file line number Diff line number Diff line change
Expand Up @@ -516,4 +516,5 @@ type Canary struct {
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"`
}
85 changes: 79 additions & 6 deletions pkg/middlewares/canary/canary.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,12 @@ package canary

import (
"context"
"crypto/sha1"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"regexp"
"strings"
Expand All @@ -14,6 +17,7 @@ import (
"github.com/containous/traefik/v2/pkg/log"
"github.com/containous/traefik/v2/pkg/middlewares"
"github.com/containous/traefik/v2/pkg/middlewares/accesslog"
"github.com/containous/traefik/v2/pkg/server/cookie"
"github.com/opentracing/opentracing-go"
"github.com/opentracing/opentracing-go/ext"
)
Expand Down Expand Up @@ -42,6 +46,7 @@ type Canary struct {
canaryResponseHeader bool
loadLabels bool
ls *LabelStore
sticky *dynamic.Sticky
next http.Handler
}

Expand Down Expand Up @@ -74,7 +79,16 @@ func New(ctx context.Context, next http.Handler, cfg dynamic.Canary, name string
loadLabels: cfg.Server != "",
addRequestID: cfg.AddRequestID,
canaryResponseHeader: cfg.CanaryResponseHeader,
sticky: cfg.Sticky,
}

if cfg.Sticky != nil {
c.sticky.Cookie.Name = cookie.GetName(cfg.Sticky.Cookie.Name, name)
if !strSliceHas(c.uidCookies, c.sticky.Cookie.Name) {
c.uidCookies = append(c.uidCookies, c.sticky.Cookie.Name)
}
}

if c.loadLabels {
c.ls = NewLabelStore(logger, cfg, expiration, cacheCleanDuration)
}
Expand Down Expand Up @@ -142,6 +156,18 @@ func (c *Canary) processCanary(rw http.ResponseWriter, req *http.Request) {
info.product = c.product
info.uid = extractUserID(req, c.uidCookies)

if info.uid == "" && c.sticky != nil {
addr := req.Header.Get("X-Real-Ip")
if addr == "" {
addr = req.Header.Get("X-Forwarded-For")
}
if addr == "" {
addr, _, _ = net.SplitHostPort(req.RemoteAddr)
}
info.uid = anonymousID(addr, req.Header.Get(headerUA), req.Header.Get("Cookie"), time.Now().Format(time.RFC822))
c.addSticky(info.uid, rw)
}

if info.label == "" && info.uid != "" {
labels := c.ls.MustLoadLabels(req.Context(), info.uid, req.Header.Get(headerXRequestID))
for _, l := range labels {
Expand All @@ -167,13 +193,27 @@ func (c *Canary) processCanary(rw http.ResponseWriter, req *http.Request) {
}
}

func (c *Canary) addSticky(id string, rw http.ResponseWriter) {
if data, err := json.Marshal(userInfo{UID5: id}); err == nil {
http.SetCookie(rw, &http.Cookie{
Name: c.sticky.Cookie.Name,
Value: base64.RawURLEncoding.EncodeToString(data),
Path: "/",
MaxAge: 60 * 60 * 24 * 7,
Secure: c.sticky.Cookie.Secure,
HttpOnly: c.sticky.Cookie.HTTPOnly,
SameSite: convertSameSite(c.sticky.Cookie.SameSite),
})
}
}

type userInfo struct {
UID0 string `json:"uid"`
UID1 string `json:"_userId"`
UID2 string `json:"userId"`
UID3 string `json:"user_id"`
UID4 string `json:"sub"`
UID5 string `json:"id"`
UID0 string `json:"uid,omitempty"`
UID1 string `json:"_userId,omitempty"`
UID2 string `json:"userId,omitempty"`
UID3 string `json:"user_id,omitempty"`
UID4 string `json:"sub,omitempty"`
UID5 string `json:"id,omitempty"`
}

func extractUserID(req *http.Request, uidCookies []string) string {
Expand Down Expand Up @@ -323,6 +363,9 @@ func (ch *canaryHeader) feed(vals []string, trust bool) {
}
}
}
if ch.testing && ch.label == "" {
ch.label = "testing"
}
}

// label should not be empty
Expand Down Expand Up @@ -358,3 +401,33 @@ func (ch *canaryHeader) String() string {
}
return strings.Join(vals, ",")
}

func convertSameSite(sameSite string) http.SameSite {
switch sameSite {
case "none":
return http.SameSiteNoneMode
case "lax":
return http.SameSiteLaxMode
case "strict":
return http.SameSiteStrictMode
default:
return 0
}
}

func anonymousID(feeds ...string) string {
h := sha1.New()
for _, v := range feeds {
io.WriteString(h, v)
}
return fmt.Sprintf("anon-%x", h.Sum(nil))
}

func strSliceHas(s []string, t string) bool {
for _, v := range s {
if v == t {
return true
}
}
return false
}
36 changes: 36 additions & 0 deletions pkg/middlewares/canary/canary_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -302,4 +302,40 @@ func TestCanary(t *testing.T) {
a.Equal("iOS", ch.client)
a.Equal("someuid", ch.uid)
})

t.Run("sticky should work", func(t *testing.T) {
a := assert.New(t)

cfg := dynamic.Canary{MaxCacheSize: 3, Server: "localhost", Product: "Urbs", Sticky: &dynamic.Sticky{
Cookie: &dynamic.Cookie{Name: "_urbs_"},
}}
c, err := New(context.Background(), next, cfg, "test")
c.ls.mustFetchLabels = func(ctx context.Context, uid, requestID string) ([]Label, int64) {
return []Label{{Label: uid}}, time.Now().Unix()
}
a.Nil(err)

req := httptest.NewRequest("GET", "http://example.com/foo", nil)
rw := httptest.NewRecorder()
c.processCanary(rw, req)
ch := &canaryHeader{}
ch.fromHeader(req.Header, true)
a.NotEqual("", ch.label)
a.Equal(ch.uid, ch.label)
a.Contains(rw.Header().Get("Set-Cookie"), "_urbs_=ey")

uid := ch.uid
cookies := rw.Result().Cookies()
a.Equal(1, len(cookies))
a.Equal("_urbs_", cookies[0].Name)

req = httptest.NewRequest("GET", "http://example.com/foo", nil)
req.AddCookie(cookies[0])
rw = httptest.NewRecorder()
c.processCanary(rw, req)
ch = &canaryHeader{}
ch.fromHeader(req.Header, true)
a.Equal(uid, ch.label)
a.Equal(ch.uid, ch.label)
})
}
2 changes: 1 addition & 1 deletion pkg/provider/kubernetes/crd/traefik/v1alpha1/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ type WeightedRoundRobin struct {
// +k8s:deepcopy-gen=true

// LabeledRoundRobin defines a labeled load-balancer of services, which select service by label.
// Label will be extract from request header or cookie, with key `X-Canary-Label`.
// Label will be extract from request header or cookie, with key `X-Canary`.
// services should be named as `{defaultService}-{label}`. Ex. "myservice-stable", "myservice-beta", "myservice-dev"
type LabeledRoundRobin struct {
Service
Expand Down

0 comments on commit d8afb49

Please sign in to comment.