From 96c09b47521a0a028e95d0d6f19b372814503400 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 5 Mar 2024 13:11:27 +0200
Subject: [PATCH] Send acks before sending set active request

---
 libgm/client.go          | 16 +++++++---------
 libgm/longpoll.go        |  6 +++++-
 libgm/pair.go            |  2 +-
 libgm/pair_google.go     |  7 +++++--
 libgm/session_handler.go |  2 +-
 5 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/libgm/client.go b/libgm/client.go
index f6f6476..32804db 100644
--- a/libgm/client.go
+++ b/libgm/client.go
@@ -8,7 +8,6 @@ import (
 	"io"
 	"net/http"
 	"net/url"
-	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -75,8 +74,6 @@ type Client struct {
 
 	conversationsFetchedOnce bool
 
-	hackyDelaySetActive atomic.Bool
-
 	AuthData *AuthData
 	cfg      *gmproto.Config
 
@@ -156,18 +153,19 @@ func (c *Client) Connect() error {
 	//	return fmt.Errorf("failed to get web encryption key: %w", err)
 	//}
 	//c.updateWebEncryptionKey(webEncryptionKeyResponse.GetKey())
-	go c.doLongPoll(true)
+	go c.doLongPoll(true, c.postConnect)
 	c.sessionHandler.startAckInterval()
-	go c.postConnect()
 	return nil
 }
 
 func (c *Client) postConnect() {
-	// For some reason SetActiveSession fails if it's called immediately after reconnecting after a google login,
-	// so hackily delay it a few seconds to make it work.
-	if c.hackyDelaySetActive.CompareAndSwap(true, false) {
-		time.Sleep(3 * time.Second)
+	time.Sleep(2 * time.Second)
+	c.Logger.Debug().Msg("Sending acks before get updates request")
+	if c.skipCount > 0 {
+		c.Logger.Warn().Int("skip_count", c.skipCount).Msg("Skip count is still non-zero")
 	}
+	c.sessionHandler.sendAckRequest()
+	time.Sleep(1 * time.Second)
 	c.Logger.Debug().Msg("Sending get updates request")
 	err := c.SetActiveSession()
 	if err != nil {
diff --git a/libgm/longpoll.go b/libgm/longpoll.go
index e2facec..8685266 100644
--- a/libgm/longpoll.go
+++ b/libgm/longpoll.go
@@ -205,7 +205,7 @@ func tryReadBody(resp io.ReadCloser) []byte {
 	return data
 }
 
-func (c *Client) doLongPoll(loggedIn bool) {
+func (c *Client) doLongPoll(loggedIn bool, onFirstConnect func()) {
 	c.listenID++
 	listenID := c.listenID
 	listenReqID := uuid.NewString()
@@ -304,6 +304,10 @@ func (c *Client) doLongPoll(loggedIn bool) {
 				log.Debug().Msg("Ditto pinger is still waiting for previous ping, skipping new ping")
 			}
 		}
+		if onFirstConnect != nil {
+			go onFirstConnect()
+			onFirstConnect = nil
+		}
 		c.readLongPoll(&log, resp.Body)
 		c.longPollingConn = nil
 	}
diff --git a/libgm/pair.go b/libgm/pair.go
index 81ab02f..7478011 100644
--- a/libgm/pair.go
+++ b/libgm/pair.go
@@ -20,7 +20,7 @@ func (c *Client) StartLogin() (string, error) {
 		return "", err
 	}
 	c.updateTachyonAuthToken(registered.GetAuthKeyData())
-	go c.doLongPoll(false)
+	go c.doLongPoll(false, nil)
 	qr, err := c.GenerateQRCodeData(registered.GetPairingKey())
 	if err != nil {
 		return "", fmt.Errorf("failed to generate QR code: %w", err)
diff --git a/libgm/pair_google.go b/libgm/pair_google.go
index 4f59037..260047c 100644
--- a/libgm/pair_google.go
+++ b/libgm/pair_google.go
@@ -29,6 +29,7 @@ import (
 	"fmt"
 	"io"
 	"math/big"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -267,7 +268,10 @@ func (c *Client) DoGaiaPairing(ctx context.Context, emojiCallback func(string))
 		return fmt.Errorf("failed to parse destination UUID: %w", err)
 	}
 	c.AuthData.DestRegID = destRegUUID
-	go c.doLongPoll(false)
+	var longPollConnectWait sync.WaitGroup
+	longPollConnectWait.Add(1)
+	go c.doLongPoll(false, longPollConnectWait.Done)
+	longPollConnectWait.Wait()
 	ps := NewPairingSession()
 	clientInit, clientFinish, err := ps.PreparePayloads()
 	if err != nil {
@@ -300,7 +304,6 @@ func (c *Client) DoGaiaPairing(ctx context.Context, emojiCallback func(string))
 	c.AuthData.PairingID = ps.UUID
 	c.triggerEvent(&events.PairSuccessful{PhoneID: c.AuthData.Mobile.GetSourceID()})
 
-	c.hackyDelaySetActive.Store(true)
 	go func() {
 		// Sleep for a bit to let the phone save the pair data. If we reconnect too quickly,
 		// the phone won't recognize the session the bridge will get unpaired.
diff --git a/libgm/session_handler.go b/libgm/session_handler.go
index 6b21142..449f52c 100644
--- a/libgm/session_handler.go
+++ b/libgm/session_handler.go
@@ -254,7 +254,7 @@ func (s *SessionHandler) queueMessageAck(messageID string) {
 
 func (s *SessionHandler) startAckInterval() {
 	if s.ackTicker != nil {
-		s.ackTicker.Stop()
+		return
 	}
 	ticker := time.NewTicker(5 * time.Second)
 	s.ackTicker = ticker