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 send_query(self, neighbor: MeshAddress):
self.debug(f"Querying {neighbor} for link states")
known_link_states = dict()
for node, link_states in self.valid_link_states().items():
known_link_states[node] = self.link_state_epochs[node]
query = LinkStateQueryHeader(node=self.our_address,
epoch=self.our_link_state_epoch,
link_nodes=list(known_link_states.keys()),
link_epochs=list(
known_link_states.values()))
network_header = NetworkHeader(
version=0,
qos=QoS.Lower,
protocol=Protocol.LINK_STATE_QUERY,
ttl=1,
identity=self.next_sequence(),
length=0,
source=self.our_address,
destination=neighbor,
)
stream = BytesIO()
network_header.encode(stream)
query.encode(stream)
stream.seek(0)
buffer = stream.read()
self.send(network_header, buffer)
def send_advertisement(self):
link_states = self.generate_our_adverts()
advertisement = LinkStateAdvertisementHeader(
node=self.our_address,
name=self.host_name,
epoch=self.our_link_state_epoch,
link_states=link_states)
self.debug("Sending Advertisement {}".format(advertisement))
network_header = NetworkHeader(
version=0,
qos=QoS.Lower,
protocol=Protocol.LINK_STATE,
ttl=7,
identity=self.next_sequence(),
length=0,
source=self.our_address,
destination=self.BroadcastAddress,
)
stream = BytesIO()
network_header.encode(stream)
advertisement.encode(stream)
stream.seek(0)
buffer = stream.read()
self.send(network_header, buffer)
def encode_packet(network_header: NetworkHeader,
additional_headers: List[Header], payload: bytes) -> bytes:
stream = BytesIO()
network_header.encode(stream)
for header in additional_headers:
header.encode(stream)
stream.write(payload)
stream.seek(0)
return stream.read()
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):
reliable_header = ReliableHeader.decode(stream)
self.debug(f"Handling reliable {reliable_header}")
if reliable_header.flags & ReliableFlags.ACK:
# We received an ACK
self.reliable_manager.handle_ack(network_header.source,
reliable_header.sequence)
else:
# An ACK is being requested
if network_header.source not in self.network.alive_neighbors():
# Ignore if we haven't seen this neighbor before, we need discovery first
# TODO send discovery to this neighbor
self.debug(
f"Ignoring {reliable_header} since we don't know this neighbor."
)
return
outgoing_network_header = NetworkHeader(
version=0,
qos=QoS.Default,
protocol=Protocol.RELIABLE,
ttl=1,
identity=self.network.next_sequence(),
length=0,
source=self.network.our_address,
destination=network_header.source,
)
ack_header = ReliableHeader(
protocol=Protocol.NONE,
flags=ReliableFlags.ACK,
sequence=reliable_header.sequence,
)
response = BytesIO()
outgoing_network_header.encode(response)
ack_header.encode(response)
response.seek(0)
buffer = response.read()
self.info(
f"Sending Ack for sequence {reliable_header.sequence} to address {outgoing_network_header.source}"
)
self.network.send(outgoing_network_header, buffer)
self.handlers.handle_l4(network_header, reliable_header.protocol,
stream)
def send_hello(self):
self.debug("Sending Hello")
hello = HelloHeader(self.config.get("host.name"),
self.alive_neighbors())
network_header = NetworkHeader(
version=0,
qos=QoS.Lower,
protocol=Protocol.HELLO,
ttl=1,
identity=self.next_sequence(),
length=0,
source=self.our_address,
destination=self.BroadcastAddress,
)
stream = BytesIO()
network_header.encode(stream)
hello.encode(stream)
stream.seek(0)
buffer = stream.read()
self.send(network_header, buffer)
def send_ping(self, node: MeshAddress) -> int:
ctrl = ControlHeader(True, ControlType.PING, 1,
bytes([secure_random_byte()]))
network_header = NetworkHeader(
version=0,
qos=QoS.Higher,
protocol=Protocol.CONTROL,
ttl=7,
identity=self.network.next_sequence(),
length=0,
source=self.network.our_address,
destination=node,
)
stream = BytesIO()
network_header.encode(stream)
ctrl.encode(stream)
stream.seek(0)
buffer = stream.read()
self.stats[node].results.append(
PingResult(ctrl.extra[0], time.time_ns(), None,
threading.Condition(self.mutex)))
self.network.send(network_header, buffer)
return ctrl.extra[0]
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 handle_query(self, link_id: int, network_header: NetworkHeader,
query: LinkStateQueryHeader):
self.debug(f"Handling {query}")
dest = network_header.source
adverts = []
# Check our local cache of link states
for node, link_states in self.valid_link_states().items():
if node == dest:
continue
if node not in query.link_nodes:
# Sender is missing this node's data
header = LinkStateAdvertisementHeader(
node=node,
name=self.host_names.get(node, "unknown"),
epoch=self.link_state_epochs.get(node),
link_states=list(self.link_states.get(node)))
adverts.append(header)
else:
idx = query.link_nodes.index(node)
epoch = query.link_epochs[idx]
if lollipop_compare(epoch, self.link_state_epochs[node]) > -1:
header = LinkStateAdvertisementHeader(
node=node,
name=self.host_names.get(node, "unknown"),
epoch=self.link_state_epochs.get(node),
link_states=list(self.link_states.get(node)))
adverts.append(header)
# Include our latest state if they don't have it
if self.our_address in query.link_nodes:
idx = query.link_nodes.index(self.our_address)
epoch = query.link_epochs[idx]
if lollipop_compare(epoch, self.our_link_state_epoch) > -1:
link_states = self.generate_our_adverts()
our_advert = LinkStateAdvertisementHeader(
node=self.our_address,
name=self.host_name,
epoch=self.our_link_state_epoch,
link_states=link_states)
adverts.append(our_advert)
dest = network_header.source
for advert in adverts:
self.debug(f"Sending {advert} advert to {network_header.source}")
resp_header = NetworkHeader(
version=0,
qos=QoS.Lower,
protocol=Protocol.LINK_STATE,
ttl=1,
identity=self.next_sequence(),
length=0,
source=self.our_address,
destination=dest,
)
stream = BytesIO()
resp_header.encode(stream)
advert.encode(stream)
stream.seek(0)
buffer = stream.read()
self.send(resp_header, buffer)