fixup! fix(core/thp): keep reading when writes are blocked

Author: romanz

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

@@ -436,49 +436,46 @@ async def write_encrypted_payload(self, ctrl_byte: int, payload: AnyBytes) -> No
 
         header = PacketHeader(ctrl_byte, self.get_channel_id_int(), payload_len)
 
+        ack_latency_ms = self.channel_cache.get_int(CHANNEL_ACK_LATENCY_MS) or 0
+
+        # 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()
+
         async def _write_loop() -> None:
             """Send the payload and wait for an ACK with retransmissions."""
+            from trezor.loop import sleep
 
-            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)
-
-            # allows preempting this channel, if another channel becomes active
-            self.last_write_ms = utime.ticks_ms()
-
             for i in range(_MAX_RETRANSMISSION_COUNT):
                 await 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
+                await sleep(timeout_ms)
+                if __debug__:
+                    log.warning(__name__, "Retransmit after %d ms", timeout_ms)
 
             # restart event loop due to unresponsive channel
             raise Timeout("THP retransmission timeout")
 
+        async def _wait_for_ack() -> None:
+            while not ABP.is_sending_allowed(self.channel_cache):
+                # `ABP.set_sending_allowed()` will be called after a valid ACK
+                await self.recv_payload(expected_ctrl_byte=None)
+
+        # wait and return after receiving an ACK, or raise in case of an unexpected message / retransmission timeout.
         try:
-            return await _write_loop()
+            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)