Cancel login if websocket is closed

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

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=

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

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