From fec43431f780baccdedff02a1a822779ad6dfb67 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 7 Feb 2021 22:14:13 +0200
Subject: [PATCH] Cancel login if websocket is closed

---
 go.mod          |  2 +-
 go.sum          |  2 ++
 provisioning.go | 19 ++++++++++++++++++-
 user.go         |  2 +-
 4 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index fc526c8..6c16fbc 100644
--- a/go.mod
+++ b/go.mod
@@ -16,4 +16,4 @@ require (
 	maunium.net/go/mautrix v0.8.0
 )
 
-replace github.com/Rhymen/go-whatsapp => github.com/tulir/go-whatsapp v0.3.17
+replace github.com/Rhymen/go-whatsapp => github.com/tulir/go-whatsapp v0.3.18
diff --git a/go.sum b/go.sum
index 18bc1e4..adfc23d 100644
--- a/go.sum
+++ b/go.sum
@@ -137,6 +137,8 @@ github.com/tulir/go-whatsapp v0.3.16 h1:NfcXC2DQXwls3qkAjbFqSeoMX+rUbbpBBGGvCXI3
 github.com/tulir/go-whatsapp v0.3.16/go.mod h1:U5+sm33vrv3wz62YyRM/VS7q2ObXkxU4Xqj/3KOmN9o=
 github.com/tulir/go-whatsapp v0.3.17 h1:HMRT6HzP1seUt5P0waD8CxThB2bfBgKX2uVjOoXCaf8=
 github.com/tulir/go-whatsapp v0.3.17/go.mod h1:U5+sm33vrv3wz62YyRM/VS7q2ObXkxU4Xqj/3KOmN9o=
+github.com/tulir/go-whatsapp v0.3.18 h1:45pkdjEnAp6yV4RTSWCZn2Eenep+MN7Kndf/rRXQtyY=
+github.com/tulir/go-whatsapp v0.3.18/go.mod h1:U5+sm33vrv3wz62YyRM/VS7q2ObXkxU4Xqj/3KOmN9o=
 golang.org/x/crypto v0.0.0-20170930174604-9419663f5a44/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
diff --git a/provisioning.go b/provisioning.go
index 0c6dfb1..3a6c72a 100644
--- a/provisioning.go
+++ b/provisioning.go
@@ -360,8 +360,25 @@ func (prov *ProvisioningAPI) Login(w http.ResponseWriter, r *http.Request) {
 			})
 		}
 	}()
+
+	go func() {
+		// Read everything so SetCloseHandler() works
+		for {
+			_, _, err = c.ReadMessage()
+			if err != nil {
+				break
+			}
+		}
+	}()
+	ctx, cancel := context.WithCancel(context.Background())
+	c.SetCloseHandler(func(code int, text string) error {
+		user.log.Debugfln("Login websocket closed (%d), cancelling login", code)
+		cancel()
+		return nil
+	})
+
 	user.log.Debugln("Starting login via provisioning API")
-	session, err := user.Conn.LoginWithRetry(qrChan, user.bridge.Config.Bridge.LoginQRRegenCount)
+	session, err := user.Conn.LoginWithRetry(qrChan, ctx, user.bridge.Config.Bridge.LoginQRRegenCount)
 	qrChan <- "stop"
 	if err != nil {
 		var msg string
diff --git a/user.go b/user.go
index b07b28e..3c27dfa 100644
--- a/user.go
+++ b/user.go
@@ -384,7 +384,7 @@ func (user *User) Login(ce *CommandEvent) {
 	qrChan := make(chan string, 3)
 	eventIDChan := make(chan id.EventID, 1)
 	go user.loginQrChannel(ce, qrChan, eventIDChan)
-	session, err := user.Conn.LoginWithRetry(qrChan, user.bridge.Config.Bridge.LoginQRRegenCount)
+	session, err := user.Conn.LoginWithRetry(qrChan, nil, user.bridge.Config.Bridge.LoginQRRegenCount)
 	qrChan <- "stop"
 	if err != nil {
 		var eventID id.EventID