diff --git a/src/wormhole/test/test_wormhole.py b/src/wormhole/test/test_wormhole.py index 7a701df..a8d59c9 100644 --- a/src/wormhole/test/test_wormhole.py +++ b/src/wormhole/test/test_wormhole.py @@ -522,11 +522,13 @@ class Basic(unittest.TestCase): self.assertEqual(len(pieces), 3) # nameplate plus two words self.assert_(re.search(r'^\d+-\w+-\w+$', code), code) - def _test_establish_key_hook(self, established): - + def _test_establish_key_hook(self, established, before): timing = DebugTiming() w = wormhole._Wormhole(APPID, "relay_url", reactor, None, timing) + if before: + d = w.establish_key() + if established is True: w._key = b"key" elif established is False: @@ -535,7 +537,10 @@ class Basic(unittest.TestCase): w._key = b"key" w._error = WelcomeError() - d = w.establish_key() + if not before: + d = w.establish_key() + else: + w._maybe_notify_key() if w._key is not None and established is True: self.successResultOf(d) @@ -546,7 +551,8 @@ class Basic(unittest.TestCase): def test_establish_key_hook(self): for established in (True, False, "error"): - self._test_establish_key_hook(established) + for before in (True, False): + self._test_establish_key_hook(established, before) # make sure verify() can be called both before and after the verifier is # computed diff --git a/src/wormhole/wormhole.py b/src/wormhole/wormhole.py index c3b0c6f..d783e14 100644 --- a/src/wormhole/wormhole.py +++ b/src/wormhole/wormhole.py @@ -237,6 +237,7 @@ class _Wormhole: self._flag_need_to_see_mailbox_used = True self._flag_need_to_build_msg1 = True self._flag_need_to_send_PAKE = True + self._key_waiter = None self._key = None self._version_message = None @@ -283,6 +284,15 @@ class _Wormhole: # todo: restore-saved-state entry points + def establish_key(self): + """ + returns a Deferred that fires when we've established the shared key. + When successful, the Deferred fires with a simple `True`, otherwise + it fails. + + """ + return self._API_establish_key() + def verify(self): """Returns a Deferred that fires when we've heard back from the other side, and have confirmed that they used the right wormhole code. When @@ -559,6 +569,7 @@ class _Wormhole: def _event_established_key(self): self._timing.add("key established") + self._maybe_notify_key() # both sides send different (random) version messages self._send_version_message() @@ -569,6 +580,23 @@ class _Wormhole: self._maybe_check_version() self._maybe_send_phase_messages() + def _API_establish_key(self): + if self._error: return defer.fail(self._error) + if not self._key is None: + return defer.succeed(True) + self._key_waiter = defer.Deferred() + return self._key_waiter + + def _maybe_notify_key(self): + if self._key is None: + return + if self._error: + result = failure.Failure(self._error) + else: + result = True + if self._key_waiter and not self._key_waiter.called: + self._key_waiter.callback(result) + def _send_version_message(self): # this is encrypted like a normal phase message, and includes a # dictionary of version flags to let the other Wormhole know what