async def test_update_token(self):
loop = asyncio.get_event_loop()
with dht_mocks.mock_network_loop(loop):
node_id1 = constants.generate_id()
peer1 = KademliaProtocol(loop, PeerManager(loop), node_id1,
'1.2.3.4', 4444, 3333)
peer2 = KademliaProtocol(loop, PeerManager(loop),
constants.generate_id(), '1.2.3.5', 4444,
3333)
await loop.create_datagram_endpoint(lambda: peer1,
('1.2.3.4', 4444))
await loop.create_datagram_endpoint(lambda: peer2,
('1.2.3.5', 4444))
peer = peer2.peer_manager.get_kademlia_peer(node_id1,
'1.2.3.4',
udp_port=4444)
self.assertEqual(None,
peer2.peer_manager.get_node_token(peer.node_id))
await peer2.get_rpc_peer(peer).find_value(b'1' * 48)
self.assertNotEqual(
None, peer2.peer_manager.get_node_token(peer.node_id))
peer1.stop()
peer2.stop()
peer1.disconnect()
peer2.disconnect()
def __init__(self,
loop: asyncio.BaseEventLoop,
peer_manager: 'PeerManager',
node_id: bytes,
external_ip: str,
udp_port: int,
peer_port: int,
rpc_timeout: float = 5.0):
self.peer_manager = peer_manager
self.loop = loop
self.node_id = node_id
self.external_ip = external_ip
self.udp_port = udp_port
self.peer_port = peer_port
self.is_seed_node = False
self.partial_messages: typing.Dict[bytes, typing.Dict[bytes,
bytes]] = {}
self.sent_messages: typing.Dict[bytes,
typing.Tuple['KademliaPeer',
asyncio.Future,
RequestDatagram]] = {}
self.protocol_version = constants.protocol_version
self.started_listening_time = 0
self.transport: DatagramTransport = None
self.old_token_secret = constants.generate_id()
self.token_secret = constants.generate_id()
self.routing_table = TreeRoutingTable(self.loop, self.peer_manager,
self.node_id)
self.data_store = DictDataStore(self.loop, self.peer_manager)
self.ping_queue = PingQueue(self.loop, self)
self.node_rpc = KademliaRPC(self, self.loop, self.peer_port)
self.lock = asyncio.Lock(loop=self.loop)
self.rpc_timeout = rpc_timeout
self._split_lock = asyncio.Lock(loop=self.loop)
async def test_split_buckets(self):
loop = asyncio.get_event_loop()
peer_addresses = [
(constants.generate_id(1), '1.2.3.1'),
]
for i in range(2, 200):
peer_addresses.append((constants.generate_id(i), f'1.2.3.{i}'))
with dht_mocks.mock_network_loop(loop):
nodes = {
i: Node(loop, PeerManager(loop), node_id, 4444, 4444, 3333, address)
for i, (node_id, address) in enumerate(peer_addresses)
}
node_1 = nodes[0]
for i in range(1, len(peer_addresses)):
node = nodes[i]
peer = node_1.protocol.peer_manager.get_kademlia_peer(
node.protocol.node_id, node.protocol.external_ip,
udp_port=node.protocol.udp_port
)
# set all of the peers to good (as to not attempt pinging stale ones during split)
node_1.protocol.peer_manager.report_last_replied(peer.address, peer.udp_port)
node_1.protocol.peer_manager.report_last_replied(peer.address, peer.udp_port)
await node_1.protocol._add_peer(peer)
# check that bucket 0 is always the one covering the local node id
self.assertEqual(True, node_1.protocol.routing_table.buckets[0].key_in_range(node_1.protocol.node_id))
self.assertEqual(40, len(node_1.protocol.routing_table.get_peers()))
self.assertEqual(len(expected_ranges), len(node_1.protocol.routing_table.buckets))
covered = 0
for (expected_min, expected_max), bucket in zip(expected_ranges, node_1.protocol.routing_table.buckets):
self.assertEqual(expected_min, bucket.range_min)
self.assertEqual(expected_max, bucket.range_max)
covered += bucket.range_max - bucket.range_min
self.assertEqual(2**384, covered)
for node in nodes.values():
node.stop()
async def test_add_peer_after_handle_request(self):
with dht_mocks.mock_network_loop(self.loop):
node_id1 = constants.generate_id()
node_id2 = constants.generate_id()
node_id3 = constants.generate_id()
node_id4 = constants.generate_id()
peer1 = KademliaProtocol(
self.loop, PeerManager(self.loop), node_id1, '1.2.3.4', 4444, 3333
)
await self.loop.create_datagram_endpoint(lambda: peer1, ('1.2.3.4', 4444))
peer1.start()
peer2, peer_2_from_peer_1 = await self._make_protocol(peer1, node_id2, '1.2.3.5', 4444, 3333)
peer3, peer_3_from_peer_1 = await self._make_protocol(peer1, node_id3, '1.2.3.6', 4444, 3333)
peer4, peer_4_from_peer_1 = await self._make_protocol(peer1, node_id4, '1.2.3.7', 4444, 3333)
# peers who reply should be added
await peer1.get_rpc_peer(peer_2_from_peer_1).ping()
await asyncio.sleep(0.5)
self.assertListEqual([peer_2_from_peer_1], peer1.routing_table.get_peers())
peer1.routing_table.remove_peer(peer_2_from_peer_1)
# peers not known by be good/bad should be enqueued to maybe-ping
peer1_from_peer3 = peer3.get_rpc_peer(peer3.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444))
self.assertEqual(0, len(peer1.ping_queue._pending_contacts))
pong = await peer1_from_peer3.ping()
self.assertEqual(b'pong', pong)
self.assertEqual(1, len(peer1.ping_queue._pending_contacts))
peer1.ping_queue._pending_contacts.clear()
# peers who are already good should be added
peer1_from_peer4 = peer4.get_rpc_peer(peer4.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444))
peer1.peer_manager.update_contact_triple(node_id4,'1.2.3.7', 4444)
peer1.peer_manager.report_last_replied('1.2.3.7', 4444)
self.assertEqual(0, len(peer1.ping_queue._pending_contacts))
pong = await peer1_from_peer4.ping()
self.assertEqual(b'pong', pong)
await asyncio.sleep(0.5)
self.assertEqual(1, len(peer1.routing_table.get_peers()))
self.assertEqual(0, len(peer1.ping_queue._pending_contacts))
peer1.routing_table.buckets[0].peers.clear()
# peers who are known to be bad recently should not be added or maybe-pinged
peer1_from_peer4 = peer4.get_rpc_peer(peer4.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', 4444))
peer1.peer_manager.update_contact_triple(node_id4,'1.2.3.7', 4444)
peer1.peer_manager.report_failure('1.2.3.7', 4444)
peer1.peer_manager.report_failure('1.2.3.7', 4444)
self.assertEqual(0, len(peer1.ping_queue._pending_contacts))
pong = await peer1_from_peer4.ping()
self.assertEqual(b'pong', pong)
self.assertEqual(0, len(peer1.routing_table.get_peers()))
self.assertEqual(0, len(peer1.ping_queue._pending_contacts))
for p in [peer1, peer2, peer3, peer4]:
p.stop()
p.disconnect()
async def test_fill_one_bucket(self):
loop = asyncio.get_event_loop()
peer_addresses = [
(constants.generate_id(1), '1.2.3.1'),
(constants.generate_id(2), '1.2.3.2'),
(constants.generate_id(3), '1.2.3.3'),
(constants.generate_id(4), '1.2.3.4'),
(constants.generate_id(5), '1.2.3.5'),
(constants.generate_id(6), '1.2.3.6'),
(constants.generate_id(7), '1.2.3.7'),
(constants.generate_id(8), '1.2.3.8'),
(constants.generate_id(9), '1.2.3.9'),
]
with dht_mocks.mock_network_loop(loop):
nodes = {
i: Node(loop, PeerManager(loop), node_id, 4444, 4444, 3333, address)
for i, (node_id, address) in enumerate(peer_addresses)
}
node_1 = nodes[0]
contact_cnt = 0
for i in range(1, len(peer_addresses)):
self.assertEqual(len(node_1.protocol.routing_table.get_peers()), contact_cnt)
node = nodes[i]
peer = node_1.protocol.peer_manager.get_kademlia_peer(
node.protocol.node_id, node.protocol.external_ip,
udp_port=node.protocol.udp_port
)
added = await node_1.protocol._add_peer(peer)
self.assertEqual(True, added)
contact_cnt += 1
self.assertEqual(len(node_1.protocol.routing_table.get_peers()), 8)
self.assertEqual(node_1.protocol.routing_table.buckets_with_contacts(), 1)
for node in nodes.values():
node.protocol.stop()
def test_add_peer(self):
peer = KademliaPeer(None,
'1.2.3.4',
constants.generate_id(2),
udp_port=4444)
peer_update2 = KademliaPeer(None,
'1.2.3.4',
constants.generate_id(2),
udp_port=4445)
self.assertListEqual([], self.kbucket.peers)
# add the peer
self.kbucket.add_peer(peer)
self.assertListEqual([peer], self.kbucket.peers)
# re-add it
self.kbucket.add_peer(peer)
self.assertListEqual([peer], self.kbucket.peers)
self.assertEqual(self.kbucket.peers[0].udp_port, 4444)
# add a new peer object with the same id and address but a different port
self.kbucket.add_peer(peer_update2)
self.assertListEqual([peer_update2], self.kbucket.peers)
self.assertEqual(self.kbucket.peers[0].udp_port, 4445)
# modify the peer object to have a different port
peer_update2.udp_port = 4444
self.kbucket.add_peer(peer_update2)
self.assertListEqual([peer_update2], self.kbucket.peers)
self.assertEqual(self.kbucket.peers[0].udp_port, 4444)
self.kbucket.peers.clear()
# Test if contacts can be added to empty list
# Add k contacts to bucket
for i in range(constants.k):
peer = self.peer_manager.get_kademlia_peer(
generate_id(), next(self.address_generator), 4444)
self.assertTrue(self.kbucket.add_peer(peer))
self.assertEqual(peer, self.kbucket.peers[i])
# Test if contact is not added to full list
peer = self.peer_manager.get_kademlia_peer(
generate_id(), next(self.address_generator), 4444)
self.assertFalse(self.kbucket.add_peer(peer))
# Test if an existing contact is updated correctly if added again
existing_peer = self.kbucket.peers[0]
self.assertTrue(self.kbucket.add_peer(existing_peer))
self.assertEqual(existing_peer, self.kbucket.peers[-1])
async def test_store_to_peer(self):
loop = asyncio.get_event_loop()
with dht_mocks.mock_network_loop(loop):
node_id1 = constants.generate_id()
peer1 = KademliaProtocol(loop, PeerManager(loop), node_id1,
'1.2.3.4', 4444, 3333)
peer2 = KademliaProtocol(loop, PeerManager(loop),
constants.generate_id(), '1.2.3.5', 4444,
3333)
await loop.create_datagram_endpoint(lambda: peer1,
('1.2.3.4', 4444))
await loop.create_datagram_endpoint(lambda: peer2,
('1.2.3.5', 4444))
peer = peer2.peer_manager.get_kademlia_peer(node_id1,
'1.2.3.4',
udp_port=4444)
peer2_from_peer1 = peer1.peer_manager.get_kademlia_peer(
peer2.node_id, peer2.external_ip, udp_port=peer2.udp_port)
peer2_from_peer1.update_tcp_port(3333)
peer3 = peer1.peer_manager.get_kademlia_peer(
constants.generate_id(), '1.2.3.6', udp_port=4444)
store_result = await peer2.store_to_peer(b'2' * 48, peer)
self.assertEqual(store_result[0], peer.node_id)
self.assertEqual(True, store_result[1])
self.assertEqual(True,
peer1.data_store.has_peers_for_blob(b'2' * 48))
self.assertEqual(False,
peer1.data_store.has_peers_for_blob(b'3' * 48))
self.assertListEqual([peer2_from_peer1],
peer1.data_store.get_storing_contacts())
peer1.data_store.completed_blobs.add(
binascii.hexlify(b'2' * 48).decode())
find_value_response = peer1.node_rpc.find_value(peer3, b'2' * 48)
self.assertEqual(len(find_value_response[b'contacts']), 0)
self.assertSetEqual(
{b'2' * 48, b'token', b'protocolVersion', b'contacts'},
set(find_value_response.keys()))
self.assertEqual(2, len(find_value_response[b'2' * 48]))
self.assertEqual(find_value_response[b'2' * 48][0],
peer2_from_peer1.compact_address_tcp())
self.assertDictEqual(bdecode(bencode(find_value_response)),
find_value_response)
peer1.stop()
peer2.stop()
peer1.disconnect()
peer2.disconnect()
def make_ping(cls, from_node_id: bytes, rpc_id: typing.Optional[bytes] = None) -> 'RequestDatagram':
if rpc_id and len(rpc_id) != constants.rpc_id_length:
raise ValueError("invalid rpc id length")
elif not rpc_id:
rpc_id = constants.generate_id()[:constants.rpc_id_length]
if len(from_node_id) != constants.hash_bits // 8:
raise ValueError("invalid node id")
return cls(REQUEST_TYPE, rpc_id, from_node_id, b'ping')
async def test_announce_blobs(self):
blob1 = binascii.hexlify(b'1' * 48).decode()
blob2 = binascii.hexlify(b'2' * 48).decode()
async with self._test_network_context():
await self.storage.add_blobs((blob1, 1024), (blob2, 1024),
finished=True)
await self.storage.db.execute(
"update blob set next_announce_time=0, should_announce=1 where blob_hash in (?, ?)",
(blob1, blob2))
to_announce = await self.storage.get_blobs_to_announce()
self.assertEqual(2, len(to_announce))
self.blob_announcer.start(
batch_size=1) # so it covers batching logic
# takes 60 seconds to start, but we advance 120 to ensure it processed all batches
await self.advance(60.0 * 2)
to_announce = await self.storage.get_blobs_to_announce()
self.assertEqual(0, len(to_announce))
self.blob_announcer.stop()
# test that we can route from a poorly connected peer all the way to the announced blob
await self.chain_peer(constants.generate_id(10), '1.2.3.10')
await self.chain_peer(constants.generate_id(11), '1.2.3.11')
await self.chain_peer(constants.generate_id(12), '1.2.3.12')
await self.chain_peer(constants.generate_id(13), '1.2.3.13')
await self.chain_peer(constants.generate_id(14), '1.2.3.14')
await self.advance(61.0)
last = self.nodes[len(self.nodes) - 1]
search_q, peer_q = asyncio.Queue(loop=self.loop), asyncio.Queue(
loop=self.loop)
search_q.put_nowait(blob1)
_, task = last.accumulate_peers(search_q, peer_q)
found_peers = await peer_q.get()
task.cancel()
self.assertEqual(1, len(found_peers))
self.assertEqual(self.node.protocol.node_id,
found_peers[0].node_id)
self.assertEqual(self.node.protocol.external_ip,
found_peers[0].address)
self.assertEqual(self.node.protocol.peer_port,
found_peers[0].tcp_port)
def __init__(self,
protocol: 'KademliaProtocol',
loop: asyncio.BaseEventLoop,
peer_port: int = 3333):
self.protocol = protocol
self.loop = loop
self.peer_port = peer_port
self.old_token_secret: bytes = None
self.token_secret = constants.generate_id()
def make_store(cls, from_node_id: bytes, blob_hash: bytes, token: bytes, port: int,
rpc_id: typing.Optional[bytes] = None) -> 'RequestDatagram':
if rpc_id and len(rpc_id) != constants.rpc_id_length:
raise ValueError("invalid rpc id length")
if not rpc_id:
rpc_id = constants.generate_id()[:constants.rpc_id_length]
if len(from_node_id) != constants.hash_bits // 8:
raise ValueError("invalid node id")
store_args = [blob_hash, token, port, from_node_id, 0]
return cls(REQUEST_TYPE, rpc_id, from_node_id, b'store', store_args)
async def test_ping(self):
loop = asyncio.get_event_loop()
with dht_mocks.mock_network_loop(loop):
node_id1 = constants.generate_id()
peer1 = KademliaProtocol(
loop, PeerManager(loop), node_id1, '1.2.3.4', 4444, 3333
)
peer2 = KademliaProtocol(
loop, PeerManager(loop), constants.generate_id(), '1.2.3.5', 4444, 3333
)
await loop.create_datagram_endpoint(lambda: peer1, ('1.2.3.4', 4444))
await loop.create_datagram_endpoint(lambda: peer2, ('1.2.3.5', 4444))
peer = peer2.peer_manager.get_kademlia_peer(node_id1, '1.2.3.4', udp_port=4444)
result = await peer2.get_rpc_peer(peer).ping()
self.assertEqual(result, b'pong')
peer1.stop()
peer2.stop()
peer1.disconnect()
peer2.disconnect()
async def test_get_token_on_announce(self):
await self.setup_network(2, seed_nodes=2)
node1, node2 = self.nodes
node1.protocol.peer_manager.clear_token(node2.protocol.node_id)
blob_hash = hexlify(constants.generate_id(1337)).decode()
node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids)
node2.protocol.node_rpc.refresh_token()
node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids)
node2.protocol.node_rpc.refresh_token()
node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids)
async def _test_network_context(self, peer_addresses=None):
self.peer_addresses = peer_addresses or [
(constants.generate_id(2), '1.2.3.2'),
(constants.generate_id(3), '1.2.3.3'),
(constants.generate_id(4), '1.2.3.4'),
(constants.generate_id(5), '1.2.3.5'),
(constants.generate_id(6), '1.2.3.6'),
(constants.generate_id(7), '1.2.3.7'),
(constants.generate_id(8), '1.2.3.8'),
(constants.generate_id(9), '1.2.3.9'),
]
try:
with dht_mocks.mock_network_loop(self.loop):
await self.setup_node(self.peer_addresses, '1.2.3.1', constants.generate_id(1))
yield
finally:
self.blob_announcer.stop()
self.node.stop()
for n in self.nodes.values():
n.stop()
async def setup_network(self, size: int, start_port=40000, seed_nodes=1):
for i in range(size):
node_port = start_port + i
node = Node(self.loop, PeerManager(self.loop), node_id=constants.generate_id(i),
udp_port=node_port, internal_udp_port=node_port,
peer_port=3333, external_ip='127.0.0.1')
self.nodes.append(node)
self.known_node_addresses.append(('127.0.0.1', node_port))
await node.start_listening('127.0.0.1')
self.addCleanup(node.stop)
for node in self.nodes:
node.protocol.rpc_timeout = .5
node.protocol.ping_queue._default_delay = .5
node.start('127.0.0.1', self.known_node_addresses[:seed_nodes])
await asyncio.gather(*[node.joined.wait() for node in self.nodes])
def change_token(self):
self.old_token_secret = self.token_secret
self.token_secret = constants.generate_id()
async def test_ping_queue_discover(self):
loop = asyncio.get_event_loop()
loop.set_debug(False)
peer_addresses = [
(constants.generate_id(1), '1.2.3.1'),
(constants.generate_id(2), '1.2.3.2'),
(constants.generate_id(3), '1.2.3.3'),
(constants.generate_id(4), '1.2.3.4'),
(constants.generate_id(5), '1.2.3.5'),
(constants.generate_id(6), '1.2.3.6'),
(constants.generate_id(7), '1.2.3.7'),
(constants.generate_id(8), '1.2.3.8'),
(constants.generate_id(9), '1.2.3.9'),
]
with dht_mocks.mock_network_loop(loop):
advance = dht_mocks.get_time_accelerator(loop, loop.time())
# start the nodes
nodes: typing.Dict[int, Node] = {
i: Node(loop, PeerManager(loop), node_id, 4444, 4444, 3333,
address)
for i, (node_id, address) in enumerate(peer_addresses)
}
for i, n in nodes.items():
n.start(peer_addresses[i][1], [])
await advance(1)
node_1 = nodes[0]
# ping 8 nodes from node_1, this will result in a delayed return ping
futs = []
for i in range(1, len(peer_addresses)):
node = nodes[i]
assert node.protocol.node_id != node_1.protocol.node_id
peer = node_1.protocol.peer_manager.get_kademlia_peer(
node.protocol.node_id,
node.protocol.external_ip,
udp_port=node.protocol.udp_port)
futs.append(node_1.protocol.get_rpc_peer(peer).ping())
await advance(3)
replies = await asyncio.gather(*tuple(futs))
self.assertTrue(all(map(lambda reply: reply == b"pong", replies)))
# run for long enough for the delayed pings to have been sent by node 1
await advance(1000)
# verify all of the previously pinged peers have node_1 in their routing tables
for n in nodes.values():
peers = n.protocol.routing_table.get_peers()
if n is node_1:
self.assertEqual(8, len(peers))
else:
self.assertEqual(1, len(peers))
self.assertEqual(
(peers[0].node_id, peers[0].address,
peers[0].udp_port),
(node_1.protocol.node_id, node_1.protocol.external_ip,
node_1.protocol.udp_port))
# run long enough for the refresh loop to run
await advance(3600)
# verify all the nodes know about each other
for n in nodes.values():
if n is node_1:
continue
peers = n.protocol.routing_table.get_peers()
self.assertEqual(8, len(peers))
self.assertSetEqual(
{
n_id[0]
for n_id in peer_addresses
if n_id[0] != n.protocol.node_id
}, {c.node_id
for c in peers})
self.assertSetEqual(
{
n_addr[1]
for n_addr in peer_addresses
if n_addr[1] != n.protocol.external_ip
}, {c.address
for c in peers})
# teardown
for n in nodes.values():
n.stop()
def refresh_token(self): # TODO: this needs to be called periodically
self.old_token_secret = self.token_secret
self.token_secret = constants.generate_id()
async def test_announce_no_peers(self):
await self.setup_network(1)
node = self.nodes[0]
blob_hash = hexlify(constants.generate_id(1337)).decode()
peers = await node.announce_blob(blob_hash)
self.assertEqual(len(peers), 0)
