From 59a4315c283aeac5232c805b8143ab03ff72abbd Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Sun, 17 Sep 2017 14:25:53 +0200
Subject: [PATCH] Avoid calling InitConnectionRequest when it's not needed

---
 telethon/telegram_bare_client.py | 71 ++++++++++++++++++--------------
 telethon/tl/session.py           | 19 +++++----
 2 files changed, 53 insertions(+), 37 deletions(-)

diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index e49be945..3473b5c4 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -125,46 +125,42 @@ class TelegramBareClient:
 
         try:
             if not self.session.auth_key:
+                # New key, we need to tell the server we're going to use
+                # the latest layer
                 self.session.auth_key, self.session.time_offset = \
                     authenticator.do_authentication(connection)
 
+                self.session.layer = layer
                 self.session.save()
+                init_connection = True
+            else:
+                init_connection = self.session.layer != layer
 
             self._sender = MtProtoSender(connection, self.session)
             self._sender.connect()
 
-            # Now it's time to send an InitConnectionRequest
-            # This must always be invoked with the layer we'll be using
-            if exported_auth is None:
-                query = initial_query if initial_query else GetConfigRequest()
-            else:
-                query = ImportAuthorizationRequest(
-                    exported_auth.id, exported_auth.bytes)
-
-            request = InitConnectionRequest(
-                api_id=self.api_id,
-                device_model=self.session.device_model,
-                system_version=self.session.system_version,
-                app_version=self.session.app_version,
-                lang_code=self.session.lang_code,
-                system_lang_code=self.session.system_lang_code,
-                lang_pack='',  # "langPacks are for official apps only"
-                query=query)
-
-            result = self(InvokeWithLayerRequest(
-                layer=layer, query=request
-            ))
-
-            if initial_query is None:
+            if init_connection:
                 if exported_auth is not None:
-                    result = self(GetConfigRequest())
-
-                # We're only interested in the DC options,
-                # although many other options are available!
-                self.dc_options = result.dc_options
-                return True
+                    self._init_connection(ImportAuthorizationRequest(
+                        exported_auth.id, exported_auth.bytes
+                    ))
+                elif initial_query:
+                    return self._init_connection(initial_query)
+                else:
+                    self.dc_options = \
+                        self._init_connection(GetConfigRequest()).dc_options
             else:
-                return result
+                # TODO Avoid duplicated code
+                if exported_auth is not None:
+                    self(ImportAuthorizationRequest(
+                        exported_auth.id, exported_auth.bytes
+                    ))
+                elif initial_query:
+                    return self(initial_query)
+                if not self.dc_options:
+                    self.dc_options = self(GetConfigRequest()).dc_options
+
+            return True
 
         except TypeNotFoundError as e:
             # This is fine, probably layer migration
@@ -181,6 +177,21 @@ class TelegramBareClient:
             )
             return None if initial_query else False
 
+    def _init_connection(self, query=None):
+        result = self(InvokeWithLayerRequest(layer, InitConnectionRequest(
+            api_id=self.api_id,
+            device_model=self.session.device_model,
+            system_version=self.session.system_version,
+            app_version=self.session.app_version,
+            lang_code=self.session.lang_code,
+            system_lang_code=self.session.system_lang_code,
+            lang_pack='',  # "langPacks are for official apps only"
+            query=query
+        )))
+        self.session.layer = layer
+        self.session.save()
+        return result
+
     def disconnect(self):
         """Disconnects from the Telegram server"""
         if self._sender:
diff --git a/telethon/tl/session.py b/telethon/tl/session.py
index efaca3d7..e72c2a21 100644
--- a/telethon/tl/session.py
+++ b/telethon/tl/session.py
@@ -51,15 +51,17 @@ class Session:
         # Cross-thread safety
         self._lock = Lock()
 
+        self.id = utils.generate_random_long(signed=False)
+        self._sequence = 0
+        self.time_offset = 0
+        self._last_msg_id = 0  # Long
+
         # These values will be saved
         self.server_address = '91.108.56.165'
         self.port = 443
         self.auth_key = None
-        self.id = utils.generate_random_long(signed=False)
-        self._sequence = 0
+        self.layer = 0
         self.salt = 0  # Unsigned long
-        self.time_offset = 0
-        self._last_msg_id = 0  # Long
 
     def save(self):
         """Saves the current session object as session_user_id.session"""
@@ -68,6 +70,7 @@ class Session:
                 json.dump({
                     'port': self.port,
                     'salt': self.salt,
+                    'layer': self.layer,
                     'server_address': self.server_address,
                     'auth_key_data':
                         b64encode(self.auth_key.key).decode('ascii')
@@ -106,9 +109,11 @@ class Session:
             try:
                 with open(path, 'r') as file:
                     data = json.load(file)
-                    result.port = data['port']
-                    result.salt = data['salt']
-                    result.server_address = data['server_address']
+                    result.port = data.get('port', result.port)
+                    result.salt = data.get('salt', result.salt)
+                    result.layer = data.get('layer', result.layer)
+                    result.server_address = \
+                        data.get('server_address', result.server_address)
 
                     # FIXME We need to import the AuthKey here or otherwise
                     # we get cyclic dependencies.