def test_topic_table():
table = discovery.TopicTable(logging.getLogger("test"))
topic = b'topic'
node = random_node()
table.add_node(node, topic)
assert len(table.get_nodes(topic)) == 1
assert table.get_nodes(topic)[0] == node
node2 = random_node()
table.add_node(node2, topic)
assert len(table.get_nodes(topic)) == 2
assert table.get_nodes(topic)[1] == node2
# Adding the same node again won't cause a duplicate entry to be added.
table.add_node(node, topic)
assert len(table.get_nodes(topic)) == 2
# When we reach the max number of entries for a given topic, the first added items are evicted
# to make room for the new ones.
for _ in range(discovery.MAX_ENTRIES_PER_TOPIC + 2):
table.add_node(random_node(), topic)
assert node not in table.get_nodes(topic)
assert node2 not in table.get_nodes(topic)
async def test_wait_neighbours():
proto = MockDiscoveryProtocol([])
node = random_node()
# Schedule a call to proto.recv_neighbours_v4() simulating a neighbours response from the node
# we expect.
neighbours = (random_node(), random_node(), random_node())
neighbours_msg_payload = [
[n.address.to_endpoint() + [n.pubkey.to_bytes()] for n in neighbours],
discovery._get_msg_expiration()]
recv_neighbours_coroutine = asyncio.coroutine(
lambda: proto.recv_neighbours_v4(node, neighbours_msg_payload, b''))
asyncio.ensure_future(recv_neighbours_coroutine())
received_neighbours = await proto.wait_neighbours(node)
assert neighbours == received_neighbours
# Ensure wait_neighbours() cleaned up after itself.
assert node not in proto.neighbours_callbacks
# If wait_neighbours() times out, we get an empty list of neighbours.
received_neighbours = await proto.wait_neighbours(node)
assert received_neighbours == tuple()
assert node not in proto.neighbours_callbacks
async def test_wait_ping(echo):
proto = MockDiscoveryProtocol([])
node = random_node()
# 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 = random_node()
got_ping = await proto.wait_ping(node2)
assert not got_ping
assert node2 not in proto.ping_callbacks
def _test_find_node_neighbours(use_v5):
alice = get_discovery_protocol(b"alice")
bob = get_discovery_protocol(b"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(kademlia.k_bucket_size * 2):
bob.update_routing_table(random_node())
# 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) == kademlia.k_bucket_size
def test_update_routing_table():
proto = MockDiscoveryProtocol([])
node = random_node()
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 = random_node()
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_protocol_bootstrap():
node1, node2 = [random_node(), random_node()]
proto = MockDiscoveryProtocol([node1, node2])
async def bond(node):
assert proto.routing.add_node(node) is None
return True
# Pretend we bonded successfully with our bootstrap nodes.
proto.bond = bond
await proto.bootstrap()
assert len(proto.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 proto.messages
]) == sorted([(node1, 'find_node'), (node2, 'find_node')])
async def test_bond():
proto = MockDiscoveryProtocol([])
node = random_node()
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 = random_node()
bonded = await proto.bond(node2)
assert not bonded
async def test_topic_query(event_loop):
bob = await get_listening_discovery_protocol(event_loop)
les_nodes = [random_node() for _ in range(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 test_update_routing_table_triggers_bond_if_eviction_candidate():
proto = MockDiscoveryProtocol([])
old_node, new_node = random_node(), random_node()
bond_called = False
def bond(node):
nonlocal bond_called
bond_called = True
assert node == old_node
proto.bond = asyncio.coroutine(bond)
# Pretend our routing table failed to add the new node by returning the least recently seen
# node for an eviction check.
proto.routing.add_node = lambda n: old_node
proto.update_routing_table(new_node)
assert new_node not in proto.routing
# The update_routing_table() call above will have scheduled a future call to proto.bond() so
# we need to yield here to give it a chance to run.
await asyncio.sleep(0.001)
assert bond_called
