diff --git a/pkg/gotd/telegram/auth/qrlogin/qrlogin.go b/pkg/gotd/telegram/auth/qrlogin/qrlogin.go index 28347025..64c44b04 100644 --- a/pkg/gotd/telegram/auth/qrlogin/qrlogin.go +++ b/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)) diff --git a/pkg/gotd/telegram/auth/qrlogin/qrlogin_test.go b/pkg/gotd/telegram/auth/qrlogin/qrlogin_test.go index 95f282fe..31453fc1 100644 --- a/pkg/gotd/telegram/auth/qrlogin/qrlogin_test.go +++ b/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()) +} diff --git a/pkg/gotd/telegram/auth/qrlogin/token.go b/pkg/gotd/telegram/auth/qrlogin/token.go index 0107526f..f79b4af9 100644 --- a/pkg/gotd/telegram/auth/qrlogin/token.go +++ b/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.