reverseproxy: Implement cookie hash selection policy (#3809)

* add CookieHashSelection for session affinity

* add CookieHashSelection for session affinity

* register module

* reverse_proxy: Add and fix cookie lb_policy

* reverse_proxy: Manage hmac.write error on cookie hash selection

* reverse_proxy: fix some comments

* reverse_proxy: variable `cookieValue` is inside the else block

* reverse_proxy: Abstract duplicate nuanced logic of reservoir sampling into a function

* reverse_proxy: Set a default secret is indeed useless

* reverse_proxy: add configuration syntax for cookie lb_policy

* reverse_proxy: doc typo and improvement

Co-authored-by: utick <123liuqingdong@163.com>

1 files changed, 83 insertions, 35 deletions

diff --git a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
index 49585da..5368a1a 100644
--- a/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
+++ b/modules/caddyhttp/reverseproxy/selectionpolicies_test.go
@@ -33,37 +33,37 @@ func TestRoundRobinPolicy(t *testing.T) {
 	rrPolicy := new(RoundRobinSelection)
 	req, _ := http.NewRequest("GET", "/", nil)
 
-	h := rrPolicy.Select(pool, req)
+	h := rrPolicy.Select(pool, req, nil)
 	// First selected host is 1, because counter starts at 0
 	// and increments before host is selected
 	if h != pool[1] {
 		t.Error("Expected first round robin host to be second host in the pool.")
 	}
-	h = rrPolicy.Select(pool, req)
+	h = rrPolicy.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected second round robin host to be third host in the pool.")
 	}
-	h = rrPolicy.Select(pool, req)
+	h = rrPolicy.Select(pool, req, nil)
 	if h != pool[0] {
 		t.Error("Expected third round robin host to be first host in the pool.")
 	}
 	// mark host as down
 	pool[1].SetHealthy(false)
-	h = rrPolicy.Select(pool, req)
+	h = rrPolicy.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected to skip down host.")
 	}
 	// mark host as up
 	pool[1].SetHealthy(true)
 
-	h = rrPolicy.Select(pool, req)
+	h = rrPolicy.Select(pool, req, nil)
 	if h == pool[2] {
 		t.Error("Expected to balance evenly among healthy hosts")
 	}
 	// mark host as full
 	pool[1].CountRequest(1)
 	pool[1].MaxRequests = 1
-	h = rrPolicy.Select(pool, req)
+	h = rrPolicy.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected to skip full host.")
 	}
@@ -76,12 +76,12 @@ func TestLeastConnPolicy(t *testing.T) {
 
 	pool[0].CountRequest(10)
 	pool[1].CountRequest(10)
-	h := lcPolicy.Select(pool, req)
+	h := lcPolicy.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected least connection host to be third host.")
 	}
 	pool[2].CountRequest(100)
-	h = lcPolicy.Select(pool, req)
+	h = lcPolicy.Select(pool, req, nil)
 	if h != pool[0] && h != pool[1] {
 		t.Error("Expected least connection host to be first or second host.")
 	}
@@ -94,44 +94,44 @@ func TestIPHashPolicy(t *testing.T) {
 
 	// We should be able to predict where every request is routed.
 	req.RemoteAddr = "172.0.0.1:80"
-	h := ipHash.Select(pool, req)
+	h := ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 	req.RemoteAddr = "172.0.0.2:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 	req.RemoteAddr = "172.0.0.3:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected ip hash policy host to be the third host.")
 	}
 	req.RemoteAddr = "172.0.0.4:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 
 	// we should get the same results without a port
 	req.RemoteAddr = "172.0.0.1"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 	req.RemoteAddr = "172.0.0.2"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 	req.RemoteAddr = "172.0.0.3"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected ip hash policy host to be the third host.")
 	}
 	req.RemoteAddr = "172.0.0.4"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
@@ -140,13 +140,13 @@ func TestIPHashPolicy(t *testing.T) {
 	// healthy host is available
 	req.RemoteAddr = "172.0.0.1"
 	pool[1].SetHealthy(false)
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected ip hash policy host to be the third host.")
 	}
 
 	req.RemoteAddr = "172.0.0.2"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[2] {
 		t.Error("Expected ip hash policy host to be the third host.")
 	}
@@ -154,12 +154,12 @@ func TestIPHashPolicy(t *testing.T) {
 
 	req.RemoteAddr = "172.0.0.3"
 	pool[2].SetHealthy(false)
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[0] {
 		t.Error("Expected ip hash policy host to be the first host.")
 	}
 	req.RemoteAddr = "172.0.0.4"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
@@ -171,22 +171,22 @@ func TestIPHashPolicy(t *testing.T) {
 		{Host: new(upstreamHost)},
 	}
 	req.RemoteAddr = "172.0.0.1:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[0] {
 		t.Error("Expected ip hash policy host to be the first host.")
 	}
 	req.RemoteAddr = "172.0.0.2:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
 	req.RemoteAddr = "172.0.0.3:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[0] {
 		t.Error("Expected ip hash policy host to be the first host.")
 	}
 	req.RemoteAddr = "172.0.0.4:80"
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected ip hash policy host to be the second host.")
 	}
@@ -194,7 +194,7 @@ func TestIPHashPolicy(t *testing.T) {
 	// We should get nil when there are no healthy hosts
 	pool[0].SetHealthy(false)
 	pool[1].SetHealthy(false)
-	h = ipHash.Select(pool, req)
+	h = ipHash.Select(pool, req, nil)
 	if h != nil {
 		t.Error("Expected ip hash policy host to be nil.")
 	}
@@ -205,13 +205,13 @@ func TestFirstPolicy(t *testing.T) {
 	firstPolicy := new(FirstSelection)
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
 
-	h := firstPolicy.Select(pool, req)
+	h := firstPolicy.Select(pool, req, nil)
 	if h != pool[0] {
 		t.Error("Expected first policy host to be the first host.")
 	}
 
 	pool[0].SetHealthy(false)
-	h = firstPolicy.Select(pool, req)
+	h = firstPolicy.Select(pool, req, nil)
 	if h != pool[1] {
 		t.Error("Expected first policy host to be the second host.")
 	}
@@ -222,19 +222,19 @@ func TestURIHashPolicy(t *testing.T) {
 	uriPolicy := new(URIHashSelection)
 
 	request := httptest.NewRequest(http.MethodGet, "/test", nil)
-	h := uriPolicy.Select(pool, request)
+	h := uriPolicy.Select(pool, request, nil)
 	if h != pool[0] {
 		t.Error("Expected uri policy host to be the first host.")
 	}
 
 	pool[0].SetHealthy(false)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != pool[1] {
 		t.Error("Expected uri policy host to be the first host.")
 	}
 
 	request = httptest.NewRequest(http.MethodGet, "/test_2", nil)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != pool[1] {
 		t.Error("Expected uri policy host to be the second host.")
 	}
@@ -247,26 +247,26 @@ func TestURIHashPolicy(t *testing.T) {
 	}
 
 	request = httptest.NewRequest(http.MethodGet, "/test", nil)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != pool[0] {
 		t.Error("Expected uri policy host to be the first host.")
 	}
 
 	pool[0].SetHealthy(false)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != pool[1] {
 		t.Error("Expected uri policy host to be the first host.")
 	}
 
 	request = httptest.NewRequest(http.MethodGet, "/test_2", nil)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != pool[1] {
 		t.Error("Expected uri policy host to be the second host.")
 	}
 
 	pool[0].SetHealthy(false)
 	pool[1].SetHealthy(false)
-	h = uriPolicy.Select(pool, request)
+	h = uriPolicy.Select(pool, request, nil)
 	if h != nil {
 		t.Error("Expected uri policy policy host to be nil.")
 	}
@@ -311,7 +311,7 @@ func TestRandomChoicePolicy(t *testing.T) {
 	randomChoicePolicy := new(RandomChoiceSelection)
 	randomChoicePolicy.Choose = 2
 
-	h := randomChoicePolicy.Select(pool, request)
+	h := randomChoicePolicy.Select(pool, request, nil)
 
 	if h == nil {
 		t.Error("RandomChoicePolicy should not return nil")
@@ -322,3 +322,51 @@ func TestRandomChoicePolicy(t *testing.T) {
 	}
 
 }
+
+func TestCookieHashPolicy(t *testing.T) {
+	pool := testPool()
+	pool[0].Dial = "localhost:8080"
+	pool[1].Dial = "localhost:8081"
+	pool[2].Dial = "localhost:8082"
+	pool[0].SetHealthy(true)
+	pool[1].SetHealthy(false)
+	pool[2].SetHealthy(false)
+	request := httptest.NewRequest(http.MethodGet, "/test", nil)
+	w := httptest.NewRecorder()
+	cookieHashPolicy := new(CookieHashSelection)
+	h := cookieHashPolicy.Select(pool, request, w)
+	cookie_server1 := w.Result().Cookies()[0]
+	if cookie_server1 == nil {
+		t.Error("cookieHashPolicy should set a cookie")
+	}
+	if cookie_server1.Name != "lb" {
+		t.Error("cookieHashPolicy should set a cookie with name lb")
+	}
+	if h != pool[0] {
+		t.Error("Expected cookieHashPolicy host to be the first only available host.")
+	}
+	pool[1].SetHealthy(true)
+	pool[2].SetHealthy(true)
+	request = httptest.NewRequest(http.MethodGet, "/test", nil)
+	w = httptest.NewRecorder()
+	request.AddCookie(cookie_server1)
+	h = cookieHashPolicy.Select(pool, request, w)
+	if h != pool[0] {
+		t.Error("Expected cookieHashPolicy host to stick to the first host (matching cookie).")
+	}
+	s := w.Result().Cookies()
+	if len(s) != 0 {
+		t.Error("Expected cookieHashPolicy to not set a new cookie.")
+	}
+	pool[0].SetHealthy(false)
+	request = httptest.NewRequest(http.MethodGet, "/test", nil)
+	w = httptest.NewRecorder()
+	request.AddCookie(cookie_server1)
+	h = cookieHashPolicy.Select(pool, request, w)
+	if h == pool[0] {
+		t.Error("Expected cookieHashPolicy to select a new host.")
+	}
+	if w.Result().Cookies() == nil {
+		t.Error("Expected cookieHashPolicy to set a new cookie.")
+	}
+}