From 538f2a2ec036fdeba11a4b22255fc181ee89e08c Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Wed, 9 Apr 2025 10:35:08 +0300
Subject: [PATCH] client: Wait before returning from disconnect

---
 pkg/connector/client.go     | 16 +++++++++++-----
 pkg/connector/loginphone.go |  6 +++++-
 pkg/connector/loginqr.go    |  6 +++++-
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/pkg/connector/client.go b/pkg/connector/client.go
index 217198d4..5b5f1a30 100644
--- a/pkg/connector/client.go
+++ b/pkg/connector/client.go
@@ -66,6 +66,7 @@ type TelegramClient struct {
 	updatesManager *updates.Manager
 	clientCtx      context.Context
 	clientCancel   context.CancelFunc
+	clientCloseC   <-chan struct{}
 
 	appConfigLock sync.Mutex
 	appConfig     map[string]any
@@ -370,11 +371,13 @@ func NewTelegramClient(ctx context.Context, tc *TelegramConnector, login *bridge
 // connectTelegramClient blocks until client is connected, calling Run
 // internally.
 // Technique from: https://github.com/gotd/contrib/blob/master/bg/connect.go
-func connectTelegramClient(ctx context.Context, cancel context.CancelFunc, client *telegram.Client) error {
+func connectTelegramClient(ctx context.Context, cancel context.CancelFunc, client *telegram.Client) (<-chan struct{}, error) {
 	errC := make(chan error, 1)
 	initDone := make(chan struct{})
+	closeC := make(chan struct{})
 	go func() {
 		defer close(errC)
+		defer close(closeC)
 		errC <- client.Run(ctx, func(ctx context.Context) error {
 			close(initDone)
 			<-ctx.Done()
@@ -388,13 +391,13 @@ func connectTelegramClient(ctx context.Context, cancel context.CancelFunc, clien
 	select {
 	case <-ctx.Done(): // context canceled
 		cancel()
-		return fmt.Errorf("context cancelled before init done: %w", ctx.Err())
+		return nil, fmt.Errorf("context cancelled before init done: %w", ctx.Err())
 	case err := <-errC: // startup timeout
 		cancel()
-		return fmt.Errorf("client connection timeout: %w", err)
+		return nil, fmt.Errorf("client connection timeout: %w", err)
 	case <-initDone: // init done
 	}
-	return nil
+	return closeC, nil
 }
 
 func (t *TelegramClient) onDead() {
@@ -497,7 +500,7 @@ func (t *TelegramClient) Connect(ctx context.Context) {
 
 	var err error
 	t.clientCtx, t.clientCancel = context.WithCancel(ctx)
-	if err = connectTelegramClient(t.clientCtx, t.clientCancel, t.client); err != nil {
+	if t.clientCloseC, err = connectTelegramClient(t.clientCtx, t.clientCancel, t.client); err != nil {
 		t.sendBadCredentialsOrUnknownError(err)
 		return
 	}
@@ -536,6 +539,9 @@ func (t *TelegramClient) Disconnect() {
 	if t.clientCancel != nil {
 		t.clientCancel()
 	}
+	if t.clientCloseC != nil {
+		<-t.clientCloseC
+	}
 }
 
 func (t *TelegramClient) getInputUser(ctx context.Context, id int64) (*tg.InputUser, error) {
diff --git a/pkg/connector/loginphone.go b/pkg/connector/loginphone.go
index c3ed10c5..a079d017 100644
--- a/pkg/connector/loginphone.go
+++ b/pkg/connector/loginphone.go
@@ -44,6 +44,7 @@ type PhoneLogin struct {
 	authClient       *telegram.Client
 	authClientCtx    context.Context
 	authClientCancel context.CancelFunc
+	authClientCloseC <-chan struct{}
 
 	phone string
 	hash  string
@@ -55,6 +56,9 @@ func (p *PhoneLogin) Cancel() {
 	if p.authClientCancel != nil {
 		p.authClientCancel()
 	}
+	if p.authClientCloseC != nil {
+		<-p.authClientCloseC
+	}
 }
 
 func (p *PhoneLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
@@ -85,7 +89,7 @@ func (p *PhoneLogin) SubmitUserInput(ctx context.Context, input map[string]strin
 		})
 		var err error
 		p.authClientCtx, p.authClientCancel = context.WithTimeoutCause(log.WithContext(context.Background()), time.Hour, errors.New("phone login took over one hour"))
-		if err = connectTelegramClient(p.authClientCtx, p.authClientCancel, p.authClient); err != nil {
+		if p.authClientCloseC, err = connectTelegramClient(p.authClientCtx, p.authClientCancel, p.authClient); err != nil {
 			return nil, err
 		}
 		sentCode, err := p.authClient.Auth().SendCode(p.authClientCtx, p.phone, auth.SendCodeOptions{})
diff --git a/pkg/connector/loginqr.go b/pkg/connector/loginqr.go
index 823b8b24..299464c5 100644
--- a/pkg/connector/loginqr.go
+++ b/pkg/connector/loginqr.go
@@ -47,6 +47,7 @@ type QRLogin struct {
 
 	authClientCtx    context.Context
 	authClientCancel context.CancelFunc
+	authClientCloseC <-chan struct{}
 
 	auth    chan qrAuthResult
 	qrToken chan qrlogin.Token
@@ -61,6 +62,9 @@ func (q *QRLogin) Cancel() {
 	if q.authClientCancel != nil {
 		q.authClientCancel()
 	}
+	if q.authClientCloseC != nil {
+		<-q.authClientCloseC
+	}
 }
 
 func (q *QRLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
@@ -85,7 +89,7 @@ func (q *QRLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
 
 	var err error
 	q.authClientCtx, q.authClientCancel = context.WithTimeoutCause(log.WithContext(context.Background()), time.Hour, errors.New("phone login took over one hour"))
-	if err = connectTelegramClient(q.authClientCtx, q.authClientCancel, q.authClient); err != nil {
+	if q.authClientCloseC, err = connectTelegramClient(q.authClientCtx, q.authClientCancel, q.authClient); err != nil {
 		return nil, err
 	}