From 097e0333476653801abfb18745577d2f6a0ee534 Mon Sep 17 00:00:00 2001
From: Monet Lee <monet_lee@163.com>
Date: Mon, 28 Jul 2025 10:57:44 +0800
Subject: [PATCH] fix: fix incorrect kicked logic. (#3480)

* fix: fix incorrect ws check logic.

* optimize checkSameToken logic.
---
 internal/msggateway/hub_server.go |  1 +
 internal/msggateway/ws_server.go  | 37 +++++++++++++++++++++++++++----
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/internal/msggateway/hub_server.go b/internal/msggateway/hub_server.go
index 3d0d1e3f9..8374af06d 100644
--- a/internal/msggateway/hub_server.go
+++ b/internal/msggateway/hub_server.go
@@ -249,6 +249,7 @@ func (s *Server) MultiTerminalLoginCheck(ctx context.Context, req *msggateway.Mu
 		tempUserCtx.SetOperationID(mcontext.GetOperationID(ctx))
 		client := &Client{}
 		client.ctx = tempUserCtx
+		client.token = req.Token
 		client.UserID = req.UserID
 		client.PlatformID = int(req.PlatformID)
 		i := &kickHandler{
diff --git a/internal/msggateway/ws_server.go b/internal/msggateway/ws_server.go
index 9b7343938..bc7a2fa5f 100644
--- a/internal/msggateway/ws_server.go
+++ b/internal/msggateway/ws_server.go
@@ -334,6 +334,27 @@ func (ws *WsServer) multiTerminalLoginChecker(clientOK bool, oldClients []*Clien
 		}
 	}
 
+	// If reconnect: When multiple msgGateway instances are deployed, a client may disconnect from instance A and reconnect to instance B.
+	// During this process, instance A might still be executing, resulting in two clients with the same token existing simultaneously.
+	// This situation needs to be filtered to prevent duplicate clients.
+	checkSameTokenFunc := func(oldClients []*Client) []*Client {
+		var clientsNeedToKick []*Client
+
+		for _, c := range oldClients {
+			if c.token == newClient.token {
+				log.ZDebug(newClient.ctx, "token is same, not kick",
+					"userID", newClient.UserID,
+					"platformID", newClient.PlatformID,
+					"token", newClient.token)
+				continue
+			}
+
+			clientsNeedToKick = append(clientsNeedToKick, c)
+		}
+
+		return clientsNeedToKick
+	}
+
 	switch ws.msgGatewayConfig.Share.MultiLogin.Policy {
 	case constant.DefalutNotKick:
 	case constant.PCAndOther:
@@ -349,11 +370,15 @@ func (ws *WsServer) multiTerminalLoginChecker(clientOK bool, oldClients []*Clien
 			}
 			oldClients = append(oldClients, c)
 		}
+
 		fallthrough
 	case constant.AllLoginButSameTermKick:
 		if !clientOK {
 			return
 		}
+
+		oldClients = checkSameTokenFunc(oldClients)
+
 		ws.clients.DeleteClients(newClient.UserID, oldClients)
 		for _, c := range oldClients {
 			err := c.KickOnlineMessage()
@@ -361,6 +386,7 @@ func (ws *WsServer) multiTerminalLoginChecker(clientOK bool, oldClients []*Clien
 				log.ZWarn(c.ctx, "KickOnlineMessage", err)
 			}
 		}
+
 		ctx := mcontext.WithMustInfoCtx(
 			[]string{newClient.ctx.GetOperationID(), newClient.ctx.GetUserID(),
 				constant.PlatformIDToName(newClient.PlatformID), newClient.ctx.GetConnID()},
@@ -379,14 +405,17 @@ func (ws *WsServer) multiTerminalLoginChecker(clientOK bool, oldClients []*Clien
 		if !ok {
 			return
 		}
-		var (
-			kickClients []*Client
-		)
+
+		var kickClients []*Client
 		for _, client := range clients {
 			if constant.PlatformIDToClass(client.PlatformID) == constant.PlatformIDToClass(newClient.PlatformID) {
-				kickClients = append(kickClients, client)
+				{
+					kickClients = append(kickClients, client)
+				}
 			}
 		}
+		kickClients = checkSameTokenFunc(kickClients)
+
 		kickTokenFunc(kickClients)
 	}
 }