-
Notifications
You must be signed in to change notification settings - Fork 0
/
simple_wan_optimizer.py
executable file
·98 lines (81 loc) · 4.05 KB
/
simple_wan_optimizer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import tcp_packet
import utils
import wan_optimizer
class WanOptimizer(wan_optimizer.BaseWanOptimizer):
""" WAN Optimizer that divides data into fixed-size blocks.
This WAN optimizer should implement part 1 of project 4.
"""
# Size of blocks to store, and send only the hash when the block has been
# sent previously
BLOCK_SIZE = 8000
def __init__(self):
wan_optimizer.BaseWanOptimizer.__init__(self)
# Add any code that you like here (but do not add any constructor arguments).
self.dest_packetSize = {}
self.hash_packetList = {}
self.dest_packetList = {}
self.dest_block = {}
def initializer(self, packet, flag):
if flag or packet.dest not in self.dest_packetSize:
self.dest_packetSize[packet.dest] = 0
if flag or packet.dest not in self.dest_packetList:
self.dest_packetList[packet.dest] = []
if flag or packet.dest not in self.dest_block:
self.dest_block[packet.dest] = ""
def receive(self, packet):
self.initializer(packet, False)
if packet.is_raw_data:
self.handle_raw_data(packet)
else:
self.send_code(packet, self.hash_packetList[packet.payload])
self.initializer(packet, True)
def handle_raw_data(self, packet):
diff = self.dest_packetSize[packet.dest] + packet.size() - self.BLOCK_SIZE
if diff >= 0:
self.handle_full_block(packet, diff)
else:
self.handle_partial_block(packet)
def handle_full_block(self, packet, diff):
dif = self.BLOCK_SIZE - self.dest_packetSize[packet.dest]
if dif == packet.size():
self.dest_packetList[packet.dest].append(tcp_packet.Packet(packet.src, packet.dest, True, packet.is_fin, packet.payload[:dif]))
else:
self.dest_packetList[packet.dest].append(tcp_packet.Packet(packet.src, packet.dest, True, False, packet.payload[:dif]))
self.dest_block[packet.dest] = self.dest_block[packet.dest] + packet.payload[:dif]
hash = utils.get_hash(self.dest_block[packet.dest])
if hash in self.hash_packetList.keys():
comp_packet = tcp_packet.Packet(packet.src, packet.dest, False, packet.is_fin, hash)
self.send(comp_packet, self.wan_port)
else:
self.hash_packetList[hash] = self.dest_packetList[packet.dest]
self.send_code(packet, self.dest_packetList[packet.dest])
self.dest_packetSize[packet.dest] = diff
self.dest_block[packet.dest] = ""
self.dest_packetList[packet.dest] = []
if diff != 0:
self.dest_block[packet.dest] = packet.payload[dif:]
self.dest_packetList[packet.dest].append(tcp_packet.Packet(packet.src, packet.dest, True, packet.is_fin, packet.payload[dif:]))
def handle_partial_block(self, packet):
self.dest_packetSize[packet.dest] += packet.size()
self.dest_block[packet.dest] += packet.payload
self.dest_packetList[packet.dest].append(packet)
if packet.is_fin:
# lets send what we have so far
hash = utils.get_hash(self.dest_block[packet.dest])
if hash in self.hash_packetList:
comp_packet = tcp_packet.Packet(packet.src, packet.dest, False, False, hash)
self.send(comp_packet, self.wan_port)
else:
self.hash_packetList[hash] = self.dest_packetList[packet.dest]
self.send_code(packet, self.dest_packetList[packet.dest])
self.initializer(packet, True)
def send_code(self, packet, packet_lst):
# Send the packets in the block
for p in packet_lst:
comp_packet = tcp_packet.Packet(packet.src, packet.dest, p.is_raw_data, p.is_fin, p.payload)
if packet.dest in self.address_to_port:
# Sending to the clients connected to this middlebox
self.send(comp_packet, self.address_to_port[comp_packet.dest])
else:
# sending thru the WAN.
self.send(comp_packet, self.wan_port)