From 94048d91025aa3fd190018ecce581006d565bbe2 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Fri, 7 Jun 2024 22:14:40 +0200
Subject: [PATCH] finalize does not need to return last_msg_id

---
 client/src/telethon/_impl/mtproto/mtp/encrypted.py | 14 +++++---------
 client/tests/mtproto_test.py                       |  7 +++++--
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/client/src/telethon/_impl/mtproto/mtp/encrypted.py b/client/src/telethon/_impl/mtproto/mtp/encrypted.py
index 6f2f65dc..5195b950 100644
--- a/client/src/telethon/_impl/mtproto/mtp/encrypted.py
+++ b/client/src/telethon/_impl/mtproto/mtp/encrypted.py
@@ -196,7 +196,7 @@ class Encrypted(Mtp):
     def _get_current_salt(self) -> int:
         return self._salts[-1].salt if self._salts else 0
 
-    def _finalize_plain(self) -> Optional[tuple[MsgId, bytes]]:
+    def _finalize_plain(self) -> Optional[bytes]:
         if not self._msg_count:
             return None
 
@@ -207,13 +207,10 @@ class Encrypted(Mtp):
             "<qq", self._get_current_salt(), self._client_id
         )
 
-        if self._msg_count == 1:
-            container_msg_id = self._last_msg_id
-        else:
-            container_msg_id = self._get_new_msg_id()
+        if self._msg_count != 1:
             self._buffer[HEADER_LEN : HEADER_LEN + CONTAINER_HEADER_LEN] = struct.pack(
                 "<qiiIi",
-                container_msg_id,
+                self._get_new_msg_id(),
                 self._get_seq_no(False),
                 len(self._buffer) - HEADER_LEN - CONTAINER_HEADER_LEN + 8,
                 MsgContainer.constructor_id(),
@@ -223,7 +220,7 @@ class Encrypted(Mtp):
         self._msg_count = 0
         result = bytes(self._buffer)
         self._buffer.clear()
-        return MsgId(container_msg_id), result
+        return result
 
     def _process_message(self, message: Message) -> None:
         if message_requires_ack(message):
@@ -436,8 +433,7 @@ class Encrypted(Mtp):
         if not result:
             return None
 
-        msg_id, buffer = result
-        return msg_id, encrypt_data_v2(buffer, self._auth_key)
+        return MsgId(self._last_msg_id), encrypt_data_v2(result, self._auth_key)
 
     def deserialize(
         self, payload: bytes | bytearray | memoryview
diff --git a/client/tests/mtproto_test.py b/client/tests/mtproto_test.py
index a0734a3d..6b184ca8 100644
--- a/client/tests/mtproto_test.py
+++ b/client/tests/mtproto_test.py
@@ -49,9 +49,12 @@ def test_rpc_error_parsing() -> None:
 PLAIN_REQUEST = b"Hey!"
 
 
-def unwrap_finalize(finalized: Optional[tuple[MsgId, bytes]]) -> bytes:
+def unwrap_finalize(finalized: Optional[tuple[MsgId, bytes] | bytes]) -> bytes:
     assert finalized is not None
-    _, buffer = finalized
+    if isinstance(finalized, tuple):
+        _, buffer = finalized
+    else:
+        buffer = finalized
     return buffer