From 39dcab3aa137b31abfa2a85728cd223d7449cdd8 Mon Sep 17 00:00:00 2001 From: Bogdan Opanchuk Date: Thu, 2 Sep 2021 23:53:34 -0700 Subject: [PATCH] Ensure enough Ursulas are available before retrieval --- nucypher/characters/lawful.py | 60 ------------ nucypher/network/nodes.py | 13 --- nucypher/network/retrieval.py | 40 +++++++- .../characters/test_bob_handles_frags.py | 95 +------------------ 4 files changed, 37 insertions(+), 171 deletions(-) diff --git a/nucypher/characters/lawful.py b/nucypher/characters/lawful.py index f4c65413f..2b8c8435c 100644 --- a/nucypher/characters/lawful.py +++ b/nucypher/characters/lawful.py @@ -541,66 +541,6 @@ class Bob(Character): card = Card.from_character(self) return card - def peek_at_treasure_map(self, treasure_map): - """ - Take a quick gander at the TreasureMap to see which - nodes are already known to us. - - Don't do any learning, pinging, or anything other than just seeing - whether we know or don't know the nodes. - - Return two sets: nodes that are unknown to us, nodes that are known to us. - """ - - # The intersection of the map and our known nodes will be the known Ursulas... - known_treasure_ursulas = treasure_map.destinations.keys() & self.known_nodes.addresses() - - # while the difference will be the unknown Ursulas. - unknown_treasure_ursulas = treasure_map.destinations.keys() - self.known_nodes.addresses() - - return unknown_treasure_ursulas, known_treasure_ursulas - - def follow_treasure_map(self, - treasure_map=None, - block=False, - new_thread=False, - timeout=10, - allow_missing=0): - """ - Follows a known TreasureMap. - - Determines which Ursulas are known and which are unknown. - - If block, will block until either unknown nodes are discovered or until timeout seconds have elapsed. - After timeout seconds, if more than allow_missing nodes are still unknown, raises NotEnoughUrsulas. - - If block and new_thread, does the same thing but on a different thread, returning a Deferred which - fires after the blocking has concluded. - - Otherwise, returns (unknown_nodes, known_nodes). - - # TODO: Check if nodes are up, declare them phantom if not. 567 - """ - unknown_ursulas, known_ursulas = self.peek_at_treasure_map(treasure_map=treasure_map) - - if unknown_ursulas: - self.learn_about_specific_nodes(unknown_ursulas) - - # TODO: what does this even do? - self._push_certain_newly_discovered_nodes_here(known_ursulas, unknown_ursulas) - - if block: - if new_thread: - return threads.deferToThread(self.block_until_specific_nodes_are_known, - unknown_ursulas, - timeout=timeout, - allow_missing=allow_missing) - else: - self.block_until_specific_nodes_are_known(unknown_ursulas, - timeout=timeout, - allow_missing=allow_missing, - learn_on_this_thread=True) - def _decrypt_treasure_map(self, encrypted_treasure_map: EncryptedTreasureMap) -> TreasureMap: publisher = Alice.from_public_keys(verifying_key=encrypted_treasure_map.publisher_verifying_key) diff --git a/nucypher/network/nodes.py b/nucypher/network/nodes.py index 3af2384f7..cc186d72b 100644 --- a/nucypher/network/nodes.py +++ b/nucypher/network/nodes.py @@ -255,7 +255,6 @@ class Learner: self.learn_on_same_thread = learn_on_same_thread self._abort_on_learning_error = abort_on_learning_error - self._learning_listeners = defaultdict(list) self._node_ids_to_learn_about_immediately = set() self.__known_nodes = self.tracker_class(domain=domain, this_node=self if include_self_in_the_state else None) @@ -463,10 +462,6 @@ class Learner: # TODO: What about InvalidNode? (for that matter, any SuspiciousActivity) 1714, 567 too really - listeners = self._learning_listeners.pop(node.checksum_address, tuple()) - - for listener in listeners: - listener.add(node.checksum_address) self._node_ids_to_learn_about_immediately.discard(node.checksum_address) if record_fleet_state: @@ -709,14 +704,6 @@ class Learner: self._LONG_LEARNING_DELAY)) self._learning_task.interval = self._LONG_LEARNING_DELAY - def _push_certain_newly_discovered_nodes_here(self, queue_to_push, node_addresses): - """ - If any node_addresses are discovered, push them to queue_to_push. - """ - for node_address in node_addresses: - self.log.info("Adding listener for {}".format(node_address)) - self._learning_listeners[node_address].append(queue_to_push) - def network_bootstrap(self, node_list: list) -> None: for node_addr, port in node_list: new_nodes = self.learn_about_nodes_now(node_addr, port) diff --git a/nucypher/network/retrieval.py b/nucypher/network/retrieval.py index 2f2a1e71d..1aead86a3 100644 --- a/nucypher/network/retrieval.py +++ b/nucypher/network/retrieval.py @@ -251,6 +251,40 @@ class RetrievalClient: self._learner = learner self.log = Logger(self.__class__.__name__) + def _ensure_ursula_availability(self, treasure_map: TreasureMap, timeout=10): + """ + Make sure we know enough nodes from the treasure map to decrypt; + otherwise block and wait for them to come online. + """ + + # OK, so we're going to need to do some network activity for this retrieval. + # Let's make sure we've seeded. + if not self._learner.done_seeding: + self._learner.learn_from_teacher_node() + + ursulas_in_map = treasure_map.destinations.keys() + all_known_ursulas = self._learner.known_nodes.addresses() + + # Push all unknown Ursulas from the map in the queue for learning + unknown_ursulas = ursulas_in_map - all_known_ursulas + if unknown_ursulas: + self._learner.learn_about_specific_nodes(unknown_ursulas) + + # How many nodes over the threshold we want to know (just in case) + redundancy = 0 + + # If we know enough to decrypt, we can proceed. + known_ursulas = ursulas_in_map & all_known_ursulas + + if len(known_ursulas) - redundancy >= treasure_map.threshold: + return + + allow_missing = max(len(unknown_ursulas) - redundancy, 0) + self._learner.block_until_specific_nodes_are_known(unknown_ursulas, + timeout=timeout, + allow_missing=allow_missing, + learn_on_this_thread=True) + def _request_reencryption(self, ursula: 'Ursula', reencryption_request: 'ReencryptionRequest', @@ -340,11 +374,7 @@ class RetrievalClient: policy_encrypting_key: PublicKey, # Key used to create the policy ) -> List[RetrievalResult]: - # TODO: why is it here? This code shouldn't know about these details. - # OK, so we're going to need to do some network activity for this retrieval. - # Let's make sure we've seeded. - if not self._learner.done_seeding: - self._learner.learn_from_teacher_node() + self._ensure_ursula_availability(treasure_map) retrieval_plan = RetrievalPlan(treasure_map=treasure_map, retrieval_kits=retrieval_kits) diff --git a/tests/integration/characters/test_bob_handles_frags.py b/tests/integration/characters/test_bob_handles_frags.py index 2902886f1..732724f42 100644 --- a/tests/integration/characters/test_bob_handles_frags.py +++ b/tests/integration/characters/test_bob_handles_frags.py @@ -19,105 +19,14 @@ import pytest import pytest_twisted from twisted.internet import threads -from nucypher.characters.lawful import Enrico +from nucypher.characters.lawful import Enrico, Bob from nucypher.config.constants import TEMPORARY_DOMAIN +from nucypher.network.retrieval import RetrievalClient from nucypher.policy.kits import RetrievalKit from tests.utils.middleware import MockRestMiddleware, NodeIsDownMiddleware -def test_bob_cannot_follow_the_treasure_map_in_isolation(federated_treasure_map, federated_bob): - # Assume for the moment that Bob has already received a TreasureMap, perhaps via a side channel. - - # Bob knows of no Ursulas. - assert len(federated_bob.known_nodes) == 0 - - # He can't successfully follow the TreasureMap until he learns of a node to ask. - unknown, known = federated_bob.peek_at_treasure_map(treasure_map=federated_treasure_map) - assert len(known) == 0 - - # TODO: Show that even with learning loop going, nothing happens here. - # Probably use Clock? - federated_bob.follow_treasure_map(treasure_map=federated_treasure_map) - assert len(known) == 0 - - -def test_bob_already_knows_all_nodes_in_treasure_map(enacted_federated_policy, - federated_ursulas, - federated_bob, - federated_alice): - # Bob knows of no Ursulas. - assert len(federated_bob.known_nodes) == 0 - - # Now we'll inform Bob of some Ursulas. - for ursula in federated_ursulas: - federated_bob.remember_node(ursula) - - # Now, Bob can get the TreasureMap all by himself, and doesn't need a side channel. - the_map = federated_bob._decrypt_treasure_map(enacted_federated_policy.treasure_map) - unknown, known = federated_bob.peek_at_treasure_map(treasure_map=the_map) - - # He finds that he didn't need to discover any new nodes... - assert len(unknown) == 0 - - # ...because he already knew of all the Ursulas on the map. - assert len(known) == enacted_federated_policy.shares - - -@pytest_twisted.inlineCallbacks -def test_bob_can_follow_treasure_map_even_if_he_only_knows_of_one_node(federated_treasure_map, - federated_ursulas, - certificates_tempdir): - """ - Similar to above, but this time, we'll show that if Bob can connect to a single node, he can - learn enough to follow the TreasureMap. - - Also, we'll get the TreasureMap from the hrac alone (ie, not via a side channel). - """ - - from nucypher.characters.lawful import Bob - - bob = Bob(network_middleware=MockRestMiddleware(), - domain=TEMPORARY_DOMAIN, - start_learning_now=False, - abort_on_learning_error=True, - federated_only=True) - - # Again, let's assume that he received the TreasureMap via a side channel. - - # Now, let's create a scenario in which Bob knows of only one node. - assert len(bob.known_nodes) == 0 - first_ursula = list(federated_ursulas).pop(0) - bob.remember_node(first_ursula) - assert len(bob.known_nodes) == 1 - - # This time, when he follows the TreasureMap... - unknown_nodes, known_nodes = bob.peek_at_treasure_map(treasure_map=federated_treasure_map) - - # Bob already knew about one node; the rest are unknown. - assert len(unknown_nodes) == len(federated_treasure_map) - 1 - - # He needs to actually follow the treasure map to get the rest. - bob.follow_treasure_map(treasure_map=federated_treasure_map) - - # The nodes in the learning loop are now his top target, but he's not learning yet. - assert not bob._learning_task.running - - # ...so he hasn't learned anything (ie, Bob still knows of just one node). - assert len(bob.known_nodes) == 1 - - # Now, we'll start his learning loop. - bob.start_learning_loop() - - # ...and block until the unknown_nodes have all been found. - d = threads.deferToThread(bob.block_until_specific_nodes_are_known, unknown_nodes) - yield d - - # ...and he now has no more unknown_nodes. - assert len(bob.known_nodes) == len(federated_treasure_map) - bob.disenchant() - - def _policy_info_kwargs(enacted_policy): return dict( encrypted_treasure_map=enacted_policy.treasure_map,