From 86089fa5ad2d5fd19e8a0cf6957d2cae59af5516 Mon Sep 17 00:00:00 2001 From: matejcik Date: Wed, 26 May 2021 12:01:47 +0200 Subject: [PATCH] feat(core): avoid restarting session for select messages (fixes #1631) --- .../management/recovery_device/homescreen.py | 3 +- core/src/session.py | 5 +- core/src/trezor/wire/__init__.py | 4 +- .../test_shamir_persistence.py | 68 +++++++++++++++++++ 4 files changed, 77 insertions(+), 3 deletions(-) diff --git a/core/src/apps/management/recovery_device/homescreen.py b/core/src/apps/management/recovery_device/homescreen.py index b4c0db88c..1da176fe1 100644 --- a/core/src/apps/management/recovery_device/homescreen.py +++ b/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: diff --git a/core/src/session.py b/core/src/session.py index b51b8c36a..2e5ab5dd9 100644 --- a/core/src/session.py +++ b/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") diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py index b1e56e1a1..040ff5095 100644 --- a/core/src/trezor/wire/__init__.py +++ b/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: diff --git a/tests/persistence_tests/test_shamir_persistence.py b/tests/persistence_tests/test_shamir_persistence.py index 7118ebc7c..c60e1bd59 100644 --- a/tests/persistence_tests/test_shamir_persistence.py +++ b/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):