async def get_local_enr(boot_info: BootInfo, node_db: NodeDBAPI,
local_private_key: PrivateKey) -> ENR:
minimal_enr = UnsignedENR(
sequence_number=1,
kv_pairs={
b"id": b"v4",
b"secp256k1": local_private_key.public_key.to_compressed_bytes(),
b"udp": boot_info.args.discovery_port,
},
identity_scheme_registry=default_identity_scheme_registry,
).to_signed_enr(local_private_key.to_bytes())
node_id = minimal_enr.node_id
try:
base_enr = node_db.get_enr(node_id)
except KeyError:
logger.info(
f"No Node for {encode_hex(node_id)} found, creating new one")
return minimal_enr
else:
if any(base_enr[key] != value for key, value in minimal_enr.items()):
logger.debug(f"Updating local ENR")
return UnsignedENR(
sequence_number=base_enr.sequence_number + 1,
kv_pairs=merge(dict(base_enr), dict(minimal_enr)),
identity_scheme_registry=default_identity_scheme_registry,
).to_signed_enr(local_private_key.to_bytes())
else:
return base_enr
async def handle_vote(self, vote: EndpointVote) -> None:
self.logger.debug(
"Received vote for %s from %s",
vote.endpoint,
encode_hex(vote.node_id),
)
current_enr = self.node_db.get_enr(self.local_node_id)
# TODO: majority voting, discard old votes
are_endpoint_keys_present = (IP_V4_ADDRESS_ENR_KEY in current_enr
and UDP_PORT_ENR_KEY in current_enr)
enr_needs_update = not are_endpoint_keys_present or (
vote.endpoint.ip_address != current_enr[IP_V4_ADDRESS_ENR_KEY]
or vote.endpoint.port != current_enr[UDP_PORT_ENR_KEY])
if enr_needs_update:
kv_pairs = merge(
current_enr, {
IP_V4_ADDRESS_ENR_KEY: vote.endpoint.ip_address,
UDP_PORT_ENR_KEY: vote.endpoint.port,
})
new_unsigned_enr = UnsignedENR(
kv_pairs=kv_pairs,
sequence_number=current_enr.sequence_number + 1,
identity_scheme_registry=self.identity_scheme_registry,
)
signed_enr = new_unsigned_enr.to_signed_enr(self.local_private_key)
self.logger.info(
f"Updating local endpoint to %s (new ENR sequence number: %d)",
vote.endpoint,
signed_enr.sequence_number,
)
self.node_db.set_enr(signed_enr)
def test_v4_structure_validation(invalid_kv_pairs, identity_scheme_registry):
with pytest.raises(ValidationError):
UnsignedENR(
sequence_number=0,
kv_pairs=invalid_kv_pairs,
identity_scheme_registry=identity_scheme_registry,
)
class ENRFactory(factory.Factory):
class Meta:
model = ENR
sequence_number = factory.Faker("pyint", min_value=0, max_value=100)
kv_pairs = factory.LazyAttribute(lambda o: merge(
{
b"id":
b"v4",
b"secp256k1":
keys.PrivateKey(o.private_key).public_key.to_compressed_bytes(),
b"ip":
o.address.ip_packed,
b"udp":
o.address.udp_port,
b"tcp":
o.address.tcp_port,
}, o.custom_kv_pairs))
signature = factory.LazyAttribute(lambda o: UnsignedENR(
o.sequence_number,
o.kv_pairs,
).to_signed_enr(o.private_key).signature)
class Params:
private_key = factory.Faker("binary",
length=V4IdentityScheme.private_key_size)
address = factory.SubFactory(AddressFactory)
custom_kv_pairs: Dict[bytes, Any] = {}
def test_signing(mock_identity_scheme, identity_scheme_registry):
unsigned_enr = UnsignedENR(
sequence_number=0,
kv_pairs={b"id": b"mock"},
identity_scheme_registry=identity_scheme_registry)
private_key = b"\x00" * 32
enr = unsigned_enr.to_signed_enr(private_key)
assert enr.signature == mock_identity_scheme.create_enr_signature(
enr, private_key)
def test_inititialization(identity_scheme_registry):
valid_sequence_number = 0
valid_kv_pairs = {b"id": b"mock"}
valid_signature = b"" # signature is not validated during initialization
assert UnsignedENR(
sequence_number=valid_sequence_number,
kv_pairs=valid_kv_pairs,
identity_scheme_registry=identity_scheme_registry,
)
assert ENR(
sequence_number=valid_sequence_number,
kv_pairs=valid_kv_pairs,
signature=valid_signature,
identity_scheme_registry=identity_scheme_registry,
)
with pytest.raises(ValidationError):
UnsignedENR(
sequence_number=valid_sequence_number,
kv_pairs={b"no-id": b""},
identity_scheme_registry=identity_scheme_registry,
)
with pytest.raises(ValidationError):
ENR(
sequence_number=valid_sequence_number,
kv_pairs={b"no-id": b""},
signature=valid_signature,
identity_scheme_registry=identity_scheme_registry,
)
with pytest.raises(ValidationError):
UnsignedENR(
sequence_number=-1,
kv_pairs=valid_kv_pairs,
identity_scheme_registry=identity_scheme_registry,
)
with pytest.raises(ValidationError):
ENR(
sequence_number=-1,
kv_pairs=valid_kv_pairs,
signature=valid_signature,
identity_scheme_registry=identity_scheme_registry,
)
def test_enr_v4_compat_signing():
private_key = PrivateKey(b"\x11" * 32)
unsigned_enr = UnsignedENR(
0, {
b"id": b"v4-compat",
b"secp256k1": private_key.public_key.to_compressed_bytes(),
b"key1": b"value1",
})
with pytest.raises(NotImplementedError):
V4CompatIdentityScheme.create_enr_signature(unsigned_enr, b'')
def test_repr(mock_identity_scheme, identity_scheme_registry):
unsigned_enr = UnsignedENR(0, {b"id": b"mock"}, identity_scheme_registry)
enr = unsigned_enr.to_signed_enr(b"\x00" * 32)
base64_encoded_enr = base64.urlsafe_b64encode(rlp.encode(enr))
represented_enr = repr(enr)
assert represented_enr.startswith("enr:")
assert base64_encoded_enr.rstrip(b"=").decode() == represented_enr[4:]
assert ENR.from_repr(represented_enr, identity_scheme_registry) == enr
def test_enr_node_id():
private_key = PrivateKey(b"\x11" * 32)
unsigned_enr = UnsignedENR(
0, {
b"id": b"v4",
b"secp256k1": private_key.public_key.to_compressed_bytes(),
b"key1": b"value1",
})
enr = unsigned_enr.to_signed_enr(private_key.to_bytes())
node_id = V4IdentityScheme.extract_node_id(enr)
assert node_id == keccak(private_key.public_key.to_bytes())
def test_enr_public_key():
private_key = PrivateKey(b"\x11" * 32)
public_key = private_key.public_key.to_compressed_bytes()
unsigned_enr = UnsignedENR(0, {
b"id": b"v4",
b"secp256k1": public_key,
b"key1": b"value1",
})
enr = unsigned_enr.to_signed_enr(private_key.to_bytes())
assert V4IdentityScheme.extract_public_key(unsigned_enr) == public_key
assert V4IdentityScheme.extract_public_key(enr) == public_key
def test_enr_signing():
private_key = PrivateKey(b"\x11" * 32)
unsigned_enr = UnsignedENR(
0, {
b"id": b"v4",
b"secp256k1": private_key.public_key.to_compressed_bytes(),
b"key1": b"value1",
})
signature = V4IdentityScheme.create_enr_signature(unsigned_enr,
private_key.to_bytes())
message_hash = keccak(unsigned_enr.get_signing_message())
assert private_key.public_key.verify_msg_hash(
message_hash, NonRecoverableSignature(signature))
def test_official_test_vector():
enr = ENR.from_repr(
OFFICIAL_TEST_DATA["repr"]) # use default identity scheme registry
assert enr.sequence_number == OFFICIAL_TEST_DATA["sequence_number"]
assert dict(enr) == OFFICIAL_TEST_DATA["kv_pairs"]
assert enr.public_key == OFFICIAL_TEST_DATA["public_key"]
assert enr.node_id == OFFICIAL_TEST_DATA["node_id"]
assert enr.identity_scheme is OFFICIAL_TEST_DATA["identity_scheme"]
assert repr(enr) == OFFICIAL_TEST_DATA["repr"]
unsigned_enr = UnsignedENR(enr.sequence_number, dict(enr))
reconstructed_enr = unsigned_enr.to_signed_enr(
OFFICIAL_TEST_DATA["private_key"])
assert reconstructed_enr == enr
def test_enr_signature_validation():
private_key = PrivateKey(b"\x11" * 32)
unsigned_enr = UnsignedENR(
0, {
b"id": b"v4",
b"secp256k1": private_key.public_key.to_compressed_bytes(),
b"key1": b"value1",
})
enr = unsigned_enr.to_signed_enr(private_key.to_bytes())
V4IdentityScheme.validate_enr_signature(enr)
forged_enr = ENR(enr.sequence_number, dict(enr), b"\x00" * 64)
with pytest.raises(ValidationError):
V4IdentityScheme.validate_enr_signature(forged_enr)
def test_signature_validation(mock_identity_scheme, identity_scheme_registry):
unsigned_enr = UnsignedENR(0, {b"id": b"mock"}, identity_scheme_registry)
private_key = b"\x00" * 32
enr = unsigned_enr.to_signed_enr(private_key)
enr.validate_signature()
invalid_signature = b"\xff" * 64
invalid_enr = ENR(enr.sequence_number,
dict(enr),
invalid_signature,
identity_scheme_registry=identity_scheme_registry)
with pytest.raises(ValidationError):
invalid_enr.validate_signature()
with pytest.raises(ValidationError):
ENR(
0,
{b"id": b"unknown"},
b"",
identity_scheme_registry=identity_scheme_registry,
)
async def test_request_enr(nursery, manually_driven_discovery_pair):
alice, bob = manually_driven_discovery_pair
# Pretend that bob and alice have already bonded, otherwise bob will ignore alice's ENR
# request.
bob.node_db.set_last_pong_time(alice.this_node.id, int(time.monotonic()))
# Add a copy of Bob's node with a stub ENR to alice's RT as later we're going to check that it
# gets updated with the received ENR.
bobs_node_with_stub_enr = Node.from_pubkey_and_addr(
bob.this_node.pubkey, bob.this_node.address)
alice.node_db.set_last_pong_time(bob.this_node.id, int(time.monotonic()))
alice.node_db.delete_enr(bobs_node_with_stub_enr.id)
alice.node_db.set_enr(bobs_node_with_stub_enr.enr)
assert alice.node_db.get_enr(
bobs_node_with_stub_enr.id).sequence_number == 0
received_enr = None
got_enr = trio.Event()
async def fetch_enr(event):
nonlocal received_enr
received_enr = await alice.request_enr(bobs_node_with_stub_enr)
event.set()
# Start a task in the background that requests an ENR to bob and then waits for it.
nursery.start_soon(fetch_enr, got_enr)
# Bob will now consume one datagram containing the ENR_REQUEST from alice, and as part of that
# will send an ENR_RESPONSE, which will then be consumed by alice, and as part of that it will
# be fed into the request_enr() task we're running the background.
with trio.fail_after(0.5):
await bob.consume_datagram()
await alice.consume_datagram()
with trio.fail_after(1):
await got_enr.wait()
validate_node_enr(bob.this_node, received_enr, sequence_number=1)
assert alice.node_db.get_enr(bob.this_node.id) == received_enr
# Now, if Bob later sends us a new ENR with no endpoint information, we'll evict him from both
# our DB and RT.
sequence_number = bob.this_node.enr.sequence_number + 1
new_unsigned_enr = UnsignedENR(sequence_number,
kv_pairs={
IDENTITY_SCHEME_ENR_KEY:
V4IdentityScheme.id,
V4IdentityScheme.public_key_enr_key:
bob.pubkey.to_compressed_bytes(),
})
bob.this_node = Node(new_unsigned_enr.to_signed_enr(
bob.privkey.to_bytes()))
received_enr = None
got_new_enr = trio.Event()
nursery.start_soon(fetch_enr, got_new_enr)
with trio.fail_after(0.5):
await bob.consume_datagram()
await alice.consume_datagram()
with trio.fail_after(1):
await got_new_enr.wait()
assert Node(received_enr).address is None
assert not alice.routing._contains(bob.this_node.id,
include_replacement_cache=True)
with pytest.raises(KeyError):
alice.node_db.get_enr(bob.this_node.id)
def test_node_id(mock_identity_scheme, identity_scheme_registry):
unsigned_enr = UnsignedENR(0, {b"id": b"mock"}, identity_scheme_registry)
private_key = b"\x00" * 32
enr = unsigned_enr.to_signed_enr(private_key)
assert enr.node_id == private_key
def test_public_key(mock_identity_scheme, identity_scheme_registry):
unsigned_enr = UnsignedENR(0, {b"id": b"mock"}, identity_scheme_registry)
private_key = b"\x00" * 32
enr = unsigned_enr.to_signed_enr(private_key)
assert enr.public_key == mock_identity_scheme.extract_public_key(enr)