From cdbd1f61935ab3d9c97cd5cc98231791f6b866ed Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sat, 3 Nov 2018 18:53:26 +0100
Subject: [PATCH] Fix valid auth_key never being saved after switching DC

---
 telethon/client/telegrambaseclient.py |  5 +++--
 telethon/network/mtprotosender.py     | 12 ++++++++++++
 telethon/version.py                   |  2 +-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 33ebfa6f..7fd60dd2 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -13,7 +13,6 @@ from ..network import MTProtoSender, ConnectionTcpFull
 from ..sessions import Session, SQLiteSession, MemorySession
 from ..tl import TLObject, functions, types
 from ..tl.alltlobjects import LAYER
-from ..crypto import AuthKey
 
 DEFAULT_DC_ID = 4
 DEFAULT_IPV4_IP = '149.154.167.51'
@@ -245,6 +244,7 @@ class TelegramBaseClient(abc.ABC):
             delay=self._retry_delay,
             auto_reconnect=self._auto_reconnect,
             connect_timeout=self._timeout,
+            auth_key_callback=self._auth_key_callback,
             update_callback=self._handle_update,
             auto_reconnect_callback=self._handle_auto_reconnect
         )
@@ -386,12 +386,13 @@ class TelegramBaseClient(abc.ABC):
         self.session.set_dc(dc.id, dc.ip_address, dc.port)
         # auth_key's are associated with a server, which has now changed
         # so it's not valid anymore. Set to None to force recreating it.
+        self._sender.auth_key.key = None
         self.session.auth_key = None
         self.session.save()
         self._disconnect()
         return await self.connect()
 
-    async def _auth_key_callback(self, auth_key):
+    def _auth_key_callback(self, auth_key):
         """
         Callback from the sender whenever it needed to generate a
         new authorization key. This means we are not authorized.
diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py
index b626abe6..d35bb4ea 100644
--- a/telethon/network/mtprotosender.py
+++ b/telethon/network/mtprotosender.py
@@ -61,6 +61,7 @@ class MTProtoSender:
     """
     def __init__(self, auth_key, loop, *,
                  retries=5, delay=1, auto_reconnect=True, connect_timeout=None,
+                 auth_key_callback=None,
                  update_callback=None, auto_reconnect_callback=None):
         self._connection = None
         self._loop = loop
@@ -68,6 +69,7 @@ class MTProtoSender:
         self._delay = delay
         self._auto_reconnect = auto_reconnect
         self._connect_timeout = connect_timeout
+        self._auth_key_callback = auth_key_callback
         self._update_callback = update_callback
         self._auto_reconnect_callback = auto_reconnect_callback
 
@@ -232,6 +234,13 @@ class MTProtoSender:
                     self.auth_key.key, self._state.time_offset =\
                         await authenticator.do_authentication(plain)
 
+                    # This is *EXTREMELY* important since we don't control
+                    # external references to the authorization key, we must
+                    # notify whenever we change it. This is crucial when we
+                    # switch to different data centers.
+                    if self._auth_key_callback:
+                        self._auth_key_callback(self.auth_key)
+
                     break
                 except (SecurityError, AssertionError) as e:
                     __log__.warning('Attempt {} at new auth_key failed: {}'
@@ -420,6 +429,9 @@ class MTProtoSender:
                     __log__.warning('Invalid buffer %s', e)
 
                 self.auth_key.key = None
+                if self._auth_key_callback:
+                    self._auth_key_callback(None)
+
                 self._start_reconnect()
                 return
             except Exception:
diff --git a/telethon/version.py b/telethon/version.py
index 1d06faa9..cb7653c8 100644
--- a/telethon/version.py
+++ b/telethon/version.py
@@ -1,3 +1,3 @@
 # Versions should comply with PEP440.
 # This line is parsed in setup.py:
-__version__ = '1.4'
+__version__ = '1.4.1'