From b4811261e9af6b1e78611f50312c03acffd93656 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Tue, 29 Aug 2017 20:25:49 +0200 Subject: [PATCH] Set a known buffer size if possible on BinaryWriter()s --- telethon/extensions/binary_writer.py | 11 ++++++++--- telethon/network/authenticator.py | 2 +- telethon/network/connection.py | 9 +++++---- telethon/network/mtproto_plain_sender.py | 2 +- 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/telethon/extensions/binary_writer.py b/telethon/extensions/binary_writer.py index 5b363297..4a24616a 100644 --- a/telethon/extensions/binary_writer.py +++ b/telethon/extensions/binary_writer.py @@ -1,4 +1,4 @@ -from io import BufferedWriter, BytesIO +from io import BufferedWriter, BytesIO, DEFAULT_BUFFER_SIZE from struct import pack @@ -8,11 +8,16 @@ class BinaryWriter: Also creates a "Memory Stream" if necessary """ - def __init__(self, stream=None): + def __init__(self, stream=None, known_length=None): if not stream: stream = BytesIO() - self.writer = BufferedWriter(stream) + if known_length is None: + # On some systems, DEFAULT_BUFFER_SIZE defaults to 8192 + # That's over 16 times as big as necessary for most messages + known_length = max(DEFAULT_BUFFER_SIZE, 1024) + + self.writer = BufferedWriter(stream, buffer_size=known_length) self.written_count = 0 # region Writing diff --git a/telethon/network/authenticator.py b/telethon/network/authenticator.py index 3ea19acf..6b065b9a 100644 --- a/telethon/network/authenticator.py +++ b/telethon/network/authenticator.py @@ -19,7 +19,7 @@ def do_authentication(connection): # Step 1 sending: PQ Request nonce = os.urandom(16) - with BinaryWriter() as writer: + with BinaryWriter(known_length=20) as writer: writer.write_int(0x60469778, signed=False) # Constructor number writer.write(nonce) sender.send(writer.get_bytes()) diff --git a/telethon/network/connection.py b/telethon/network/connection.py index a4db9417..af3ab817 100644 --- a/telethon/network/connection.py +++ b/telethon/network/connection.py @@ -151,8 +151,9 @@ class Connection: def _send_tcp_full(self, message): # https://core.telegram.org/mtproto#tcp-transport # total length, sequence number, packet and checksum (CRC32) - with BinaryWriter() as writer: - writer.write_int(len(message) + 12) + length = len(message) + 12 + with BinaryWriter(known_length=length) as writer: + writer.write_int(length) writer.write_int(self._send_counter) writer.write(message) writer.write_int(crc32(writer.get_bytes()), signed=False) @@ -160,13 +161,13 @@ class Connection: self.write(writer.get_bytes()) def _send_intermediate(self, message): - with BinaryWriter() as writer: + with BinaryWriter(known_length=len(message) + 4) as writer: writer.write_int(len(message)) writer.write(message) self.write(writer.get_bytes()) def _send_abridged(self, message): - with BinaryWriter() as writer: + with BinaryWriter(known_length=len(message) + 4) as writer: length = len(message) >> 2 if length < 127: writer.write_byte(length) diff --git a/telethon/network/mtproto_plain_sender.py b/telethon/network/mtproto_plain_sender.py index a453ce72..c387287e 100644 --- a/telethon/network/mtproto_plain_sender.py +++ b/telethon/network/mtproto_plain_sender.py @@ -20,7 +20,7 @@ class MtProtoPlainSender: def send(self, data): """Sends a plain packet (auth_key_id = 0) containing the given message body (data)""" - with BinaryWriter() as writer: + with BinaryWriter(known_length=len(data) + 20) as writer: writer.write_long(0) writer.write_long(self._get_new_msg_id()) writer.write_int(len(data))