gotd: handle all login token response types in QR login

Cherry-picked from https://github.com/gotd/td/commit/4c22747e9a0299b457f45084c3dbcbcbb5a7a5e7

pkg/gotd/telegram/auth/qrlogin/qrlogin.go

@@ -47,11 +47,21 @@ func (q QR) Export(ctx context.Context, exceptIDs ...int64) (Token, error) {
 		return Token{}, errors.Wrap(err, "export")
 	}
 
-	t, ok := result.(*tg.AuthLoginToken)
-	if !ok {
+	switch t := result.(type) {
+	case *tg.AuthLoginToken:
+		return NewToken(t.Token, t.Expires), nil
+	case *tg.AuthLoginTokenSuccess:
+		// Token was already accepted, authentication successful
+		// Return empty token since no new token is needed
+		return Token{}, nil
+	case *tg.AuthLoginTokenMigrateTo:
+		// Migration needed
+		return Token{}, &MigrationNeededError{
+			MigrateTo: t,
+		}
+	default:
 		return Token{}, errors.Errorf("unexpected type %T", result)
 	}
-	return NewToken(t.Token, t.Expires), nil
 }
 
 // Accept accepts given token.
@@ -147,6 +157,18 @@ func (q QR) Auth(
 	if err != nil {
 		return nil, err
 	}
+
+	// If token is empty, it means AuthLoginTokenSuccess was returned
+	// and authentication is already complete, but we should wait for the signal.
+	if token.Empty() {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-loggedIn:
+			return q.Import(ctx)
+		}
+	}
+
 	timer := q.clock.Timer(until(token))
 	defer clock.StopTimer(timer)
 
@@ -163,6 +185,13 @@ func (q QR) Auth(
 			if err != nil {
 				return nil, err
 			}
+
+			if t.Empty() {
+				// If empty token, it means AuthLoginTokenSuccess was returned.
+				// QR was scanned and accepted, break to import.
+				break
+			}
+
 			token = t
 			timer.Reset(until(token))
 

pkg/gotd/telegram/auth/qrlogin/qrlogin_test.go

@@ -2,12 +2,13 @@ package qrlogin
 
 import (
 	"context"
+	"runtime"
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/require"
-
 	"github.com/gotd/neo"
+	"github.com/stretchr/testify/require"
+	"rsc.io/qr"
 
 	"go.mau.fi/mautrix-telegram/pkg/gotd/constant"
 	"go.mau.fi/mautrix-telegram/pkg/gotd/testutil"
@@ -185,3 +186,140 @@ func TestQR_Auth(t *testing.T) {
 
 	a.NoError(<-done)
 }
+
+func TestMigrationNeededError_Error(t *testing.T) {
+	a := require.New(t)
+	err := &MigrationNeededError{
+		MigrateTo: &tg.AuthLoginTokenMigrateTo{
+			DCID: 2,
+		},
+	}
+	a.Equal("migration to 2 needed", err.Error())
+}
+
+// Mock dispatcher that implements the required interface.
+type mockDispatcher struct {
+	handler tg.LoginTokenHandler
+}
+
+func (m *mockDispatcher) OnLoginToken(h tg.LoginTokenHandler) {
+	m.handler = h
+}
+
+func TestOnLoginToken(t *testing.T) {
+	if runtime.GOOS == "darwin" {
+		t.Skip("skipping on macOS")
+	}
+
+	a := require.New(t)
+
+	dispatcher := &mockDispatcher{}
+	loggedIn := OnLoginToken(dispatcher)
+
+	// Verify that handler was set.
+	a.NotNil(dispatcher.handler)
+
+	// Test the handler
+	ctx := context.Background()
+	entities := tg.Entities{}
+	update := &tg.UpdateLoginToken{}
+
+	// First call should send to channel.
+	done := make(chan error, 1)
+	go func() {
+		done <- dispatcher.handler(ctx, entities, update)
+	}()
+
+	// Should receive signal.
+	select {
+	case <-loggedIn:
+		// Good
+	case <-time.After(time.Second * 5):
+		t.Fatal("should receive signal")
+	}
+
+	// Handler should return nil.
+	a.NoError(<-done)
+
+	// Second call when channel is full should not block.
+	err := dispatcher.handler(ctx, entities, update)
+	a.NoError(err)
+}
+
+func TestToken_Image(t *testing.T) {
+	a := require.New(t)
+	token := NewToken([]byte("test_token"), int(time.Now().Unix()))
+
+	// Test with valid QR level.
+	img, err := token.Image(qr.L)
+	a.NoError(err)
+	a.NotNil(img)
+
+	// Test with different QR levels.
+	levels := []qr.Level{qr.L, qr.M, qr.Q, qr.H}
+
+	for _, level := range levels {
+		img, err := token.Image(level)
+		a.NoError(err)
+		a.NotNil(img)
+	}
+}
+
+func TestQR_Import_WithMigration(t *testing.T) {
+	ctx := context.Background()
+	a := require.New(t)
+
+	// Test with migration function.
+	migrateCalled := false
+	migrate := func(ctx context.Context, dcID int) error {
+		migrateCalled = true
+		a.Equal(2, dcID)
+		return nil
+	}
+
+	mock, qr := testQR(t, migrate)
+
+	auth := &tg.AuthAuthorization{
+		User: &tg.User{ID: 10},
+	}
+
+	// First call returns migration needed.
+	mock.ExpectCall(&tg.AuthExportLoginTokenRequest{
+		APIID:   constant.TestAppID,
+		APIHash: constant.TestAppHash,
+	}).ThenResult(&tg.AuthLoginTokenMigrateTo{
+		DCID:  2,
+		Token: testToken.token,
+	}).ExpectCall(&tg.AuthImportLoginTokenRequest{
+		Token: testToken.token,
+	}).ThenResult(&tg.AuthLoginTokenSuccess{
+		Authorization: auth,
+	})
+
+	result, err := qr.Import(ctx)
+	a.NoError(err)
+	a.Equal(auth, result)
+	a.True(migrateCalled)
+}
+
+func TestQR_Import_MigrationError(t *testing.T) {
+	ctx := context.Background()
+	a := require.New(t)
+
+	// Test with migration function that returns error,
+	migrate := func(ctx context.Context, dcID int) error {
+		return testutil.TestError()
+	}
+
+	mock, qr := testQR(t, migrate)
+
+	mock.ExpectCall(&tg.AuthExportLoginTokenRequest{
+		APIID:   constant.TestAppID,
+		APIHash: constant.TestAppHash,
+	}).ThenResult(&tg.AuthLoginTokenMigrateTo{
+		DCID: 2,
+	})
+
+	_, err := qr.Import(ctx)
+	a.ErrorIs(err, testutil.TestError())
+}

pkg/gotd/telegram/auth/qrlogin/token.go

@@ -59,6 +59,11 @@ func (t Token) String() string {
 	return base64.URLEncoding.EncodeToString(t.token)
 }
 
+// Empty reports whether token is empty.
+func (t Token) Empty() bool {
+	return len(t.token) == 0
+}
+
 // URL returns login URL.
 //
 // See https://core.telegram.org/api/qr-login#exporting-a-login-token.