Fix needing to reconnect after initial login

commands.go

@@ -230,6 +230,13 @@ func (handler *CommandHandler) CommandLogout(ce *CommandEvent) {
 		ce.Reply("You're not logged in.")
 		return
 	}
+	puppet := handler.bridge.GetPuppetByJID(ce.User.JID)
+	if puppet.CustomMXID != "" {
+		err := puppet.SwitchCustomMXID("", "")
+		if err != nil {
+			ce.User.log.Warnln("Failed to logout-matrix while logging out of WhatsApp:", err)
+		}
+	}
 	err := ce.User.Conn.Logout()
 	if err != nil {
 		ce.User.log.Warnln("Error while logging out:", err)
@@ -242,6 +249,9 @@ func (handler *CommandHandler) CommandLogout(ce *CommandEvent) {
 	}
 	ce.User.Conn.RemoveHandlers()
 	ce.User.Conn = nil
+	ce.User.removeFromJIDMap()
+	// TODO this causes a foreign key violation, which should be fixed
+	//ce.User.JID = ""
 	ce.User.SetSession(nil)
 	ce.Reply("Logged out successfully.")
 }

provisioning.go

@@ -292,6 +292,9 @@ func (prov *ProvisioningAPI) Logout(w http.ResponseWriter, r *http.Request) {
 	}
 	user.Conn.RemoveHandlers()
 	user.Conn = nil
+	user.removeFromJIDMap()
+	// TODO this causes a foreign key violation, which should be fixed
+	//ce.User.JID = ""
 	user.SetSession(nil)
 	jsonResponse(w, http.StatusOK, Response{true, "Logged out successfully."})
 }
@@ -351,6 +354,7 @@ func (prov *ProvisioningAPI) Login(w http.ResponseWriter, r *http.Request) {
 	}
 	user.ConnectionErrors = 0
 	user.JID = strings.Replace(user.Conn.Info.Wid, whatsappExt.OldUserSuffix, whatsappExt.NewUserSuffix, 1)
+	user.addToJIDMap()
 	user.SetSession(&session)
 	_ = c.WriteJSON(map[string]interface{}{
 		"success": true,

user.go

@@ -90,6 +90,18 @@ func (bridge *Bridge) GetUserByJID(userID types.WhatsAppID) *User {
 	return user
 }
 
+func (user *User) addToJIDMap() {
+	user.bridge.usersLock.Lock()
+	user.bridge.usersByJID[user.JID] = user
+	user.bridge.usersLock.Unlock()
+}
+
+func (user *User) removeFromJIDMap() {
+	user.bridge.usersLock.Lock()
+	delete(user.bridge.usersByJID, user.JID)
+	user.bridge.usersLock.Unlock()
+}
+
 func (bridge *Bridge) GetAllUsers() []*User {
 	bridge.usersLock.Lock()
 	defer bridge.usersLock.Unlock()
@@ -332,8 +344,11 @@ func (user *User) Login(ce *CommandEvent) {
 		_, _ = ce.Bot.SendMessageEvent(ce.RoomID, event.EventMessage, &msg)
 		return
 	}
+	// TODO there's a bit of duplication between this and the provisioning API login method
+	//      Also between the two logout methods (commands.go and provisioning.go)
 	user.ConnectionErrors = 0
 	user.JID = strings.Replace(user.Conn.Info.Wid, whatsappExt.OldUserSuffix, whatsappExt.NewUserSuffix, 1)
+	user.addToJIDMap()
 	user.SetSession(&session)
 	ce.Reply("Successfully logged in, synchronizing chats...")
 	user.PostLogin()