From e6706080ad2ffeac9f3db53e393952eda6902ad8 Mon Sep 17 00:00:00 2001
From: Lonami <totufals@hotmail.com>
Date: Wed, 7 Sep 2016 19:48:49 +0200
Subject: [PATCH] Improved TcpClient receive

---
 main.py                |  2 +-
 network/tcp_client.py  | 16 +++++++++-------
 utils/binary_writer.py | 13 +++++++++++++
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/main.py b/main.py
index a036ef10..331590ae 100755
--- a/main.py
+++ b/main.py
@@ -9,7 +9,7 @@ if __name__ == '__main__':
 
     else:
         print('Loading interactive example...')
-        
+
         # First, initialize our TelegramClient and connect
         settings = load_settings()
         client = TelegramClient(session_user_id=settings.get('session_name', 'anonymous'),
diff --git a/network/tcp_client.py b/network/tcp_client.py
index abbc59b4..7287f58a 100755
--- a/network/tcp_client.py
+++ b/network/tcp_client.py
@@ -1,5 +1,6 @@
 # Python rough implementation of a C# TCP client
 import socket
+from utils import BinaryWriter
 
 
 class TcpClient:
@@ -23,11 +24,12 @@ class TcpClient:
 
     def read(self, buffer_size):
         """Reads (receives) the specified bytes from the connected peer"""
-        # TODO improve (don't cast to list, use a mutable byte list instead or similar, see recv_into)
-        result = []
-        while len(result) < buffer_size:
-            left_data = buffer_size - len(result)
-            partial = self.socket.recv(left_data)
-            result.extend(list(partial))
+        with BinaryWriter() as writer:
+            while writer.written_count < buffer_size:
+                # When receiving from the socket, we may not receive all the data at once
+                # This is why we need to keep checking to make sure that we receive it all
+                left_count = buffer_size - writer.written_count
+                partial = self.socket.recv(left_count)
+                writer.write(partial)
 
-        return bytes(result)
+            return writer.get_bytes()
diff --git a/utils/binary_writer.py b/utils/binary_writer.py
index 8bb6023e..db2aa35c 100755
--- a/utils/binary_writer.py
+++ b/utils/binary_writer.py
@@ -14,6 +14,7 @@ class BinaryWriter:
 
         self.stream = stream
         self.writer = BufferedWriter(self.stream)
+        self.written_count = 0
 
     # region Writing
 
@@ -21,30 +22,37 @@ class BinaryWriter:
     def write_byte(self, value):
         """Writes a single byte value"""
         self.writer.write(pack('B', value))
+        self.written_count += 1
 
     def write_int(self, value, signed=True):
         """Writes an integer value (4 bytes), which can or cannot be signed"""
         self.writer.write(int.to_bytes(value, length=4, byteorder='little', signed=signed))
+        self.written_count += 4
 
     def write_long(self, value, signed=True):
         """Writes a long integer value (8 bytes), which can or cannot be signed"""
         self.writer.write(int.to_bytes(value, length=8, byteorder='little', signed=signed))
+        self.written_count += 8
 
     def write_float(self, value):
         """Writes a floating point value (4 bytes)"""
         self.writer.write(pack('<f', value))
+        self.written_count += 4
 
     def write_double(self, value):
         """Writes a floating point value (8 bytes)"""
         self.writer.write(pack('<d', value))
+        self.written_count += 8
 
     def write_large_int(self, value, bits, signed=True):
         """Writes a n-bits long integer value"""
         self.writer.write(int.to_bytes(value, length=bits // 8, byteorder='little', signed=signed))
+        self.written_count += bits // 8
 
     def write(self, data):
         """Writes the given bytes array"""
         self.writer.write(data)
+        self.written_count += len(data)
 
     # endregion
 
@@ -98,6 +106,11 @@ class BinaryWriter:
             self.writer.flush()
         return self.stream.getvalue()
 
+    def get_written_bytes_count(self):
+        """Gets the count of bytes written in the buffer.
+           This may NOT be equal to the stream length if one was provided when initializing the writer"""
+        return self.written_count
+
     # with block
     def __enter__(self):
         return self