From 68ea4080df1ec7b22304bc39c1fbb142c7a17e26 Mon Sep 17 00:00:00 2001
From: teidesu <me@tei.su>
Date: Sat, 5 Jun 2021 18:56:43 +0300
Subject: [PATCH] feat(core): handle AUTH_KEY_UNREGISTERED error

---
 .../src/methods/files/download-iterable.ts    |  8 +--
 packages/core/src/base-client.ts              | 56 +++++++++++++++----
 2 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/packages/client/src/methods/files/download-iterable.ts b/packages/client/src/methods/files/download-iterable.ts
index e263f767..e8fd288f 100644
--- a/packages/client/src/methods/files/download-iterable.ts
+++ b/packages/client/src/methods/files/download-iterable.ts
@@ -86,21 +86,21 @@ export async function* downloadAsIterable(
     const requestCurrent = async (): Promise<Buffer> => {
         let result: tl.RpcCallReturn['upload.getFile'] | tl.RpcCallReturn['upload.getWebFile']
         try {
-            result = await connection.sendForResult({
+            result = await this.call({
                 _: isWeb ? 'upload.getWebFile' : 'upload.getFile',
                 location: location as any,
                 offset,
                 limit: chunkSize
-            })
+            }, { connection })
         } catch (e) {
-            if (e instanceof FileMigrateError) {
+            if (e.constructor === FileMigrateError) {
                 connection = this._downloadConnections[e.newDc]
                 if (!connection) {
                     connection = await this.createAdditionalConnection(e.newDc)
                     this._downloadConnections[e.newDc] = connection
                 }
                 return requestCurrent()
-            } else if (e instanceof FilerefUpgradeNeededError) {
+            } else if (e.constructor === FilerefUpgradeNeededError) {
                 // todo: implement someday
                 // see: https://github.com/LonamiWebs/Telethon/blob/0e8bd8248cc649637b7c392616887c50986427a0/telethon/client/downloads.py#L99
                 throw new MtCuteUnsupportedError('File ref expired!')
diff --git a/packages/core/src/base-client.ts b/packages/core/src/base-client.ts
index 16fe353e..f7714b55 100644
--- a/packages/core/src/base-client.ts
+++ b/packages/core/src/base-client.ts
@@ -17,6 +17,7 @@ import {
     defaultProductionIpv6Dc,
 } from './utils/default-dcs'
 import {
+    AuthKeyUnregisteredError,
     FloodTestPhoneWaitError,
     FloodWaitError,
     InternalError,
@@ -347,7 +348,9 @@ export class BaseTelegramClient {
             this.storage.setAuthKeyFor(this._primaryDc.id, key)
             await this._saveStorage()
         })
-        this.primaryConnection.on('error', (err) => this._emitError(err, this.primaryConnection))
+        this.primaryConnection.on('error', (err) =>
+            this._emitError(err, this.primaryConnection)
+        )
     }
 
     /**
@@ -488,7 +491,11 @@ export class BaseTelegramClient {
 
         for (let i = 0; i < this._rpcRetryCount; i++) {
             try {
-                const res = await connection.sendForResult(message, stack, params?.timeout)
+                const res = await connection.sendForResult(
+                    message,
+                    stack,
+                    params?.timeout
+                )
                 await this._cachePeersFrom(res)
 
                 return res
@@ -530,16 +537,37 @@ export class BaseTelegramClient {
                         continue
                     }
                 }
-                if (
-                    connection === this.primaryConnection &&
-                    (e.constructor === PhoneMigrateError ||
+
+                if (connection.params.dc.id === this._primaryDc.id) {
+                    if (
+                        e.constructor === PhoneMigrateError ||
                         e.constructor === UserMigrateError ||
-                        e.constructor === NetworkMigrateError)
-                ) {
-                    debug('Migrate error, new dc = %d', e.newDc)
-                    await this.changeDc(e.newDc)
-                    continue
+                        e.constructor === NetworkMigrateError
+                    ) {
+                        debug('Migrate error, new dc = %d', e.newDc)
+                        await this.changeDc(e.newDc)
+                        continue
+                    }
+                } else {
+                    if (e.constructor === AuthKeyUnregisteredError) {
+                        // we can try re-exporting auth from the primary connection
+                        debug('exported auth key error, re-exporting..')
+
+                        const auth = await this.call({
+                            _: 'auth.exportAuthorization',
+                            dcId: connection.params.dc.id,
+                        })
+
+                        await connection.sendForResult({
+                            _: 'auth.importAuthorization',
+                            id: auth.id,
+                            bytes: auth.bytes,
+                        })
+
+                        continue
+                    }
                 }
+
                 throw e
             }
         }
@@ -635,7 +663,9 @@ export class BaseTelegramClient {
     changeTransport(factory: TransportFactory): void {
         this.primaryConnection.changeTransport(factory)
 
-        this._additionalConnections.forEach((conn) => conn.changeTransport(factory))
+        this._additionalConnections.forEach((conn) =>
+            conn.changeTransport(factory)
+        )
     }
 
     /**
@@ -647,7 +677,9 @@ export class BaseTelegramClient {
      *     the connection in which the error has occurred, in case
      *     this was connection-related error.
      */
-    onError(handler: (err: Error, connection?: TelegramConnection) => void): void {
+    onError(
+        handler: (err: Error, connection?: TelegramConnection) => void
+    ): void {
         this._onError = handler
     }