core: rework wait_layout()

The original wait_layout was unreliable, because there are no guarantees
re order of arrival of the respective events. Still, TT's event handling
is basically deterministic, so as long as the host sent its messages
close enough to each other, the order worked out.

This is no longer the case with the introduction of loop.spawn: TT's
behavior is still deterministic, but now ButtonAck is processed *before*
the corresponding wait_layout, so the waiting side waits forever.

In the new process, the host must first register to receive layout
events, and then receives all of them (so the number of calls to
wait_layout must match the number of layout changes).

DebugLinkWatchLayout message must be version-gated, because of an
unfortunate collection of bugs in previous versions wrt unknown message
handling; and this interests us because upgrade-tests are using
wait_layout feature.

common/protob/messages-debug.proto

@@ -185,3 +185,14 @@ message DebugLinkEraseSdCard {
     optional bool format = 1;  // if true, the card will be formatted to FAT32.
                                // if false, it will be all 0xFF bytes.
 }
+
+
+/**
+ * Request: Start or stop tracking layout changes
+ * @start
+ * @next Success
+ */
+message DebugLinkWatchLayout {
+    optional bool watch = 1;  // if true, start watching layout.
+                              // if false, stop.
+}

common/protob/messages.proto

@@ -115,6 +115,7 @@ enum MessageType {
     MessageType_DebugLinkRecordScreen = 9003 [(wire_debug_in) = true];
     MessageType_DebugLinkShowText = 9004 [(wire_debug_in) = true];
     MessageType_DebugLinkEraseSdCard = 9005 [(wire_debug_in) = true];
+    MessageType_DebugLinkWatchLayout = 9006 [(wire_debug_in) = true];
 
     // Ethereum
     MessageType_EthereumGetPublicKey = 450 [(wire_in) = true];

core/src/apps/debug/__init__.py

@@ -18,6 +18,7 @@ if __debug__:
         from trezor.messages.DebugLinkReseedRandom import DebugLinkReseedRandom
         from trezor.messages.DebugLinkState import DebugLinkState
         from trezor.messages.DebugLinkEraseSdCard import DebugLinkEraseSdCard
+        from trezor.messages.DebugLinkWatchLayout import DebugLinkWatchLayout
 
     save_screen = False
     save_screen_directory = "."
@@ -37,6 +38,7 @@ if __debug__:
 
     layout_change_chan = loop.chan()
     current_content = None  # type: Optional[List[str]]
+    watch_layout_changes = False
 
     def screenshot() -> bool:
         if save_screen:
@@ -47,7 +49,7 @@ if __debug__:
     def notify_layout_change(layout: ui.Layout) -> None:
         global current_content
         current_content = layout.read_content()
-        if layout_change_chan.takers:
+        if watch_layout_changes:
             layout_change_chan.publish(current_content)
 
     async def debuglink_decision_dispatcher() -> None:
@@ -77,6 +79,15 @@ if __debug__:
         content = await layout_change_chan.take()
         await ctx.write(DebugLinkLayout(lines=content))
 
+    async def dispatch_DebugLinkWatchLayout(  # type: ignore
+        ctx: wire.Context, msg: DebugLinkWatchLayout
+    ) -> Success:
+        global watch_layout_changes
+        layout_change_chan.putters.clear()
+        watch_layout_changes = msg.watch
+        log.debug(__name__, "Watch layout changes: {}".format(watch_layout_changes))
+        return Success()
+
     async def dispatch_DebugLinkDecision(
         ctx: wire.Context, msg: DebugLinkDecision
     ) -> None:
@@ -171,3 +182,4 @@ if __debug__:
         wire.register(MessageType.DebugLinkReseedRandom, dispatch_DebugLinkReseedRandom)
         wire.register(MessageType.DebugLinkRecordScreen, dispatch_DebugLinkRecordScreen)
         wire.register(MessageType.DebugLinkEraseSdCard, dispatch_DebugLinkEraseSdCard)
+        wire.register(MessageType.DebugLinkWatchLayout, dispatch_DebugLinkWatchLayout)

core/src/trezor/messages/DebugLinkWatchLayout.py

@@ -0,0 +1,26 @@
+# Automatically generated by pb2py
+# fmt: off
+import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class DebugLinkWatchLayout(p.MessageType):
+    MESSAGE_WIRE_TYPE = 9006
+
+    def __init__(
+        self,
+        watch: bool = None,
+    ) -> None:
+        self.watch = watch
+
+    @classmethod
+    def get_fields(cls) -> Dict:
+        return {
+            1: ('watch', p.BoolType, 0),
+        }

core/src/trezor/messages/MessageType.py

@@ -79,6 +79,7 @@ DebugLinkReseedRandom = 9002  # type: Literal[9002]
 DebugLinkRecordScreen = 9003  # type: Literal[9003]
 DebugLinkShowText = 9004  # type: Literal[9004]
 DebugLinkEraseSdCard = 9005  # type: Literal[9005]
+DebugLinkWatchLayout = 9006  # type: Literal[9006]
 if not utils.BITCOIN_ONLY:
     EthereumGetPublicKey = 450  # type: Literal[450]
     EthereumPublicKey = 451  # type: Literal[451]

legacy/firmware/protob/Makefile

@@ -3,7 +3,7 @@ Q := @
 endif
 
 SKIPPED_MESSAGES := Binance Cardano DebugMonero Eos Monero Ontology Ripple SdProtect Tezos WebAuthn \
-	DebugLinkRecordScreen DebugLinkReseedRandom DebugLinkShowText DebugLinkEraseSdCard
+	DebugLinkRecordScreen DebugLinkReseedRandom DebugLinkShowText DebugLinkEraseSdCard DebugLinkWatchLayout
 
 ifeq ($(BITCOIN_ONLY), 1)
 SKIPPED_MESSAGES += Ethereum Lisk NEM Stellar

python/src/trezorlib/debuglink.py

@@ -87,6 +87,15 @@ class DebugLink:
         obj = self._call(messages.DebugLinkGetState(wait_layout=True))
         return layout_lines(obj.layout_lines)
 
+    def watch_layout(self, watch: bool) -> None:
+        """Enable or disable watching layouts.
+        If disabled, wait_layout will not work.
+
+        The message is missing on T1. Use `TrezorClientDebugLink.watch_layout` for
+        cross-version compatibility.
+        """
+        self._call(messages.DebugLinkWatchLayout(watch=watch))
+
     def encode_pin(self, pin, matrix=None):
         """Transform correct PIN according to the displayed matrix."""
         if matrix is None:
@@ -335,9 +344,25 @@ class TrezorClientDebugLink(TrezorClient):
         self.ui.input_flow = input_flow
         input_flow.send(None)  # start the generator
 
+    def watch_layout(self, watch: bool) -> None:
+        """Enable or disable watching layout changes.
+
+        Happens implicitly in a `with client` block.
+
+        Since trezor-core v2.3.2, it is necessary to call `watch_layout()` before
+        using `debug.wait_layout()`, otherwise layout changes are not reported.
+        """
+        if self.version >= (2, 3, 2):
+            # version check is necessary because otherwise we cannot reliably detect
+            # whether and where to wait for reply:
+            # - T1 reports unknown debuglink messages on the wirelink
+            # - TT < 2.3.0 does not reply to unknown debuglink messages due to a bug
+            self.debug.watch_layout(watch)
+
     def __enter__(self):
         # For usage in with/expected_responses
         self.in_with_statement += 1
+        self.watch_layout(True)
         return self
 
     def __exit__(self, _type, value, traceback):
@@ -362,6 +387,7 @@ class TrezorClientDebugLink(TrezorClient):
             self.expected_responses = None
             self.current_response = None
             self.ui.clear()
+            self.watch_layout(False)
 
         return False
 

python/src/trezorlib/messages/DebugLinkWatchLayout.py

@@ -0,0 +1,26 @@
+# Automatically generated by pb2py
+# fmt: off
+from .. import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class DebugLinkWatchLayout(p.MessageType):
+    MESSAGE_WIRE_TYPE = 9006
+
+    def __init__(
+        self,
+        watch: bool = None,
+    ) -> None:
+        self.watch = watch
+
+    @classmethod
+    def get_fields(cls) -> Dict:
+        return {
+            1: ('watch', p.BoolType, 0),
+        }

python/src/trezorlib/messages/MessageType.py

@@ -77,6 +77,7 @@ DebugLinkReseedRandom = 9002  # type: Literal[9002]
 DebugLinkRecordScreen = 9003  # type: Literal[9003]
 DebugLinkShowText = 9004  # type: Literal[9004]
 DebugLinkEraseSdCard = 9005  # type: Literal[9005]
+DebugLinkWatchLayout = 9006  # type: Literal[9006]
 EthereumGetPublicKey = 450  # type: Literal[450]
 EthereumPublicKey = 451  # type: Literal[451]
 EthereumGetAddress = 56  # type: Literal[56]

python/src/trezorlib/messages/__init__.py

@@ -53,6 +53,7 @@ from .DebugLinkShowText import DebugLinkShowText
 from .DebugLinkShowTextItem import DebugLinkShowTextItem
 from .DebugLinkState import DebugLinkState
 from .DebugLinkStop import DebugLinkStop
+from .DebugLinkWatchLayout import DebugLinkWatchLayout
 from .DebugMoneroDiagAck import DebugMoneroDiagAck
 from .DebugMoneroDiagRequest import DebugMoneroDiagRequest
 from .Deprecated_PassphraseStateAck import Deprecated_PassphraseStateAck