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