fix(telegram): restore DC-only sessions and migrate QR on the same cl…

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
 			}
 

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)

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 {

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.

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

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")
+}