Merge pull request #560 from trezor/matejcik/word-count

core/recovery: allow changing word count again (fixes #554)

core/src/apps/management/recovery_device/homescreen.py

@@ -46,13 +46,23 @@ async def _continue_recovery_process(ctx: wire.Context) -> Success:
     # gather the current recovery state from storage
     dry_run = storage.recovery.is_dry_run()
     word_count, backup_type = recover.load_slip39_state()
-    if word_count:
+
+    # Both word_count and backup_type are derived from the same data. Both will be
+    # either set or unset. We use 'backup_type is None' to detect status of both.
+    # The following variable indicates that we are (re)starting the first recovery step,
+    # which includes word count selection.
+    is_first_step = backup_type is None
+
+    if not is_first_step:
+        # If we continue recovery, show starting screen with word count immediately.
         await _request_share_first_screen(ctx, word_count)
 
     secret = None
     while secret is None:
-        if not word_count:  # the first run, prompt word count from the user
+        if is_first_step:
+            # If we are starting recovery, ask for word count first...
             word_count = await _request_word_count(ctx, dry_run)
+            # ...and only then show the starting screen with word count.
             await _request_share_first_screen(ctx, word_count)
 
         # ask for mnemonic words one by one
@@ -63,7 +73,11 @@ async def _continue_recovery_process(ctx: wire.Context) -> Success:
             continue
 
         try:
-            secret, word_count, backup_type = await _process_words(ctx, words)
+            secret, backup_type = await _process_words(ctx, words)
+            # If _process_words succeeded, we now have both backup_type (from
+            # its result) and word_count (from _request_word_count earlier), which means
+            # that the first step is complete.
+            is_first_step = False
         except MnemonicError:
             await layout.show_invalid_mnemonic(ctx, word_count)
 
@@ -142,8 +156,7 @@ async def _request_word_count(ctx: wire.Context, dry_run: bool) -> int:
 
 async def _process_words(
     ctx: wire.Context, words: str
-) -> Tuple[Optional[bytes], EnumTypeBackupType, int]:
-
+) -> Tuple[Optional[bytes], EnumTypeBackupType]:
     word_count = len(words.split(" "))
     is_slip39 = backup_types.is_slip39_word_count(word_count)
 
@@ -159,7 +172,7 @@ async def _process_words(
             await layout.show_group_share_success(ctx, share.index, share.group_index)
         await _request_share_next_screen(ctx)
 
-    return secret, word_count, backup_type
+    return secret, backup_type
 
 
 async def _request_share_first_screen(ctx: wire.Context, word_count: int) -> None:

tests/device_tests/test_msg_recoverydevice_slip39_basic.py

@@ -133,6 +133,97 @@ def test_noabort(client):
         assert client.features.initialized is True
 
 
+@pytest.mark.setup_client(uninitialized=True)
+def test_ask_word_number(client):
+    debug = client.debug
+
+    def input_flow_retry_first():
+        yield  # Confirm Recovery
+        debug.press_yes()
+        yield  # Homescreen - start process
+        debug.press_yes()
+        yield  # Enter number of words
+        debug.input("20")
+        yield  # Homescreen - proceed to share entry
+        debug.press_yes()
+        yield  # Enter first share
+        for _ in range(20):
+            debug.input("slush")
+
+        code = yield  # Invalid share
+        assert code == messages.ButtonRequestType.Warning
+        debug.press_yes()
+
+        yield  # Homescreen - start process
+        debug.press_yes()
+        yield  # Enter number of words
+        debug.input("33")
+        yield  # Homescreen - proceed to share entry
+        debug.press_yes()
+        yield  # Enter first share
+        for _ in range(33):
+            debug.input("slush")
+
+        code = yield  # Invalid share
+        assert code == messages.ButtonRequestType.Warning
+        debug.press_yes()
+
+        yield  # Homescreen
+        debug.press_no()
+        yield  # Confirm abort
+        debug.press_yes()
+
+    with client:
+        client.set_input_flow(input_flow_retry_first)
+        with pytest.raises(exceptions.Cancelled):
+            device.recover(client, pin_protection=False, label="label")
+        client.init_device()
+        assert client.features.initialized is False
+
+    def input_flow_retry_second():
+        yield  # Confirm Recovery
+        debug.press_yes()
+        yield  # Homescreen - start process
+        debug.press_yes()
+        yield  # Enter number of words
+        debug.input("20")
+        yield  # Homescreen - proceed to share entry
+        debug.press_yes()
+        yield  # Enter first share
+        share = MNEMONIC_SLIP39_BASIC_20_3of6[0].split(" ")
+        for word in share:
+            debug.input(word)
+
+        yield  # More shares needed
+        debug.press_yes()
+
+        yield  # Enter another share
+        share = share[:3] + ["slush"] * 17
+        for word in share:
+            debug.input(word)
+
+        code = yield  # Invalid share
+        assert code == messages.ButtonRequestType.Warning
+        debug.press_yes()
+
+        yield  # Proceed to next share
+        share = MNEMONIC_SLIP39_BASIC_20_3of6[1].split(" ")
+        for word in share:
+            debug.input(word)
+
+        yield  # More shares needed
+        debug.press_no()
+        yield  # Confirm abort
+        debug.press_yes()
+
+    with client:
+        client.set_input_flow(input_flow_retry_second)
+        with pytest.raises(exceptions.Cancelled):
+            device.recover(client, pin_protection=False, label="label")
+        client.init_device()
+        assert client.features.initialized is False
+
+
 @pytest.mark.setup_client(uninitialized=True)
 @pytest.mark.parametrize("nth_word", range(3))
 def test_wrong_nth_word(client, nth_word):