login: wait for 515 event

mautrix · Sep 16, 2024 · 7264c7f · 7264c7f
1 parent e6f9c8e
commit 7264c7f
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/go.mod b/go.mod
@@ -15,7 +15,7 @@ require (
 	github.com/tidwall/gjson v1.17.3
 	go.mau.fi/util v0.7.1-0.20240913091524-7617daa66719
 	go.mau.fi/webp v0.1.0
-	go.mau.fi/whatsmeow v0.0.0-20240911102933-bb3364aa3986
+	go.mau.fi/whatsmeow v0.0.0-20240916205343-ea8c175b2e2c
 	golang.org/x/exp v0.0.0-20240909161429-701f63a606c0
 	golang.org/x/image v0.20.0
 	golang.org/x/net v0.29.0

diff --git a/go.sum b/go.sum
@@ -82,8 +82,8 @@ go.mau.fi/util v0.7.1-0.20240913091524-7617daa66719 h1:sg1P/f4RHY1JuAwsPOjTCsZr8
 go.mau.fi/util v0.7.1-0.20240913091524-7617daa66719/go.mod h1:1Ixb8HWoVbl3rT6nAX6nV4iMkzn7KU/KXwE0Rn5RmsQ=
 go.mau.fi/webp v0.1.0 h1:BHObH/DcFntT9KYun5pDr0Ot4eUZO8k2C7eP7vF4ueA=
 go.mau.fi/webp v0.1.0/go.mod h1:e42Z+VMFrUMS9cpEwGRIor+lQWO8oUAyPyMtcL+NMt8=
-go.mau.fi/whatsmeow v0.0.0-20240911102933-bb3364aa3986 h1:7X+3826qoRBHPCtxY89tqMcYEsi9+OuWE6hHZfRc0qI=
-go.mau.fi/whatsmeow v0.0.0-20240911102933-bb3364aa3986/go.mod h1:BhHKalSq0qNtSCuGIUIvoJyU5KbT4a7k8DQ5yw1Ssk4=
+go.mau.fi/whatsmeow v0.0.0-20240916205343-ea8c175b2e2c h1:0pNAbeBNHdDmLbHG2bj+tnQnwE5YZVE83/QAfMlaYs4=
+go.mau.fi/whatsmeow v0.0.0-20240916205343-ea8c175b2e2c/go.mod h1:BhHKalSq0qNtSCuGIUIvoJyU5KbT4a7k8DQ5yw1Ssk4=
 go.mau.fi/zeroconfig v0.1.3 h1:As9wYDKmktjmNZW5i1vn8zvJlmGKHeVxHVIBMXsm4kM=
 go.mau.fi/zeroconfig v0.1.3/go.mod h1:NcSJkf180JT+1IId76PcMuLTNa1CzsFFZ0nBygIQM70=
 golang.org/x/crypto v0.27.0 h1:GXm2NjJrPaiv/h1tb2UH8QfgC/hOf/+z0p6PT8o1w7A=

diff --git a/pkg/connector/login.go b/pkg/connector/login.go
@@ -76,6 +76,7 @@ func (wa *WhatsAppConnector) CreateLogin(_ context.Context, user *bridgev2.User,
 
 		WaitForQRs:    exsync.NewEvent(),
 		LoginComplete: exsync.NewEvent(),
+		Received515:   exsync.NewEvent(),
 	}, nil
 }
 
@@ -92,6 +93,7 @@ type WALogin struct {
 	LoginSuccess  *events.PairSuccess
 	WaitForQRs    *exsync.Event
 	LoginComplete *exsync.Event
+	Received515   *exsync.Event
 	PrevQRIndex   atomic.Int32
 
 	Closed         atomic.Bool
@@ -109,6 +111,7 @@ func (wl *WALogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
 	device := wl.Main.DeviceStore.NewDevice()
 	wl.Client = whatsmeow.NewClient(device, waLog.Zerolog(wl.Log))
 	wl.Client.EnableAutoReconnect = false
+	wl.Client.DisableLoginAutoReconnect = true
 	wl.EventHandlerID = wl.Client.AddEventHandler(wl.handleEvent)
 	if err := wl.Main.updateProxy(ctx, wl.Client, true); err != nil {
 		return nil, err
@@ -234,6 +237,8 @@ func (wl *WALogin) handleEvent(rawEvt any) {
 	case *events.Connected, *events.ConnectFailure, *events.LoggedOut, *events.TemporaryBan:
 		wl.Log.Warn().Any("event_data", evt).Type("event_type", evt).Msg("Got unexpected disconnect event")
 		wl.LoginError = ErrUnexpectedDisconnect
+	case *events.ManualLoginReconnect:
+		wl.Received515.Set()
 	default:
 		wl.Log.Warn().Type("event_type", evt).Msg("Got unexpected event")
 		return
@@ -274,13 +279,21 @@ func (wl *WALogin) Wait(ctx context.Context) (*bridgev2.LoginStep, error) {
 			wl.Cancel()
 			return nil, ctx.Err()
 		case <-wl.LoginComplete.GetChan():
-			wl.Cancel()
 		}
 	}
-	wl.Log.Debug().Err(wl.LoginError).Msg("Login completed")
 	if wl.LoginError != nil {
+		wl.Log.Debug().Err(wl.LoginError).Msg("Login completed with error")
+		wl.Cancel()
 		return nil, wl.LoginError
 	}
+	wl.Log.Debug().Msg("Login completed without error, waiting for 515 event")
+	ctxTimeout, cancel := context.WithTimeout(ctx, 10*time.Second)
+	err := wl.Received515.Wait(ctxTimeout)
+	cancel()
+	wl.Cancel()
+	if err != nil {
+		return nil, fmt.Errorf("timed out waiting for 515: %w", err)
+	}
 
 	newLoginID := waid.MakeUserLoginID(wl.LoginSuccess.ID)
 	ul, err := wl.User.NewLogin(ctx, &database.UserLogin{