core: consider lockscreen to be a separate homescreen

this involves some changes to the workflow defaults:

* workflow.start_default() takes no arguments
* workflow.set_default() (originally replace_default) configures the
  default that will be started by next call to start_default().
  The intended usecase is to set_default() first and then start it
  separately.
* apps.base.set_homescreen() factors out the logic originally in
  main.py, that decides which homescreen should be launched. This uses
  set_default() call. start_default() is then used explicitly in main.py

core/src/apps/base.py

@@ -112,14 +112,33 @@ async def handle_Ping(ctx: wire.Context, msg: Ping) -> Success:
 ALLOW_WHILE_LOCKED = (MessageType.Initialize, MessageType.GetFeatures)
 
 
+def set_homescreen() -> None:
+    if not config.is_unlocked():
+        from apps.homescreen.lockscreen import lockscreen
+
+        workflow.set_default(lockscreen)
+
+    elif storage.recovery.is_in_progress():
+        from apps.management.recovery_device.homescreen import recovery_homescreen
+
+        workflow.set_default(recovery_homescreen)
+
+    else:
+        from apps.homescreen.homescreen import homescreen
+
+        workflow.set_default(homescreen)
+
+
 def lock_device() -> None:
     config.lock()
+    set_homescreen()
     wire.find_handler = get_pinlocked_handler
 
 
 async def unlock_device(ctx: wire.GenericContext = wire.DUMMY_CONTEXT) -> None:
     await verify_user_pin(ctx)
     # verify_user_pin will raise if the PIN was invalid
+    set_homescreen()
     wire.find_handler = wire.find_registered_workflow_handler
 
 

core/src/apps/homescreen/__init__.py

@@ -6,10 +6,9 @@ if False:
 
 
 class HomescreenBase(ui.Layout):
-    def __init__(self, lock_label = "Locked") -> None:
+    def __init__(self) -> None:
         self.repaint = True
 
-        self.lock_label = lock_label
         self.label = storage.device.get_label() or "My Trezor"
         self.image = storage.device.get_homescreen() or res.load(
             "apps/homescreen/res/bg.toif"
@@ -19,18 +18,6 @@ class HomescreenBase(ui.Layout):
         ui.display.avatar(48, 48 - 10, self.image, ui.WHITE, ui.BLACK)
         ui.display.text_center(ui.WIDTH // 2, 220, self.label, ui.BOLD, ui.FG, ui.BG)
 
-    def render_lock(self) -> None:
-        ui.display.bar_radius(40, 100, 160, 40, ui.TITLE_GREY, ui.BG, 4)
-        ui.display.bar_radius(42, 102, 156, 36, ui.BG, ui.TITLE_GREY, 4)
-        ui.display.text_center(
-            ui.WIDTH // 2, 128, self.lock_label, ui.BOLD, ui.TITLE_GREY, ui.BG
-        )
-
-        ui.display.text_center(
-            ui.WIDTH // 2 + 10, 220, "Tap to unlock", ui.BOLD, ui.TITLE_GREY, ui.BG
-        )
-        ui.display.icon(45, 202, res.load(ui.ICON_CLICK), ui.TITLE_GREY, ui.BG)
-
     def dispatch(self, event: int, x: int, y: int) -> None:
         if event is ui.RENDER and self.repaint:
             self.repaint = False

core/src/apps/homescreen/homescreen.py

@@ -12,7 +12,7 @@ async def homescreen() -> None:
 class Homescreen(HomescreenBase):
     def __init__(self) -> None:
         super().__init__()
-        if config.is_unlocked() and not storage.is_initialized():
+        if not storage.is_initialized():
             self.label = "Go to trezor.io/start"
 
     def render_warning(self) -> None:
@@ -30,5 +30,3 @@ class Homescreen(HomescreenBase):
     def on_render(self) -> None:
         self.render_warning()
         self.render_homescreen()
-        if not config.is_unlocked():
-            self.render_lock()

core/src/apps/homescreen/lockscreen.py

@@ -1,9 +1,35 @@
-from trezor import ui
+from trezor import res, ui
 
 from . import HomescreenBase
 
 
+async def lockscreen() -> None:
+    from apps.common.request_pin import can_lock_device
+    from apps.base import unlock_device
+
+    if can_lock_device():
+        await Lockscreen()
+
+    await unlock_device()
+
+
 class Lockscreen(HomescreenBase):
+    def __init__(self, lock_label: str = "Locked") -> None:
+        self.lock_label = lock_label
+        super().__init__()
+
+    def render_lock(self) -> None:
+        ui.display.bar_radius(40, 100, 160, 40, ui.TITLE_GREY, ui.BG, 4)
+        ui.display.bar_radius(42, 102, 156, 36, ui.BG, ui.TITLE_GREY, 4)
+        ui.display.text_center(
+            ui.WIDTH // 2, 128, self.lock_label, ui.BOLD, ui.TITLE_GREY, ui.BG
+        )
+
+        ui.display.text_center(
+            ui.WIDTH // 2 + 10, 220, "Tap to unlock", ui.BOLD, ui.TITLE_GREY, ui.BG
+        )
+        ui.display.icon(45, 202, res.load(ui.ICON_CLICK), ui.TITLE_GREY, ui.BG)
+
     def on_render(self) -> None:
         self.render_homescreen()
         self.render_lock()

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

@@ -63,7 +63,7 @@ async def recovery_device(ctx: wire.Context, msg: RecoveryDevice) -> Success:
     storage.recovery.set_in_progress(True)
     storage.recovery.set_dry_run(bool(msg.dry_run))
 
-    workflow.replace_default(recovery_homescreen)
+    workflow.set_default(recovery_homescreen)
     return await recovery_process(ctx)
 
 

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

@@ -24,7 +24,7 @@ if False:
 
 async def recovery_homescreen() -> None:
     if not storage.recovery.is_in_progress():
-        workflow.replace_default(homescreen)
+        workflow.set_default(homescreen)
         return
 
     # recovery process does not communicate on the wire

core/src/main.py

@@ -6,7 +6,6 @@ import boot  # noqa: F401
 # prepare the USB interfaces, but do not connect to the host yet
 import usb
 
-import storage.recovery
 from trezor import loop, utils, wire, workflow
 
 # start the USB
@@ -57,14 +56,8 @@ def _boot_apps() -> None:
         apps.debug.boot()
 
     # run main event loop and specify which screen is the default
-    if storage.recovery.is_in_progress():
+    apps.base.set_homescreen()
-        from apps.management.recovery_device.homescreen import recovery_homescreen
+    workflow.start_default()
-
-        workflow.start_default(recovery_homescreen)
-    else:
-        from apps.homescreen.homescreen import homescreen
-
-        workflow.start_default(homescreen)
 
 
 # initialize the wire codec

core/src/trezor/workflow.py

@@ -42,7 +42,7 @@ def on_close(workflow: loop.Task) -> None:
     if not tasks and default_constructor:
         # If no workflows are running, we should create a new default workflow
         # and run it.
-        start_default(default_constructor)
+        start_default()
     if __debug__:
         # In debug builds, we dump a memory info right after a workflow is
         # finished.
@@ -50,18 +50,19 @@ def on_close(workflow: loop.Task) -> None:
             micropython.mem_info()
 
 
-def start_default(constructor: Callable[[], loop.Task]) -> None:
+def start_default() -> None:
-    """Start a default workflow, created from `constructor`.
+    """Start a default workflow.
 
-    If a default task is already running, nothing will happen. Use `replace_default`
+    Use `set_default` to set the default workflow constructor.
-    to set up a new default task for the next run.
+    If a default task is already running, nothing will happen.
     """
     global default_task
     global default_constructor
 
+    assert default_constructor is not None
+
     if not default_task:
-        default_constructor = constructor
+        default_task = default_constructor()
-        default_task = constructor()
         if __debug__:
             log.debug(__name__, "start default: %s", default_task)
         # Schedule the default task.  Because the task can complete on its own,
@@ -72,7 +73,7 @@ def start_default(constructor: Callable[[], loop.Task]) -> None:
             log.debug(__name__, "default already started")
 
 
-def replace_default(constructor: Callable[[], loop.Task]) -> None:
+def set_default(constructor: Callable[[], loop.Task]) -> None:
     """Configure a default workflow, which will be started next time it is needed."""
     global default_constructor
     if __debug__:
@@ -111,7 +112,7 @@ def _finalize_default(task: loop.Task, value: Any) -> None:
             # finalizer, so when this function finished, nothing will be running.
             # We must schedule a new instance of the default now.
             if default_constructor is not None:
-                start_default(default_constructor)
+                start_default()
             else:
                 raise RuntimeError  # no tasks and no default constructor