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)
