feat(core): handle power button in FW loop

- power button is handled before the event is passed to rust layouts

[no changelog]

Author: obrusvit

core/src/apps/base.py

@@ -2,7 +2,7 @@
 
 import storage.device as storage_device
 from storage.cache_common import APP_COMMON_BUSY_DEADLINE_MS, APP_COMMON_SEED
-from trezor import TR, config, utils, wire, workflow
+from trezor import TR, config, io, utils, wire, workflow
 from trezor.enums import HomescreenFormat, MessageType
 from trezor.messages import Success, UnlockPath
 from trezor.ui.layouts import confirm_action
@@ -396,6 +396,18 @@ def set_homescreen() -> None:
         set_default(homescreen)
 
 
+if utils.USE_POWER_MANAGER:
+
+    def suspend_device() -> None:
+        """Suspends the device when the power button is pressed."""
+        if config.has_pin() and config.is_unlocked():
+            lock_device(interrupt_workflow=True)
+        else:
+            set_homescreen()
+            workflow.close_others()
+        io.pm.suspend()
+
+
 def lock_device(interrupt_workflow: bool = True) -> None:
     if config.has_pin():
         config.lock()

core/src/trezor/ui/__init__.py

@@ -374,34 +374,59 @@ def create_tasks(self) -> Iterator[loop.Task]:
         """Set up background tasks for a layout.
 
         Called from `start()`. Creates and yields a list of background tasks, typically
-        event handlers for different interfaces.
+        event handlers for different interfaces. Event handlers are enabled based on build options to prevent keeping stale events in the event queue.
 
         Override and then `yield from super().create_tasks()` to add more tasks."""
         if utils.USE_BUTTON:
-            yield self._handle_input_iface(io.BUTTON, self.layout.button_event)
+            yield self._handle_button_events()
         if utils.USE_TOUCH:
-            yield self._handle_input_iface(io.TOUCH, self.layout.touch_event)
+            yield self._handle_touch_events()
         if utils.USE_BLE:
-            # most layouts don't care but we don't want to keep stale events in the queue
             yield self._handle_ble_events()
         if utils.USE_POWER_MANAGER:
             yield self._handle_power_manager()
 
-    def _handle_input_iface(
-        self, iface: int, event_call: Callable[..., LayoutState | None]
-    ) -> Generator:
-        """Task that is waiting for the user input."""
-        touch = loop.wait(iface)
-        try:
-            while True:
-                # Using `yield` instead of `await` to avoid allocations.
-                event = yield touch
-                workflow.idle_timer.touch()
-                self._event(event_call, *event)
-        except Shutdown:
-            return
-        finally:
-            touch.close()
+    if utils.USE_BUTTON:
+
+        def _handle_button_events(self) -> Generator:
+            """Task that is waiting for the user button input."""
+            button = loop.wait(io.BUTTON)
+            try:
+                while True:
+                    # Using `yield` instead of `await` to avoid allocations.
+                    event = yield button
+                    if utils.USE_POWER_MANAGER:
+                        event_type, event_button = event
+                        # check for POWER_BUTTON (2), BUTTON_UP (0)
+                        if event_button == 2 and event_type == 0:
+                            from apps.base import suspend_device
+
+                            if __debug__:
+                                log.info(__name__, "suspend_device")
+                            suspend_device()
+                            break
+                    workflow.idle_timer.touch()
+                    self._event(self.layout.button_event, *event)
+            except Shutdown:
+                return
+            finally:
+                button.close()
+
+    if utils.USE_TOUCH:
+
+        def _handle_touch_events(self) -> Generator:
+            """Task that is waiting for the user touch input."""
+            touch = loop.wait(io.TOUCH)
+            try:
+                while True:
+                    # Using `yield` instead of `await` to avoid allocations.
+                    event = yield touch
+                    workflow.idle_timer.touch()
+                    self._event(self.layout.touch_event, *event)
+            except Shutdown:
+                return
+            finally:
+                touch.close()
 
     async def _handle_usb_iface(self) -> None:
         if self.context is None: