fix(core): handle ACK bit on THP data messages

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

@@ -10,10 +10,7 @@ def is_ack_valid(cache: ChannelCache, ack_bit: int) -> bool:
     if not _is_ack_expected(cache):
         return False
 
-    if not _has_ack_correct_sync_bit(cache, ack_bit):
-        return False
-
-    return True
+    return _has_ack_correct_sync_bit(cache, ack_bit)
 
 
 def _is_ack_expected(cache: ChannelCache) -> bool:

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

@@ -217,7 +217,8 @@ def is_channel_to_replace(self) -> bool:
 
     async def recv_payload(
         self,
-        expected_ctrl_byte: Callable[[int], bool] | None,
+        expected_ctrl_byte: Callable[[int], bool],
+        return_ack: bool = False,
         timeout_ms: int | None = None,
     ) -> memoryview:
         """
@@ -226,7 +227,7 @@ async def recv_payload(
 
         Raise if the received control byte is an unexpected one.
 
-        If `expected_ctrl_byte` is `None`, returns after the first received ACK.
+        If `return_ack` is set, return after receiving an expected ACK.
         """
 
         while True:
@@ -236,17 +237,32 @@ async def recv_payload(
 
             # Synchronization process
             ctrl_byte = msg[0]
+            if not (control_byte.is_ack(ctrl_byte) or control_byte.is_data(ctrl_byte)):
+                if __debug__:
+                    self._log(
+                        "Unrelated control byte - ignoring ",
+                        utils.hexlify_if_bytes(msg),
+                        logger=log.warning,
+                    )
+                continue
+
             payload = msg[PacketHeader.INIT_LENGTH : -CHECKSUM_LENGTH]
             seq_bit = control_byte.get_seq_bit(ctrl_byte)
+            ack_bit = control_byte.get_ack_bit(ctrl_byte)
 
-            # 1: Handle ACKs
-            if control_byte.is_ack(ctrl_byte):
-                handle_ack(self, control_byte.get_ack_bit(ctrl_byte))
-                if expected_ctrl_byte is None:
-                    return payload
-                continue
+            # 0: Handle ACKs
+            if received_expected_ack(self, ack_bit):
+                if control_byte.is_ack(ctrl_byte):
+                    # ACK message is handled
+                    self.reassembler.reset()
 
-            if expected_ctrl_byte is None or not expected_ctrl_byte(ctrl_byte):
+                if return_ack:
+                    # keep the reassembled payload (for future handling)
+                    return payload[:0]
+
+            # 1: Check expected control byte
+            self.reassembler.reset()
+            if not expected_ctrl_byte(ctrl_byte) or return_ack:
                 if __debug__:
                     self._log(
                         "Unexpected control byte - ignoring ",
@@ -400,8 +416,12 @@ async def write_encrypted_payload(self, ctrl_byte: int, payload: bytes) -> None:
             # starting from 100ms till ~3.42s
             timeout_ms = round(10200 - 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)
+                # wait and return after receiving an expected ACK.
+                await self.recv_payload(
+                    expected_ctrl_byte=control_byte.is_encrypted_transport,
+                    return_ack=True,
+                    timeout_ms=timeout_ms,
+                )
             except Timeout:
                 if __debug__:
                     log.warning(__name__, "Retransmit after %d ms", timeout_ms)
@@ -462,9 +482,9 @@ def send_ack(channel: Channel, ack_bit: int) -> Awaitable[None]:
     return channel.ctx.write_payload(header, b"")
 
 
-def handle_ack(ctx: Channel, ack_bit: int) -> None:
+def received_expected_ack(ctx: Channel, ack_bit: int) -> bool:
     if not ABP.is_ack_valid(ctx.channel_cache, ack_bit):
-        return
+        return False
     # ACK is expected and it has correct sync bit
     if __debug__:
         log.debug(
@@ -473,3 +493,4 @@ def handle_ack(ctx: Channel, ack_bit: int) -> None:
             iface=ctx.iface,
         )
     ABP.set_sending_allowed(ctx.channel_cache, True)
+    return True

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

@@ -41,6 +41,11 @@ def is_ack(ctrl_byte: int) -> bool:
     return ctrl_byte & _ACK_MASK == ACK_MESSAGE
 
 
+def is_data(ctrl_byte: int) -> bool:
+    """Return `True` for ACK, handshake and encrypted transport messages."""
+    return ctrl_byte & 0xE0 == 0
+
+
 def is_continuation(ctrl_byte: int) -> bool:
     return ctrl_byte & _CONTINUATION_PACKET_MASK == CONTINUATION_PACKET