feat(core): reimplement stale THP channel interruption

It has been removed by #5546.

[no changelog]

Author: romanz

core/src/trezor/wire/__init__.py

@@ -104,8 +104,7 @@ async def handle_session(iface: WireInterface) -> None:
         if __debug__:
             _THP_CHANNELS.append(ctx._channels)
         try:
-            channel = await ctx.get_next_message()
-            while await received_message_handler.handle_received_message(channel):
+            while await received_message_handler.handle_received_message(ctx):
                 pass
         finally:
             # Wait for all active workflows to finish.

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

@@ -1,4 +1,5 @@
 import ustruct
+import utime
 from micropython import const
 from typing import TYPE_CHECKING
 
@@ -58,6 +59,8 @@
 _MAX_RETRANSMISSION_COUNT = const(50)
 _MIN_RETRANSMISSION_COUNT = const(2)
 
+_STALE_CHANNEL_TIMEOUT_MS = const(500)
+
 
 class Reassembler:
     def __init__(self, cid: int, read_buf: ThpBuffer) -> None:
@@ -129,6 +132,11 @@ def verify_checksum(buffer: memoryview) -> memoryview | None:
     return None
 
 
+# Raised when a new channel is being allocated, in order to close an existing (unresponsive) workflow.
+class StaleChannelInterrupt(Exception):
+    pass
+
+
 class Channel:
     """
     THP protocol encrypted communication channel.
@@ -270,9 +278,23 @@ async def _get_reassembled_message(
         self, timeout_ms: int | None = None
     ) -> memoryview:
         """Doesn't block if a message has been already reassembled."""
+        start_ms = utime.ticks_ms()
         while self.reassembler.message is None:
             # receive and reassemble a new message from this channel
             channel = await self.ctx.get_next_message(timeout_ms=timeout_ms)
+            if channel is None:
+                # interrupted by a new channel allocation request
+                elapsed = utime.ticks_diff(utime.ticks_ms(), start_ms)
+                if elapsed < _STALE_CHANNEL_TIMEOUT_MS:
+                    continue
+                if __debug__:
+                    self._log(
+                        f"Stale channel interrupted after {elapsed} ms",
+                        logger=log.warning,
+                    )
+                # Will close stale workflow and restart the event loop, in order to allow new channels to be handled
+                raise StaleChannelInterrupt
+
             if channel is self:
                 break
 

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

@@ -47,11 +47,12 @@ def __init__(self, iface: WireInterface) -> None:
         self._write = loop.wait(iface.iface_num() | io.POLL_WRITE)
         self._channels: dict[int, Channel] = {}
 
-    async def get_next_message(self, timeout_ms: int | None = None) -> Channel:
+    async def get_next_message(self, timeout_ms: int | None = None) -> Channel | None:
         """
         Reassemble a valid THP payload and return its channel.
 
         Also handle THP channel allocation.
+        If buffer allocation fails, return `None` to allow stale channel interruption.
         """
         from .. import THP_BUFFERS_PROVIDER
 
@@ -80,9 +81,10 @@ async def get_next_message(self, timeout_ms: int | None = None) -> Channel:
 
             if (channel := self._channels.get(cid)) is None:
                 if (buffers := THP_BUFFERS_PROVIDER.take()) is None:
-                    # concurrent payload reassembly is not supported
+                    # send TRANSPORT_BUSY to the new channel, which should retry later
                     await self.write_error(cid, ThpErrorType.TRANSPORT_BUSY)
-                    continue
+                    return None
+
                 channel = self._channels[cid] = Channel(cache, self, buffers)
 
             if channel.reassemble(packet):

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

@@ -35,6 +35,7 @@
     from trezor.messages import ThpHandshakeCompletionReqNoisePayload
 
     from .channel import Channel
+    from .interface_context import ThpContext
 
 if __debug__:
     from trezor import log
@@ -45,12 +46,18 @@
 _TREZOR_STATE_PAIRED_AUTOCONNECT = b"\x02"
 
 
-async def handle_received_message(channel: Channel) -> bool:
+async def handle_received_message(ctx: ThpContext) -> bool:
     """
     Handle a message received from the channel.
 
     Returns False if we can restart the event loop.
     """
+    channel = await ctx.get_next_message()
+    if channel is None:
+        # Channel buffers' allocation failed (due to another active THP interface).
+        # Continue handling THP messages until the event loop is restarted.
+        return True
+
     try:
         state = channel.get_channel_state()
         if state is ChannelState.ENCRYPTED_TRANSPORT:

python/src/trezorlib/transport/thp/protocol_v2.py

@@ -278,8 +278,8 @@ def _read_handshake_completion_response(self) -> int:
         self._send_ack_bit(bit=1)
         return int.from_bytes(trezor_state, "big")
 
-    def _read_ack(self):
-        header, payload = self._read_until_valid_crc_check()
+    def _read_ack(self, timeout: float | None = None):
+        header, payload = self._read_until_valid_crc_check(timeout=timeout)
         if not header.is_ack() or len(payload) > 0:
             LOG.error("Received message is not a valid ACK")
 

tests/device_tests/thp/test_multiple_hosts.py

@@ -1,6 +1,8 @@
+import time
+
 import pytest
 
-from trezorlib import exceptions
+from trezorlib import exceptions, transport
 from trezorlib.client import ProtocolV2Channel
 from trezorlib.debuglink import TrezorClientDebugLink as Client
 
@@ -42,3 +44,26 @@ def test_concurrent_handshakes(client: Client) -> None:
 
     # Now the second handshake can be done
     protocol_2._do_handshake()
+
+
+def test_channel_interrupt(client: Client) -> None:
+    protocol_1 = _new_channel(client)
+    protocol_2 = _new_channel(client)
+
+    # The first host starts handshake
+    protocol_1._send_handshake_init_request()
+    protocol_1._read_ack()
+    protocol_1._read_handshake_init_response()
+
+    # The second host can interrupt the first host's handshake after a while
+    time.sleep(1)
+    with pytest.raises(exceptions.TransportBusy):
+        protocol_2._do_handshake()
+
+    # The first host has been interrupted, retrying handshake should succeed now
+    protocol_2._do_handshake()
+
+    # The first handshake was interrupted, so no response will be sent
+    protocol_1._send_handshake_completion_request()
+    with pytest.raises(transport.Timeout):
+        protocol_1._read_ack(timeout=1)