async def test_advertisement_collector_handle_existing_valid_remote_advertisement( alice, bob, alice_alexandria_network, bob_alexandria_network, autojump_clock, ): content = ContentFactory(2048) proof = compute_proof(content, sedes=content_sedes) ad_collector = alice_alexandria_network.advertisement_collector advertisement = AdvertisementFactory( private_key=bob.private_key, hash_tree_root=proof.get_hash_tree_root(), ) bob_alexandria_network.commons_content_storage.set_content( advertisement.content_key, content, ) alice_alexandria_network.remote_advertisement_db.add(advertisement) assert alice_alexandria_network.remote_advertisement_db.exists( advertisement) async with ad_collector.new_advertisement.subscribe() as subscription: with trio.fail_after(5): await ad_collector.handle_advertisement(advertisement) with pytest.raises(trio.TooSlowError): with trio.fail_after(5): await subscription.receive() assert alice_alexandria_network.remote_advertisement_db.exists( advertisement)
async def test_advertisement_collector_acks_false_if_advertisements_already_known( alice, bob, alice_alexandria_client, bob_alexandria_network, autojump_clock, ): content = ContentFactory(2048) proof = compute_proof(content, sedes=content_sedes) advertisement = AdvertisementFactory( private_key=bob.private_key, hash_tree_root=proof.get_hash_tree_root(), ) bob_alexandria_network.commons_content_storage.set_content( advertisement.content_key, content, ) bob_alexandria_network.local_advertisement_db.add(advertisement) with pytest.raises(trio.TooSlowError): with trio.fail_after(10): async with bob_alexandria_network.advertisement_collector.new_advertisement.subscribe_and_wait( ): # noqa: E501 ack_message = await alice_alexandria_client.advertise( bob.node_id, bob.endpoint, advertisements=(advertisement, ), ) assert len(ack_message.payload.acked) == 1 assert ack_message.payload.acked[0] is False
async def test_content_provider_restricts_max_chunks( alice, bob, alice_alexandria_network, bob_alexandria_client, ): content = ContentFactory(length=1024 * 10) content_key = b"test-content-key" content_storage = MemoryContentStorage({content_key: content}) proof = compute_proof(content, sedes=content_sedes) content_provider = ContentProvider(bob_alexandria_client, (content_storage, ), max_chunks_per_request=16) async with background_trio_service(content_provider): # this ensures that the subscription is in place. await content_provider.ready() with trio.fail_after(2): proof = await alice_alexandria_network.get_content_proof( bob.node_id, hash_tree_root=proof.get_hash_tree_root(), content_key=content_key, start_chunk_index=0, max_chunks=100, endpoint=bob.endpoint, ) validate_proof(proof) num_leaf_elements = len( tuple(element for element in proof.elements if len(element.path) == proof.path_bit_length)) assert num_leaf_elements == 16
async def test_advertisement_collector_handle_new_valid_remote_advertisement( alice, bob, alice_alexandria_network, bob_alexandria_network, ): content = ContentFactory(2048) proof = compute_proof(content, sedes=content_sedes) ad_collector = alice_alexandria_network.advertisement_collector advertisement = AdvertisementFactory( private_key=bob.private_key, hash_tree_root=proof.get_hash_tree_root(), ) bob_alexandria_network.commons_content_storage.set_content( advertisement.content_key, content, ) assert not alice_alexandria_network.remote_advertisement_db.exists( advertisement) with trio.fail_after(5): async with ad_collector.new_advertisement.subscribe_and_wait(): await ad_collector.handle_advertisement(advertisement) assert alice_alexandria_network.remote_advertisement_db.exists( advertisement)
async def test_content_provider_serves_large_content( alice, bob, alice_alexandria_network, bob_alexandria_client, ): content = ContentFactory(length=1024 * 10) content_key = b"test-content-key" content_storage = MemoryContentStorage({content_key: content}) proof = compute_proof(content, sedes=content_sedes) content_provider = ContentProvider(bob_alexandria_client, (content_storage, )) async with background_trio_service(content_provider): async with alice_alexandria_network.client.subscribe( ContentMessage) as subscription: # this ensures that the subscription is in place. await content_provider.ready() with trio.fail_after(2): content_retrieval_ctx = alice_alexandria_network.retrieve_content( content_key, proof.get_hash_tree_root(), ) async with content_retrieval_ctx as content_retrieval: await content_retrieval.node_queue.add(bob.node_id) result = await content_retrieval.wait_content_proof() validate_proof(result) result_data = result.get_proven_data() assert result_data[0:len(content)] == content response = await subscription.receive() assert response.message.payload.is_proof is True
async def test_advertisement_collector_validate_local_hash_tree_root_mismatch( alice, alice_alexandria_network, ): ad_collector = alice_alexandria_network.advertisement_collector advertisement = AdvertisementFactory(private_key=alice.private_key) content_storage = alice_alexandria_network.commons_content_storage content_storage.set_content(advertisement.content_key, ContentFactory()) with pytest.raises(ValidationError, match="Mismatched roots"): await ad_collector.validate_advertisement(advertisement)
def test_content_storage_set_content_exists_ok(content_storage): content_key = b"\x00test-key" content = ContentFactory(128) assert content_storage.has_content(content_key) is False content_storage.set_content(content_key, content) with pytest.raises(ContentAlreadyExists): content_storage.set_content(content_key, content) with pytest.raises(ContentAlreadyExists): content_storage.set_content(content_key, content, exists_ok=False) # sanity check assert content_storage.get_content(content_key) == content new_content = ContentFactory(256) assert new_content != content content_storage.set_content(content_key, new_content, exists_ok=True) # sanity check assert content_storage.get_content(content_key) == new_content
async def test_content_retrieval_with_griefing_peer_sending_tiny_chunks( alice, bob, alice_alexandria_network, bob_alexandria_client, ): content = ContentFactory(length=1024 * 10) proof = compute_proof(content, sedes=content_sedes) content_retrieval = ContentRetrieval( alice_alexandria_network, content_key=b"test-key", hash_tree_root=proof.get_hash_tree_root(), ) async with bob_alexandria_client.subscribe( GetContentMessage) as subscription: async with trio.open_nursery() as nursery: async def _serve(): async for request in subscription: start_at = request.message.payload.start_chunk_index * 32 # We only every return a proof for 1 chunk of data end_at = min(len(content), start_at + 32) partial = proof.to_partial(start_at, end_at - start_at) payload = partial.serialize() await bob_alexandria_client.send_content( request.sender_node_id, request.sender_endpoint, is_proof=True, payload=payload, request_id=request.request_id, ) nursery.start_soon(_serve) await content_retrieval.node_queue.add(bob.node_id) with trio.fail_after(10): async with background_trio_service(content_retrieval): with trio.fail_after(5): result = await content_retrieval.wait_content_proof() validate_proof(result) result_data = result.get_proven_data() assert result_data[0:len(content)] == content nursery.cancel_scope.cancel()
async def test_alexandria_network_get_content_proof_api( alice, bob, alice_alexandria_network, bob_alexandria_client, content_size, ): content = ContentFactory(length=content_size) proof = compute_proof(content, sedes=content_sedes) async with bob_alexandria_client.subscribe( GetContentMessage) as subscription: async with trio.open_nursery() as nursery: async def _serve(): request = await subscription.receive() if content_size > 1024: partial = proof.to_partial( request.message.payload.start_chunk_index * 32, request.message.payload.max_chunks * 32, ) payload = partial.serialize() is_proof = True else: payload = content is_proof = False await bob_alexandria_client.send_content( request.sender_node_id, request.sender_endpoint, is_proof=is_proof, payload=payload, request_id=request.request_id, ) nursery.start_soon(_serve) with trio.fail_after(2): partial = await alice_alexandria_network.get_content_proof( bob.node_id, hash_tree_root=proof.get_hash_tree_root(), content_key=b"test-content-key", start_chunk_index=0, max_chunks=16, ) validate_proof(partial) partial_data = partial.get_proven_data() assert partial_data[0:16 * 32] == content[0:16 * 32]
def test_ssz_proof_get_missing_segments_last_chunk_not_full(): content = ContentFactory(180) # 12 bytes short full_proof = compute_proof(content, sedes=short_content_sedes) assert tuple(full_proof.get_missing_segments()) == () partial = full_proof.to_partial(64, 64) missing_segments = tuple(partial.get_missing_segments()) assert len(missing_segments) == 2 segment_a, segment_b = missing_segments assert segment_a.start_at == 0 assert segment_a.length == 64 assert segment_b.start_at == 128 assert segment_b.length == 52
def test_ssz_proof_get_missing_segments_only_head(): content = ContentFactory(512) full_proof = compute_proof(content, sedes=short_content_sedes) assert tuple(full_proof.get_missing_segments()) == () partial = full_proof.to_partial(160, 352) missing_segments = tuple(partial.get_missing_segments()) assert len(missing_segments) == 1 segment = missing_segments[0] assert segment.start_at == 0 # The length is 128 instead of 160 because the segment just before the # partial boundary gets included as part of the proof since the proof # starts in the middle of two sibling leaf nodes. assert segment.length == 128
async def test_alexandria_network_get_content( tester, alice, ): content = ContentFactory(4096) proof = compute_proof(content, sedes=content_sedes) hash_tree_root = proof.get_hash_tree_root() content_key = b"test-key" async with AsyncExitStack() as stack: networks = await stack.enter_async_context( tester.alexandria.network_group(4)) for network in networks: advertisement = Advertisement.create( content_key=content_key, hash_tree_root=hash_tree_root, private_key=network.client.local_private_key, ) network.local_advertisement_db.add(advertisement) network.pinned_content_storage.set_content(content_key, content) # give the the network some time to interconnect. with trio.fail_after(30): for _ in range(1000): await trio.lowlevel.checkpoint() bootnodes = tuple(network.enr_manager.enr for network in networks[:2]) alice_alexandria_network = await stack.enter_async_context( alice.alexandria.network(bootnodes=bootnodes)) # give alice some time to interconnect too with trio.fail_after(30): for _ in range(1000): await trio.lowlevel.checkpoint() with trio.fail_after(60): result = await alice_alexandria_network.get_content( content_key, hash_tree_root=hash_tree_root) assert result == proof assert result.get_content() == content
def test_ssz_proof_get_missing_segments_only_middle(): content = ContentFactory(512) full_proof = compute_proof(content, sedes=short_content_sedes) assert tuple(full_proof.get_missing_segments()) == () head_proof = full_proof.to_partial(0, 160) tail_proof = full_proof.to_partial(352, 160) partial = head_proof.merge(tail_proof) missing_segments = tuple(partial.get_missing_segments()) assert len(missing_segments) == 1 segment = missing_segments[0] # The segmeent starts at 192 because the starting proof ends between two # sibling nodes, causing an extra leaf node to be included. assert segment.start_at == 192 # The length is 128 instead of 192 because the tail segment # starts in the middle of two sibling leaves which causes one extra leaf # to be included. assert segment.length == 128
async def test_advertisement_collector_validate_remote_fail_custody_proof_check( bob, alice_alexandria_network, bob_alexandria_network, autojump_clock, ): content = ContentFactory(2048) proof = compute_proof(content, sedes=content_sedes) ad_collector = alice_alexandria_network.advertisement_collector advertisement = AdvertisementFactory( private_key=bob.private_key, hash_tree_root=proof.get_hash_tree_root(), ) async with bob_alexandria_network.client.subscribe( GetContentMessage) as subscription: did_serve_initial_proof = False async with trio.open_nursery() as nursery: did_serve_initial_proof = True async def _respond(): request = await subscription.receive() partial = proof.to_partial(0, 64) await bob_alexandria_network.client.send_content( request.sender_node_id, request.sender_endpoint, is_proof=True, payload=partial.serialize(), request_id=request.request_id, ) nursery.start_soon(_respond) with pytest.raises(ValidationError, match="Proof of custody check failed"): await ad_collector.validate_advertisement(advertisement) assert did_serve_initial_proof is True
def test_content_storage_read_write_exists_delete(content_storage): content_key = b"\x00test-key" content = ContentFactory(128) assert content_storage.has_content(content_key) is False with pytest.raises(ContentNotFound): content_storage.get_content(content_key) with pytest.raises(ContentNotFound): content_storage.delete_content(content_key) content_storage.set_content(content_key, content) assert content_storage.has_content(content_key) is True assert content_storage.get_content(content_key) == content content_storage.delete_content(content_key) assert content_storage.has_content(content_key) is False with pytest.raises(ContentNotFound): content_storage.get_content(content_key) with pytest.raises(ContentNotFound): content_storage.delete_content(content_key)
async def test_content_manager_enumerates_and_broadcasts_content( alice, bob, ): # # Test Setup: # # 15 pieces of content # 5 that don't already have advertisements to test lazy creation contents = tuple((b"\xfftest-content-" + bytes([idx]), ContentFactory(256 + 25)[idx:idx + 256]) for idx in range(12)) advertisement_db = alice.alexandria.local_advertisement_db pinned_storage = alice.alexandria.pinned_content_storage for content_key, content in contents[0:13:3]: hash_tree_root = ssz.get_hash_tree_root(content, sedes=content_sedes) advertisement = Advertisement.create( content_key=content_key, hash_tree_root=hash_tree_root, private_key=alice.private_key, ) advertisement_db.add(advertisement) content_keys, content_payloads = zip(*contents) for content_key, content in contents: pinned_storage.set_content(content_key, content) received_advertisements = [] async with AsyncExitStack() as stack: bob_alexandria_client = await stack.enter_async_context( bob.alexandria.client()) ad_subscription = await stack.enter_async_context( bob_alexandria_client.subscribe(AdvertiseMessage)) ping_subscription = await stack.enter_async_context( bob_alexandria_client.subscribe(PingMessage)) find_nodes_subscription = await stack.enter_async_context( bob_alexandria_client.subscribe(FindNodesMessage)) done = trio.Event() async def _do_found_nodes(): async for request in find_nodes_subscription: await bob_alexandria_client.send_found_nodes( request.sender_node_id, request.sender_endpoint, enrs=(), request_id=request.request_id, ) async def _do_pong(): async for request in ping_subscription: await bob_alexandria_client.send_pong( request.sender_node_id, request.sender_endpoint, advertisement_radius=2**256 - 1, enr_seq=bob.enr.sequence_number, request_id=request.request_id, ) async def _do_ack(): async for request in ad_subscription: received_advertisements.extend(request.message.payload) await bob_alexandria_client.send_ack( request.sender_node_id, request.sender_endpoint, advertisement_radius=2**256 - 1, acked=tuple(True for _ in request.message.payload), request_id=request.request_id, ) if len(received_advertisements) >= 12: done.set() break async with trio.open_nursery() as nursery: nursery.start_soon(_do_ack) nursery.start_soon(_do_pong) nursery.start_soon(_do_found_nodes) async with alice.alexandria.network() as alice_alexandria_network: await alice_alexandria_network.bond(bob.node_id) with trio.fail_after(90): await done.wait() nursery.cancel_scope.cancel() # 1. All 12 pieces of content should have been advertised received_keys = { advertisement.content_key for advertisement in received_advertisements } assert len(received_keys) == 12 assert received_keys == set(content_keys) # 2. The 2 contents that didn't have advertisements should have been lazily created. assert all( advertisement_db.exists(advertisement) for advertisement in received_advertisements)