From 6e0849d4c2855375b66974f85eb9840d05197694 Mon Sep 17 00:00:00 2001 From: Dimitri Masson <30894448+d-masson@users.noreply.github.com> Date: Fri, 20 Nov 2020 20:39:26 +0100 Subject: 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> --- .../reverseproxy/selectionpolicies_test.go | 118 +++++++++++++++------ 1 file changed, 83 insertions(+), 35 deletions(-) (limited to 'modules/caddyhttp/reverseproxy/selectionpolicies_test.go') 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.") + } +} -- cgit v1.2.3