refactor(session): harden mtproto auth and optimize session handling

Author: SychO3

pyrogram/session/auth.py

@@ -60,11 +60,12 @@ def __init__(
 
     @staticmethod
     def pack(data: TLObject, server_time: float) -> bytes:
+        payload = data.write()
         return (
             bytes(8)
             + Long(int(server_time * (2**32)) & ~0b11)
-            + Int(len(data.write()))
-            + data.write()
+            + Int(len(payload))
+            + payload
         )
 
     @staticmethod
@@ -150,7 +151,7 @@ async def create(self):
 
                 log.debug("Done encrypt data with RSA")
 
-                # Step 5. TODO: Handle "server_DH_params_fail". Code assumes response is ok
+                # Step 5
                 log.debug("Send req_DH_params")
                 server_dh_params = await self.invoke(
                     raw.functions.ReqDHParams(
@@ -163,6 +164,9 @@ async def create(self):
                     )
                 )
 
+                if isinstance(server_dh_params, raw.types.ServerDhParamsFail):
+                    raise Exception("Server DH params generation failed")
+
                 encrypted_answer = server_dh_params.encrypted_answer
 
                 server_nonce = server_nonce.to_bytes(16, "little", signed=True)
@@ -189,21 +193,49 @@ async def create(self):
 
                 dh_prime = int.from_bytes(server_dh_inner_data.dh_prime, "big")
                 delta_time = server_dh_inner_data.server_time - time.time()
-
                 log.debug("Delta time: %s", round(delta_time, 3))
 
-                # Step 6
                 g = server_dh_inner_data.g
+                g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
+
+                # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
+                answer = server_dh_inner_data.write()
+                SecurityCheckMismatch.check(
+                    answer_with_hash[:20] == sha1(answer).digest(),
+                    "answer_with_hash[:20] == sha1(answer).digest()"
+                )
+                log.debug("SHA1 hash values check: OK")
+
+                # Validate DH parameters BEFORE expensive modular exponentiation
+                SecurityCheckMismatch.check(dh_prime == prime.CURRENT_DH_PRIME, "dh_prime == prime.CURRENT_DH_PRIME")
+                SecurityCheckMismatch.check(1 < g < dh_prime - 1, "1 < g < dh_prime - 1")
+                SecurityCheckMismatch.check(1 < g_a < dh_prime - 1, "1 < g_a < dh_prime - 1")
+                SecurityCheckMismatch.check(
+                    2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64),
+                    "2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)"
+                )
+                log.debug("DH parameters and g_a validation: OK")
+
+                # Step 6 — now safe to compute
                 b = int.from_bytes(urandom(256), "big")
-                g_b = pow(g, b, dh_prime).to_bytes(256, "big")
+                g_b = pow(g, b, dh_prime)
+
+                SecurityCheckMismatch.check(1 < g_b < dh_prime - 1, "1 < g_b < dh_prime - 1")
+                SecurityCheckMismatch.check(
+                    2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64),
+                    "2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)"
+                )
+                log.debug("g_b validation: OK")
+
+                g_b_bytes = g_b.to_bytes(256, "big")
 
                 retry_id = 0
 
                 data = raw.types.ClientDHInnerData(
                     nonce=nonce,
                     server_nonce=server_nonce,
                     retry_id=retry_id,
-                    g_b=g_b
+                    g_b=g_b_bytes
                 ).write()
 
                 sha = sha1(data).digest()
@@ -220,56 +252,25 @@ async def create(self):
                     )
                 )
 
-                # TODO: Handle "auth_key_aux_hash" if the previous step fails
+                if isinstance(set_client_dh_params_answer, raw.types.DhGenFail):
+                    raise Exception("DH key generation failed (dh_gen_fail)")
+
+                if isinstance(set_client_dh_params_answer, raw.types.DhGenRetry):
+                    raise Exception("DH key generation requires retry (dh_gen_retry)")
 
                 # Step 7; Step 8
-                g_a = int.from_bytes(server_dh_inner_data.g_a, "big")
                 auth_key = pow(g_a, b, dh_prime).to_bytes(256, "big")
                 server_nonce = server_nonce.to_bytes(16, "little", signed=True)
 
-                # TODO: Handle errors
-
-                #######################
-                # Security checks
-                #######################
-
-                SecurityCheckMismatch.check(dh_prime == prime.CURRENT_DH_PRIME, "dh_prime == prime.CURRENT_DH_PRIME")
-                log.debug("DH parameters check: OK")
-
-                # https://core.telegram.org/mtproto/security_guidelines#g-a-and-g-b-validation
-                g_b = int.from_bytes(g_b, "big")
-                SecurityCheckMismatch.check(1 < g < dh_prime - 1, "1 < g < dh_prime - 1")
-                SecurityCheckMismatch.check(1 < g_a < dh_prime - 1, "1 < g_a < dh_prime - 1")
-                SecurityCheckMismatch.check(1 < g_b < dh_prime - 1, "1 < g_b < dh_prime - 1")
-                SecurityCheckMismatch.check(
-                    2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64),
-                    "2 ** (2048 - 64) < g_a < dh_prime - 2 ** (2048 - 64)"
-                )
-                SecurityCheckMismatch.check(
-                    2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64),
-                    "2 ** (2048 - 64) < g_b < dh_prime - 2 ** (2048 - 64)"
-                )
-                log.debug("g_a and g_b validation: OK")
-
-                # https://core.telegram.org/mtproto/security_guidelines#checking-sha1-hash-values
-                answer = server_dh_inner_data.write()  # Call .write() to remove padding
-                SecurityCheckMismatch.check(
-                    answer_with_hash[:20] == sha1(answer).digest(),
-                    "answer_with_hash[:20] == sha1(answer).digest()"
-                )
-                log.debug("SHA1 hash values check: OK")
-
                 # https://core.telegram.org/mtproto/security_guidelines#checking-nonce-server-nonce-and-new-nonce-fields
-                # 1st message
                 SecurityCheckMismatch.check(nonce == res_pq.nonce, "nonce == res_pq.nonce")
-                # 2nd message
+
                 server_nonce = int.from_bytes(server_nonce, "little", signed=True)
                 SecurityCheckMismatch.check(nonce == server_dh_params.nonce, "nonce == server_dh_params.nonce")
                 SecurityCheckMismatch.check(
                     server_nonce == server_dh_params.server_nonce,
                     "server_nonce == server_dh_params.server_nonce"
                 )
-                # 3rd message
                 SecurityCheckMismatch.check(
                     nonce == set_client_dh_params_answer.nonce,
                     "nonce == set_client_dh_params_answer.nonce"
@@ -287,13 +288,12 @@ async def create(self):
                 log.debug("Server salt: %s", int.from_bytes(server_salt, "little"))
 
                 log.info("Done auth key exchange: %s", set_client_dh_params_answer.__class__.__name__)
-            except ConnectionError as e:
+            except ConnectionError:
                 log.info("Unable to connect due to network issues. Retrying...")
-                # Treat like transient network error: retry according to MAX_RETRIES
                 if retries_left:
                     retries_left -= 1
                 else:
-                    raise e
+                    raise
                 await asyncio.sleep(1)
                 continue
             except Exception as e:
@@ -302,7 +302,7 @@ async def create(self):
                 if retries_left > 0:
                     retries_left -= 1
                 elif retries_left == 0:
-                    raise e
+                    raise
 
                 await asyncio.sleep(1)
                 continue

pyrogram/session/internals/msg_factory.py

@@ -30,14 +30,11 @@ def __init__(self, client: "pyrogram.Client"):
         self.client = client
 
         self._last_msg_id = 0
-
-        self._msg_id_lock = asyncio.Lock()
-        self._seq_no_lock = asyncio.Lock()
-
+        self._lock = asyncio.Lock()
         self._content_related_messages_sent = 0
 
     async def allocate_message_identity(self) -> int:
-        async with self._msg_id_lock:
+        async with self._lock:
             base_msg_id = int(self.client.server_time * (2**32)) & ~0b11
 
             if base_msg_id <= self._last_msg_id:
@@ -47,19 +44,19 @@ async def allocate_message_identity(self) -> int:
 
             return base_msg_id
 
-    async def allocate_message_sequence(self, is_content_related: bool) -> int:
-        async with self._seq_no_lock:
-            seq_no = (self._content_related_messages_sent * 2) + (1 if is_content_related else 0)
+    async def create(self, body: TLObject) -> Message:
+        async with self._lock:
+            base_msg_id = int(self.client.server_time * (2**32)) & ~0b11
 
-            if is_content_related:
-                self._content_related_messages_sent += 1
+            if base_msg_id <= self._last_msg_id:
+                base_msg_id = self._last_msg_id + 4
 
-            return seq_no
+            self._last_msg_id = base_msg_id
 
-    async def create(self, body: TLObject) -> Message:
-        msg_id = await self.allocate_message_identity()
+            is_content_related = not isinstance(body, (Ping, HttpWait, MsgsAck, MsgContainer))
+            seq_no = (self._content_related_messages_sent * 2) + (1 if is_content_related else 0)
 
-        is_content_related = not isinstance(body, (Ping, HttpWait, MsgsAck, MsgContainer))
-        seq_no = await self.allocate_message_sequence(is_content_related)
+            if is_content_related:
+                self._content_related_messages_sent += 1
 
-        return Message(body, msg_id, seq_no, len(body))
+            return Message(body, base_msg_id, seq_no, len(body))