From 4c65194dbc8c5a4ccbe8d1198f978ae66b3cc0d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=BE=D0=BC=D0=B0=D0=BD=20=D0=91=D0=BE=D1=80=D0=BE?=
 =?UTF-8?q?=D0=B4=D0=B8=D0=BD?= <roman@ukamnya.ru>
Date: Mon, 8 Apr 2019 17:48:13 +0300
Subject: [PATCH] .

---
 mods/utcp.py | 38 +++++++++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/mods/utcp.py b/mods/utcp.py
index 65f39f8..b13acd1 100644
--- a/mods/utcp.py
+++ b/mods/utcp.py
@@ -35,6 +35,13 @@ class Connection:
         self.peer_pub = None
         self.recv_lock = threading.Lock()
         self.send_lock = threading.Lock()
+        self.packet_buffer = {
+            'SYN': [],
+            'ACK': [],
+            'SYN-ACK': [],
+            'DATA': [],
+            'FIN-ACK': []
+            }
     @staticmethod
     def gen_starting_seq_num():
         return random.randint(Connection.SMALLEST_STARTING_SEQ, Connection.HIGHEST_STARTING_SEQ)
@@ -46,15 +53,19 @@ class Connection:
         return ack
 
 class Ack:
+    type = 'ACK'
     def __init__(self, id_):
         self.id = id_
 class Fin:
+    type = 'FIN'
     def __init__(self):
         self.id = uuid.uuid4().bytes
 class FinAck:
+    type = 'FIN-ACK'
     def __init__(self, id_):
         self.id = id_
 class Syn:
+    type = 'SYN'
     checksum = None
     def __init__(self):
         self.id = uuid.uuid4().bytes
@@ -62,6 +73,7 @@ class Syn:
         self.checksum = TCP.checksum(pubkey)
         self.pubkey = pubkey
 class SynAck:
+    type = 'SYN-ACK'
     checksum = None
     def __init__(self, id_):
         self.id = id_
@@ -87,7 +99,7 @@ class Packet:
         self.checksum = TCP.checksum(data)
         self.data = data
 
-    
+
 
 class RestrictedUnpickler(pickle.Unpickler):
     def find_class(self, module, name):
@@ -233,7 +245,7 @@ class TCP(object):
                 packet = Packet()
                 packet.set_data(data_chunk)
                 packet_to_send = pickle.dumps(packet)
-                answer = self.retransmit(connection, packet_to_send, wnat_id=packet.id)
+                answer = self.retransmit(connection, packet_to_send, want_id=packet.id)
             return len(data)
         except socket.error as error:
             raise EOFError('Socket was closed before executing command. Error is: %s.' % error)
@@ -281,7 +293,7 @@ class TCP(object):
                             conn = Connection(address, self.encrypted)
                             if self.encrypted:
                                 try:
-                                    conn.peer_pub = simplecrypto.RsaPublicKey(answer.data)
+                                    conn.peer_pub = simplecrypto.RsaPublicKey(answer.pubkey)
                                 except:
                                     raise socket.error('Init peer public key error')
                             self.connection_queue.append((answer, conn))
@@ -341,7 +353,7 @@ class TCP(object):
                 raise socket.error('Server cant receive any connections right now.')
             if self.encrypted:
                 try:
-                    peer_pub = conn.my_key.decrypt_raw(answer.data)
+                    peer_pub = conn.my_key.decrypt_raw(answer.pubkey)
                     conn.peer_pub = simplecrypto.RsaPublicKey(peer_pub)
                 except:
                     raise socket.error('Decrypt peer public key error')
@@ -363,6 +375,8 @@ class TCP(object):
         with self.connection_lock:
             if len(self.connections):
                 self.connections.pop(connection)
+            for k in list(self.packets_received.keys()):
+                self.packets_received[k].pop(connection)
     def close(self, connection=None):
         try:
             if connection not in list(self.connections.keys()):
@@ -432,7 +446,8 @@ class TCP(object):
                         if not len(self.packets_received[condition][address]):
                             del self.packets_received[condition][address]
                 else:
-                    packet = self.packets_received[condition].pop(address)
+                    with conn.recv_lock:
+                        packet = self.packets_received[condition].pop(address)
                     if want_id and packet.id != want_id:
                         raise KeyError
                 return packet
@@ -448,6 +463,9 @@ class TCP(object):
     def sort_answers(self, packet, address):
         if address not in self.connections and not isinstance(packet, Syn):
             return
+        if isinstance(packet, Syn):
+            with self.queue_lock:
+                pass
         if isinstance(packet, Fin):
             self.disconnect(address, packet.id)
         elif isinstance(packet, Packet):
@@ -462,11 +480,13 @@ class TCP(object):
                 ack = Ack(packet.id)
                 self.send_ack(address, ack)
             self.blink_incoming_packet_event()
-        elif packet.packet_type() == '':
-            #print('redundant packet found', packet)
-            pass
         else:
-            self.packets_received[packet.packet_type()][address] = packet
+            if address in self.packets_received[packet.type]:
+                conn = self.connections[address]
+                with conn.recv_lock:
+                    if packet.id == self.packets_received[packet.type][address].id:
+                        return
+            self.packets_received[packet.type][address] = packet
             self.blink_incoming_packet_event()
 
     def central_receive_handler(self):