From df1b2338b1a6d69e8ff6f82d41d2c20600e0a103 Mon Sep 17 00:00:00 2001
From: Brian Warner <warner@lothar.com>
Date: Thu, 6 Apr 2017 10:44:04 -0700
Subject: [PATCH] tests: exercise Key receiving PAKE before set_code

---
 src/wormhole/_key.py               |  2 ++
 src/wormhole/test/common.py        |  9 ++++++++-
 src/wormhole/test/test_machines.py | 28 ++++++++++++++++++++++++++++
 src/wormhole/test/test_wormhole.py | 24 +++++++++++++++++++++++-
 4 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/src/wormhole/_key.py b/src/wormhole/_key.py
index 67433ec..622af65 100644
--- a/src/wormhole/_key.py
+++ b/src/wormhole/_key.py
@@ -69,6 +69,7 @@ class Key(object):
     def __attrs_post_init__(self):
         self._SK = _SortedKey(self._appid, self._versions, self._side,
                               self._timing)
+        self._debug_pake_stashed = False # for tests
 
     def wire(self, boss, mailbox, receive):
         self._SK.wire(boss, mailbox, receive)
@@ -90,6 +91,7 @@ class Key(object):
     @m.output()
     def stash_pake(self, body):
         self._pake = body
+        self._debug_pake_stashed = True
     @m.output()
     def deliver_code(self, code):
         self._SK.got_code(code)
diff --git a/src/wormhole/test/common.py b/src/wormhole/test/common.py
index f3b57ad..dd1d490 100644
--- a/src/wormhole/test/common.py
+++ b/src/wormhole/test/common.py
@@ -1,6 +1,6 @@
 # no unicode_literals untill twisted update
 from twisted.application import service
-from twisted.internet import defer, task
+from twisted.internet import defer, task, reactor
 from twisted.python import log
 from click.testing import CliRunner
 import mock
@@ -84,3 +84,10 @@ def config(*argv):
         cfg = go.call_args[0][1]
     return cfg
 
+@defer.inlineCallbacks
+def poll_until(predicate):
+    # return a Deferred that won't fire until the predicate is True
+    while not predicate():
+        d = defer.Deferred()
+        reactor.callLater(0.001, d.callback, None)
+        yield d
diff --git a/src/wormhole/test/test_machines.py b/src/wormhole/test/test_machines.py
index 3e45ab7..df002b7 100644
--- a/src/wormhole/test/test_machines.py
+++ b/src/wormhole/test/test_machines.py
@@ -257,6 +257,34 @@ class Key(unittest.TestCase):
         k.got_pake(dict_to_bytes(bad_pake_d))
         self.assertEqual(events, [("b.scared",)])
 
+    def test_reversed(self):
+        # A receiver using input_code() will choose the nameplate first, then
+        # the rest of the code. Once the nameplate is selected, we'll claim
+        # it and open the mailbox, which will cause the senders PAKE to
+        # arrive before the code has been set. Key() is supposed to stash the
+        # PAKE message until the code is set (allowing the PAKE computation
+        # to finish). This test exercises that PAKE-then-code sequence.
+        k, b, m, r, events = self.build()
+        code = u"1-foo"
+
+        sp = SPAKE2_Symmetric(to_bytes(code), idSymmetric=to_bytes(u"appid"))
+        msg2_bytes = sp.start()
+        msg2 = dict_to_bytes({"pake_v1": bytes_to_hexstr(msg2_bytes)})
+        k.got_pake(msg2)
+        self.assertEqual(len(events), 0)
+
+        k.got_code(code)
+        self.assertEqual(len(events), 5)
+        self.assertEqual(events[0][:2], ("m.add_message", "pake"))
+        msg1_json = events[0][2].decode("utf-8")
+        msg1 = json.loads(msg1_json)
+        msg1_bytes = hexstr_to_bytes(msg1["pake_v1"])
+        key2 = sp.finish(msg1_bytes)
+        self.assertEqual(events[1], ("b.got_key", key2))
+        self.assertEqual(events[2][0], "b.got_verifier")
+        self.assertEqual(events[3][:2], ("m.add_message", "version"))
+        self.assertEqual(events[4], ("r.got_key", key2))
+
 class Code(unittest.TestCase):
     def build(self):
         events = []
diff --git a/src/wormhole/test/test_wormhole.py b/src/wormhole/test/test_wormhole.py
index fc329d6..dd2a552 100644
--- a/src/wormhole/test/test_wormhole.py
+++ b/src/wormhole/test/test_wormhole.py
@@ -4,7 +4,7 @@ import mock
 from twisted.trial import unittest
 from twisted.internet import reactor
 from twisted.internet.defer import gatherResults, inlineCallbacks
-from .common import ServerBase
+from .common import ServerBase, poll_until
 from .. import wormhole, _rendezvous
 from ..errors import (WrongPasswordError,
                       KeyFormatError, WormholeClosed, LonelyError,
@@ -234,6 +234,28 @@ class Wormholes(ServerBase, unittest.TestCase):
         yield w1.close()
         yield w2.close()
 
+    @inlineCallbacks
+    def test_input_code(self):
+        w1 = wormhole.create(APPID, self.relayurl, reactor)
+        w2 = wormhole.create(APPID, self.relayurl, reactor)
+        w1.set_code("123-purple-elephant")
+        h = w2.input_code()
+        h.choose_nameplate("123")
+        # Pause to allow some messages to get delivered. Specifically we want
+        # to wait until w2 claims the nameplate, opens the mailbox, and
+        # receives the PAKE message, to exercise the PAKE-before-CODE path in
+        # Key.
+        yield poll_until(lambda: w2._boss._K._debug_pake_stashed)
+        h.choose_words("purple-elephant")
+
+        w1.send(b"data1"), w2.send(b"data2")
+        dl = yield self.doBoth(w1.when_received(), w2.when_received())
+        (dataX, dataY) = dl
+        self.assertEqual(dataX, b"data2")
+        self.assertEqual(dataY, b"data1")
+        yield w1.close()
+        yield w2.close()
+
 
     @inlineCallbacks
     def test_multiple_messages(self):