class PeerTest(AsyncioTestCase): def setUp(self): self.loop = asyncio.get_event_loop() self.peer_manager = PeerManager(self.loop) self.node_ids = [generate_id(), generate_id(), generate_id()] self.first_contact = self.peer_manager.get_kademlia_peer( self.node_ids[1], '127.0.0.1', udp_port=1000) self.second_contact = self.peer_manager.get_kademlia_peer( self.node_ids[0], '192.168.0.1', udp_port=1000) def test_make_contact_error_cases(self): self.assertRaises(ValueError, self.peer_manager.get_kademlia_peer, self.node_ids[1], '192.168.1.20', 100000) self.assertRaises(ValueError, self.peer_manager.get_kademlia_peer, self.node_ids[1], '192.168.1.20.1', 1000) self.assertRaises(ValueError, self.peer_manager.get_kademlia_peer, self.node_ids[1], 'this is not an ip', 1000) self.assertRaises(ValueError, self.peer_manager.get_kademlia_peer, self.node_ids[1], '192.168.1.20', -1000) self.assertRaises(ValueError, self.peer_manager.get_kademlia_peer, b'not valid node id', '192.168.1.20', 1000) def test_boolean(self): self.assertNotEqual(self.first_contact, self.second_contact) self.assertEqual( self.second_contact, self.peer_manager.get_kademlia_peer(self.node_ids[0], '192.168.0.1', udp_port=1000)) def test_compact_ip(self): self.assertEqual(self.first_contact.compact_ip(), b'\x7f\x00\x00\x01') self.assertEqual(self.second_contact.compact_ip(), b'\xc0\xa8\x00\x01')
class TestBlobAnnouncer(AsyncioTestCase): async def setup_node(self, peer_addresses, address, node_id): self.nodes: typing.Dict[int, Node] = {} self.advance = dht_mocks.get_time_accelerator(self.loop, self.loop.time()) self.conf = Config() self.storage = SQLiteStorage(self.conf, ":memory:", self.loop, self.loop.time) await self.storage.open() self.peer_manager = PeerManager(self.loop) self.node = Node(self.loop, self.peer_manager, node_id, 4444, 4444, 3333, address) await self.node.start_listening(address) self.blob_announcer = BlobAnnouncer(self.loop, self.node, self.storage) for node_id, address in peer_addresses: await self.add_peer(node_id, address) self.node.joined.set() async def add_peer(self, node_id, address, add_to_routing_table=True): n = Node(self.loop, PeerManager(self.loop), node_id, 4444, 4444, 3333, address) await n.start_listening(address) self.nodes.update({len(self.nodes): n}) if add_to_routing_table: await self.node.protocol.add_peer( self.peer_manager.get_kademlia_peer( n.protocol.node_id, n.protocol.external_ip, n.protocol.udp_port ) ) @contextlib.asynccontextmanager 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 chain_peer(self, node_id, address): previous_last_node = self.nodes[len(self.nodes) - 1] await self.add_peer(node_id, address, False) last_node = self.nodes[len(self.nodes) - 1] peer = last_node.protocol.get_rpc_peer( last_node.protocol.peer_manager.get_kademlia_peer( previous_last_node.protocol.node_id, previous_last_node.protocol.external_ip, previous_last_node.protocol.udp_port ) ) await peer.ping() return peer 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_completed_blob(blob1, 1024) await self.storage.add_completed_blob(blob2, 1024) 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() await self.advance(61.0) 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') 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)
class DataStoreTests(TestCase): def setUp(self): self.loop = mock.Mock(spec=asyncio.BaseEventLoop) self.loop.time = lambda: 0.0 self.peer_manager = PeerManager(self.loop) self.data_store = DictDataStore(self.loop, self.peer_manager) def _test_add_peer_to_blob(self, blob=b'2' * 48, node_id=b'1' * 48, address='1.2.3.4', tcp_port=3333, udp_port=4444): peer = self.peer_manager.get_kademlia_peer(node_id, address, udp_port) peer.update_tcp_port(tcp_port) before = self.data_store.get_peers_for_blob(blob) self.data_store.add_peer_to_blob(peer, blob) self.assertListEqual(before + [peer], self.data_store.get_peers_for_blob(blob)) return peer def test_refresh_peer_to_blob(self): blob = b'f' * 48 self.assertListEqual([], self.data_store.get_peers_for_blob(blob)) peer = self._test_add_peer_to_blob(blob=blob, node_id=b'a' * 48, address='1.2.3.4') self.assertTrue(self.data_store.has_peers_for_blob(blob)) self.assertEqual(len(self.data_store.get_peers_for_blob(blob)), 1) self.assertEqual(self.data_store._data_store[blob][0][1], 0) self.loop.time = lambda: 100.0 self.assertEqual(self.data_store._data_store[blob][0][1], 0) self.data_store.add_peer_to_blob(peer, blob) self.assertEqual(self.data_store._data_store[blob][0][1], 100) def test_add_peer_to_blob(self, blob=b'f' * 48, peers=None): peers = peers or [ (b'a' * 48, '1.2.3.4'), (b'b' * 48, '1.2.3.5'), (b'c' * 48, '1.2.3.6'), ] self.assertListEqual([], self.data_store.get_peers_for_blob(blob)) peer_objects = [] for (node_id, address) in peers: peer_objects.append( self._test_add_peer_to_blob(blob=blob, node_id=node_id, address=address)) self.assertTrue(self.data_store.has_peers_for_blob(blob)) self.assertEqual(len(self.data_store.get_peers_for_blob(blob)), len(peers)) return peer_objects def test_get_storing_contacts(self, peers=None, blob1=b'd' * 48, blob2=b'e' * 48): peers = peers or [ (b'a' * 48, '1.2.3.4'), (b'b' * 48, '1.2.3.5'), (b'c' * 48, '1.2.3.6'), ] peer_objs1 = self.test_add_peer_to_blob(blob=blob1, peers=peers) self.assertEqual(len(peers), len(peer_objs1)) self.assertEqual(len(peers), len(self.data_store.get_storing_contacts())) peer_objs2 = self.test_add_peer_to_blob(blob=blob2, peers=peers) self.assertEqual(len(peers), len(peer_objs2)) self.assertEqual(len(peers), len(self.data_store.get_storing_contacts())) for o1, o2 in zip(peer_objs1, peer_objs2): self.assertIs(o1, o2) def test_remove_expired_peers(self): peers = [ (b'a' * 48, '1.2.3.4'), (b'b' * 48, '1.2.3.5'), (b'c' * 48, '1.2.3.6'), ] blob1 = b'd' * 48 blob2 = b'e' * 48 self.data_store.removed_expired_peers() # nothing should happen self.test_get_storing_contacts(peers, blob1, blob2) self.assertEqual(len(self.data_store.get_peers_for_blob(blob1)), len(peers)) self.assertEqual(len(self.data_store.get_peers_for_blob(blob2)), len(peers)) self.assertEqual(len(self.data_store.get_storing_contacts()), len(peers)) # expire the first peer from blob1 first = self.data_store._data_store[blob1][0][0] self.data_store._data_store[blob1][0] = (first, -86401) self.assertEqual(len(self.data_store.get_storing_contacts()), len(peers)) self.data_store.removed_expired_peers() self.assertEqual(len(self.data_store.get_peers_for_blob(blob1)), len(peers) - 1) self.assertEqual(len(self.data_store.get_peers_for_blob(blob2)), len(peers)) self.assertEqual(len(self.data_store.get_storing_contacts()), len(peers)) # expire the first peer from blob2 first = self.data_store._data_store[blob2][0][0] self.data_store._data_store[blob2][0] = (first, -86401) self.data_store.removed_expired_peers() self.assertEqual(len(self.data_store.get_peers_for_blob(blob1)), len(peers) - 1) self.assertEqual(len(self.data_store.get_peers_for_blob(blob2)), len(peers) - 1) self.assertEqual(len(self.data_store.get_storing_contacts()), len(peers) - 1) # expire the second and third peers from blob1 first = self.data_store._data_store[blob2][0][0] self.data_store._data_store[blob1][0] = (first, -86401) second = self.data_store._data_store[blob2][1][0] self.data_store._data_store[blob1][1] = (second, -86401) self.data_store.removed_expired_peers() self.assertEqual(len(self.data_store.get_peers_for_blob(blob1)), 0) self.assertEqual(len(self.data_store.get_peers_for_blob(blob2)), len(peers) - 1) self.assertEqual(len(self.data_store.get_storing_contacts()), len(peers) - 1)
class TestKBucket(AsyncioTestCase): def setUp(self): self.loop = asyncio.get_event_loop() self.address_generator = address_generator() self.peer_manager = PeerManager(self.loop) self.kbucket = KBucket(self.peer_manager, 0, 2**constants.hash_bits, generate_id()) def test_add_peer(self): # 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]) # def testGetContacts(self): # # try and get 2 contacts from empty list # result = self.kbucket.getContacts(2) # self.assertFalse(len(result) != 0, "Returned list should be empty; returned list length: %d" % # (len(result))) # # # Add k-2 contacts # node_ids = [] # if constants.k >= 2: # for i in range(constants.k-2): # node_ids.append(generate_id()) # tmpContact = self.contact_manager.make_contact(node_ids[-1], next(self.address_generator), 4444, 0, # None) # self.kbucket.addContact(tmpContact) # else: # # add k contacts # for i in range(constants.k): # node_ids.append(generate_id()) # tmpContact = self.contact_manager.make_contact(node_ids[-1], next(self.address_generator), 4444, 0, # None) # self.kbucket.addContact(tmpContact) # # # try to get too many contacts # # requested count greater than bucket size; should return at most k contacts # contacts = self.kbucket.getContacts(constants.k+3) # self.assertTrue(len(contacts) <= constants.k, # 'Returned list should not have more than k entries!') # # # verify returned contacts in list # for node_id, i in zip(node_ids, range(constants.k-2)): # self.assertFalse(self.kbucket._contacts[i].id != node_id, # "Contact in position %s not same as added contact" % (str(i))) # # # try to get too many contacts # # requested count one greater than number of contacts # if constants.k >= 2: # result = self.kbucket.getContacts(constants.k-1) # self.assertFalse(len(result) != constants.k-2, # "Too many contacts in returned list %s - should be %s" % # (len(result), constants.k-2)) # else: # result = self.kbucket.getContacts(constants.k-1) # # if the count is <= 0, it should return all of it's contats # self.assertFalse(len(result) != constants.k, # "Too many contacts in returned list %s - should be %s" % # (len(result), constants.k-2)) # result = self.kbucket.getContacts(constants.k-3) # self.assertFalse(len(result) != constants.k-3, # "Too many contacts in returned list %s - should be %s" % # (len(result), constants.k-3)) def test_remove_peer(self): # try remove contact from empty list peer = self.peer_manager.get_kademlia_peer( generate_id(), next(self.address_generator), 4444) self.assertRaises(ValueError, self.kbucket.remove_peer, peer) added = [] # Add couple contacts for i in range(constants.k - 2): peer = self.peer_manager.get_kademlia_peer( generate_id(), next(self.address_generator), 4444) self.assertTrue(self.kbucket.add_peer(peer)) added.append(peer) while added: peer = added.pop() self.assertIn(peer, self.kbucket.peers) self.kbucket.remove_peer(peer) self.assertNotIn(peer, self.kbucket.peers)