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_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()
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()
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(make_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(make_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(make_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_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 = make_kademlia_peer(node_id1, '1.2.3.4', udp_port=4444) peer2_from_peer1 = make_kademlia_peer(peer2.node_id, peer2.external_ip, udp_port=peer2.udp_port) peer2_from_peer1.update_tcp_port(3333) peer3 = make_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', b'p'}, 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) find_value_page_above_pages_response = peer1.node_rpc.find_value( peer3, b'2' * 48, page=10) self.assertNotIn(b'2' * 48, find_value_page_above_pages_response) peer1.stop() peer2.stop() peer1.disconnect() peer2.disconnect()
async def _test_network_context(self, peer_count=200): self.peer_addresses = [ (constants.generate_id(i), socket.inet_ntoa(int(i + 0x01000001).to_bytes(length=4, byteorder='big'))) for i in range(1, peer_count + 1) ] try: with dht_mocks.mock_network_loop(self.loop): await self.setup_node(self.peer_addresses, '1.2.3.1', constants.generate_id(1000)) yield finally: self.blob_announcer.stop() self.node.stop() for n in self.nodes.values(): n.stop()
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 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 = make_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_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) # 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 = make_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)) # TODO: figure out why this breaks # 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()
async def test_losing_connection(self): async def wait_for(check_ok, insist, timeout=20): start = time.time() while time.time() - start < timeout: if check_ok(): break await asyncio.sleep(0) else: insist() loop = self.loop loop.set_debug(False) peer_addresses = [('1.2.3.4', 40000 + i) for i in range(10)] node_ids = [constants.generate_id(i) for i in range(10)] nodes = [ Node(loop, PeerManager(loop), node_id, udp_port, udp_port, 3333, address, storage=SQLiteStorage(Config(), ":memory:", self.loop, self.loop.time)) for node_id, (address, udp_port) in zip(node_ids, peer_addresses) ] dht_network = { peer_addresses[i]: node.protocol for i, node in enumerate(nodes) } num_seeds = 3 with dht_mocks.mock_network_loop(loop, dht_network): for i, n in enumerate(nodes): await n._storage.open() self.addCleanup(n.stop) n.start(peer_addresses[i][0], peer_addresses[:num_seeds]) await asyncio.gather(*[n.joined.wait() for n in nodes]) node = nodes[-1] advance = dht_mocks.get_time_accelerator(loop) await advance(500) # Join the network, assert that at least the known peers are in RT self.assertTrue(node.joined.is_set()) self.assertTrue( len(node.protocol.routing_table.get_peers()) >= num_seeds) # Refresh, so that the peers are persisted self.assertFalse( len(await node._storage.get_persisted_kademlia_peers()) > num_seeds) await advance(4000) self.assertTrue( len(await node._storage.get_persisted_kademlia_peers()) > num_seeds) # We lost internet connection - all the peers stop responding dht_network.pop( (node.protocol.external_ip, node.protocol.udp_port)) # The peers are cleared on refresh from RT and storage await advance(4000) self.assertListEqual([], await node._storage.get_persisted_kademlia_peers()) await wait_for( lambda: len(node.protocol.routing_table.get_peers()) == 0, lambda: self.assertListEqual( node.protocol.routing_table.get_peers(), [])) # Reconnect dht_network[(node.protocol.external_ip, node.protocol.udp_port)] = node.protocol # Check that node reconnects at least to them await advance(1000) await wait_for( lambda: len(node.protocol.routing_table.get_peers()) >= num_seeds, lambda: self.assertGreaterEqual( len(node.protocol.routing_table.get_peers()), num_seeds))