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()))