From 5cf5c140f888f425be312a278596d5765501733c Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Fri, 24 Dec 2021 01:25:51 +0100
Subject: [PATCH] GramJs: Properly handle `AUTH_KEY_UNREGISTERED` (again)

---
 src/lib/gramjs/network/MTProtoSender.js | 51 ++++++++++++++++++-------
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/lib/gramjs/network/MTProtoSender.js b/src/lib/gramjs/network/MTProtoSender.js
index 809ff639f..636dd4f1a 100644
--- a/src/lib/gramjs/network/MTProtoSender.js
+++ b/src/lib/gramjs/network/MTProtoSender.js
@@ -1,3 +1,5 @@
+const { RPCError } = require('../errors');
+
 const MtProtoPlainSender = require('./MTProtoPlainSender');
 const MTProtoState = require('./MTProtoState');
 const Helpers = require('../Helpers');
@@ -16,7 +18,10 @@ const {
 } = require('../tl').constructors;
 const MessagePacker = require('../extensions/MessagePacker');
 const BinaryReader = require('../extensions/BinaryReader');
-const { UpdateConnectionState, UpdateServerTimeOffset } = require('./index');
+const {
+    UpdateConnectionState,
+    UpdateServerTimeOffset,
+} = require('./index');
 const { BadMessageError } = require('../errors/Common');
 const {
     BadServerSalt,
@@ -420,13 +425,7 @@ class MTProtoSender {
                 } else if (e instanceof InvalidBufferError) {
                     // 404 means that the server has "forgotten" our auth key and we need to create a new one.
                     if (e.code === 404) {
-                        this._log.warn(`Broken authorization key for dc ${this._dcId}; resetting`);
-                        if (this._updateCallback && this._isMainSender) {
-                            this._updateCallback(new UpdateConnectionState(UpdateConnectionState.broken));
-                        } else if (this._onConnectionBreak && !this._isMainSender) {
-                            // Deletes the current sender from the object
-                            this._onConnectionBreak(this._dcId);
-                        }
+                        this._handleBadAuthKey();
                     } else {
                         // this happens sometimes when telegram is having some internal issues.
                         // reconnecting should be enough usually
@@ -445,12 +444,34 @@ class MTProtoSender {
             try {
                 await this._processMessage(message);
             } catch (e) {
-                this._log.error('Unhandled error while receiving data');
-                this._log.error(e);
+                // `RPCError` errors except for 'AUTH_KEY_UNREGISTERED' should be handled by the client
+                if (e instanceof RPCError) {
+                    if (e.message === 'AUTH_KEY_UNREGISTERED') {
+                        // 'AUTH_KEY_UNREGISTERED' for the main sender is thrown when unauthorized and should be ignored
+                        this._handleBadAuthKey(true);
+                    }
+                } else {
+                    this._log.error('Unhandled error while receiving data');
+                    this._log.error(e);
+                }
             }
         }
     }
 
+    _handleBadAuthKey(shouldSkipForMain) {
+        if (shouldSkipForMain && this._isMainSender) {
+            return;
+        }
+
+        this._log.warn(`Broken authorization key for dc ${this._dcId}, resetting...`);
+
+        if (this._isMainSender && this._updateCallback) {
+            this._updateCallback(new UpdateConnectionState(UpdateConnectionState.broken));
+        } else if (!this._isMainSender && this._onConnectionBreak) {
+            this._onConnectionBreak(this._dcId);
+        }
+    }
+
     // Response Handlers
 
     /**
@@ -557,6 +578,7 @@ class MTProtoSender {
             const error = RPCMessageToError(result.error, state.request);
             this._send_queue.append(new RequestState(new MsgsAck({ msgIds: [state.msgId] })));
             state.reject(error);
+            throw error;
         } else {
             try {
                 const reader = new BinaryReader(result.body);
@@ -826,10 +848,11 @@ class MTProtoSender {
             // this._user_connected = false
             // we want to wait a second between each reconnect try to not flood the server with reconnects
             // in case of internal server issues.
-            Helpers.sleep(1000).then(() => {
-                this._log.info('Started reconnecting');
-                this._reconnect();
-            });
+            Helpers.sleep(1000)
+                .then(() => {
+                    this._log.info('Started reconnecting');
+                    this._reconnect();
+                });
         }
     }