From 5d8bec7f246ccfbcc54ae168e56dd11e695a40f5 Mon Sep 17 00:00:00 2001
From: Toby Chui <tobychui@users.noreply.github.com>
Date: Fri, 14 Feb 2025 22:53:29 +0800
Subject: [PATCH] Fixed sticky session bug

- Fixed sticky session bug in new active fallback lb implementation
---
 .../dynamicproxy/loadbalance/loadbalance.go   |  5 +++
 .../dynamicproxy/loadbalance/onlineStatus.go  |  1 +
 .../dynamicproxy/loadbalance/originPicker.go  | 39 +++++++++++++++----
 src/reverseproxy.go                           |  1 +
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/mod/dynamicproxy/loadbalance/loadbalance.go b/src/mod/dynamicproxy/loadbalance/loadbalance.go
index b1be8e8..2da878e 100644
--- a/src/mod/dynamicproxy/loadbalance/loadbalance.go
+++ b/src/mod/dynamicproxy/loadbalance/loadbalance.go
@@ -109,6 +109,11 @@ func GetUpstreamsAsString(upstreams []*Upstream) string {
 	return strings.Join(targets, ", ")
 }
 
+// Reset the current session store and clear all previous sessions
+func (m *RouteManager) ResetSessions() {
+	m.SessionStore = sessions.NewCookieStore([]byte(m.Options.SystemUUID))
+}
+
 func (m *RouteManager) Close() {
 	//Close the session store
 	m.SessionStore.MaxAge(0)
diff --git a/src/mod/dynamicproxy/loadbalance/onlineStatus.go b/src/mod/dynamicproxy/loadbalance/onlineStatus.go
index 860c266..6d2d591 100644
--- a/src/mod/dynamicproxy/loadbalance/onlineStatus.go
+++ b/src/mod/dynamicproxy/loadbalance/onlineStatus.go
@@ -67,5 +67,6 @@ func (m *RouteManager) FilterOfflineOrigins(origins []*Upstream) []*Upstream {
 			onlineOrigins = append(onlineOrigins, origin)
 		}
 	}
+
 	return onlineOrigins
 }
diff --git a/src/mod/dynamicproxy/loadbalance/originPicker.go b/src/mod/dynamicproxy/loadbalance/originPicker.go
index 92d0cfd..4c4ac75 100644
--- a/src/mod/dynamicproxy/loadbalance/originPicker.go
+++ b/src/mod/dynamicproxy/loadbalance/originPicker.go
@@ -24,7 +24,7 @@ func (m *RouteManager) GetRequestUpstreamTarget(w http.ResponseWriter, r *http.R
 	if useStickySession {
 		//Use stick session, check which origins this request previously used
 		targetOriginId, err := m.getSessionHandler(r, origins)
-		if err != nil || !m.IsTargetOnline(origins[targetOriginId].OriginIpOrDomain) {
+		if err != nil {
 			// No valid session found or origin is offline
 			// Filter the offline origins
 			origins = m.FilterOfflineOrigins(origins)
@@ -39,14 +39,18 @@ func (m *RouteManager) GetRequestUpstreamTarget(w http.ResponseWriter, r *http.R
 				targetOrigin = origins[0]
 				index = 0
 			}
+
+			//fmt.Println("DEBUG: (Sticky Session) Registering session origin " + origins[index].OriginIpOrDomain)
 			m.setSessionHandler(w, r, targetOrigin.OriginIpOrDomain, index)
 			return targetOrigin, nil
 		}
 
 		//Valid session found and origin is online
+		//fmt.Println("DEBUG: (Sticky Session) Picking origin " + origins[targetOriginId].OriginIpOrDomain)
 		return origins[targetOriginId], nil
 	}
 	//No sticky session, get a random origin
+	m.clearSessionHandler(w, r) //Clear the session
 
 	//Filter the offline origins
 	origins = m.FilterOfflineOrigins(origins)
@@ -89,6 +93,20 @@ func (m *RouteManager) setSessionHandler(w http.ResponseWriter, r *http.Request,
 	return nil
 }
 
+func (m *RouteManager) clearSessionHandler(w http.ResponseWriter, r *http.Request) error {
+	session, err := m.SessionStore.Get(r, "STICKYSESSION")
+	if err != nil {
+		return err
+	}
+	session.Options.MaxAge = -1
+	session.Options.Path = "/"
+	err = session.Save(r, w)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // Get the previous connected origin from session
 func (m *RouteManager) getSessionHandler(r *http.Request, upstreams []*Upstream) (int, error) {
 	// Get existing session
@@ -105,15 +123,22 @@ func (m *RouteManager) getSessionHandler(r *http.Request, upstreams []*Upstream)
 		return -1, errors.New("no session has been set")
 	}
 	originDomain := originDomainRaw.(string)
-	originID := originIDRaw.(int)
+	//originID := originIDRaw.(int)
 
-	//Check if it has been modified
-	if len(upstreams) < originID || upstreams[originID].OriginIpOrDomain != originDomain {
-		//Mismatch or upstreams has been updated
-		return -1, errors.New("upstreams has been changed")
+	//Check if the upstream still exists
+	for i, upstream := range upstreams {
+		if upstream.OriginIpOrDomain == originDomain {
+			if !m.IsTargetOnline(originDomain) {
+				//Origin is offline
+				return -1, errors.New("origin is offline")
+			}
+
+			//Ok, the origin is still online
+			return i, nil
+		}
 	}
 
-	return originID, nil
+	return -1, errors.New("origin is no longer exists")
 }
 
 /* Functions related to random upstream picking */
diff --git a/src/reverseproxy.go b/src/reverseproxy.go
index 62c8c3a..e25e405 100644
--- a/src/reverseproxy.go
+++ b/src/reverseproxy.go
@@ -595,6 +595,7 @@ func ReverseProxyHandleEditEndpoint(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	targetProxyEntry.Remove()
+	loadBalancer.ResetSessions()
 	dynamicProxyRouter.AddProxyRouteToRuntime(readyRoutingRule)
 
 	//Save it to file