From b760f0de503c3f8038c1b815cb4461b847a8b196 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 8 Jun 2017 16:51:20 +0200
Subject: [PATCH] Fix reconnection (InvalidDC did not occur on .connect())*

* Also, use the parameters given for the initial connection
instead ignoring them.
---
 telethon/telegram_bare_client.py | 37 ++++++++++++++------------------
 telethon/telegram_client.py      | 30 +++++++++++++++++++++-----
 2 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index 9a8f1407..423d10de 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -109,22 +109,6 @@ class TelegramBareClient:
             self.dc_options = result.dc_options
             return True
 
-        except InvalidDCError as e:
-            # If an InvalidDCError occurs on CONNECTION then it for sure
-            # means that the user is associated with a different DC,
-            # which is the only case in which we need to reconnect.
-            self._logger.info('DC error on initial connection, '
-                              'attempting to reconnect at DC {}'
-                              .format(e.new_dc))
-
-            self.disconnect()
-            self.session.auth_key = None  # Force creating new auth_key
-            dc = self._get_dc(e.new_dc)
-            self.session.server_address = dc.ip_address
-            self.session.port = dc.port
-            self.session.save()
-            self.connect(device_model, system_version, app_version, lang_code)
-
         except (RPCError, ConnectionError) as error:
             # Probably errors from the previous session, ignore them
             self.disconnect()
@@ -138,12 +122,23 @@ class TelegramBareClient:
             self.sender.disconnect()
             self.sender = None
 
-    def reconnect(self):
-        """Disconnects and connects again (effectively reconnecting)"""
+    def reconnect(self, device_model, system_version, app_version, lang_code,
+                  new_dc=None):
+        """Disconnects and connects again (effectively reconnecting).
+
+           If 'new_dc' is not None, the current authorization key is
+           removed, the DC used is switched, and a new connection is made.
+        """
         self.disconnect()
-        # TODO Don't use these parameters
-        self.connect(
-            platform.node(), platform.system(), self.__version__, 'en')
+
+        if new_dc is not None:
+            self.session.auth_key = None  # Force creating new auth_key
+            dc = self._get_dc(new_dc)
+            self.session.server_address = dc.ip_address
+            self.session.port = dc.port
+            self.session.save()
+
+        self.connect(device_model, system_version, app_version, lang_code)
 
     # endregion
 
diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 19ee48dc..18839553 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -125,7 +125,8 @@ class TelegramClient(TelegramBareClient):
 
     # region Connecting
 
-    def connect(self):
+    def connect(self, a=None, b=None, c=None, d=None):
+        # TODO Use **kwargs instead 4 dummy parameters
         return super(TelegramClient, self).connect(
             device_model=self.device_model,
             system_version=self.system_version,
@@ -145,6 +146,15 @@ class TelegramClient(TelegramBareClient):
         self._cached_senders.clear()
         self._cached_sessions.clear()
 
+    def reconnect(self, new_dc=None, a=None, b=None, c=None, d=None):
+        # TODO Use **kwargs instead 4 dummy parameters
+        super(TelegramClient, self).reconnect(
+            device_model=self.device_model,
+            system_version=self.system_version,
+            app_version=self.app_version,
+            lang_code=self.lang_code,
+            new_dc=new_dc)
+
     # endregion
 
     # region Working with different Data Centers
@@ -248,11 +258,21 @@ class TelegramClient(TelegramBareClient):
             return result
 
         except InvalidDCError as e:
-            self._logger.info('DC error when invoking request, '
-                              'attempting to send it on DC {}'
-                              .format(e.new_dc))
+            if (e.message.startswith('PHONE_MIGRATE_') or
+                    e.message.startswith('USER_MIGRATE_')):
+                self._logger.info('DC error when invoking request, '
+                                  'attempting to reconnect at DC {}'
+                                  .format(e.new_dc))
 
-            return self.invoke_on_dc(request, e.new_dc, timeout=timeout)
+                self.reconnect(new_dc=e.new_dc)
+                return self.invoke(request, timeout=timeout)
+
+            else:
+                self._logger.info('DC error when invoking request, '
+                                  'attempting to send it on DC {}'
+                                  .format(e.new_dc))
+
+                return self.invoke_on_dc(request, e.new_dc, timeout=timeout)
 
         finally:
             self._lock.release()