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 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):
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])
# 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)
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()
self.node._refresh_task = self.loop.create_task(
self.node.refresh_node())
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:
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_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)
async def test_popular_blob(self):
peer_count = 150
addresses = [
(constants.generate_id(i + 1),
socket.inet_ntoa(int(i + 1).to_bytes(length=4, byteorder='big')))
for i in range(peer_count)
]
blob_hash = b'1' * 48
async with self._test_network_context(peer_addresses=addresses):
total_seen = set()
announced_to = self.nodes[0]
for i in range(1, peer_count):
node = self.nodes[i]
kad_peer = announced_to.protocol.peer_manager.get_kademlia_peer(
node.protocol.node_id, node.protocol.external_ip,
node.protocol.udp_port)
await announced_to.protocol._add_peer(kad_peer)
peer = node.protocol.get_rpc_peer(
node.protocol.peer_manager.get_kademlia_peer(
announced_to.protocol.node_id,
announced_to.protocol.external_ip,
announced_to.protocol.udp_port))
response = await peer.store(blob_hash)
self.assertEqual(response, b'OK')
peers_for_blob = await peer.find_value(blob_hash, 0)
if i == 1:
self.assertTrue(blob_hash not in peers_for_blob)
self.assertEqual(peers_for_blob[b'p'], 0)
else:
self.assertEqual(len(peers_for_blob[blob_hash]),
min(i - 1, constants.k))
self.assertEqual(
len(
announced_to.protocol.data_store.
get_peers_for_blob(blob_hash)), i)
if i - 1 > constants.k:
self.assertEqual(len(peers_for_blob[b'contacts']),
constants.k)
self.assertEqual(peers_for_blob[b'p'],
((i - 1) // (constants.k + 1)) + 1)
seen = set(peers_for_blob[blob_hash])
self.assertEqual(len(seen), constants.k)
self.assertEqual(len(peers_for_blob[blob_hash]), len(seen))
for pg in range(1, peers_for_blob[b'p']):
page_x = await peer.find_value(blob_hash, pg)
self.assertNotIn(b'contacts', page_x)
page_x_set = set(page_x[blob_hash])
self.assertEqual(len(page_x[blob_hash]),
len(page_x_set))
self.assertTrue(len(page_x_set) > 0)
self.assertSetEqual(seen.intersection(page_x_set),
set())
seen.intersection_update(page_x_set)
total_seen.update(page_x_set)
else:
self.assertEqual(len(peers_for_blob[b'contacts']), i - 1)
self.assertEqual(len(total_seen), peer_count - 2)
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)
