fix(core): keep reading when writes are blocked on THP

Author: romanz

core/.changelog.d/6506.fixed

@@ -0,0 +1 @@
+[T3W1] Avoid THP deadlock over USB.

core/src/trezor/wire/thp/channel.py

@@ -41,7 +41,7 @@
 
 if TYPE_CHECKING:
     from buffer_types import AnyBuffer, AnyBytes
-    from typing import Any, Awaitable, Callable
+    from typing import Any, Awaitable, Callable, Generator
 
     from trezor.messages import ThpPairingCredential
     from trezor.wire import WireInterface
@@ -436,49 +436,56 @@ async def write_encrypted_payload(self, ctrl_byte: int, payload: AnyBytes) -> No
 
         header = PacketHeader(ctrl_byte, self.get_channel_id_int(), payload_len)
 
-        async def _write_loop() -> None:
-            """Send the payload and wait for an ACK with retransmissions."""
+        ack_latency_ms = self.channel_cache.get_int(CHANNEL_ACK_LATENCY_MS) or 0
 
-            ack_latency_ms = self.channel_cache.get_int(CHANNEL_ACK_LATENCY_MS) or 0
-            if __debug__:
-                self._log(f"Sending {len(payload)} bytes, latency: {ack_latency_ms} ms")
+        # ACK is needed before sending more data
+        ABP.set_sending_allowed(self.channel_cache, False)
 
-            # ACK is needed before sending more data
-            ABP.set_sending_allowed(self.channel_cache, False)
+        # allows preempting this channel, if another channel becomes active
+        self.last_write_ms = utime.ticks_ms()
 
-            # allows preempting this channel, if another channel becomes active
-            self.last_write_ms = utime.ticks_ms()
+        def _write_loop() -> Generator[Any, Any, None]:
+            """
+            Retransmit the payload (with increasing delay), raising `Timeout` in the end.
 
-            for i in range(_MAX_RETRANSMISSION_COUNT):
-                await self._write_payload_once(header, payload)
+            This task is spawned concurrently with `_wait_for_ack()` using `loop.race()`,
+            so it will be cancelled when the expected ACK is received.
+            """
+            if __debug__:
+                self._log(f"Sending {len(payload)} bytes, latency: {ack_latency_ms} ms")
 
+            for i in range(_MAX_RETRANSMISSION_COUNT):
+                # Try to send the payload (split into packets), or raise if transport is blocked
+                yield from self._write_payload_once(header, payload)
                 # Channel's estimated latency + a variable delay (from 200ms till ~3.52s)
-                timeout_ms = ack_latency_ms + round(10300 - 1010000 / (100 + i))
-                try:
-                    # wait and return after receiving an ACK, or raise in case of an unexpected message.
-                    await self.recv_payload(
-                        expected_ctrl_byte=None, timeout_ms=timeout_ms
-                    )
-                except Timeout:
-                    if __debug__:
-                        log.warning(__name__, "Retransmit after %d ms", timeout_ms)
-                    continue
-
-                ack_latency_ms = utime.ticks_diff(utime.ticks_ms(), self.last_write_ms)
-                # Limit estimated latency to avoid integer overflows and too long delays
-                ack_latency_ms = max(0, min(800, ack_latency_ms))
-                self.channel_cache.set_int(CHANNEL_ACK_LATENCY_MS, ack_latency_ms)
-
-                # `ABP.set_sending_allowed()` will be called after a valid ACK
-                if ABP.is_sending_allowed(self.channel_cache):
-                    return
+                delay_ms = ack_latency_ms + round(10300 - 1010000 / (100 + i))
+                yield from sleep(delay_ms)
+                if __debug__:
+                    log.warning(__name__, "Retransmit after %d ms", delay_ms)
 
             # restart event loop due to unresponsive channel
             raise Timeout("THP retransmission timeout")
 
+        def _wait_for_ack() -> Generator[Any, Any, None]:
+            """
+            Wait for the expected ACK to be received.
+
+            This task is spawned concurrently with `_write_loop()` using `loop.race()`,
+            so it will be cancelled when retransmission loop is over.
+            """
+            while not ABP.is_sending_allowed(self.channel_cache):
+                # `ABP.set_sending_allowed()` will be called after a valid ACK
+                yield from self.recv_payload(expected_ctrl_byte=None)
+
         try:
-            return await _write_loop()
+            # wait and return after receiving an ACK, or raise in case of an unexpected message / retransmission timeout.
+            await race(_wait_for_ack(), _write_loop())
         finally:
+            ack_latency_ms = utime.ticks_diff(utime.ticks_ms(), self.last_write_ms)
+            # Limit estimated latency to avoid integer overflows and too long delays
+            ack_latency_ms = max(0, min(800, ack_latency_ms))
+            self.channel_cache.set_int(CHANNEL_ACK_LATENCY_MS, ack_latency_ms)
+
             # Make sure to use the next `seq_bit` for the next payload
             ABP.set_send_seq_bit_to_opposite(self.channel_cache)
 

core/tests/test_trezor.wire.thp.writer.py

@@ -6,7 +6,7 @@
 if utils.USE_THP:
     import thp_common
     from mock_wire_interface import MockHID
-    from trezor.loop import Timeout
+    from trezor.loop import Timeout, race
     from trezor.wire.thp import ENCRYPTED, PacketHeader
     from trezor.wire.thp import alternating_bit_protocol as ABP
     from trezor.wire.thp.channel import _MAX_RETRANSMISSION_COUNT
@@ -134,15 +134,21 @@ def test_write_timeout(self):
         seq_bit = ABP.get_send_seq_bit(channel.channel_cache)
 
         task = channel.write_encrypted_payload(ENCRYPTED, b"PAYLOAD")
-        task.send(None)  # start the generator
+        race_obj = task.send(None)  # start the generator
+        assert isinstance(race_obj, race)
+        _wait_for_ack, write_loop = race_obj.children
+        write_loop.send(None)  # start the generator
 
         for _ in range(_MAX_RETRANSMISSION_COUNT - 1):
-            task.send(None)  # complete write
-            task.throw(Timeout())  # no ACK is received
+            write_loop.send(None)  # complete write
+            write_loop.send(None)  # complete sleep
+
+        write_loop.send(None)  # complete write last time
+        with self.assertRaises(Timeout) as ctx:
+            write_loop.send(None)  # complete sleep & raise Timeout
 
-        task.send(None)  # complete write last time
         with self.assertRaises(Timeout):
-            task.throw(Timeout())  # no ACK is received
+            task.throw(ctx.value)  # re-raise timeout in `write_encrypted_payload`
 
         # next write should use the next `seq_bit` (see #6138)
         self.assertNotEqual(ABP.get_send_seq_bit(channel.channel_cache), seq_bit)
@@ -152,16 +158,22 @@ def test_write_blocked(self):
         seq_bit = ABP.get_send_seq_bit(channel.channel_cache)
 
         task = channel.write_encrypted_payload(ENCRYPTED, b"PAYLOAD")
-        task.send(None)  # start the generator
+        race_obj = task.send(None)  # start the generator
+        assert isinstance(race_obj, race)
+        _wait_for_ack, write_loop = race_obj.children
+        write_loop.send(None)  # start the generator
 
         # Re-transmit a few times
         for _ in range(3):
-            task.send(None)  # complete write
-            task.throw(Timeout())  # no ACK is received
+            write_loop.send(None)  # complete write
+            write_loop.send(None)  # complete sleep
+
+        with self.assertRaises(Timeout) as ctx:
+            # timeout `_write_payload_once()` (as if `loop.sleep` has completed) using dummy "ticks" integer value
+            write_loop.send(12345)
 
         with self.assertRaises(Timeout):
-            # timeout write (as if `loop.sleep` has completed) using dummy "ticks" integer value
-            task.send(12345)
+            task.throw(ctx.value)  # re-raise timeout in `write_encrypted_payload`
 
         # next write should use the next `seq_bit` (see #6138)
         self.assertNotEqual(ABP.get_send_seq_bit(channel.channel_cache), seq_bit)

tests/device_tests/thp/test_abp.py

@@ -1,9 +1,12 @@
+import functools
 import time
+import typing as t
+from pathlib import Path
 from unittest.mock import Mock
 
 import pytest
 
-from trezorlib import messages
+from trezorlib import device, messages
 from trezorlib.debuglink import DebugSession
 from trezorlib.thp.client import TrezorClientThp
 
@@ -33,3 +36,31 @@ def test_abp(session: DebugSession) -> None:
     # we should now successfully read the response
     resp = session.read()
     messages.Success.ensure_isinstance(resp)
+
+
+def delay_call(func: t.Callable, seconds: float) -> t.Callable:
+
+    @functools.wraps(func)
+    def wrapper(*args, **kw):
+        time.sleep(seconds)
+        return func(*args, **kw)
+
+    return wrapper
+
+
+HERE = Path(__file__).parent.resolve()
+
+
+def test_delay_acks_from_host(session: DebugSession) -> None:
+    assert isinstance(session.client, TrezorClientThp)
+    channel = session.client.channel
+
+    # delay THP ACK sending, to trigger retransmits
+    channel._send_ack = delay_call(channel._send_ack, seconds=0.6)
+    session.client.ping("Should succeed after some retransmits")
+    session.client.ping("ButtonRequest should be retransmitted", button_protection=True)
+
+    file_name = "test_bg_eckhart.jpg"
+    with open(HERE.parent / file_name, "rb") as f:
+        # Multiple requests and responses
+        device.apply_settings(session, homescreen=f.read())

tests/ui_tests/fixtures.json

@@ -32351,6 +32351,7 @@
 "T3W1_cs_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "c66aacc0efb0ef974a67a0a0b838257f3eafa0018438317824700fbf84c8c433",
 "T3W1_cs_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "a6aa76610b2f4c09d32c9f63e8a831cb561710ec6f5e02d165218fa766c1fd23",
 "T3W1_cs_thp-test_abp.py::test_abp": "c24521e569c08e3605b164212c876f8ac57c5eef6cca6f2ca53a635a883ebc4b",
+"T3W1_cs_thp-test_abp.py::test_delay_acks_from_host": "8b42717a5490020e02da5b481bb9db9f5f333dceb15813e29cac09252f85c8cd",
 "T3W1_cs_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_cs_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_cs_thp-test_handshake.py::test_no_unlock": "05416515d2a63d94fc0ba3b6afdb9b0eb139ed403b7c20b62203c776ffe5d6ee",
@@ -33931,6 +33932,7 @@
 "T3W1_de_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "253a29728bc2b0273d5b09ee7ea84c1053d66bfaa1237ce389c97ca2ad8a6381",
 "T3W1_de_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "04932ccd7cdd4bbfa6f3665038d81f68c2822c8431b36bc9668c0aa2e6a418ab",
 "T3W1_de_thp-test_abp.py::test_abp": "b3ce10e1620297758d6f2b235127f1dd0d61804c8cc32bc6664db5c3a9d81dde",
+"T3W1_de_thp-test_abp.py::test_delay_acks_from_host": "89f5aeb7233b3ca22ccee2f08cc9caee38f7dec1b72bc49d9c1a4224f4ca1b27",
 "T3W1_de_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_de_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_de_thp-test_handshake.py::test_no_unlock": "9be3988c015d9a5260eb1537ab5cd398f4ddd36e2378157fe88f8ee04f533b92",
@@ -35511,6 +35513,7 @@
 "T3W1_en_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "5559494cb76edbe2de07923e9002e0ccf580144e2f17ac879d3d465e42189879",
 "T3W1_en_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "8d2d3d0da8e522be34cdd398cc9b49c2cf860a0687a9ac2c1a37c7255800651a",
 "T3W1_en_thp-test_abp.py::test_abp": "931d9afceb0ba1e4faae891775819277242d889644d5c0c5863fc8c9fcf859b1",
+"T3W1_en_thp-test_abp.py::test_delay_acks_from_host": "f6e49fc7fcf36d8f2eaec89587656ebd905b9d843b1960f0684c51c5b775e4ff",
 "T3W1_en_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_en_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_en_thp-test_handshake.py::test_no_unlock": "a60a24172d0470601707825e0287a91a7388e8aeb00b5e49ec9b2643de35c1ca",
@@ -37091,6 +37094,7 @@
 "T3W1_es_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "f485f09bb09e2717d91380a05c6ca9a3162a7068e665a1038574bde88271d9c0",
 "T3W1_es_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "ff6418a3e46044c19372760f8162def1f7a37d5276add41d7df4ffe409ae4b2b",
 "T3W1_es_thp-test_abp.py::test_abp": "56536ae9cd7c4ff8022def4ec3350031d3614064d961f53a2850f2be425af201",
+"T3W1_es_thp-test_abp.py::test_delay_acks_from_host": "e86b0905548c3f025189b94122aacee1032a65b49b7de4d1ee880d27ce2d7895",
 "T3W1_es_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_es_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_es_thp-test_handshake.py::test_no_unlock": "d362078d046200bf6a6421a18a3ad8c3e0eba86f88512fc934c15c850ae4148d",
@@ -38671,6 +38675,7 @@
 "T3W1_fr_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "af27005f1a60e90ac80560815ad2a458e9787ad399fd1cf8ff46743a435a12b2",
 "T3W1_fr_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "f107171d02df2b7ac91c04505f3f8143016f372c82891806c31da43165ae89e0",
 "T3W1_fr_thp-test_abp.py::test_abp": "e8156cf4eda1d29060f05b4a18d50032b0b344230609b7de7cababfe0a86b20b",
+"T3W1_fr_thp-test_abp.py::test_delay_acks_from_host": "477d766a3dddbadcaf240d3702e471260b39c79cc5a63e04832f104d8e6b11a7",
 "T3W1_fr_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_fr_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_fr_thp-test_handshake.py::test_no_unlock": "56132ede5cabeccec0514da02d7c88fb8ebd78d8fb14a2b75ec8c124a7ccbf2d",
@@ -40256,6 +40261,7 @@
 "T3W1_pt_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer": "3495d5e1353fbbbe3bff511512f74500dc1d29c0be40548c5f052ebd303102c6",
 "T3W1_pt_tezos-test_sign_tx.py::test_tezos_smart_contract_transfer_to_contract": "a3b2ff8deb91d0e7d8a5e13087d4b848bad23111de73a1123c3c3f3cfd08cfd5",
 "T3W1_pt_thp-test_abp.py::test_abp": "1c74667c078e25e7e0d37c7b2aa35f7c8ab02cd88e74695ecb355c82a293b0cc",
+"T3W1_pt_thp-test_abp.py::test_delay_acks_from_host": "da50bba056af80a3881763b986575d4fa4ed40402dcb37aec75992f583629ef0",
 "T3W1_pt_thp-test_basic.py::test_v1": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_pt_thp-test_basic.py::test_v2_unallocated": "987d1c62e6576b9cc24373e00df522efdc9736af978d6032f7d319af8daaef5d",
 "T3W1_pt_thp-test_handshake.py::test_no_unlock": "d362078d046200bf6a6421a18a3ad8c3e0eba86f88512fc934c15c850ae4148d",