From 7a6d7d003798cf26450b347c6d4c9743309fc47f Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sat, 9 Jun 2018 19:36:23 +0200
Subject: [PATCH] Implement async TCP protocol

---
 pyrogram/connection/transport/tcp/tcp.py | 86 ++++++++++++++++++++----
 1 file changed, 73 insertions(+), 13 deletions(-)

diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
index 5df8aacb..f006cadd 100644
--- a/pyrogram/connection/transport/tcp/tcp.py
+++ b/pyrogram/connection/transport/tcp/tcp.py
@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+import asyncio
 import logging
 import socket
 
@@ -32,14 +33,18 @@ except ImportError as e:
 log = logging.getLogger(__name__)
 
 
-class TCP(socks.socksocket):
+class TCP:
     def __init__(self, proxy: dict):
-        super().__init__()
-        self.settimeout(10)
+        self.proxy = proxy
+
+        self.socket = socks.socksocket()
+        self.reader = None  # type: asyncio.StreamReader
+        self.writer = None  # type: asyncio.StreamWriter
+
         self.proxy_enabled = proxy.get("enabled", False)
 
         if proxy and self.proxy_enabled:
-            self.set_proxy(
+            self.socket.set_proxy(
                 proxy_type=socks.SOCKS5,
                 addr=proxy.get("hostname", None),
                 port=proxy.get("port", None),
@@ -52,26 +57,81 @@ class TCP(socks.socksocket):
                 proxy.get("port", None)
             ))
 
+    async def connect(self, address: tuple):
+        self.socket.connect(address)
+        self.reader, self.writer = await asyncio.open_connection(sock=self.socket)
+
     def close(self):
         try:
-            self.shutdown(socket.SHUT_RDWR)
-        except OSError:
-            pass
-        finally:
-            super().close()
+            self.writer.close()
+        except AttributeError:
+            try:
+                self.socket.shutdown(socket.SHUT_RDWR)
+            except OSError:
+                pass
+            finally:
+                self.socket.close()
 
-    def recvall(self, length: int) -> bytes or None:
+    async def send(self, data: bytes):
+        self.writer.write(data)
+        await self.writer.drain()
+
+    async def recv(self, length: int = 0):
         data = b""
 
         while len(data) < length:
             try:
-                packet = super().recv(length - len(data))
+                chunk = await self.reader.read(length - len(data))
             except OSError:
                 return None
             else:
-                if packet:
-                    data += packet
+                if chunk:
+                    data += chunk
                 else:
                     return None
 
         return data
+
+# class TCP(socks.socksocket):
+#     def __init__(self, proxy: dict):
+#         super().__init__()
+#         self.settimeout(10)
+#         self.proxy_enabled = proxy.get("enabled", False)
+#
+#         if proxy and self.proxy_enabled:
+#             self.set_proxy(
+#                 proxy_type=socks.SOCKS5,
+#                 addr=proxy.get("hostname", None),
+#                 port=proxy.get("port", None),
+#                 username=proxy.get("username", None),
+#                 password=proxy.get("password", None)
+#             )
+#
+#             log.info("Using proxy {}:{}".format(
+#                 proxy.get("hostname", None),
+#                 proxy.get("port", None)
+#             ))
+#
+#     def close(self):
+#         try:
+#             self.shutdown(socket.SHUT_RDWR)
+#         except OSError:
+#             pass
+#         finally:
+#             super().close()
+#
+#     def recvall(self, length: int) -> bytes or None:
+#         data = b""
+#
+#         while len(data) < length:
+#             try:
+#                 packet = super().recv(length - len(data))
+#             except OSError:
+#                 return None
+#             else:
+#                 if packet:
+#                     data += packet
+#                 else:
+#                     return None
+#
+#         return data