From d6457e6cbb44b708bed40c684de93427c601125e Mon Sep 17 00:00:00 2001
From: Yakun Sun <sunyakun2016@outlook.com>
Date: Tue, 8 Aug 2023 21:12:06 +0800
Subject: [PATCH] Set sameSite field for wrr load balancer sticky cookie

---
 pkg/server/service/loadbalancer/wrr/wrr.go    | 24 ++++++++++++++++++-
 .../service/loadbalancer/wrr/wrr_test.go      | 20 ++++++++++++++--
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/pkg/server/service/loadbalancer/wrr/wrr.go b/pkg/server/service/loadbalancer/wrr/wrr.go
index 7442fa3a1..f3a38b4c9 100644
--- a/pkg/server/service/loadbalancer/wrr/wrr.go
+++ b/pkg/server/service/loadbalancer/wrr/wrr.go
@@ -22,6 +22,20 @@ type stickyCookie struct {
 	name     string
 	secure   bool
 	httpOnly bool
+	sameSite string
+}
+
+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 http.SameSiteDefaultMode
+	}
 }
 
 // Balancer is a WeightedRoundRobin load balancer based on Earliest Deadline First (EDF).
@@ -57,6 +71,7 @@ func New(sticky *dynamic.Sticky, wantHealthCheck bool) *Balancer {
 			name:     sticky.Cookie.Name,
 			secure:   sticky.Cookie.Secure,
 			httpOnly: sticky.Cookie.HTTPOnly,
+			sameSite: sticky.Cookie.SameSite,
 		}
 	}
 	return balancer
@@ -214,7 +229,14 @@ func (b *Balancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	}
 
 	if b.stickyCookie != nil {
-		cookie := &http.Cookie{Name: b.stickyCookie.name, Value: server.name, Path: "/", HttpOnly: b.stickyCookie.httpOnly, Secure: b.stickyCookie.secure}
+		cookie := &http.Cookie{
+			Name:     b.stickyCookie.name,
+			Value:    server.name,
+			Path:     "/",
+			HttpOnly: b.stickyCookie.httpOnly,
+			Secure:   b.stickyCookie.secure,
+			SameSite: convertSameSite(b.stickyCookie.sameSite),
+		}
 		http.SetCookie(w, cookie)
 	}
 
diff --git a/pkg/server/service/loadbalancer/wrr/wrr_test.go b/pkg/server/service/loadbalancer/wrr/wrr_test.go
index e0ca68c04..b905e55dc 100644
--- a/pkg/server/service/loadbalancer/wrr/wrr_test.go
+++ b/pkg/server/service/loadbalancer/wrr/wrr_test.go
@@ -220,7 +220,12 @@ func TestBalancerAllServersZeroWeight(t *testing.T) {
 
 func TestSticky(t *testing.T) {
 	balancer := New(&dynamic.Sticky{
-		Cookie: &dynamic.Cookie{Name: "test"},
+		Cookie: &dynamic.Cookie{
+			Name:     "test",
+			Secure:   true,
+			HTTPOnly: true,
+			SameSite: "none",
+		},
 	}, false)
 
 	balancer.Add("first", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@@ -233,7 +238,11 @@ func TestSticky(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	}), Int(2))
 
-	recorder := &responseRecorder{ResponseRecorder: httptest.NewRecorder(), save: map[string]int{}}
+	recorder := &responseRecorder{
+		ResponseRecorder: httptest.NewRecorder(),
+		save:             map[string]int{},
+		cookies:          make(map[string]*http.Cookie),
+	}
 
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
 	for i := 0; i < 3; i++ {
@@ -247,6 +256,9 @@ func TestSticky(t *testing.T) {
 
 	assert.Equal(t, 0, recorder.save["first"])
 	assert.Equal(t, 3, recorder.save["second"])
+	assert.Equal(t, true, recorder.cookies["test"].HttpOnly)
+	assert.Equal(t, true, recorder.cookies["test"].Secure)
+	assert.Equal(t, http.SameSiteNoneMode, recorder.cookies["test"].SameSite)
 }
 
 // TestBalancerBias makes sure that the WRR algorithm spreads elements evenly right from the start,
@@ -282,11 +294,15 @@ type responseRecorder struct {
 	save     map[string]int
 	sequence []string
 	status   []int
+	cookies  map[string]*http.Cookie
 }
 
 func (r *responseRecorder) WriteHeader(statusCode int) {
 	r.save[r.Header().Get("server")]++
 	r.sequence = append(r.sequence, r.Header().Get("server"))
 	r.status = append(r.status, statusCode)
+	for _, cookie := range r.Result().Cookies() {
+		r.cookies[cookie.Name] = cookie
+	}
 	r.ResponseRecorder.WriteHeader(statusCode)
 }