def fragment_datagram(self, source: MeshAddress, destination: MeshAddress, datagram_header: DatagramHeader, data: bytes): stream = BytesIO() datagram_header.encode(stream) stream.write(data) stream.seek(0) fragment_size = self.network.mtu() - FragmentHeader.size() fragments = list(chunks(stream.read(), fragment_size)) with self.send_seq: sequences = range(self.send_seq, len(fragments)) self.send_seq += len(fragments) for i, fragment in zip(range(len(fragments)), fragments): if i < len(fragments) - 1: flags = FragmentFlags.FRAGMENT else: flags = FragmentFlags.NONE fragment_header = FragmentHeader(protocol=Protocol.DATAGRAM, flags=flags, fragment=i, sequence=sequences[i]) network_header = NetworkHeader( version=0, qos=QoS.Default, protocol=Protocol.FRAGMENT, ttl=MeshProtocol.DefaultTTL, identity=struct.unpack('<H', secrets.token_bytes(2))[0], length=len(fragment), source=source, destination=destination, ) buffer = encode_packet(network_header, [fragment_header], fragment) self.network.send(network_header, buffer)
def test_fragment(self): msg = """ Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum""".replace("\r", "").encode("utf-8") datagram = DatagramHeader(100, 100, len(msg), 0)
def test_encode_decode_fragment(self): msg = "Hello, World!".encode("utf-8") datagram1 = DatagramHeader(source=100, destination=100, length=len(msg), checksum=crc_b(msg)) fragment1 = FragmentHeader(Protocol.DATAGRAM, FragmentFlags.NONE, 0, 99) header1 = NetworkHeader(version=0, protocol=Protocol.FRAGMENT, qos=QoS.Default, ttl=4, identity=42, length=fragment1.size() + datagram1.size() + len(msg), source=MeshAddress(1), destination=MeshAddress(2)) stream = BytesIO() header1.encode(stream) fragment1.encode(stream) datagram1.encode(stream) stream.write(msg) stream.seek(0) header2 = NetworkHeader.decode(stream) fragment = FragmentHeader.decode(stream) datagram = DatagramHeader.decode(stream) msg2 = stream.read() self.assertEqual(header1, header2) self.assertEqual(datagram1, datagram) self.assertEqual(msg2, msg)
def test_encode_decode_datagram(self): msg = "Hello, World!".encode("utf-8") datagram_header_1 = DatagramHeader(source=100, destination=100, length=len(msg), checksum=crc_b(msg)) header1 = NetworkHeader(version=0, protocol=Protocol.DATAGRAM, qos=QoS.Default, ttl=4, identity=42, length=datagram_header_1.size() + len(msg), source=MeshAddress(1), destination=MeshAddress(2)) data = encode_packet(header1, [datagram_header_1], msg) stream = BytesIO(data) header2 = NetworkHeader.decode(stream) datagram2 = DatagramHeader.decode(stream) self.assertEqual(header1, header2) self.assertEqual(datagram_header_1, datagram2)
def test_send_receive(self): network_header = NetworkHeader( version=0, qos=QoS.Lower, protocol=Protocol.DATAGRAM, ttl=3, identity=10, length=0, source=MeshAddress(1), destination=MeshAddress(4), ) msg = "Hello, Node 4".encode("utf-8") datagram_header = DatagramHeader(source=100, destination=100, length=len(msg), checksum=0) stream = BytesIO() network_header.encode(stream) datagram_header.encode(stream) stream.write(msg) stream.seek(0) captured = None class MockTransportManager(L4Handler): def handle_l4(self, network_header: NetworkHeader, stream: BytesIO): nonlocal captured DatagramHeader.decode(stream) captured = stream.read() self.node_4.l4_handlers.handlers[ Protocol.DATAGRAM] = MockTransportManager() self.node_1.send(network_header, stream.read()) self.for_each_node(self.drain_queue) self.assertEqual(captured.decode("utf-8"), "Hello, Node 4")
def send_datagram(self, source: MeshAddress, destination: MeshAddress, datagram_header: DatagramHeader, data: bytes, reliable: bool = False): if reliable: network_header = NetworkHeader( version=0, qos=QoS.Default, protocol=Protocol.RELIABLE, ttl=MeshProtocol.DefaultTTL, identity=struct.unpack('<H', secrets.token_bytes(2))[0], length=0, source=source, destination=destination, ) reliable_header = ReliableHeader(protocol=Protocol.DATAGRAM, flags=ReliableFlags.ACK, sequence=network_header.identity, acknowledged=[]) if len(data) + DatagramHeader.size() + reliable_header.size( ) > self.network.mtu(): raise RuntimeError( "Fragmentation not supported for reliable protocol") buffer = BytesIO() network_header.encode(buffer) reliable_header.encode(buffer) datagram_header.encode(buffer) buffer.write(data) buffer.seek(0) self.reliable_protocol.send(network_header, reliable_header, buffer.read()) else: if len(data) + DatagramHeader.size() > self.network.mtu(): self.fragment_protocol.fragment_datagram( source, destination, datagram_header, data) else: network_header = NetworkHeader( version=0, qos=QoS.Default, protocol=Protocol.DATAGRAM, ttl=MeshProtocol.DefaultTTL, identity=struct.unpack('<H', secrets.token_bytes(2))[0], length=0, source=source, destination=destination, ) buffer = BytesIO() network_header.encode(buffer) datagram_header.encode(buffer) buffer.write(data) buffer.seek(0) self.network.send(network_header, buffer.read())
def handle_l4(self, network_header: NetworkHeader, stream: BytesIO): fragment_header = FragmentHeader.decode(stream) fragment = Fragment(network_header, fragment_header, stream.read()) source = fragment.network_header.source dest = fragment.network_header.destination protocol = fragment.fragment_header.protocol # Only Datagram allowed inside fragments (for now) if protocol != Protocol.DATAGRAM: self.warning( f"Dropping fragment for unsupport protocol {protocol}") # The sequence minus fragment is same for all fragments in a given PDU base_seq = ( MeshProtocol.WindowSize + fragment.fragment_header.sequence - fragment.fragment_header.fragment) % MeshProtocol.WindowSize # Buffer the incoming frame and see if it completes a whole segment self.buffer[source][base_seq].append(fragment) fragments = self.buffer[source][base_seq] sorted_fragments = sorted( fragments, key=operator.attrgetter("fragment_header.fragment")) has_more = sorted_fragments[ -1].fragment_header.flags & FragmentFlags.FRAGMENT have_all = len( fragments) == sorted_fragments[-1].fragment_header.fragment + 1 if not has_more and have_all: del self.buffer[source][base_seq] joined = BytesIO() for fragment in sorted_fragments: if fragment.network_header.source == source and \ fragment.network_header.destination == dest and \ fragment.fragment_header.protocol == protocol: joined.write(fragment.payload) else: # Fragment consistency error raise RuntimeError( "Fragment consistency error, header fields do not match" ) joined.seek(0) datagram_header = DatagramHeader.decode(joined) payload = joined.read() self.datagram_manager.handle_datagram( Datagram(fragment.network_header, datagram_header, payload))
def write_to(self, address: str, data: Any) -> None: dest = MeshAddress.parse(address) can_route, mtu = self.network.route_packet(dest) if can_route: if isinstance(data, str): encoded_data = data.encode("utf-8") elif isinstance(data, (bytes, bytearray)): encoded_data = data else: raise ValueError( "DatagramTransport.write only supports bytes and strings") if len(encoded_data) <= mtu: header = DatagramHeader(source=self.port, destination=self.port, length=len(encoded_data), checksum=crc_b(encoded_data)) self.datagram_protocol.send_datagram(self.local, self.remote, header, encoded_data) #self.broadcast_protocol.send_broadcast(self.port, encoded_data) else: raise RuntimeError(f"Message too large, maximum size is {mtu}") else: raise RuntimeError(f"Cannot route to {dest}!")
def handle_l4(self, network_header: NetworkHeader, stream: BytesIO): nonlocal captured DatagramHeader.decode(stream) captured = stream.read()
def handle_l4(self, network_header: NetworkHeader, stream: BytesIO): datagram_header = DatagramHeader.decode(stream) self.datagram_manager.handle_datagram( Datagram(network_header, datagram_header, stream.read()))