def test_successful_handshake_with_enr_update():
initiator_private_key = PrivateKeyFactory().to_bytes()
recipient_private_key = PrivateKeyFactory().to_bytes()
old_initiator_enr = ENRFactory(private_key=initiator_private_key)
new_initiator_enr = ENRFactory(
private_key=initiator_private_key,
sequence_number=old_initiator_enr.sequence_number + 1,
)
initiator = HandshakeInitiatorFactory(
local_private_key=initiator_private_key,
local_enr=new_initiator_enr,
remote_private_key=recipient_private_key,
)
recipient = HandshakeRecipientFactory(
local_private_key=recipient_private_key,
remote_enr=old_initiator_enr,
initiating_packet_auth_tag=initiator.first_packet_to_send.auth_tag)
initiator_result = initiator.complete_handshake(
recipient.first_packet_to_send)
recipient_result = recipient.complete_handshake(
initiator_result.auth_header_packet)
assert initiator_result.enr is None
assert recipient_result.enr == new_initiator_enr
def test_successful_handshake():
initiator_private_key = PrivateKeyFactory().to_bytes()
recipient_private_key = PrivateKeyFactory().to_bytes()
initiator_enr = ENRFactory(private_key=initiator_private_key)
recipient_enr = ENRFactory(private_key=recipient_private_key)
initial_message = PingMessageFactory()
initiator = HandshakeInitiatorFactory(
local_private_key=initiator_private_key,
local_enr=initiator_enr,
remote_enr=recipient_enr,
initial_message=initial_message,
)
recipient = HandshakeRecipientFactory(
local_private_key=recipient_private_key,
local_enr=recipient_enr,
remote_enr=initiator_enr,
initiating_packet_auth_tag=initiator.first_packet_to_send.auth_tag)
initiator_result = initiator.complete_handshake(
recipient.first_packet_to_send)
recipient_result = recipient.complete_handshake(
initiator_result.auth_header_packet)
assert_session_keys_equal(initiator_result.session_keys,
recipient_result.session_keys)
assert initiator_result.message is None
assert initiator_result.enr is None
assert initiator_result.auth_header_packet is not None
assert recipient_result.message == initial_message
assert recipient_result.enr is None
assert recipient_result.auth_header_packet is None
def test_successful_handshake_with_new_enr():
initiator_private_key = PrivateKeyFactory().to_bytes()
recipient_private_key = PrivateKeyFactory().to_bytes()
initiator_enr = ENRFactory(private_key=initiator_private_key)
recipient_enr = ENRFactory(private_key=recipient_private_key)
initiator = HandshakeInitiatorFactory(
local_private_key=initiator_private_key,
local_enr=initiator_enr,
remote_enr=recipient_enr,
)
recipient = HandshakeRecipientFactory(
local_private_key=recipient_private_key,
local_enr=recipient_enr,
remote_enr=None,
remote_node_id=initiator_enr.node_id,
initiating_packet_auth_tag=initiator.first_packet_to_send.auth_tag)
initiator_result = initiator.complete_handshake(
recipient.first_packet_to_send)
recipient_result = recipient.complete_handshake(
initiator_result.auth_header_packet)
assert initiator_result.enr is None
assert recipient_result.enr == initiator_enr
def test_wait_time_full_queue_and_table(topic_table, max_queue_size,
target_ad_lifetime):
# fill one queue
topic = TopicFactory()
reg_time = 0
oldest_queue_eol = reg_time + target_ad_lifetime
while not topic_table.is_queue_full(topic):
topic_table.register(topic, ENRFactory(), reg_time)
reg_time += 1
# fill the rest of the table
oldest_table_eol = reg_time + target_ad_lifetime
while not topic_table.is_full:
topic_table.register(TopicFactory(), ENRFactory(), reg_time)
reg_time += 1
assert topic_table.get_wait_time(topic, 0) == oldest_queue_eol
assert topic_table.get_wait_time(TopicFactory(), 0) == oldest_queue_eol
# refill queue
oldest_queue_eol = reg_time + target_ad_lifetime
for _ in range(max_queue_size):
topic_table.register(topic, ENRFactory(), reg_time)
reg_time += 1
assert topic_table.get_wait_time(topic, 0) == oldest_queue_eol
assert topic_table.get_wait_time(TopicFactory(), 0) == oldest_table_eol
def test_registration_single_queue(topic_table, max_queue_size):
topic = TopicFactory()
enr = ENRFactory()
other_enr = ENRFactory()
topic_table.get_enrs_for_topic(topic) == ()
topic_table.register(topic, enr, 0)
assert topic_table.get_enrs_for_topic(topic) == (enr, )
topic_table.register(topic, other_enr, 0)
assert topic_table.get_enrs_for_topic(topic) == (other_enr, enr)
with pytest.raises(ValueError):
topic_table.register(topic, enr, 0)
while not topic_table.is_queue_full(topic):
topic_table.register(topic, ENRFactory(), 0)
with pytest.raises(ValueError):
topic_table.register(topic, ENRFactory(), 0)
enrs_before = topic_table.get_enrs_for_topic(topic)
new_enr = ENRFactory()
topic_table.register(topic, new_enr, topic_table.get_wait_time(topic, 0))
enrs_after = topic_table.get_enrs_for_topic(topic)
assert enrs_after == (new_enr, ) + enrs_before[:-1]
def test_wait_time_full_table(topic_table, target_ad_lifetime):
# fill one queue
reg_time = 0
oldest_table_eol = reg_time + target_ad_lifetime
while not topic_table.is_full:
assert topic_table.get_wait_time(TopicFactory(), 0) == 0
topic_table.register(TopicFactory(), ENRFactory(), reg_time)
reg_time += 1
assert topic_table.get_wait_time(TopicFactory(), 0) == oldest_table_eol
topic_table.register(TopicFactory(), ENRFactory(), reg_time)
assert topic_table.get_wait_time(TopicFactory(), 0) == oldest_table_eol + 1
async def test_insert_or_update(db):
private_key = PrivateKeyFactory().to_bytes()
enr = ENRFactory(private_key=private_key)
await db.insert_or_update(enr)
assert await db.get(enr.node_id) == enr
await db.insert_or_update(enr)
assert await db.get(enr.node_id) == enr
updated_enr = ENRFactory(private_key=private_key,
sequence_number=enr.sequence_number + 1)
await db.insert_or_update(updated_enr)
assert await db.get(enr.node_id) == updated_enr
def test_registration_two_queues(topic_table, max_queue_size):
topic1 = TopicFactory()
topic2 = TopicFactory()
enr = ENRFactory()
topic_table.register(topic1, enr, 0)
while not topic_table.is_queue_full(topic1):
topic_table.register(topic1, ENRFactory(), 0)
topic_table.register(topic2, enr, 1)
while not topic_table.is_queue_full(topic2):
topic_table.register(topic2, ENRFactory(), 1)
with pytest.raises(ValueError):
topic_table.register(topic1, ENRFactory(), 1)
with pytest.raises(ValueError):
topic_table.register(topic2, ENRFactory(), 1)
with pytest.raises(ValueError):
topic_table.register(topic2, ENRFactory(),
topic_table.get_wait_time(topic1, 0))
enrs_topic1_before = topic_table.get_enrs_for_topic(topic1)
enrs_topic2_before = topic_table.get_enrs_for_topic(topic2)
new_enr_topic1 = ENRFactory()
new_enr_topic2 = ENRFactory()
topic_table.register(topic1, new_enr_topic1,
topic_table.get_wait_time(topic1, 0))
topic_table.register(topic2, new_enr_topic2,
topic_table.get_wait_time(topic2, 0))
enrs_topic1_after = topic_table.get_enrs_for_topic(topic1)
enrs_topic2_after = topic_table.get_enrs_for_topic(topic2)
assert enrs_topic1_after == (new_enr_topic1, ) + enrs_topic1_before[:-1]
assert enrs_topic2_after == (new_enr_topic2, ) + enrs_topic2_before[:-1]
async def test_insert(db):
private_key = PrivateKeyFactory().to_bytes()
enr = ENRFactory(private_key=private_key)
await db.insert(enr)
assert await db.contains(enr.node_id)
assert await db.get(enr.node_id) == enr
with pytest.raises(ValueError):
await db.insert(enr)
updated_enr = ENRFactory(private_key=private_key,
sequence_number=enr.sequence_number + 1)
with pytest.raises(ValueError):
await db.insert(updated_enr)
def remote_enr(remote_endpoint):
return ENRFactory(
custom_kv_pairs={
b"ip": remote_endpoint.ip_address,
b"udp": remote_endpoint.port,
},
)
async def test_update(db):
private_key = PrivateKeyFactory().to_bytes()
enr = ENRFactory(private_key=private_key)
with pytest.raises(KeyError):
await db.update(enr)
await db.insert(enr)
await db.update(enr)
assert await db.get(enr.node_id) == enr
updated_enr = ENRFactory(private_key=private_key,
sequence_number=enr.sequence_number + 1)
await db.update(updated_enr)
assert await db.get(enr.node_id) == updated_enr
async def test_file_db_saves_enrs(file_db_dir, file_db):
enr = ENRFactory()
await file_db.insert(enr)
filename = get_enr_filename(enr)
assert (file_db_dir / filename).exists()
assert (file_db_dir / filename).is_file()
assert ENR.from_repr((file_db_dir / filename).read_text()) == enr
async def filled_routing_table(routing_table, node_db):
# add entries until the first bucket is full
while len(routing_table.get_nodes_at_log_distance(255)) < routing_table.bucket_size:
enr = ENRFactory()
routing_table.update(enr.node_id)
node_db.set_enr(enr)
return routing_table
def enr(private_key, endpoint):
return ENRFactory(
private_key=private_key,
custom_kv_pairs={
b"ip": endpoint.ip_address,
b"udp": endpoint.port,
}
)
def remote_enr(remote_private_key, remote_endpoint):
return ENRFactory(
private_key=remote_private_key,
custom_kv_pairs={
b"ip": remote_endpoint.ip_address,
b"udp": remote_endpoint.port,
}
)
def test_wait_time_full_queue(topic_table, max_total_size, target_ad_lifetime):
topic = TopicFactory()
different_topic = TopicFactory()
reg_time = 0
oldest_queue_eol = reg_time + target_ad_lifetime
while not topic_table.is_queue_full(topic):
assert topic_table.get_wait_time(topic, 0) == 0
assert topic_table.get_wait_time(different_topic, 0) == 0
topic_table.register(topic, ENRFactory(), reg_time)
reg_time += 1
assert topic_table.get_wait_time(topic, 0) == oldest_queue_eol
assert topic_table.get_wait_time(different_topic, 0) == 0
topic_table.register(topic, ENRFactory(), reg_time)
assert topic_table.get_wait_time(topic, 0) == oldest_queue_eol + 1
assert topic_table.get_wait_time(different_topic, 0) == 0
async def test_get(db):
enr = ENRFactory()
with pytest.raises(KeyError):
await db.get(enr.node_id)
await db.insert(enr)
assert await db.get(enr.node_id) == enr
def test_get_and_set_enr(node_db):
private_key = PrivateKeyFactory().to_bytes()
db = node_db
enr = ENRFactory(private_key=private_key)
with pytest.raises(KeyError):
db.get_enr(enr.node_id)
db.set_enr(enr)
assert db.get_enr(enr.node_id) == enr
updated_enr = ENRFactory(private_key=private_key, sequence_number=enr.sequence_number + 1)
db.set_enr(updated_enr)
assert db.get_enr(enr.node_id) == updated_enr
with pytest.raises(ValueError):
db.set_enr(enr)
async def test_memory_checks_identity_scheme():
empty_identity_scheme_registry = IdentitySchemeRegistry()
memory_db = MemoryEnrDb(empty_identity_scheme_registry)
enr = ENRFactory()
with pytest.raises(ValueError):
await memory_db.insert(enr)
with pytest.raises(ValueError):
await memory_db.insert_or_update(enr)
async def test_memory_remove(memory_db):
enr = ENRFactory()
with pytest.raises(KeyError):
await memory_db.remove(enr.node_id)
await memory_db.insert(enr)
await memory_db.remove(enr.node_id)
assert not await memory_db.contains(enr.node_id)
async def test_generate_eth_cap_enr_field():
base_db = AtomicDB()
ChainDB(base_db).persist_header(ROPSTEN_GENESIS_HEADER)
enr_field = await generate_eth_cap_enr_field(ROPSTEN_VM_CONFIGURATION,
AsyncHeaderDB(base_db))
enr = ENRFactory(custom_kv_pairs={enr_field[0]: enr_field[1]})
assert extract_forkid(enr) == ForkID(hash=to_bytes(hexstr='0x30c7ddbc'),
next=10)
def test_table_size(topic_table, max_queue_size, target_ad_lifetime):
assert len(topic_table) == 0
topic = TopicFactory()
topic_table.register(topic, ENRFactory(), 0)
assert len(topic_table) == 1
topic_table.register(topic, ENRFactory(), 0)
assert len(topic_table) == 2
topic2 = TopicFactory()
topic_table.register(topic2, ENRFactory(), 0)
assert len(topic_table) == 3
topic3 = TopicFactory()
for _ in range(max_queue_size):
topic_table.register(topic3, ENRFactory(), 0)
assert len(topic_table) == 3 + max_queue_size
topic_table.register(topic3, ENRFactory(), target_ad_lifetime)
assert len(topic_table) == 3 + max_queue_size
def test_get_and_set_last_pong_time(node_db):
db = node_db
enr = ENRFactory()
with pytest.raises(KeyError):
db.get_last_pong_time(enr.node_id)
pong_time = int(time.monotonic())
db.set_last_pong_time(enr.node_id, pong_time)
assert db.get_last_pong_time(enr.node_id) == pong_time
def test_delete_enr(node_db):
db = node_db
enr = ENRFactory()
with pytest.raises(KeyError):
db.delete_enr(enr.node_id)
db.set_enr(enr)
db.delete_enr(enr.node_id)
with pytest.raises(KeyError):
db.get_enr(enr.node_id)
def _make_node_with_enr_and_forkid(genesis_hash, head, vm_config):
fork_blocks = forkid.extract_fork_blocks(vm_config)
node_forkid = forkid.make_forkid(genesis_hash, head, fork_blocks)
ip = socket.inet_aton(IPAddressFactory.generate())
udp_port = 30304
enr = ENRFactory(
custom_kv_pairs={
b'eth': sedes.List([forkid.ForkID]).serialize([node_forkid]),
IP_V4_ADDRESS_ENR_KEY: ip,
UDP_PORT_ENR_KEY: udp_port,
TCP_PORT_ENR_KEY: udp_port,
})
return Node(enr)
def test_delete_last_pong_time(node_db):
db = node_db
enr = ENRFactory()
with pytest.raises(KeyError):
db.delete_last_pong_time(enr.node_id)
pong_time = int(time.monotonic())
db.set_last_pong_time(enr.node_id, pong_time)
db.delete_last_pong_time(enr.node_id)
with pytest.raises(KeyError):
db.get_last_pong_time(enr.node_id)
async def test_request_nodes(message_dispatcher, remote_enr, remote_endpoint,
incoming_message_channels,
outgoing_message_channels, nursery):
request_id = message_dispatcher.get_free_request_id(remote_enr.node_id)
request = FindNodeMessage(
request_id=request_id,
distance=3,
)
enrs_per_message = [[ENRFactory() for _ in range(2)] for _ in range(3)]
response_messages = [
NodesMessage(request_id=request_id,
total=len(enrs_per_message),
enrs=enrs) for enrs in enrs_per_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
for response in response_messages:
await incoming_message_channels[0].send(
IncomingMessage(
message=response,
sender_endpoint=remote_endpoint,
sender_node_id=remote_enr.node_id,
))
nursery.start_soon(handle_request_on_remote)
with trio.fail_after(3):
received_responses = await message_dispatcher.request_nodes(
remote_enr.node_id, request)
assert len(received_responses) == len(response_messages)
for received_response, expected_response_message in zip(
received_responses, response_messages):
assert received_response.sender_endpoint == remote_endpoint
assert received_response.sender_node_id == remote_enr.node_id
assert received_response.message == expected_response_message
def test_registration_full_table(topic_table, max_queue_size, max_total_size):
for _ in range(max_total_size):
topic_table.register(TopicFactory(), ENRFactory(), 0)
assert topic_table.is_full
with pytest.raises(ValueError):
topic_table.register(TopicFactory(), ENRFactory(), 0)
wait_time = topic_table.get_wait_time(TopicFactory(), 0)
topic_table.register(TopicFactory(), ENRFactory(), wait_time)
topic = TopicFactory()
assert not topic_table.is_queue_full(topic)
while not topic_table.is_queue_full(topic):
topic_table.register(topic, ENRFactory(),
topic_table.get_wait_time(topic, 0))
with pytest.raises(ValueError):
topic_table.register(TopicFactory(), ENRFactory(), 0)
wait_time = topic_table.get_wait_time(TopicFactory(), 0)
topic_table.register(TopicFactory(), ENRFactory(), wait_time)
def remote_enr(remote_private_key):
return ENRFactory(private_key=remote_private_key)
def initial_enr(private_key):
return ENRFactory(private_key=private_key, )