From 9c2475dfbb94990e1b13e5f0b53a3459f7e333ea Mon Sep 17 00:00:00 2001 From: Arash mo <75903249+arashm404@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:36:06 +0200 Subject: [PATCH] fix(telegram): restore DC-only sessions and migrate QR on the same client Persisting a session blob with DC set but no auth key used to fail with "corrupted key" because a zero key has a non-zero SHA-1 id. Treat that as an unauthenticated handshake pinned to the stored DC (same as pool.Session.Migrate in memory). Also derive AuthKeyID when it is missing, and persist the session before signalling Run ready so SessionStorage wrappers observe DC/auth key ID in the callback. QR LoginTokenMigrateTo is bound to the current connection's auth key. Auth() now migrates via Options.Migrate and importLoginToken on the same client instead of leaving callers to open a second Client with Options.DC (that path AUTH_TOKEN_EXPIRED). --- telegram/auth/qrlogin/qrlogin.go | 70 ++++++++++++++-------- telegram/auth/qrlogin/qrlogin_test.go | 38 ++++++++++++ telegram/migrate_to_dc.go | 6 ++ telegram/options.go | 8 +++ telegram/session.go | 31 ++++++++-- telegram/session_restore_test.go | 83 +++++++++++++++++++++++++++ 6 files changed, 208 insertions(+), 28 deletions(-) create mode 100644 telegram/session_restore_test.go diff --git a/telegram/auth/qrlogin/qrlogin.go b/telegram/auth/qrlogin/qrlogin.go index 09e1c9928c..25b2bc43ce 100644 --- a/telegram/auth/qrlogin/qrlogin.go +++ b/telegram/auth/qrlogin/qrlogin.go @@ -36,6 +36,12 @@ func NewQR(api *tg.Client, appID int, appHash string, opts Options) QR { // Export exports new login token. // +// If Telegram returns LoginTokenMigrateTo, Export returns MigrationNeededError. +// Callers must migrate the *same* telegram.Client (Client.MigrateTo) and then +// auth.importLoginToken with the token. Do not start a new Client with +// Options.DC: that creates a new auth key, and import fails with +// AUTH_TOKEN_EXPIRED. Client.QR() already passes MigrateTo via Options.Migrate. +// // See https://core.telegram.org/api/qr-login#exporting-a-login-token. func (q QR) Export(ctx context.Context, exceptIDs ...int64) (Token, error) { result, err := q.api.AuthExportLoginToken(ctx, &tg.AuthExportLoginTokenRequest{ @@ -71,6 +77,34 @@ func (q QR) Accept(ctx context.Context, t Token) (*tg.Authorization, error) { return AcceptQR(ctx, q.api, t) } +// importMigrated switches the same client to t.DCID and imports t.Token. +// Telegram binds the login token to the original auth key, so this must not +// be implemented by opening a new Client with Options.DC. +func (q QR) importMigrated(ctx context.Context, t *tg.AuthLoginTokenMigrateTo) (*tg.AuthAuthorization, error) { + if q.migrate == nil { + return nil, &MigrationNeededError{MigrateTo: t} + } + if err := q.migrate(ctx, t.DCID); err != nil { + return nil, errors.Wrap(err, "migrate") + } + + res, err := q.api.AuthImportLoginToken(ctx, t.Token) + if err != nil { + return nil, errors.Wrap(err, "import") + } + + success, ok := res.(*tg.AuthLoginTokenSuccess) + if !ok { + return nil, errors.Errorf("unexpected type %T", res) + } + + auth, ok := success.Authorization.(*tg.AuthAuthorization) + if !ok { + return nil, errors.Errorf("unexpected type %T", success.Authorization) + } + return auth, nil +} + // Import imports accepted token. // // See https://core.telegram.org/api/qr-login#confirming-importing-the-login-token. @@ -85,30 +119,7 @@ func (q QR) Import(ctx context.Context) (*tg.AuthAuthorization, error) { switch t := result.(type) { case *tg.AuthLoginTokenMigrateTo: - if q.migrate == nil { - return nil, &MigrationNeededError{ - MigrateTo: t, - } - } - if err := q.migrate(ctx, t.DCID); err != nil { - return nil, errors.Wrap(err, "migrate") - } - - res, err := q.api.AuthImportLoginToken(ctx, t.Token) - if err != nil { - return nil, errors.Wrap(err, "import") - } - - success, ok := res.(*tg.AuthLoginTokenSuccess) - if !ok { - return nil, errors.Errorf("unexpected type %T", res) - } - - auth, ok := success.Authorization.(*tg.AuthAuthorization) - if !ok { - return nil, errors.Errorf("unexpected type %T", success.Authorization) - } - return auth, nil + return q.importMigrated(ctx, t) case *tg.AuthLoginTokenSuccess: auth, ok := t.Authorization.(*tg.AuthAuthorization) if !ok { @@ -155,6 +166,13 @@ func (q QR) Auth( token, err := q.Export(ctx, exceptIDs...) if err != nil { + var mig *MigrationNeededError + if errors.As(err, &mig) { + // First export asked us to rotate DC. Stay on this client: + // migrate + importLoginToken. A new Client with Options.DC + // would AUTH_TOKEN_EXPIRED. + return q.importMigrated(ctx, mig.MigrateTo) + } return nil, err } @@ -183,6 +201,10 @@ func (q QR) Auth( case <-timer.C(): t, err := q.Export(ctx, exceptIDs...) if err != nil { + var mig *MigrationNeededError + if errors.As(err, &mig) { + return q.importMigrated(ctx, mig.MigrateTo) + } return nil, err } diff --git a/telegram/auth/qrlogin/qrlogin_test.go b/telegram/auth/qrlogin/qrlogin_test.go index 21af328f32..2c257129fe 100644 --- a/telegram/auth/qrlogin/qrlogin_test.go +++ b/telegram/auth/qrlogin/qrlogin_test.go @@ -299,6 +299,44 @@ func TestQR_Import_WithMigration(t *testing.T) { a.True(migrateCalled) } +func TestQR_Auth_FirstExportMigrateTo(t *testing.T) { + ctx := context.Background() + a := require.New(t) + + migrateCalled := false + migrate := func(ctx context.Context, dcID int) error { + migrateCalled = true + a.Equal(4, dcID) + return nil + } + + mock, qr := testQR(t, migrate) + auth := &tg.AuthAuthorization{ + User: &tg.User{ID: 10}, + } + + mock.ExpectCall(&tg.AuthExportLoginTokenRequest{ + APIID: constant.TestAppID, + APIHash: constant.TestAppHash, + }).ThenResult(&tg.AuthLoginTokenMigrateTo{ + DCID: 4, + Token: testToken.token, + }).ExpectCall(&tg.AuthImportLoginTokenRequest{ + Token: testToken.token, + }).ThenResult(&tg.AuthLoginTokenSuccess{ + Authorization: auth, + }) + + loggedIn := make(chan struct{}) + result, err := qr.Auth(ctx, loggedIn, func(ctx context.Context, token Token) error { + t.Fatal("show should not be called when first export is LoginTokenMigrateTo") + return nil + }) + a.NoError(err) + a.Equal(auth, result) + a.True(migrateCalled) +} + func TestQR_Import_MigrationError(t *testing.T) { ctx := context.Background() a := require.New(t) diff --git a/telegram/migrate_to_dc.go b/telegram/migrate_to_dc.go index 156331e727..6c03bf83e3 100644 --- a/telegram/migrate_to_dc.go +++ b/telegram/migrate_to_dc.go @@ -62,6 +62,12 @@ func (c *Client) migrateToDc(ctx context.Context, dcID int) error { } // MigrateTo forces client to migrate to another DC. +// +// The in-memory auth key is cleared (pool.Session.Migrate) and the primary +// connection is restarted against dcID. Use this on the same Client after +// auth.exportLoginToken returns LoginTokenMigrateTo — the exported token is +// bound to that connection. Do not open a new Client with Options.DC to +// import the token. func (c *Client) MigrateTo(ctx context.Context, dcID int) error { // Acquire or cancel. select { diff --git a/telegram/options.go b/telegram/options.go index 754c51572e..a08c1f0c84 100644 --- a/telegram/options.go +++ b/telegram/options.go @@ -34,6 +34,14 @@ type Options struct { // DC ID to connect. // // If not provided, 2 will be used by default. + // + // This pins the initial handshake DC for a new client. It is not a + // substitute for Client.MigrateTo: after auth.exportLoginToken returns + // LoginTokenMigrateTo, the login token is bound to the current + // connection's auth key. Starting a second Client with Options.DC set + // to the target DC negotiates a new key, and auth.importLoginToken + // then fails with AUTH_TOKEN_EXPIRED. Use Client.QR() (which already + // passes MigrateTo) or call Client.MigrateTo on the same client. DC int // DCList is initial list of addresses to connect. diff --git a/telegram/session.go b/telegram/session.go index 9448de9094..32746e0835 100644 --- a/telegram/session.go +++ b/telegram/session.go @@ -34,10 +34,28 @@ func (c *Client) restoreConnection(ctx context.Context) error { data.DC = prev.DC } + if len(data.AuthKey) == 0 { + // DC-only state: pin DC for a new handshake. pool.Session.Migrate + // does the same in memory (set DC, zero key). A session blob with DC + // but no auth key used to fail with "corrupted key" because a zero + // key has a non-zero SHA-1 id. + if data.DC != 0 { + c.connMux.Lock() + c.session.Store(pool.Session{DC: data.DC}) + c.replaceConn(c.createPrimaryConn(nil)) + c.connMux.Unlock() + } + return nil + } + // Restoring persisted auth key. var key crypto.AuthKey copy(key.Value[:], data.AuthKey) - copy(key.ID[:], data.AuthKeyID) + if len(data.AuthKeyID) == 8 { + copy(key.ID[:], data.AuthKeyID) + } else { + key.ID = key.Value.ID() + } if key.Value.ID() != key.ID { return errors.New("corrupted key") @@ -112,11 +130,16 @@ func (c *Client) onSession(cfg tg.Config, s mtproto.Session) error { c.connMux.Lock() c.session.Store(sessionData) c.cfg.Store(cfg) - c.onReady() c.connMux.Unlock() - if err := c.saveSession(cfg, s); err != nil { - return errors.Wrap(err, "save") + // Persist before signalling ready so SessionStorage (and wrappers that + // capture DC / auth key ID from StoreSession) is populated when Run's + // callback starts. Previously onReady ran first, so the callback could + // observe an empty storage blob. + saveErr := c.saveSession(cfg, s) + c.onReady() + if saveErr != nil { + return errors.Wrap(saveErr, "save") } return nil diff --git a/telegram/session_restore_test.go b/telegram/session_restore_test.go new file mode 100644 index 0000000000..657d172a9d --- /dev/null +++ b/telegram/session_restore_test.go @@ -0,0 +1,83 @@ +package telegram + +import ( + "context" + "crypto/rand" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/gotd/td/crypto" + "github.com/gotd/td/session" +) + +func TestRestoreConnection_DCOnly(t *testing.T) { + ctx := context.Background() + st := &session.StorageMemory{} + loader := &session.Loader{Storage: st} + require.NoError(t, loader.Save(ctx, &session.Data{DC: 4})) + + c := NewClient(1, "hash", Options{ + SessionStorage: st, + DC: 2, + NoUpdates: true, + }) + require.NoError(t, c.restoreConnection(ctx)) + + got := c.session.Load() + require.Equal(t, 4, got.DC) + require.True(t, got.AuthKey.Zero()) +} + +func TestRestoreConnection_MissingAuthKeyID(t *testing.T) { + ctx := context.Background() + var k crypto.Key + _, err := rand.Read(k[:]) + require.NoError(t, err) + ak := k.WithID() + + st := &session.StorageMemory{} + loader := &session.Loader{Storage: st} + require.NoError(t, loader.Save(ctx, &session.Data{ + DC: 3, + AuthKey: ak.Value[:], + })) + + c := NewClient(1, "hash", Options{ + SessionStorage: st, + DC: 2, + NoUpdates: true, + }) + require.NoError(t, c.restoreConnection(ctx)) + + got := c.session.Load() + require.Equal(t, 3, got.DC) + require.Equal(t, ak.ID, got.AuthKey.ID) + require.Equal(t, ak.Value, got.AuthKey.Value) +} + +func TestRestoreConnection_CorruptedKey(t *testing.T) { + ctx := context.Background() + var k crypto.Key + _, err := rand.Read(k[:]) + require.NoError(t, err) + ak := k.WithID() + wrongID := ak.ID + wrongID[0] ^= 0xff + + st := &session.StorageMemory{} + loader := &session.Loader{Storage: st} + require.NoError(t, loader.Save(ctx, &session.Data{ + DC: 3, + AuthKey: ak.Value[:], + AuthKeyID: wrongID[:], + })) + + c := NewClient(1, "hash", Options{ + SessionStorage: st, + DC: 2, + NoUpdates: true, + }) + err = c.restoreConnection(ctx) + require.EqualError(t, err, "corrupted key") +}