From d42c4722c939246bf9042838f8c8ce0b570b336c Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 8 Mar 2023 19:49:43 +0200
Subject: [PATCH] Fill usersByID properly

---
 commands.go     |  2 +-
 provisioning.go |  4 +---
 user.go         | 21 ++++++++++++++++++---
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/commands.go b/commands.go
index cd719d2..676ede3 100644
--- a/commands.go
+++ b/commands.go
@@ -281,7 +281,7 @@ var cmdLogout = &commands.FullHandler{
 
 func fnLogout(ce *WrappedCommandEvent) {
 	wasLoggedIn := ce.User.DiscordID != ""
-	ce.User.Logout()
+	ce.User.Logout(false)
 	if wasLoggedIn {
 		ce.Reply("Logged out successfully.")
 	} else {
diff --git a/provisioning.go b/provisioning.go
index fa803dc..7116873 100644
--- a/provisioning.go
+++ b/provisioning.go
@@ -227,7 +227,7 @@ func (p *ProvisioningAPI) logout(w http.ResponseWriter, r *http.Request) {
 	} else {
 		msg = "User wasn't logged in."
 	}
-	user.Logout()
+	user.Logout(false)
 	jsonResponse(w, http.StatusOK, Response{true, msg})
 }
 
@@ -328,8 +328,6 @@ func (p *ProvisioningAPI) qrLogin(w http.ResponseWriter, r *http.Request) {
 			}
 
 			log.Infofln("Logged in as %s#%s (%s)", discordUser.Username, discordUser.Discriminator, discordUser.UserID)
-			user.DiscordID = discordUser.UserID
-			user.Update()
 
 			if err = user.Login(discordUser.Token); err != nil {
 				log.Errorln("Failed to connect after logging in:", err)
diff --git a/user.go b/user.go
index 73cd82f..0bc47a0 100644
--- a/user.go
+++ b/user.go
@@ -458,7 +458,7 @@ func (user *User) IsLoggedIn() bool {
 	return user.DiscordToken != ""
 }
 
-func (user *User) Logout() {
+func (user *User) Logout(isOverwriting bool) {
 	user.Lock()
 	defer user.Unlock()
 
@@ -479,9 +479,16 @@ func (user *User) Logout() {
 	}
 
 	user.Session = nil
-	user.DiscordID = ""
 	user.DiscordToken = ""
 	user.ReadStateVersion = 0
+	if !isOverwriting {
+		user.bridge.usersLock.Lock()
+		if user.bridge.usersByID[user.DiscordID] == user {
+			delete(user.bridge.usersByID, user.DiscordID)
+		}
+		user.bridge.usersLock.Unlock()
+	}
+	user.DiscordID = ""
 	user.Update()
 	user.log.Infoln("User logged out")
 }
@@ -599,7 +606,15 @@ func (user *User) readyHandler(_ *discordgo.Session, r *discordgo.Ready) {
 	user.bridgeStateLock.Unlock()
 
 	if user.DiscordID != r.User.ID {
+		user.bridge.usersLock.Lock()
 		user.DiscordID = r.User.ID
+		if previousUser, ok := user.bridge.usersByID[user.DiscordID]; ok && previousUser != user {
+			user.log.Warnfln("Another user (%s) is logged in with same Discord ID, logging them out", previousUser.MXID)
+			// TODO send notice?
+			previousUser.Logout(true)
+		}
+		user.bridge.usersByID[user.DiscordID] = user
+		user.bridge.usersLock.Unlock()
 		user.Update()
 	}
 	user.BridgeState.Send(status.BridgeState{StateEvent: status.StateBackfilling})
@@ -833,7 +848,7 @@ func (user *User) invalidAuthHandler(_ *discordgo.Session, _ *discordgo.InvalidA
 	user.log.Infoln("Got logged out from Discord due to invalid token")
 	user.wasLoggedOut = true
 	user.BridgeState.Send(status.BridgeState{StateEvent: status.StateBadCredentials, Error: "dc-websocket-disconnect-4004", Message: "Discord access token is no longer valid, please log in again"})
-	go user.Logout()
+	go user.Logout(false)
 }
 
 func (user *User) guildCreateHandler(_ *discordgo.Session, g *discordgo.GuildCreate) {