def __init__(self, loop: asyncio.AbstractEventLoop, peer_manager: 'PeerManager', node_id: bytes, external_ip: str,
udp_port: int, peer_port: int, rpc_timeout: float = constants.rpc_timeout,
split_buckets_under_index: int = constants.split_buckets_under_index):
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, split_buckets_under_index)
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.rpc_timeout = rpc_timeout
self._split_lock = asyncio.Lock(loop=self.loop)
self._to_remove: typing.Set['KademliaPeer'] = set()
self._to_add: typing.Set['KademliaPeer'] = set()
self._wakeup_routing_task = asyncio.Event(loop=self.loop)
self.maintaing_routing_task: typing.Optional[asyncio.Task] = None
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)
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 = make_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)