async def test_packet_decoder_error():
datagram_send_channel, datagram_receive_channel = trio.open_memory_channel(
1)
packet_send_channel, packet_receive_channel = trio.open_memory_channel(1)
service = PacketDecoder(datagram_receive_channel, packet_send_channel)
async with background_trio_service(service):
# send invalid packet
await datagram_send_channel.send(
InboundDatagram(
datagram=b"not a valid packet",
sender_endpoint=EndpointFactory(),
))
# send valid packet
packet = AuthTagPacketFactory()
sender_endpoint = EndpointFactory()
await datagram_send_channel.send(
InboundDatagram(
datagram=packet.to_wire_bytes(),
sender_endpoint=sender_endpoint,
))
# ignore the invalid one, only receive the valid one
with trio.fail_after(0.5):
inbound_packet = await packet_receive_channel.receive()
assert inbound_packet.packet == packet
assert inbound_packet.sender_endpoint.ip_address == sender_endpoint.ip_address
assert inbound_packet.sender_endpoint.port == sender_endpoint.port
async def test_pool_get_session_by_endpoint(tester, initiator, pool, events):
endpoint = EndpointFactory()
# A: initiated locally, handshake incomplete
remote_a = tester.node(endpoint=endpoint)
session_a = pool.initiate_session(endpoint, remote_a.node_id)
# B: initiated locally, handshake complete
remote_b = tester.node(endpoint=endpoint)
driver_b = tester.session_pair(
initiator,
remote_b,
)
with trio.fail_after(1):
await driver_b.handshake()
session_b = driver_b.initiator.session
# C: initiated remotely, handshake incomplete
session_c = pool.receive_session(endpoint)
# D: initiated remotely, handshake complete
remote_d = tester.node(endpoint=endpoint)
driver_d = tester.session_pair(
remote_d,
initiator,
)
await driver_d.handshake()
session_d = driver_d.recipient.session
# Some other sessions with non-matching endpoints before handshake
session_e = pool.receive_session(EndpointFactory())
session_f = pool.initiate_session(EndpointFactory(), NodeIDFactory())
# Some other sessions with non-matching endpoints after handshake
driver_g = tester.session_pair(initiator, )
await driver_g.handshake()
session_g = driver_g.initiator.session
driver_h = tester.session_pair(recipient=initiator, )
await driver_h.handshake()
session_h = driver_h.recipient.session
endpoint_matches = pool.get_sessions_for_endpoint(endpoint)
assert len(endpoint_matches) == 4
assert session_a in endpoint_matches
assert session_b in endpoint_matches
assert session_c in endpoint_matches
assert session_d in endpoint_matches
assert session_e not in endpoint_matches
assert session_f not in endpoint_matches
assert session_g not in endpoint_matches
assert session_h not in endpoint_matches
async def test_send_endpoint_vote(
ping_sender_service,
routing_table,
inbound_message_channels,
outbound_message_channels,
endpoint_vote_channels,
local_enr,
remote_enr,
remote_endpoint,
):
# wait for ping
with trio.fail_after(ROUTING_TABLE_PING_INTERVAL):
outbound_message = await outbound_message_channels[1].receive()
ping = outbound_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,
)
inbound_message = InboundMessage(
message=pong,
sender_endpoint=outbound_message.receiver_endpoint,
sender_node_id=outbound_message.receiver_node_id,
)
await inbound_message_channels[0].send(inbound_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 == inbound_message.sender_node_id
async def test_endpoint_tracker_updates_enr(endpoint_tracker, initial_enr,
enr_db, vote_channels):
endpoint = EndpointFactory()
endpoint_vote = EndpointVoteFactory(endpoint=endpoint)
await vote_channels[0].send(endpoint_vote)
await wait_all_tasks_blocked() # wait until vote has been processed
updated_enr = enr_db.get_enr(initial_enr.node_id)
assert updated_enr.sequence_number == initial_enr.sequence_number + 1
assert updated_enr[IP_V4_ADDRESS_ENR_KEY] == endpoint.ip_address
assert updated_enr[UDP_PORT_ENR_KEY] == endpoint.port
async def test_peer_packer_sends_who_are_you(peer_packer,
inbound_packet_channels,
outbound_packet_channels,
nursery):
inbound_packet = InboundPacket(
AuthTagPacketFactory(),
EndpointFactory(),
)
inbound_packet_channels[0].send_nowait(inbound_packet)
with trio.fail_after(0.5):
outbound_packet = await outbound_packet_channels[1].receive()
assert peer_packer.is_during_handshake
assert outbound_packet.receiver_endpoint == inbound_packet.sender_endpoint
assert isinstance(outbound_packet.packet, WhoAreYouPacket)
assert outbound_packet.packet.token == inbound_packet.packet.auth_tag
async def test_peer_packer_initiates_handshake(peer_packer,
outbound_message_channels,
outbound_packet_channels,
nursery):
outbound_message = OutboundMessage(
PingMessageFactory(),
EndpointFactory(),
peer_packer.remote_node_id,
)
outbound_message_channels[0].send_nowait(outbound_message)
with trio.fail_after(0.5):
outbound_packet = await outbound_packet_channels[1].receive()
assert peer_packer.is_during_handshake
assert outbound_packet.receiver_endpoint == outbound_message.receiver_endpoint
assert isinstance(outbound_packet.packet, AuthTagPacket)
def node(
self,
private_key: Optional[keys.PrivateKey] = None,
endpoint: Optional[Endpoint] = None,
enr_db: Optional[ENRDatabaseAPI] = None,
events: Optional[EventsAPI] = None,
) -> Node:
if private_key is None:
private_key = PrivateKeyFactory()
if endpoint is None:
endpoint = EndpointFactory.localhost()
if enr_db is None:
enr_db = ENRDB({})
return Node(private_key=private_key,
endpoint=endpoint,
enr_db=enr_db,
events=events)
async def test_subscribe_delivers_messages():
node_id = NodeIDFactory()
endpoint = EndpointFactory()
manager = SubscriptionManager()
async with manager.subscribe(MessageForTesting) as subscription:
with pytest.raises(trio.WouldBlock):
subscription.receive_nowait()
manager.feed_subscriptions(
InboundMessage(
message=OtherMessageForTesting(1234),
sender_node_id=node_id,
sender_endpoint=endpoint,
)
)
manager.feed_subscriptions(
InboundMessage(
message=MessageForTesting(1234),
sender_node_id=node_id,
sender_endpoint=endpoint,
)
)
with trio.fail_after(1):
message = await subscription.receive()
assert isinstance(message.message, MessageForTesting)
assert message.message.id == 1234
with pytest.raises(trio.WouldBlock):
subscription.receive_nowait()
manager.feed_subscriptions(
InboundMessage(
message=MessageForTesting(1234),
sender_node_id=node_id,
sender_endpoint=endpoint,
)
)
with pytest.raises(trio.ClosedResourceError):
subscription.receive_nowait()
def node(
self,
private_key: Optional[keys.PrivateKey] = None,
endpoint: Optional[Endpoint] = None,
enr_db: Optional[QueryableENRDatabaseAPI] = None,
events: Optional[EventsAPI] = None,
name: Optional[str] = None,
) -> Node:
if private_key is None:
private_key = PrivateKeyFactory()
if endpoint is None:
endpoint = EndpointFactory.localhost()
if enr_db is None:
enr_db = QueryableENRDB(sqlite3.connect(":memory:"))
node = Node(private_key=private_key,
endpoint=endpoint,
enr_db=enr_db,
events=events)
self.logger.info("Tester[Node]: name=%s node_id=%s", name,
node.node_id.hex())
return node
async def test_packet_encoder():
packet_send_channel, packet_receive_channel = trio.open_memory_channel(1)
datagram_send_channel, datagram_receive_channel = trio.open_memory_channel(
1)
service = PacketEncoder(packet_receive_channel, datagram_send_channel)
async with background_trio_service(service):
receiver_endpoint = EndpointFactory()
outbound_packet = OutboundPacket(
packet=AuthTagPacketFactory(),
receiver_endpoint=receiver_endpoint,
)
await packet_send_channel.send(outbound_packet)
with trio.fail_after(0.5):
outbound_datagram = await datagram_receive_channel.receive()
assert outbound_datagram.datagram == outbound_packet.packet.to_wire_bytes(
)
assert (outbound_datagram.receiver_endpoint.ip_address ==
receiver_endpoint.ip_address)
assert outbound_datagram.receiver_endpoint.port == receiver_endpoint.port
def remote_endpoint():
return EndpointFactory()
def endpoint():
return EndpointFactory()
async def test_subscribe_filters_by_node_id_and_endpoint():
node_id = NodeIDFactory()
endpoint = EndpointFactory()
manager = SubscriptionManager()
async with AsyncExitStack() as stack:
subscription_a = await stack.enter_async_context(
manager.subscribe(MessageForTesting)
)
subscription_b = await stack.enter_async_context(
manager.subscribe(MessageForTesting, node_id=node_id)
)
subscription_c = await stack.enter_async_context(
manager.subscribe(MessageForTesting, node_id=node_id, endpoint=endpoint)
)
with pytest.raises(trio.WouldBlock):
subscription_a.receive_nowait()
with pytest.raises(trio.WouldBlock):
subscription_b.receive_nowait()
with pytest.raises(trio.WouldBlock):
subscription_c.receive_nowait()
# One Message that doesn't match the message type
manager.feed_subscriptions(
InboundMessage(
message=OtherMessageForTesting(1234),
sender_node_id=node_id,
sender_endpoint=endpoint,
)
)
# One Message that only matches the message type
manager.feed_subscriptions(
InboundMessage(
message=MessageForTesting(1),
sender_node_id=NodeIDFactory(),
sender_endpoint=EndpointFactory(),
)
)
# One Message that matches the message type AND node_id
manager.feed_subscriptions(
InboundMessage(
message=MessageForTesting(2),
sender_node_id=node_id,
sender_endpoint=EndpointFactory(),
)
)
# One Message that matches the message type AND node_id AND endpoint
manager.feed_subscriptions(
InboundMessage(
message=MessageForTesting(3),
sender_node_id=node_id,
sender_endpoint=endpoint,
)
)
# now grab all the messages
with trio.fail_after(1):
message_a_0 = await subscription_a.receive()
message_a_1 = await subscription_a.receive()
message_a_2 = await subscription_a.receive()
message_b_0 = await subscription_b.receive()
message_b_1 = await subscription_b.receive()
message_c_0 = await subscription_c.receive()
# all of the subscriptions should now be empty
with pytest.raises(trio.WouldBlock):
subscription_a.receive_nowait()
with pytest.raises(trio.WouldBlock):
subscription_b.receive_nowait()
with pytest.raises(trio.WouldBlock):
subscription_c.receive_nowait()
assert message_a_0.message.id == 1
assert message_a_1.message.id == 2
assert message_a_2.message.id == 3
assert message_b_0.message.id == 2
assert message_b_1.message.id == 3
assert message_c_0.message.id == 3