Files
flagger/test/verify_session_affinity.go
2025-11-18 15:16:42 -03:00

112 lines
2.8 KiB
Go

package main
import (
"fmt"
"io"
"log"
"net/http"
"os"
"strings"
"sync"
"time"
)
// channel for canary cookie
var cc = make(chan string, 1)
// channel for primary cookie
var pc = make(chan string, 1)
var timeout = time.Second * 10
func main() {
url := os.Getenv("URL")
host := os.Getenv("HOST")
canaryVersion := os.Getenv("CANARY_VERSION")
canaryCookieName := os.Getenv("CANARY_COOKIE_NAME")
primaryVersion := os.Getenv("PRIMARY_VERSION")
primaryCookieName := os.Getenv("PRIMARY_COOKIE_NAME")
go tryUntilWorkloadIsHit(url, host, canaryVersion, canaryCookieName, cc)
go tryUntilWorkloadIsHit(url, host, primaryVersion, primaryCookieName, pc)
wg := &sync.WaitGroup{}
wg.Add(2)
go verifySessionAffinity(cc, wg, url, host, primaryVersion)
go verifySessionAffinity(pc, wg, url, host, canaryVersion)
wg.Wait()
log.Println("✔ successfully verified session affinity")
}
func verifySessionAffinity(cc chan string, wg *sync.WaitGroup, url, host, wrongVersion string) {
select {
// If we receive a cookie, then try to verify that we are always routed to the
// Canary deployment based on the cookie.
case cookie := <-cc:
for i := 0; i < 5; i++ {
headers := map[string]string{
"Cookie": cookie,
}
body, _, err := sendRequest(url, host, headers)
if err != nil {
log.Fatalf("failed to send request to verify cookie based routing: %v", err)
}
if strings.Contains(body, wrongVersion) {
log.Fatalf("received response from the wrong deployment")
}
}
wg.Done()
case <-time.After(timeout):
log.Fatal("timed out waiting for workload hit")
}
}
// sendRequest sends a request to the URL with the provided host and headers.
// It returns the response body and cookies or an error.
func sendRequest(url, host string, headers map[string]string) (string, []*http.Cookie, error) {
client := http.DefaultClient
req, err := http.NewRequest("GET", url, nil)
if err != nil {
return "", nil, err
}
for key, value := range headers {
req.Header.Add(key, value)
}
req.Host = host
resp, err := client.Do(req)
if err != nil {
return "", nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", nil, err
}
return string(body), resp.Cookies(), nil
}
// tryUntilWorkloadIsHit is a recursive function that tries to send request and
// either sends the cookie back to the main thread (if received) or re-sends
// the request.
func tryUntilWorkloadIsHit(url, host, version, cookieName string, c chan string) {
body, cookies, err := sendRequest(url, host, nil)
if err != nil {
log.Printf("warning: failed to send request: %s", err)
return
}
if strings.Contains(body, version) {
if cookies[0].Name == cookieName {
c <- fmt.Sprintf("%s=%s", cookies[0].Name, cookies[0].Value)
return
}
}
tryUntilWorkloadIsHit(url, host, version, cookieName, c)
}