From d91f705e95ecdbc7ca052550bf2fc9ebd48765ad Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 30 Aug 2019 22:04:57 +0300
Subject: [PATCH] Fix most unexpected logout bugs, handle
 connection-in-progress states better and send warning message if sending to
 whatsapp times out

---
 commands.go | 28 ++++++++++++++++++++++++----
 main.go     |  2 +-
 portal.go   | 13 +++++++++++--
 user.go     |  4 ++++
 4 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/commands.go b/commands.go
index 413ba9f..e9f10bd 100644
--- a/commands.go
+++ b/commands.go
@@ -223,6 +223,13 @@ func (handler *CommandHandler) CommandReconnect(ce *CommandEvent) {
 		} else {
 			ce.Reply("Unknown error while reconnecting: %v", err)
 		}
+		ce.User.log.Debugln("Disconnecting due to failed session restore in reconnect command...")
+		sess, err := ce.User.Conn.Disconnect()
+		if err != nil {
+			ce.User.log.Errorln("Failed to disconnect after failed session restore in reconnect command:", err)
+		} else if len(sess.Wid) > 0 {
+			ce.User.SetSession(&sess)
+		}
 		return
 	}
 	ce.User.ConnectionErrors = 0
@@ -259,8 +266,9 @@ func (handler *CommandHandler) CommandDisconnect(ce *CommandEvent) {
 		ce.User.log.Warnln("Error while disconnecting:", err)
 		ce.Reply("Unknown error while disconnecting: %v", err)
 		return
+	} else if len(sess.Wid) > 0 {
+		ce.User.SetSession(&sess)
 	}
-	ce.User.SetSession(&sess)
 	ce.Reply("Successfully disconnected. Use the `reconnect` command to reconnect.")
 }
 
@@ -268,13 +276,25 @@ const cmdPingHelp = `ping - Check your connection to WhatsApp.`
 
 func (handler *CommandHandler) CommandPing(ce *CommandEvent) {
 	if ce.User.Session == nil {
-		ce.Reply("You're not logged into WhatsApp.")
+		if ce.User.IsLoginInProgress() {
+			ce.Reply("You're not logged into WhatsApp, but there's a login in progress.")
+		} else {
+			ce.Reply("You're not logged into WhatsApp.")
+		}
 	} else if ce.User.Conn == nil {
 		ce.Reply("You don't have a WhatsApp connection.")
 	} else if ok, err := ce.User.Conn.AdminTest(); err != nil {
-		ce.Reply("Connection not OK: %v", err)
+		if ce.User.IsLoginInProgress() {
+			ce.Reply("Connection not OK: %v, but login in progress", err)
+		} else {
+			ce.Reply("Connection not OK: %v", err)
+		}
 	} else if !ok {
-		ce.Reply("Connection not OK, but no error received")
+		if ce.User.IsLoginInProgress() {
+			ce.Reply("Connection not OK, but no error received and login in progress")
+		} else {
+			ce.Reply("Connection not OK, but no error received")
+		}
 	} else {
 		ce.Reply("Connection to WhatsApp OK")
 	}
diff --git a/main.go b/main.go
index b8a1951..e02adf8 100644
--- a/main.go
+++ b/main.go
@@ -282,7 +282,7 @@ func (bridge *Bridge) Stop() {
 		sess, err := user.Conn.Disconnect()
 		if err != nil {
 			bridge.Log.Errorfln("Error while disconnecting %s: %v", user.MXID, err)
-		} else {
+		} else if len(sess.Wid) > 0 {
 			user.SetSession(&sess)
 		}
 	}
diff --git a/portal.go b/portal.go
index 5809264..d799628 100644
--- a/portal.go
+++ b/portal.go
@@ -1092,8 +1092,11 @@ func (portal *Portal) sendMatrixConnectionError(sender *User, eventID string) bo
 		if portal.IsPrivateChat() {
 			inRoom = " in your management room"
 		}
-		msg := format.RenderMarkdown(fmt.Sprintf("\u26a0 You are not connected to WhatsApp, so your message was not bridged. " +
-			"Use `%s reconnect`%s to reconnect.", portal.bridge.Config.Bridge.CommandPrefix, inRoom))
+		reconnect := fmt.Sprintf("Use `%s reconnect`%s to reconnect.", portal.bridge.Config.Bridge.CommandPrefix, inRoom)
+		if sender.IsLoginInProgress() {
+			reconnect = "You have a login attempt in progress, please wait."
+		}
+		msg := format.RenderMarkdown("\u26a0 You are not connected to WhatsApp, so your message was not bridged. " + reconnect)
 		msg.MsgType = mautrix.MsgNotice
 		_, err := portal.MainIntent().SendMessageEvent(portal.MXID, mautrix.EventMessage, msg)
 		if err != nil {
@@ -1227,6 +1230,12 @@ func (portal *Portal) HandleMatrixMessage(sender *User, evt *mautrix.Event) {
 	_, err = sender.Conn.Send(info)
 	if err != nil {
 		portal.log.Errorfln("Error handling Matrix event %s: %v", evt.ID, err)
+		msg := format.RenderMarkdown(fmt.Sprintf("\u26a0 Your message may not have been bridged: %v", err))
+		msg.MsgType = mautrix.MsgNotice
+		_, err := portal.MainIntent().SendMessageEvent(portal.MXID, mautrix.EventMessage, msg)
+		if err != nil {
+			portal.log.Errorln("Failed to send bridging failure message:", err)
+		}
 	} else {
 		portal.log.Debugln("Handled Matrix event:", evt)
 	}
diff --git a/user.go b/user.go
index 1d72c2d..af7f171 100644
--- a/user.go
+++ b/user.go
@@ -233,6 +233,10 @@ func (user *User) IsConnected() bool {
 	return user.Conn != nil && user.Conn.IsConnected() && user.Conn.IsLoggedIn()
 }
 
+func (user *User) IsLoginInProgress() bool {
+	return user.Conn != nil && user.Conn.IsLoginInProgress()
+}
+
 func (user *User) loginQrChannel(ce *CommandEvent, qrChan <-chan string, eventIDChan chan<- string) {
 	var qrEventID string
 	for code := range qrChan {