def test_routingtable_split_bucket():
table = RoutingTable(NodeFactory())
assert len(table.buckets) == 1
old_bucket = table.buckets[0]
table.split_bucket(0)
assert len(table.buckets) == 2
assert old_bucket not in table.buckets
def _test_find_node_neighbours(use_v5, alice, bob):
# Add some nodes to bob's routing table so that it has something to use when replying to
# alice's find_node.
for _ in range(constants.KADEMLIA_BUCKET_SIZE * 2):
bob.update_routing_table(NodeFactory())
# Connect alice's and bob's transports directly so we don't need to deal with the complexities
# of going over the wire.
link_transports(alice, bob)
# Collect all neighbours packets received by alice in a list for later inspection.
received_neighbours = []
alice.recv_neighbours_v4 = lambda node, payload, hash_: received_neighbours.append((node, payload)) # noqa: E501
# Pretend that bob and alice have already bonded, otherwise bob will ignore alice's find_node.
bob.update_routing_table(alice.this_node)
if use_v5:
alice.send_find_node_v5(bob.this_node, alice.this_node.id)
else:
alice.send_find_node_v4(bob.this_node, alice.this_node.id)
# Bob should have sent two neighbours packets in order to keep the total packet size under the
# 1280 bytes limit.
assert len(received_neighbours) == 2
packet1, packet2 = received_neighbours
neighbours = []
for packet in [packet1, packet2]:
node, payload = packet
assert node == bob.this_node
neighbours.extend(discovery._extract_nodes_from_payload(
node.address, payload[0], bob.logger))
assert len(neighbours) == constants.KADEMLIA_BUCKET_SIZE
async def test_update_routing_table_triggers_bond_if_eviction_candidate(
manually_driven_discovery, monkeypatch):
discovery = manually_driven_discovery
old_node, new_node = NodeFactory.create_batch(2)
bond_called = False
async def bond(node_id):
nonlocal bond_called
bond_called = True
assert node_id == old_node.id
monkeypatch.setattr(discovery, 'bond', bond)
# Pretend our routing table failed to add the new node by returning the least recently seen
# node for an eviction check.
monkeypatch.setattr(discovery.routing, 'update', lambda n: old_node.id)
discovery.update_routing_table(new_node)
assert not discovery.routing._contains(new_node.id,
include_replacement_cache=False)
# The update_routing_table() call above will have scheduled a future call to discovery.bond() so
# we need to wait a bit here to give it a chance to run.
with trio.fail_after(0.5):
while not bond_called:
await trio.sleep(0.001)
async def test_protocol_bootstrap(monkeypatch):
node1, node2 = NodeFactory.create_batch(2)
discovery = MockDiscoveryService([node1, node2])
invalidated_bonds = []
def invalidate_bond(node_id):
invalidated_bonds.append(node_id)
async def bond(node_id):
assert discovery.routing.update(node_id) is None
return True
monkeypatch.setattr(discovery, 'invalidate_bond', invalidate_bond)
# Pretend we bonded successfully with our bootstrap nodes.
monkeypatch.setattr(discovery, 'bond', bond)
await discovery.bootstrap()
assert sorted(invalidated_bonds) == sorted([node.id for node in [node1, node2]])
assert len(discovery.messages) == 2
# We don't care in which order the bootstrap nodes are contacted, nor which node_id was used
# in the find_node request, so we just assert that we sent find_node msgs to both nodes.
assert sorted([(node, cmd) for (node, cmd, _) in discovery.messages]) == sorted([
(node1, 'find_node'),
(node2, 'find_node')])
def test_update_routing_table():
proto = MockDiscoveryProtocol([])
node = NodeFactory()
assert proto.update_routing_table(node) is None
assert node in proto.routing
async def test_wait_pong():
proto = MockDiscoveryProtocol([])
us = proto.this_node
node = NodeFactory()
token = b'token'
# Schedule a call to proto.recv_pong() simulating a pong from the node we expect.
pong_msg_payload = [us.address.to_endpoint(), token, discovery._get_msg_expiration()]
recv_pong_coroutine = asyncio.coroutine(lambda: proto.recv_pong_v4(node, pong_msg_payload, b''))
asyncio.ensure_future(recv_pong_coroutine())
got_pong = await proto.wait_pong_v4(node, token)
assert got_pong
# Ensure wait_pong() cleaned up after itself.
pingid = proto._mkpingid(token, node)
assert pingid not in proto.pong_callbacks
# If the remote node echoed something different than what we expected, wait_pong() would
# timeout.
wrong_token = b"foo"
pong_msg_payload = [us.address.to_endpoint(), wrong_token, discovery._get_msg_expiration()]
recv_pong_coroutine = asyncio.coroutine(lambda: proto.recv_pong_v4(node, pong_msg_payload, b''))
asyncio.ensure_future(recv_pong_coroutine())
got_pong = await proto.wait_pong_v4(node, token)
assert not got_pong
assert pingid not in proto.pong_callbacks
async def test_update_routing_table():
discovery = MockDiscoveryService([])
node = NodeFactory()
discovery.update_routing_table(node)
assert discovery.routing._contains(node.id, include_replacement_cache=False)
async def test_update_routing_table():
discovery = MockDiscoveryService([])
node = NodeFactory()
assert discovery.update_routing_table(node) is None
assert node in discovery.routing
async def test_aurora_walk(network_size, malpn, malpg, mistake_threshold,
test_runs):
""" TODO this is a non-deterministic test, should be changed"""
response_size = constants.KADEMLIA_BUCKET_SIZE
batch = NodeFactory.create_batch(network_size)
pubkey_honesty: Dict[any, Tuple[NodeAPI, bool]] = {}
honest_nodes: Set[NodeAPI] = set()
malicious_nodes: Set[NodeAPI] = set()
for index, node in enumerate(batch):
if index < network_size * malpn:
pubkey_honesty.update({node.pubkey: False})
malicious_nodes.add(node)
else:
pubkey_honesty.update({node.pubkey: True})
honest_nodes.add(node)
proto = MockDiscoveryProtocolAurora(batch, honest_nodes, malicious_nodes,
malpg)
hit_number = 0
miss_number = 0
for _ in range(test_runs):
entry_node = random.choice(tuple(malicious_nodes))
_, result_pubkey, _ = await proto.aurora_walk(entry_node, network_size,
response_size,
mistake_threshold)
if pubkey_honesty[result_pubkey]:
hit_number += 1
else:
miss_number += 1
assert hit_number > miss_number
def test_bucket_ordering():
first = KBucket(0, 50)
second = KBucket(51, 100)
third = NodeFactory()
assert first < second
with pytest.raises(TypeError):
assert first > third
async def test_aurora_pick_existing_candidates():
candidates = NodeFactory.create_batch(4)
node1, node2, *other_nodes = candidates
exclusion_candidates = {node1, node2}
result = aurora_pick(set(candidates), exclusion_candidates)
assert result in candidates
assert result not in exclusion_candidates
def test_kbucket_split():
bucket = KBucket(0, 100)
for i in range(1, bucket.size + 1):
node = NodeFactory()
# Set the IDs of half the nodes below the midpoint, so when we split we should end up with
# two buckets containing k/2 nodes.
if i % 2 == 0:
node._id_int = bucket.midpoint + i
else:
node._id_int = bucket.midpoint - i
bucket.add(node)
assert bucket.is_full
bucket1, bucket2 = bucket.split()
assert bucket1.start == 0
assert bucket1.end == 50
assert bucket2.start == 51
assert bucket2.end == 100
assert len(bucket1) == bucket.size / 2
assert len(bucket2) == bucket.size / 2
def test_kbucket_add():
bucket = KBucket(0, 100)
node = NodeFactory()
assert bucket.add(node) is None
assert bucket.nodes == [node]
node2 = NodeFactory()
assert bucket.add(node2) is None
assert bucket.nodes == [node, node2]
assert bucket.head == node
assert bucket.add(node) is None
assert bucket.nodes == [node2, node]
assert bucket.head == node2
bucket.size = 2
node3 = NodeFactory()
assert bucket.add(node3) == node2
assert bucket.nodes == [node2, node]
assert bucket.head == node2
async def test_bond():
proto = MockDiscoveryProtocol([])
node = NodeFactory()
token = b'token'
# Do not send pings, instead simply return the pingid we'd expect back together with the pong.
proto.send_ping_v4 = lambda remote: token
# Pretend we get a pong from the node we are bonding with.
proto.wait_pong_v4 = asyncio.coroutine(lambda n, t: t == token and n == node)
bonded = await proto.bond(node)
assert bonded
# If we try to bond with any other nodes we'll timeout and bond() will return False.
node2 = NodeFactory()
bonded = await proto.bond(node2)
assert not bonded
def test_multiplexer_pair_factory():
alice_remote, bob_remote = NodeFactory.create_batch(2)
alice_multiplexer, bob_multiplexer = MultiplexerPairFactory(
alice_remote=alice_remote,
bob_remote=bob_remote,
)
assert alice_multiplexer.remote == bob_remote
assert bob_multiplexer.remote == alice_remote
assert alice_multiplexer.get_base_protocol().version == DEVP2P_V5
assert bob_multiplexer.get_base_protocol().version == DEVP2P_V5
async def test_bond(nursery, monkeypatch):
discovery = MockDiscoveryService([])
us = discovery.this_node
node = NodeFactory()
discovery.node_db.set_enr(node.enr)
token = b'token'
async def send_ping(node):
return token
# Do not send pings, instead simply return the pingid we'd expect back together with the pong.
monkeypatch.setattr(discovery, 'send_ping_v4', send_ping)
# Schedule a call to service.recv_pong() simulating a pong from the node we expect.
enr_seq = 1
pong_msg_payload = [
us.address.to_endpoint(), token,
_get_msg_expiration(),
int_to_big_endian(enr_seq)
]
nursery.start_soon(discovery.recv_pong_v4, node, pong_msg_payload, b'')
bonded = await discovery.bond(node.id)
assert bonded
assert discovery.is_bond_valid_with(node.id)
# Upon successfully bonding, retrieval of the remote's ENR will be scheduled.
with trio.fail_after(1):
scheduled_enr_node_id, scheduled_enr_seq = await discovery.pending_enrs_consumer.receive(
)
assert scheduled_enr_node_id == node.id
assert scheduled_enr_seq == enr_seq
# If we try to bond with any other nodes we'll timeout and bond() will return False.
node2 = NodeFactory()
discovery.node_db.set_enr(node2.enr)
bonded = await discovery.bond(node2.id)
assert not bonded
async def test_unsolicited_neighbours(manually_driven_discovery_pair):
# Ensure our routing table cannot be poisoned by malicious nodes sending us unsolicited
# neighbours packages.
alice, bob = manually_driven_discovery_pair
node = NodeFactory()
await alice.send_neighbours_v4(bob.this_node, [node])
with trio.fail_after(1):
await bob.consume_datagram()
assert not bob.routing._contains(node.id, include_replacement_cache=True)
async def test_records_failures():
connection_tracker = MemoryConnectionTracker()
node = NodeFactory()
blacklisted_ids = await connection_tracker.get_blacklisted()
assert node.id not in blacklisted_ids
connection_tracker.record_failure(node, HandshakeFailure())
blacklisted_ids = await connection_tracker.get_blacklisted()
assert node.id in blacklisted_ids
assert connection_tracker._record_exists(node.id)
async def test_topic_query(event_loop):
bob = await get_listening_discovery_protocol(event_loop)
les_nodes = NodeFactory.create_batch(10)
topic = b'les'
for n in les_nodes:
bob.topic_table.add_node(n, topic)
alice = await get_listening_discovery_protocol(event_loop)
echo = alice.send_topic_query(bob.this_node, topic)
received_nodes = await alice.wait_topic_nodes(bob.this_node, echo)
assert len(received_nodes) == 10
assert sorted(received_nodes) == sorted(les_nodes)
async def get_directly_linked_peers_without_handshake(
alice_factory: BasePeerFactory = None,
bob_factory: BasePeerFactory = None) -> Tuple[BasePeer, BasePeer]:
"""
See get_directly_linked_peers().
Neither the P2P handshake nor the sub-protocol handshake will be performed here.
"""
cancel_token = CancelToken("get_directly_linked_peers_without_handshake")
if alice_factory is None:
alice_factory = ParagonPeerFactory(
privkey=ecies.generate_privkey(),
context=ParagonContext(),
token=cancel_token,
)
if bob_factory is None:
bob_factory = ParagonPeerFactory(
privkey=ecies.generate_privkey(),
context=ParagonContext(),
token=cancel_token,
)
alice_remote = NodeFactory(pubkey=alice_factory.privkey.public_key)
bob_remote = NodeFactory(pubkey=bob_factory.privkey.public_key)
alice_transport, bob_transport = MemoryTransportPairFactory(
alice_remote=alice_remote,
alice_private_key=alice_factory.privkey,
bob_remote=bob_remote,
bob_private_key=bob_factory.privkey,
)
alice = alice_factory.create_peer(alice_transport)
bob = bob_factory.create_peer(bob_transport)
return alice, bob
async def test_sql_does_persist(tmpdir):
db_path = Path(tmpdir.join("nodedb"))
node = NodeFactory()
connection_tracker_a = SQLiteConnectionTracker(get_tracking_database(db_path))
assert await connection_tracker_a.should_connect_to(node) is True
connection_tracker_a.record_failure(node, HandshakeFailure())
assert await connection_tracker_a.should_connect_to(node) is False
del connection_tracker_a
# open a second instance
connection_tracker_b = SQLiteConnectionTracker(get_tracking_database(db_path))
# the second instance remembers the failure
assert await connection_tracker_b.should_connect_to(node) is False
async def test_memory_does_not_persist():
node = NodeFactory()
connection_tracker_a = MemoryConnectionTracker()
assert await connection_tracker_a.should_connect_to(node) is True
connection_tracker_a.record_failure(node, HandshakeFailure())
assert await connection_tracker_a.should_connect_to(node) is False
# open a second instance
connection_tracker_b = MemoryConnectionTracker()
# the second instance has no memory of the failure
assert await connection_tracker_b.should_connect_to(node) is True
assert await connection_tracker_a.should_connect_to(node) is False
async def test_bond(nursery, monkeypatch):
discovery = MockDiscoveryService([])
node = NodeFactory()
token = b'token'
# Do not send pings, instead simply return the pingid we'd expect back together with the pong.
monkeypatch.setattr(discovery, 'send_ping_v4', lambda remote: token)
# Pretend we get a pong from the node we are bonding with.
async def wait_pong_v4(remote, t) -> bool:
return t == token and remote == node
monkeypatch.setattr(discovery, 'wait_pong_v4', wait_pong_v4)
bonded = await discovery.bond(node)
assert bonded
# If we try to bond with any other nodes we'll timeout and bond() will return False.
node2 = NodeFactory()
bonded = await discovery.bond(node2)
assert not bonded
async def test_wait_ping(echo):
proto = MockDiscoveryProtocol([])
node = NodeFactory()
# Schedule a call to proto.recv_ping() simulating a ping from the node we expect.
recv_ping_coroutine = asyncio.coroutine(lambda: proto.recv_ping_v4(node, echo, b''))
asyncio.ensure_future(recv_ping_coroutine())
got_ping = await proto.wait_ping(node)
assert got_ping
# Ensure wait_ping() cleaned up after itself.
assert node not in proto.ping_callbacks
# If we waited for a ping from a different node, wait_ping() would timeout and thus return
# false.
recv_ping_coroutine = asyncio.coroutine(lambda: proto.recv_ping_v4(node, echo, b''))
asyncio.ensure_future(recv_ping_coroutine())
node2 = NodeFactory()
got_ping = await proto.wait_ping(node2)
assert not got_ping
assert node2 not in proto.ping_callbacks
async def test_wait_ping(nursery, echo):
discovery = MockDiscoveryService([])
node = NodeFactory()
# Schedule a call to discovery.recv_ping() simulating a ping from the node we expect.
async def recv_ping() -> None:
discovery.recv_ping_v4(node, echo, b'')
nursery.start_soon(recv_ping)
got_ping = await discovery.wait_ping(node)
assert got_ping
# Ensure wait_ping() cleaned up after itself.
assert node not in discovery.ping_callbacks
# If we waited for a ping from a different node, wait_ping() would timeout and thus return
# false.
nursery.start_soon(recv_ping)
node2 = NodeFactory()
got_ping = await discovery.wait_ping(node2)
assert not got_ping
assert node2 not in discovery.ping_callbacks
async def test_wait_neighbours(nursery):
service = MockDiscoveryService([])
addresses = [
AddressFactory(ip='10.0.0.1'),
AddressFactory(ip='10.0.0.2'),
AddressFactory(ip='10.0.0.3'),
AddressFactory(ip='10.0.0.4')
]
sender = NodeFactory(address=addresses[0])
# All our nodes are on the same network as the sender to ensure they pass the
# check_relayed_addr() check, otherwise in some rare cases we may get a random IP address
# that doesn't and it will be ignored by wait_neighbours, causing the test to fail.
neighbours = tuple(
NodeFactory(address=address) for address in addresses[1:])
# Schedule a call to service.recv_neighbours_v4() simulating a neighbours response from the
# node we expect.
expiration = _get_msg_expiration()
neighbours_msg_payload = [[
n.address.to_endpoint() + [n.pubkey.to_bytes()] for n in neighbours
], expiration]
nursery.start_soon(service.recv_neighbours_v4, sender,
neighbours_msg_payload, b'')
received_neighbours = await service.wait_neighbours(sender)
assert neighbours == received_neighbours
# Ensure wait_neighbours() cleaned up after itself.
assert not service.neighbours_channels.already_waiting_for(sender)
# If wait_neighbours() times out, we get an empty list of neighbours.
received_neighbours = await service.wait_neighbours(sender)
assert received_neighbours == tuple()
assert not service.neighbours_channels.already_waiting_for(sender)
async def test_find_node_neighbours(manually_driven_discovery_pair, monkeypatch):
alice, bob = manually_driven_discovery_pair
nodes_in_rt = 0
# Ensure we have plenty of nodes in our RT's buckets so that the NEIGHBOURS response sent by
# bob is split into multiple messages.
while nodes_in_rt < (constants.KADEMLIA_BUCKET_SIZE * 2):
node = NodeFactory()
eviction_candidate = bob.routing.update(node.id)
if eviction_candidate is not None:
continue
nodes_in_rt += 1
bob.enr_db.set_enr(node.enr)
# Collect all neighbours packets received by alice in a list for later inspection.
received_neighbours = []
async def recv_neighbours(node, payload, hash_):
received_neighbours.append((node, payload))
monkeypatch.setattr(alice, 'recv_neighbours_v4', recv_neighbours)
# Pretend that bob and alice have already bonded, otherwise bob will ignore alice's find_node.
bob._last_pong_at[alice.this_node.id] = int(time.monotonic())
await alice.send_find_node_v4(bob.this_node, alice.pubkey.to_bytes())
with trio.fail_after(1):
await bob.consume_datagram()
# Alice needs to consume two datagrams here because we expect bob's response to be split
# across two packets since a single one would be bigger than protocol's byte limit.
await alice.consume_datagram()
await alice.consume_datagram()
# Bob should have sent two neighbours packets in order to keep the total packet size
# under the 1280 bytes limit. However, the two consume_datagram() calls above will have
# spawned background tasks so we take a few short naps here to wait for them to complete.
while len(received_neighbours) != 2:
await trio.sleep(0.01)
packet1, packet2 = received_neighbours
neighbours = []
for packet in [packet1, packet2]:
node, payload = packet
assert node == bob.this_node
neighbours.extend(_extract_nodes_from_payload(
node.address, payload[0], bob.logger))
assert len(neighbours) == constants.KADEMLIA_BUCKET_SIZE
async def test_bond_short_circuits(monkeypatch):
discovery = MockDiscoveryService([])
bob = NodeFactory()
discovery.enr_db.set_enr(bob.enr)
# Pretend we have a valid bond with bob.
discovery._last_pong_at[bob.id] = int(time.monotonic())
class AttemptedNewBond(Exception):
pass
async def send_ping(node):
raise AttemptedNewBond()
monkeypatch.setattr(discovery, 'send_ping_v4', send_ping)
# When we have a valid bond, we won't attempt a new one.
assert discovery.is_bond_valid_with(bob.id)
assert await discovery.bond(bob.id)
async def test_protocol_bootstrap():
node1, node2 = NodeFactory.create_batch(2)
discovery = MockDiscoveryService([node1, node2])
async def bond(node):
assert discovery.routing.add_node(node) is None
return True
# Pretend we bonded successfully with our bootstrap nodes.
discovery.bond = bond
await discovery.bootstrap()
assert len(discovery.messages) == 2
# We don't care in which order the bootstrap nodes are contacted, nor which node_id was used
# in the find_node request, so we just assert that we sent find_node msgs to both nodes.
assert sorted([(node, cmd) for (node, cmd, _) in discovery.messages
]) == sorted([(node1, 'find_node'), (node2, 'find_node')])
async def test_sql_does_persist(tmpdir):
db_path = Path(tmpdir.join("nodedb"))
node = NodeFactory()
connection_tracker_a = SQLiteConnectionTracker(
get_tracking_database(db_path))
blacklisted_ids = await connection_tracker_a.get_blacklisted()
assert node.id not in blacklisted_ids
connection_tracker_a.record_failure(node, HandshakeFailure())
blacklisted_ids = await connection_tracker_a.get_blacklisted()
assert node.id in blacklisted_ids
del connection_tracker_a
# open a second instance
connection_tracker_b = SQLiteConnectionTracker(
get_tracking_database(db_path))
blacklisted_ids = await connection_tracker_b.get_blacklisted()
# the second instance remembers the failure
assert node.id in blacklisted_ids