From ebe883e4a9920a151a7901522c6c3c183b1e35e4 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Wed, 24 Jul 2019 14:39:14 +0200
Subject: [PATCH] core: support rebooting into default mode

---
 core/src/apps/homescreen/__init__.py          |  9 +++--
 .../management/recovery_device/homescreen.py  |  7 ++--
 core/src/main.py                              | 37 ++++++++++---------
 core/src/trezor/wire/__init__.py              |  5 +++
 4 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/core/src/apps/homescreen/__init__.py b/core/src/apps/homescreen/__init__.py
index 536edd8b5..37e9fad33 100644
--- a/core/src/apps/homescreen/__init__.py
+++ b/core/src/apps/homescreen/__init__.py
@@ -73,9 +73,10 @@ async def handle_Ping(ctx: wire.Context, msg: Ping) -> Success:
     return Success(message=msg.message)
 
 
-def boot() -> None:
+def boot(features_only: bool = False) -> None:
     register(MessageType.Initialize, protobuf_workflow, handle_Initialize)
     register(MessageType.GetFeatures, protobuf_workflow, handle_GetFeatures)
-    register(MessageType.Cancel, protobuf_workflow, handle_Cancel)
-    register(MessageType.ClearSession, protobuf_workflow, handle_ClearSession)
-    register(MessageType.Ping, protobuf_workflow, handle_Ping)
+    if not features_only:
+        register(MessageType.Cancel, protobuf_workflow, handle_Cancel)
+        register(MessageType.ClearSession, protobuf_workflow, handle_ClearSession)
+        register(MessageType.Ping, protobuf_workflow, handle_Ping)
diff --git a/core/src/apps/management/recovery_device/homescreen.py b/core/src/apps/management/recovery_device/homescreen.py
index fde4c8d3d..73b98e6d1 100644
--- a/core/src/apps/management/recovery_device/homescreen.py
+++ b/core/src/apps/management/recovery_device/homescreen.py
@@ -15,10 +15,11 @@ async def recovery_homescreen() -> None:
     ctx = wire.DummyContext()
     try:
         await recovery_process(ctx)
-    except Exception:
-        # clear the loop state, so loop.run will exit and the device is soft-rebooted
+    finally:
+        # clear the loop state, so loop.run will exit
         loop.clear()
-        raise
+        # clear the registered wire handlers to avoid conflicts
+        wire.clear()
 
 
 async def recovery_process(ctx: wire.Context) -> Success:
diff --git a/core/src/main.py b/core/src/main.py
index cef80d253..80141b7bc 100644
--- a/core/src/main.py
+++ b/core/src/main.py
@@ -6,12 +6,9 @@ import boot  # noqa: F401
 # prepare the USB interfaces, but do not connect to the host yet
 import usb
 
-from trezor import wire, utils
+from trezor import utils
 
-# initialize the wire codec and start the USB
-wire.setup(usb.iface_wire)
-if __debug__:
-    wire.setup(usb.iface_debug)
+# start the USB
 usb.bus.open()
 
 # switch into unprivileged mode, as we don't need the extra permissions anymore
@@ -20,15 +17,10 @@ utils.set_mode_unprivileged()
 
 def _boot_recovery():
     # load applications
-    import apps.management
+    import apps.homescreen
 
     # boot applications
-    apps.management.boot()
-
-    if __debug__:
-        import apps.debug
-
-        apps.debug.boot()
+    apps.homescreen.boot(features_only=True)
 
     from apps.management.recovery_device.homescreen import recovery_homescreen
 
@@ -79,11 +71,20 @@ def _boot_default():
     workflow.startdefault(homescreen)
 
 
-from trezor import loop, workflow
+from trezor import loop, wire, workflow
 from apps.common.storage import recovery
 
-if recovery.is_in_progress():
-    _boot_recovery()
-else:
-    _boot_default()
-loop.run()
+while True:
+    # initialize the wire codec
+    wire.setup(usb.iface_wire)
+    if __debug__:
+        wire.setup(usb.iface_debug)
+
+    # boot either in recovery or default mode
+    if recovery.is_in_progress():
+        _boot_recovery()
+    else:
+        _boot_default()
+    loop.run()
+
+    # loop is empty, reboot
diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py
index 2631922c8..6674e6794 100644
--- a/core/src/trezor/wire/__init__.py
+++ b/core/src/trezor/wire/__init__.py
@@ -58,6 +58,11 @@ def setup(iface: WireInterface) -> None:
     loop.schedule(session_handler(iface, codec_v1.SESSION_ID))
 
 
+def clear() -> None:
+    """Remove all registered handlers."""
+    workflow_handlers.clear()
+
+
 class DummyContext:
     async def call(*argv):
         pass