chore(core) disallow other messages while backup is unlocked

core/src/apps/base.py

@@ -374,7 +374,7 @@ def set_homescreen() -> None:
 def lock_device(interrupt_workflow: bool = True) -> None:
     if config.has_pin():
         config.lock()
-        wire.find_handler = get_pinlocked_handler
+        wire.find_handler = _get_pinlocked_handler
         set_homescreen()
         if interrupt_workflow:
             workflow.close_others()
@@ -413,7 +413,7 @@ async def unlock_device() -> None:
     wire.find_handler = workflow_handlers.find_registered_handler
 
 
-def get_pinlocked_handler(
+def _get_pinlocked_handler(
     iface: wire.WireInterface, msg_type: int
 ) -> wire.Handler[wire.Msg] | None:
     orig_handler = workflow_handlers.find_registered_handler(iface, msg_type)
@@ -436,6 +436,28 @@ def get_pinlocked_handler(
     return wrapper
 
 
+def _get_backup_handler(
+    iface: wire.WireInterface, msg_type: int
+) -> wire.Handler[wire.Msg] | None:
+    orig_handler = workflow_handlers.find_registered_handler(iface, msg_type)
+    if orig_handler is None:
+        return None
+
+    if __debug__:
+        import usb
+
+        if iface is usb.iface_debug:
+            return orig_handler
+
+    if msg_type in workflow.ALLOW_WHILE_REPEATED_BACKUP_UNLOCKED:
+        return orig_handler
+
+    async def wrapper(_msg: wire.Msg) -> protobuf.MessageType:
+        raise wire.ProcessError("Operation not allowed when in repeated backup state")
+
+    return wrapper
+
+
 # this function is also called when handling ApplySettings
 def reload_settings_from_storage() -> None:
     from trezor import ui
@@ -468,7 +490,12 @@ def boot() -> None:
         workflow_handlers.register(msg_type, handler)
 
     reload_settings_from_storage()
-    if config.is_unlocked():
+    if (
+        storage_cache.get(storage_cache.APP_RECOVERY_REPEATED_BACKUP_UNLOCKED)
+        == b"\x01"
+    ):
+        wire.find_handler = _get_backup_handler
+    elif config.is_unlocked():
         wire.find_handler = workflow_handlers.find_registered_handler
     else:
-        wire.find_handler = get_pinlocked_handler
+        wire.find_handler = _get_pinlocked_handler

core/src/apps/management/backup_device.py

@@ -15,6 +15,7 @@ async def backup_device(msg: BackupDevice) -> Success:
     from trezor import wire
     from trezor.messages import Success
 
+    from apps import workflow_handlers
     from apps.common import mnemonic
 
     from .reset_device import backup_seed, backup_slip39_custom, layout
@@ -63,5 +64,6 @@ async def backup_device(msg: BackupDevice) -> Success:
     await layout.show_backup_success()
 
     storage_cache.delete(storage_cache.APP_RECOVERY_REPEATED_BACKUP_UNLOCKED)
+    wire.find_handler = workflow_handlers.find_registered_handler
 
     return Success(message="Seed successfully backed up")

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

@@ -60,6 +60,7 @@ async def _continue_repeated_backup() -> None:
     from trezor.ui.layouts import confirm_action
     from trezor.wire import ActionCancelled
 
+    from apps import workflow_handlers
     from apps.common import mnemonic
     from apps.homescreen import homescreen
     from apps.management.reset_device import backup_seed
@@ -88,6 +89,7 @@ async def _continue_repeated_backup() -> None:
         workflow.set_default(homescreen)
     finally:
         storage_cache.delete(storage_cache.APP_RECOVERY_REPEATED_BACKUP_UNLOCKED)
+        wire.find_handler = workflow_handlers.find_registered_handler
         storage_recovery.end_progress()
 
 

core/src/trezor/workflow.py

@@ -30,6 +30,15 @@ ALLOW_WHILE_LOCKED = (
 )
 
 
+ALLOW_WHILE_REPEATED_BACKUP_UNLOCKED = (
+    MessageType.Initialize,
+    MessageType.GetFeatures,
+    MessageType.EndSession,
+    MessageType.BackupDevice,
+    MessageType.WipeDevice,
+)
+
+
 # Set of workflow tasks.  Multiple workflows can be running at the same time.
 tasks: set[loop.spawn] = set()