From 005d69e1c213f1d5c5dee0a398f1becf7ee733a0 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Sat, 6 Apr 2024 15:23:36 +0200
Subject: [PATCH] feat(core/debug): better ButtonAck deadlock detection

---
 core/src/apps/debug/__init__.py | 13 +++++-
 tests/input_flows.py            | 76 ++-------------------------------
 2 files changed, 14 insertions(+), 75 deletions(-)

diff --git a/core/src/apps/debug/__init__.py b/core/src/apps/debug/__init__.py
index 163407b82..d13fee970 100644
--- a/core/src/apps/debug/__init__.py
+++ b/core/src/apps/debug/__init__.py
@@ -4,6 +4,7 @@ if not __debug__:
     halt("debug mode inactive")
 
 if __debug__:
+    from micropython import const
     from typing import TYPE_CHECKING
 
     import trezorui2
@@ -36,7 +37,9 @@ if __debug__:
 
     REFRESH_INDEX = 0
 
-    _DEADLOCK_DETECT_SLEEP = loop.sleep(2000)
+    _DEADLOCK_WAIT_TRIES = const(5)
+    _DEADLOCK_SLEEP_MS = const(2000)
+    _DEADLOCK_DETECT_SLEEP = loop.sleep(_DEADLOCK_SLEEP_MS)
 
     def screenshot() -> bool:
         if storage.save_screen:
@@ -51,8 +54,14 @@ if __debug__:
     def notify_layout_change(layout: Layout | None) -> None:
         layout_change_chan.put(layout, replace=True)
 
-    def wait_until_layout_is_running() -> Awaitable[None]:  # type: ignore [awaitable-is-generator]
+    def wait_until_layout_is_running(tries: int | None = _DEADLOCK_WAIT_TRIES) -> Awaitable[None]:  # type: ignore [awaitable-is-generator]
+        counter = 0
         while ui.CURRENT_LAYOUT is None:
+            counter += 1
+            if tries is not None and counter > tries:
+                raise wire.FirmwareError(
+                    "layout deadlock detected (did you send a ButtonAck?)"
+                )
             yield
 
     async def return_layout_change(
diff --git a/tests/input_flows.py b/tests/input_flows.py
index 70e99f5a4..accbb617d 100644
--- a/tests/input_flows.py
+++ b/tests/input_flows.py
@@ -863,22 +863,7 @@ class InputFlowSignTxInformationReplacement(InputFlowBase):
         self.debug.press_right()
         self.debug.press_right()
 
-    def input_flow_t3t1(self) -> BRGeneratorType:
-        yield  # confirm txid
-        self.debug.press_yes()
-        yield  # confirm address
-        self.debug.press_yes()
-        # go back to address
-        self.debug.press_no()
-        # confirm address
-        self.debug.press_yes()
-        yield  # confirm amount
-        self.debug.press_yes()
-
-        yield  # transaction summary, press info
-        self.debug.click(buttons.CORNER_BUTTON)
-        self.debug.click(buttons.CORNER_BUTTON)
-        self.debug.press_yes()
+    input_flow_t3t1 = input_flow_tt
 
 
 def lock_time_input_flow_tt(
@@ -1283,30 +1268,7 @@ class InputFlowSlip39BasicBackup(InputFlowBase):
         assert br.code == B.Success
         self.debug.press_yes()
 
-    def input_flow_t3t1(self) -> BRGeneratorType:
-        yield  # 1. Checklist
-        self.debug.press_yes()
-        if self.click_info:
-            yield from click_info_button_tt(self.debug)
-        yield  # 2. Number of shares (5)
-        self.debug.press_yes()
-        yield  # 3. Checklist
-        self.debug.press_yes()
-        if self.click_info:
-            yield from click_info_button_tt(self.debug)
-        yield  # 4. Threshold (3)
-        self.debug.press_yes()
-        yield  # 5. Checklist
-        self.debug.press_yes()
-        yield  # 6. Confirm show seeds
-        self.debug.press_yes()
-
-        # Mnemonic phrases
-        self.mnemonics = yield from load_5_shares(self.debug)
-
-        br = yield  # Confirm backup
-        assert br.code == B.Success
-        self.debug.press_yes()
+    input_flow_t3t1 = input_flow_tt
 
 
 class InputFlowSlip39BasicResetRecovery(InputFlowBase):
@@ -1468,39 +1430,7 @@ class InputFlowSlip39AdvancedBackup(InputFlowBase):
         assert br.code == B.Success
         self.debug.press_yes()
 
-    def input_flow_t3t1(self) -> BRGeneratorType:
-        yield  # 1. Checklist
-        self.debug.press_yes()
-        if self.click_info:
-            yield from click_info_button_tt(self.debug)
-        yield  # 2. Set and confirm group count
-        self.debug.press_yes()
-        yield  # 3. Checklist
-        self.debug.press_yes()
-        if self.click_info:
-            yield from click_info_button_tt(self.debug)
-        yield  # 4. Set and confirm group threshold
-        self.debug.press_yes()
-        yield  # 5. Checklist
-        self.debug.press_yes()
-        for _ in range(5):  # for each of 5 groups
-            if self.click_info:
-                yield from click_info_button_tt(self.debug)
-            yield  # Set & Confirm number of shares
-            self.debug.press_yes()
-            if self.click_info:
-                yield from click_info_button_tt(self.debug)
-            yield  # Set & confirm share threshold value
-            self.debug.press_yes()
-        yield  # Confirm show seeds
-        self.debug.press_yes()
-
-        # Mnemonic phrases - show & confirm shares for all groups
-        self.mnemonics = yield from load_5_groups_5_shares(self.debug)
-
-        br = yield  # Confirm backup
-        assert br.code == B.Success
-        self.debug.press_yes()
+    input_flow_t3t1 = input_flow_tt
 
 
 class InputFlowSlip39AdvancedResetRecovery(InputFlowBase):