async def test_send_endpoint_vote(ping_sender, routing_table,
incoming_message_channels,
outgoing_message_channels,
endpoint_vote_channels, local_enr,
remote_enr, remote_endpoint):
# wait for ping
with trio.fail_after(ROUTING_TABLE_PING_INTERVAL):
outgoing_message = await outgoing_message_channels[1].receive()
ping = outgoing_message.message
# respond with pong
fake_local_endpoint = EndpointFactory()
pong = PongMessage(
request_id=ping.request_id,
enr_seq=0,
packet_ip=fake_local_endpoint.ip_address,
packet_port=fake_local_endpoint.port,
)
incoming_message = IncomingMessage(
message=pong,
sender_endpoint=outgoing_message.receiver_endpoint,
sender_node_id=outgoing_message.receiver_node_id,
)
await incoming_message_channels[0].send(incoming_message)
await wait_all_tasks_blocked()
# receive endpoint vote
endpoint_vote = endpoint_vote_channels[1].receive_nowait()
assert endpoint_vote.endpoint == fake_local_endpoint
assert endpoint_vote.node_id == incoming_message.sender_node_id
async def handle_incoming_packet_post_handshake(self, incoming_packet: IncomingPacket) -> None:
if not self.is_post_handshake:
raise ValueError("Can only handle packets post handshake")
if self.session_keys is None:
raise TypeError("session_keys are None even though handshake has been completed")
if isinstance(incoming_packet.packet, AuthTagPacket):
try:
message = incoming_packet.packet.decrypt_message(
self.session_keys.decryption_key,
self.message_type_registry,
)
except DecryptionError:
self.logger.info(
"Failed to decrypt message from peer, starting another handshake as recipient"
)
self.reset_handshake_state()
await self.handle_incoming_packet_pre_handshake(incoming_packet)
except ValidationError as validation_error:
self.logger.warning("Received invalid packet: %s", validation_error)
raise # let the service fail
else:
incoming_message = IncomingMessage(
message,
incoming_packet.sender_endpoint,
self.remote_node_id,
)
self.logger.debug("Received %s", incoming_message)
await self.incoming_message_send_channel.send(incoming_message)
else:
self.logger.debug(
"Dropping %s as handshake has already been completed",
incoming_packet,
)
async def respond_with_remote_enrs(
self, incoming_message: IncomingMessage) -> None:
"""Send a Nodes message containing ENRs of peers at a given node distance."""
node_ids = self.routing_table.get_nodes_at_log_distance(
incoming_message.message.distance)
enrs = []
for node_id in node_ids:
try:
enr = self.node_db.get_enr(node_id)
except KeyError:
self.logger.warning("Missing ENR for node %s",
encode_hex(node_id))
else:
enrs.append(enr)
enr_partitions = partition_enrs(enrs,
NODES_MESSAGE_PAYLOAD_SIZE) or ((), )
self.logger.debug(
"Responding to %s with %d Nodes message containing %d ENRs at distance %d",
incoming_message.sender_endpoint,
len(enr_partitions),
len(enrs),
incoming_message.message.distance,
)
for partition in enr_partitions:
nodes_message = NodesMessage(
request_id=incoming_message.message.request_id,
total=len(enr_partitions),
enrs=partition,
)
outgoing_message = incoming_message.to_response(nodes_message)
await self.outgoing_message_send_channel.send(outgoing_message)
async def handle_request_on_remote():
async for outgoing_message in outgoing_message_channels[1]:
assert outgoing_message.message == request
assert outgoing_message.receiver_endpoint == remote_endpoint
assert outgoing_message.receiver_node_id == remote_enr.node_id
await incoming_message_channels[0].send(IncomingMessage(
message=response,
sender_endpoint=remote_endpoint,
sender_node_id=remote_enr.node_id,
))
async def respond_with_local_enr(
self, incoming_message: IncomingMessage) -> None:
"""Send a Nodes message containing the local ENR in response to an incoming message."""
local_enr = await self.get_local_enr()
nodes_message = NodesMessage(
request_id=incoming_message.message.request_id,
total=1,
enrs=(local_enr, ),
)
outgoing_message = incoming_message.to_response(nodes_message)
self.logger.debug(
"Responding to %s with Nodes message containing local ENR",
incoming_message.sender_endpoint,
)
await self.outgoing_message_send_channel.send(outgoing_message)
async def test_request_handling(message_dispatcher, incoming_message_channels,
remote_enr, remote_endpoint):
ping_send_channel, ping_receive_channel = trio.open_memory_channel(0)
async with message_dispatcher.add_request_handler(
PingMessage) as request_subscription:
incoming_message = IncomingMessage(
message=PingMessageFactory(),
sender_endpoint=remote_endpoint,
sender_node_id=remote_enr.node_id,
)
await incoming_message_channels[0].send(incoming_message)
with trio.fail_after(1):
handled_incoming_message = await request_subscription.receive()
assert handled_incoming_message == incoming_message
async def test_response_handling(message_dispatcher, remote_enr,
incoming_message_channels):
request_id = message_dispatcher.get_free_request_id(remote_enr.node_id)
async with message_dispatcher.add_response_handler(
remote_enr.node_id,
request_id,
) as response_subscription:
incoming_message = IncomingMessage(
message=PingMessageFactory(request_id=request_id, ),
sender_endpoint=remote_endpoint,
sender_node_id=remote_enr.node_id,
)
await incoming_message_channels[0].send(incoming_message)
with trio.fail_after(1):
handled_response = await response_subscription.receive()
assert handled_response == incoming_message
async def respond_with_pong(self,
incoming_message: IncomingMessage) -> None:
if not isinstance(incoming_message.message, PingMessage):
raise TypeError(f"Can only respond with Pong to Ping, not "
f"{incoming_message.message.__class__.__name__}")
local_enr = await self.get_local_enr()
pong = PongMessage(
request_id=incoming_message.message.request_id,
enr_seq=local_enr.sequence_number,
packet_ip=incoming_message.sender_endpoint.ip_address,
packet_port=incoming_message.sender_endpoint.port,
)
outgoing_message = incoming_message.to_response(pong)
self.logger.debug(
"Responding with Pong to %s",
encode_hex(outgoing_message.receiver_node_id),
)
await self.outgoing_message_send_channel.send(outgoing_message)
async def handle_incoming_packet_during_handshake(self,
incoming_packet: IncomingPacket,
) -> None:
if not self.is_during_handshake:
raise ValueError("Can only handle packets during handshake")
if self.handshake_participant is None:
raise TypeError("handshake_participant is None even though handshake is in progress")
packet = incoming_packet.packet
if self.handshake_participant.is_response_packet(packet):
self.logger.debug("Received %s as handshake response", packet.__class__.__name__)
else:
self.logger.debug("Dropping %s unexpectedly received during handshake", incoming_packet)
return
try:
handshake_result = self.handshake_participant.complete_handshake(packet)
except HandshakeFailure as handshake_failure:
self.logger.warning(
"Handshake with %s has failed: %s",
encode_hex(self.remote_node_id),
handshake_failure,
)
raise # let the service fail
else:
self.logger.info("Handshake with %s was successful", encode_hex(self.remote_node_id))
self.handshake_successful_event.set()
# copy message backlog before we reset it
outgoing_message_backlog = tuple(self.outgoing_message_backlog)
self.reset_handshake_state()
self.session_keys = handshake_result.session_keys
if not self.is_post_handshake:
raise Exception(
"Invariant: As session_keys are set now, peer packer is in post handshake state"
)
if handshake_result.enr is not None:
self.logger.debug("Updating ENR in DB with %r", handshake_result.enr)
await self.enr_db.insert_or_update(handshake_result.enr)
if handshake_result.auth_header_packet is not None:
outgoing_packet = OutgoingPacket(
handshake_result.auth_header_packet,
incoming_packet.sender_endpoint,
)
self.logger.debug(
"Sending %s packet to let peer complete handshake",
outgoing_packet,
)
await self.outgoing_packet_send_channel.send(outgoing_packet)
if handshake_result.message:
incoming_message = IncomingMessage(
handshake_result.message,
incoming_packet.sender_endpoint,
self.remote_node_id,
)
self.logger.debug("Received %s during handshake", incoming_message)
await self.incoming_message_send_channel.send(incoming_message)
self.logger.debug("Sending %d messages from backlog", len(outgoing_message_backlog))
for outgoing_message in outgoing_message_backlog:
await self.handle_outgoing_message(outgoing_message)
