introduce server match model for sticky session

doc/reference/filters.md

@@ -1128,8 +1128,10 @@ Rules to revise request header.
 
 | Name          | Type   | Description                                                                                                 | Required |
 | ------------- | ------ | ----------------------------------------------------------------------------------------------------------- | -------- |
-| mode          | string | Mode of session stickiness, only `CookieConsistentHash` is supported by now                                 | Yes      |
-| appCookieName | string | Name of the application cookie, its value will be used as the session identifier for stickiness             | Yes      |
+| mode          | string | Mode of session stickiness, support `CookieConsistentHash`,`DurationBased`,`ApplicationBased`                                 | Yes      |
+| appCookieName | string | Name of the application cookie, its value will be used as the session identifier for stickiness in `CookieConsistentHash` and `ApplicationBased` mode             | No      |
+| lbCookieName | string | Name of the cookie generated by load balancer, its value will be used as the session identifier for stickiness in `DurationBased` and `ApplicationBased` mode, default is `EG_SESSION`             | No      |
+| lbCookieExpire | string | Expire duration of the cookie generated by load balancer, its value will be used as the session expire time for stickiness in `DurationBased` and `ApplicationBased` mode, default is 2 hours             | No      |
 
 ### proxy.MemoryCacheSpec
 

pkg/filters/proxy/loadbalance.go

@@ -18,10 +18,17 @@
 package proxy
 
 import (
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/binary"
+	"encoding/hex"
 	"fmt"
 	"hash/fnv"
+	"hash/maphash"
 	"math/rand"
+	"net/http"
 	"sync/atomic"
+	"time"
 
 	"github.com/buraksezer/consistent"
 	"github.com/megaease/easegress/pkg/logger"
@@ -40,18 +47,35 @@ const (
 	LoadBalancePolicyIPHash = "ipHash"
 	// LoadBalancePolicyHeaderHash is the load balance policy of HTTP header hash.
 	LoadBalancePolicyHeaderHash = "headerHash"
+	// StickySessionModeCookieConsistentHash is the sticky session mode of consistent hash on app cookie.
+	StickySessionModeCookieConsistentHash = "CookieConsistentHash"
+	// StickySessionModeDurationBased uses a load balancer-generated cookie for stickiness.
+	StickySessionModeDurationBased = "DurationBased"
+	// StickySessionModeApplicationBased uses a load balancer-generated cookie depends on app cookie for stickiness.
+	StickySessionModeApplicationBased = "ApplicationBased"
+	// StickySessionDefaultLBCookieName is the default name of the load balancer-generated cookie.
+	StickySessionDefaultLBCookieName = "EG_SESSION"
+	// StickySessionDefaultLBCookieExpire is the default expiration duration of the load balancer-generated cookie.
+	StickySessionDefaultLBCookieExpire = time.Hour * 2
+	// KeyLen is the key length used by HMAC.
+	KeyLen = 8
 )
 
 // LoadBalancer is the interface of an HTTP load balancer.
 type LoadBalancer interface {
 	ChooseServer(req *httpprot.Request) *Server
+	ReturnServer(server *Server, req *httpprot.Request, resp *httpprot.Response)
 }
 
 // StickySessionSpec is the spec for sticky session.
 type StickySessionSpec struct {
-	Mode string `json:"mode" jsonschema:"required,enum=CookieConsistentHash"`
-	// AppCookieName need to be omitempty when we support other sticky mode.
-	AppCookieName string `json:"appCookieName" jsonschema:"required"`
+	Mode string `json:"mode" jsonschema:"required,enum=CookieConsistentHash,enum=DurationBased,enum=ApplicationBased"`
+	// AppCookieName is the user-defined cookie name in CookieConsistentHash and ApplicationBased mode.
+	AppCookieName string `json:"appCookieName" jsonschema:"omitempty"`
+	// LBCookieName is the generated cookie name in DurationBased and ApplicationBased mode.
+	LBCookieName string `json:"lbCookieName" jsonschema:"omitempty"`
+	// LBCookieExpire is the expire seconds of generated cookie in DurationBased and ApplicationBased mode.
+	LBCookieExpire string `json:"lbCookieExpire" jsonschema:"omitempty,format=duration"`
 }
 
 // LoadBalanceSpec is the spec to create a load balancer.
@@ -103,6 +127,7 @@ type BaseLoadBalancer struct {
 	spec           *LoadBalanceSpec
 	Servers        []*Server
 	consistentHash *consistent.Consistent
+	cookieExpire   time.Duration
 }
 
 func (blb *BaseLoadBalancer) init(spec *LoadBalanceSpec, servers []*Server) {
@@ -113,9 +138,18 @@ func (blb *BaseLoadBalancer) init(spec *LoadBalanceSpec, servers []*Server) {
 		return
 	}
 
-	// For now, we only support HeaderConsistentHash & CookieConsistentHash
-	members := make([]consistent.Member, len(servers))
-	for i, s := range servers {
+	switch spec.StickySession.Mode {
+	case StickySessionModeCookieConsistentHash:
+		blb.initConsistentHash()
+	case StickySessionModeDurationBased, StickySessionModeApplicationBased:
+		blb.configLBCookie()
+	}
+}
+
+// initConsistentHash initializes for consistent hash mode
+func (blb *BaseLoadBalancer) initConsistentHash() {
+	members := make([]consistent.Member, len(blb.Servers))
+	for i, s := range blb.Servers {
 		members[i] = hashMember{server: s}
 	}
 
@@ -128,12 +162,36 @@ func (blb *BaseLoadBalancer) init(spec *LoadBalanceSpec, servers []*Server) {
 	blb.consistentHash = consistent.New(members, cfg)
 }
 
+// configLBCookie configures properties for load balancer-generated cookie
+func (blb *BaseLoadBalancer) configLBCookie() {
+	if blb.spec.StickySession.LBCookieName == "" {
+		blb.spec.StickySession.LBCookieName = StickySessionDefaultLBCookieName
+	}
+
+	blb.cookieExpire, _ = time.ParseDuration(blb.spec.StickySession.LBCookieExpire)
+	if blb.cookieExpire <= 0 {
+		blb.cookieExpire = StickySessionDefaultLBCookieExpire
+	}
+}
+
 // ChooseServer chooses the sticky server if enable
 func (blb *BaseLoadBalancer) ChooseServer(req *httpprot.Request) *Server {
 	if blb.spec.StickySession == nil {
 		return nil
 	}
 
+	switch blb.spec.StickySession.Mode {
+	case StickySessionModeCookieConsistentHash:
+		return blb.chooseServerByConsistentHash(req)
+	case StickySessionModeDurationBased, StickySessionModeApplicationBased:
+		return blb.chooseServerByLBCookie(req)
+	}
+
+	return nil
+}
+
+// chooseServerByConsistentHash chooses server using consistent hash on cookie
+func (blb *BaseLoadBalancer) chooseServerByConsistentHash(req *httpprot.Request) *Server {
 	cookie, err := req.Cookie(blb.spec.StickySession.AppCookieName)
 	if err != nil {
 		return nil
@@ -147,6 +205,75 @@ func (blb *BaseLoadBalancer) ChooseServer(req *httpprot.Request) *Server {
 	return nil
 }
 
+// chooseServerByLBCookie chooses server by load balancer-generated cookie
+func (blb *BaseLoadBalancer) chooseServerByLBCookie(req *httpprot.Request) *Server {
+	cookie, err := req.Cookie(blb.spec.StickySession.LBCookieName)
+	if err != nil {
+		return nil
+	}
+
+	signed, err := hex.DecodeString(cookie.Value)
+	if err != nil || len(signed) != KeyLen+sha256.Size {
+		return nil
+	}
+
+	key := signed[:KeyLen]
+	macBytes := signed[KeyLen:]
+	for _, s := range blb.Servers {
+		mac := hmac.New(sha256.New, key)
+		mac.Write([]byte(s.ID()))
+		expected := mac.Sum(nil)
+		if hmac.Equal(expected, macBytes) {
+			return s
+		}
+	}
+
+	return nil
+}
+
+// ReturnServer does some custom work before return server
+func (blb *BaseLoadBalancer) ReturnServer(server *Server, req *httpprot.Request, resp *httpprot.Response) {
+	if blb.spec.StickySession == nil {
+		return
+	}
+
+	setCookie := false
+	switch blb.spec.StickySession.Mode {
+	case StickySessionModeDurationBased:
+		setCookie = true
+	case StickySessionModeApplicationBased:
+		for _, c := range resp.Cookies() {
+			if c.Name == blb.spec.StickySession.AppCookieName {
+				setCookie = true
+				break
+			}
+		}
+	}
+	if setCookie {
+		cookie := &http.Cookie{
+			Name:    blb.spec.StickySession.LBCookieName,
+			Value:   sign([]byte(server.ID())),
+			Expires: time.Now().Add(blb.cookieExpire),
+		}
+		resp.SetCookie(cookie)
+	}
+}
+
+// sign signs plain text byte array to encoded string
+func sign(plain []byte) string {
+	signed := make([]byte, KeyLen+sha256.Size)
+	key := signed[:KeyLen]
+	macBytes := signed[KeyLen:]
+
+	// use maphash to generate random key fast
+	binary.LittleEndian.PutUint64(key, new(maphash.Hash).Sum64())
+	mac := hmac.New(sha256.New, key)
+	mac.Write(plain)
+	mac.Sum(macBytes[:0])
+
+	return hex.EncodeToString(signed)
+}
+
 // randomLoadBalancer does load balancing in a random manner.
 type randomLoadBalancer struct {
 	BaseLoadBalancer

pkg/filters/proxy/loadbalance_test.go

@@ -193,3 +193,73 @@ func TestStickySession_ConsistentHash(t *testing.T) {
 		assert.Equal(svr1, svr)
 	}
 }
+
+func TestStickySession_DurationBased(t *testing.T) {
+	assert := assert.New(t)
+
+	servers := prepareServers(10)
+	lb := NewLoadBalancer(&LoadBalanceSpec{
+		Policy: LoadBalancePolicyRandom,
+		StickySession: &StickySessionSpec{
+			Mode: StickySessionModeDurationBased,
+		},
+	}, servers)
+
+	r, _ := httpprot.NewRequest(&http.Request{Header: http.Header{}})
+	svr1 := lb.ChooseServer(r)
+	resp, _ := httpprot.NewResponse(&http.Response{Header: http.Header{}})
+	lb.ReturnServer(svr1, r, resp)
+	c := readCookie(resp.Cookies(), StickySessionDefaultLBCookieName)
+
+	for i := 0; i < 100; i++ {
+		req := &http.Request{Header: http.Header{}}
+		req.AddCookie(&http.Cookie{Name: StickySessionDefaultLBCookieName, Value: c.Value})
+		r, _ = httpprot.NewRequest(req)
+		svr := lb.ChooseServer(r)
+		assert.Equal(svr1, svr)
+
+		resp, _ = httpprot.NewResponse(&http.Response{Header: http.Header{}})
+		lb.ReturnServer(svr, r, resp)
+		c = readCookie(resp.Cookies(), StickySessionDefaultLBCookieName)
+	}
+}
+
+func TestStickySession_ApplicationBased(t *testing.T) {
+	assert := assert.New(t)
+
+	servers := prepareServers(10)
+	appCookieName := "x-app-cookie"
+	lb := NewLoadBalancer(&LoadBalanceSpec{
+		Policy: LoadBalancePolicyRandom,
+		StickySession: &StickySessionSpec{
+			Mode:          StickySessionModeApplicationBased,
+			AppCookieName: appCookieName,
+		},
+	}, servers)
+
+	r, _ := httpprot.NewRequest(&http.Request{Header: http.Header{}})
+	svr1 := lb.ChooseServer(r)
+	resp, _ := httpprot.NewResponse(&http.Response{Header: http.Header{}})
+	resp.SetCookie(&http.Cookie{Name: appCookieName, Value: ""})
+	lb.ReturnServer(svr1, r, resp)
+	c := readCookie(resp.Cookies(), StickySessionDefaultLBCookieName)
+
+	for i := 0; i < 100; i++ {
+		req := &http.Request{Header: http.Header{}}
+		req.AddCookie(&http.Cookie{Name: StickySessionDefaultLBCookieName, Value: c.Value})
+		r, _ = httpprot.NewRequest(req)
+		svr := lb.ChooseServer(r)
+		assert.Equal(svr1, svr)
+
+		resp, _ = httpprot.NewResponse(&http.Response{Header: http.Header{}})
+		resp.SetCookie(&http.Cookie{Name: appCookieName, Value: ""})
+		lb.ReturnServer(svr, r, resp)
+		c = readCookie(resp.Cookies(), StickySessionDefaultLBCookieName)
+	}
+}
+
+func BenchmarkSign(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		sign([]byte("192.168.1.2"))
+	}
+}

pkg/filters/proxy/pool.go

@@ -438,6 +438,8 @@ func (sp *ServerPool) doHandle(stdctx stdcontext.Context, spCtx *serverPoolConte
 		return serverPoolError{http.StatusInternalServerError, resultInternalError}
 	}
 
+	sp.LoadBalancer().ReturnServer(svr, spCtx.req, spCtx.resp)
+
 	spCtx.LazyAddTag(func() string {
 		return fmt.Sprintf("status code: %d", resp.StatusCode)
 	})