From 949fbc7ba87d2500655bd675240bc7b6e6b8c944 Mon Sep 17 00:00:00 2001
From: zhaoYangguang <1163765691@qq.com>
Date: Tue, 19 Jan 2021 18:22:19 +0800
Subject: [PATCH] optimize the process that logout with double puppeting
 enabled

---
 commands.go | 18 +++++++++++++++++-
 user.go     | 24 +++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/commands.go b/commands.go
index aaf075d..f3ff562 100644
--- a/commands.go
+++ b/commands.go
@@ -251,6 +251,7 @@ func (handler *CommandHandler) CommandLogout(ce *CommandEvent) {
 	if ce.User.Conn == nil {
 		return
 	}
+	_ = ce.User.Conn.GetConversations("", ce.User.bridge.Config.Bridge.InitialChatSync)
 	ce.User.Conn.LoggedIn = false
 	username := ""
 	password := ""
@@ -269,7 +270,22 @@ func (handler *CommandHandler) CommandLogout(ce *CommandEvent) {
 		Username: username,
 		Password: password,
 	}
+
 	portals := ce.User.GetPortals()
+	newPortals := ce.User.GetPortalsNew()
+	allPortals := newPortals[0:]
+	for _, portal := range portals {
+		var newPortalsHas bool
+		for _, newPortal := range newPortals {
+			if portal.Key == newPortal.Key {
+				newPortalsHas = true
+			}
+		}
+		if !newPortalsHas {
+			allPortals = append(allPortals, portal)
+		}
+	}
+
 	leave := func(portal *Portal) {
 		if len(portal.MXID) > 0 {
 			_, _ = portal.MainIntent().KickUser(portal.MXID, &mautrix.ReqKickUser{
@@ -278,7 +294,7 @@ func (handler *CommandHandler) CommandLogout(ce *CommandEvent) {
 			})
 		}
 	}
-	for _, portal := range portals {
+	for _, portal := range allPortals {
 		leave(portal)
 	}
 	ce.Reply("Logged out successfully.")
diff --git a/user.go b/user.go
index 25949e9..21f92c5 100644
--- a/user.go
+++ b/user.go
@@ -156,6 +156,28 @@ func (user *User) GetPortals() []*Portal {
 	return portals
 }
 
+func (user *User) GetPortalsNew() []*Portal {
+	keys := make([]database.PortalKey, len(user.Conn.Store.Chats))
+	i := 0
+	for jid, _ := range user.Conn.Store.Chats {
+		keys[i] = database.NewPortalKey(jid, user.JID)
+		i++
+	}
+
+	portals := make([]*Portal, len(keys))
+
+	user.bridge.portalsLock.Lock()
+	for i, key := range keys {
+		portal, ok := user.bridge.portalsByJID[key]
+		if !ok {
+			portal = user.bridge.loadDBPortal(user.bridge.DB.Portal.GetByJID(key), &key)
+		}
+		portals[i] = portal
+	}
+	user.bridge.portalsLock.Unlock()
+	return portals
+}
+
 func (bridge *Bridge) NewUser(dbUser *database.User) *User {
 	user := &User{
 		User:   dbUser,
@@ -394,9 +416,9 @@ func (user *User) Login(ce *CommandEvent, name string, password string) (err err
 	go user.Conn.Poll()
 
 	user.ConnectionErrors = 0
-	user.SetSession(user.Conn.LoginInfo)
 	user.JID = "8:" + user.Conn.UserProfile.Username + skypeExt.NewUserSuffix
 	user.addToJIDMap()
+	user.SetSession(user.Conn.LoginInfo)
 	_ = ce.User.Conn.GetConversations("", user.bridge.Config.Bridge.InitialChatSync)
 	user.PostLogin()
 	return