Lock the send method for every tcp mode, not only for tcp_full

delivrance · delivrance · commit b23b41bc7dc6 · 2017-12-18T14:14:44.000+01:00
diff --git a/pyrogram/connection/connection.py b/pyrogram/connection/connection.py
@@ -17,6 +17,7 @@
 # along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import logging
+import threading
 import time
 
 from .transport import *
@@ -34,6 +35,7 @@ class Connection:
     def __init__(self, ipv4: str, mode: int = 1):
         self.address = (ipv4, 80)
         self.mode = self.MODES.get(mode, TCPAbridged)
+        self.lock = threading.Lock()
         self.connection = None
 
     def connect(self):
@@ -53,7 +55,8 @@ def close(self):
         self.connection.close()
 
     def send(self, data: bytes):
-        self.connection.send(data)
+        with self.lock:
+            self.connection.send(data)
 
     def recv(self) -> bytes or None:
         return self.connection.recv()
diff --git a/pyrogram/connection/transport/tcp/tcp_full.py b/pyrogram/connection/transport/tcp/tcp_full.py
@@ -19,7 +19,6 @@
 import logging
 from binascii import crc32
 from struct import pack, unpack
-from threading import Lock
 
 from .tcp import TCP
 
@@ -30,22 +29,19 @@ class TCPFull(TCP):
     def __init__(self):
         super().__init__()
         self.seq_no = None
-        self.lock = Lock()
 
     def connect(self, address: tuple):
         super().connect(address)
         self.seq_no = 0
         log.info("Connected!")
 
     def send(self, data: bytes):
-        with self.lock:
-            # 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)
-            data = pack("<II", len(data) + 12, self.seq_no) + data
-            data += pack("<I", crc32(data))
+        # 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)
+        data = pack("<II", len(data) + 12, self.seq_no) + data
+        data += pack("<I", crc32(data))
+        self.seq_no += 1
 
-            self.seq_no += 1
-
-            super().sendall(data)
+        super().sendall(data)
 
     def recv(self) -> bytes or None:
         length = self.recvall(4)
