feat(core): avoid restarting session for select messages (fixes #1631)

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

@@ -6,7 +6,7 @@ from trezor import strings, utils, wire, workflow
 from trezor.crypto import slip39
 from trezor.crypto.hashlib import sha256
 from trezor.errors import MnemonicError
-from trezor.messages import BackupType
+from trezor.messages import BackupType, MessageType
 from trezor.messages.Success import Success
 from trezor.ui.layouts import show_success
 
@@ -31,6 +31,7 @@ async def recovery_homescreen() -> None:
 
 
 async def recovery_process(ctx: wire.GenericContext) -> Success:
+    wire.AVOID_RESTARTING_FOR = (MessageType.Initialize, MessageType.GetFeatures)
     try:
         return await _continue_recovery_process(ctx)
     except recover.RecoveryAborted:

core/src/session.py

@@ -1,4 +1,4 @@
-from trezor import loop, utils, wire, workflow
+from trezor import log, loop, utils, wire, workflow
 
 import apps.base
 import usb
@@ -25,3 +25,6 @@ if __debug__:
     wire.setup(usb.iface_debug, is_debug_session=True)
 
 loop.run()
+
+if __debug__:
+    log.debug(__name__, "Restarting main loop")

core/src/trezor/wire/__init__.py

@@ -52,6 +52,7 @@ if False:
         Any,
         Awaitable,
         Callable,
+        Container,
         Coroutine,
         Iterable,
         TypeVar,
@@ -431,7 +432,7 @@ async def handle_session(
                     # workflow running on wire.
                     utils.unimport_end(modules)
 
-                    if next_msg is None:
+                    if next_msg is None and msg.type not in AVOID_RESTARTING_FOR:
                         # Shut down the loop if there is no next message waiting.
                         # Let the session be restarted from `main`.
                         loop.clear()
@@ -450,6 +451,7 @@ def _find_handler_placeholder(iface: WireInterface, msg_type: int) -> Handler |
 
 
 find_handler = _find_handler_placeholder
+AVOID_RESTARTING_FOR: Container[int] = ()
 
 
 def failure(exc: BaseException) -> Failure:

tests/persistence_tests/test_shamir_persistence.py

@@ -99,6 +99,74 @@ def test_recovery_single_reset(emulator):
     assert features.recovery_mode is False
 
 
+@core_only
+def test_recovery_on_old_wallet(emulator):
+    """Check that the recovery workflow started on a disconnected device can survive
+    handling by the old Wallet.
+
+    While Suite will send a RecoveryDevice message and hook into the running recovery
+    flow, old Wallet can't do that and instead must repeatedly ask for features (via
+    Initialize+GetFeatures). At minimum, these two messages must not interrupt the
+    running recovery.
+    """
+    device_handler = BackgroundDeviceHandler(emulator.client)
+    debug = device_handler.debuglink()
+    features = device_handler.features()
+
+    assert features.initialized is False
+    assert features.recovery_mode is False
+
+    # enter recovery mode
+    device_handler.run(device.recover, pin_protection=False)
+    recovery.confirm_recovery(debug)
+
+    # restart to get into stand-alone recovery
+    debug = _restart(device_handler, emulator)
+    features = device_handler.features()
+    assert features.recovery_mode is True
+
+    # enter number of words
+    recovery.select_number_of_words(debug)
+
+    first_share = MNEMONIC_SLIP39_BASIC_20_3of6[0]
+    words = first_share.split(" ")
+
+    # start entering first share
+    layout = debug.read_layout()
+    assert "Enter any share" in layout.text
+    debug.press_yes()
+    layout = debug.wait_layout()
+    assert layout.text == "Slip39Keyboard"
+
+    # enter first word
+    debug.input(words[0])
+    layout = debug.wait_layout()
+
+    # while keyboard is open, hit the device with Initialize/GetFeatures
+    device_handler.client.init_device()
+    device_handler.client.refresh_features()
+
+    # try entering remaining 19 words
+    for word in words[1:]:
+        assert layout.text == "Slip39Keyboard"
+        debug.input(word)
+        layout = debug.wait_layout()
+
+    # check that we entered the first share successfully
+    assert "2 more shares" in layout.text
+
+    # try entering the remaining shares
+    for share in MNEMONIC_SLIP39_BASIC_20_3of6[1:3]:
+        recovery.enter_share(debug, share)
+
+    recovery.finalize(debug)
+
+    # check that the recovery succeeded
+    features = device_handler.features()
+    assert features.initialized is True
+    assert features.recovery_mode is False
+
+
 @core_only
 def test_recovery_multiple_resets(emulator):
     def enter_shares_with_restarts(debug):