refactor(core): move checksum validation to Channel

Also, move channel usage update to `interface_context`.

Keep the last valid reassembled message at `Reassembler.message`.
It will be used to implement `ThpChannel` read functionality.

[no changelog]

Author: romanz

core/src/trezor/wire/__init__.py

@@ -105,7 +105,9 @@ async def handle_session(iface: WireInterface) -> None:
 
         while True:
             try:
-                (channel, message) = await ctx.get_next_message()
+                channel = await ctx.get_next_message()
+                message = channel.reassembler.message
+                assert message is not None
                 await received_message_handler.handle_received_message(channel, message)
             except Exception:
                 loop.clear()  # restart event loop in case of error

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

@@ -25,7 +25,7 @@
 from . import ENCRYPTED, ChannelState, PacketHeader, ThpDecryptionError, ThpError
 from . import alternating_bit_protocol as ABP
 from . import control_byte, crypto, memory_manager
-from .checksum import CHECKSUM_LENGTH
+from .checksum import CHECKSUM_LENGTH, is_valid
 from .transmission_loop import TransmissionLoop
 from .writer import MESSAGE_TYPE_LENGTH
 
@@ -52,20 +52,23 @@ def __init__(self, cid: int) -> None:
         self.reset()
 
     def reset(self) -> None:
-        self.bytes_read = 0
-        self.buffer_len = 0
+        self.bytes_read: int = 0
+        self.buffer_len: int = 0
+        self.message: memoryview | None = None
 
-    def get_next_message(self, packet: memoryview) -> memoryview | None:
+    def handle_packet(self, packet: memoryview) -> bool:
         """
-        Process current packet, returning the payload buffer on success.
+        Process current packet, returning `True` when a valid message is reassembled.
+        The parsed message can retrieved via the `message` field (if it's not `None`).
+        In case of a checksum error or if the reassembly is not over, return `False`.
 
         May raise `WireBufferError` if there is a concurrent payload reassembly in progress.
         """
         ctrl_byte = packet[0]
         if control_byte.is_continuation(ctrl_byte):
             if not self.bytes_read:
                 # ignore unexpected continuation packets
-                return None
+                return False
 
             # may raise WireBufferError
             buffer = memory_manager.get_existing_read_buffer(self.cid)
@@ -86,19 +89,36 @@ def get_next_message(self, packet: memoryview) -> memoryview | None:
 
         assert len(buffer) == self.buffer_len
         if self.bytes_read < self.buffer_len:
-            return None
-        elif self.bytes_read == self.buffer_len:
-            self.reset()
-            return buffer
-        else:
+            return False
+
+        if self.bytes_read > self.buffer_len:
             raise ThpError("read more bytes than expected")
 
+        if not verify_checksum(buffer):
+            return False
+
+        assert self.message is None
+        self.message = buffer
+        return True
+
     def _buffer_packet_data(
         self, payload_buffer: memoryview, packet: memoryview, offset: int
     ) -> None:
         self.bytes_read += utils.memcpy(payload_buffer, self.bytes_read, packet, offset)
 
 
+def verify_checksum(buffer: memoryview) -> memoryview | None:
+    """
+    Return the buffer if the checksum is valid, otherwise return `None`.
+    """
+    if is_valid(buffer[-CHECKSUM_LENGTH:], buffer[:-CHECKSUM_LENGTH]):
+        return buffer
+    # ignore invalid payloads
+    if __debug__:
+        log.warning("Invalid payload checksum: %s", utils.hexlify_if_bytes(buffer))
+    return None
+
+
 class Channel:
     """
     THP protocol encrypted communication channel.
@@ -184,11 +204,18 @@ def is_channel_to_replace(self) -> bool:
 
     # READ and DECRYPT
 
-    def handle_packet(self, packet: utils.BufferType) -> memoryview | None:
+    def reassemble(self, packet: utils.BufferType) -> bool:
+        """
+        Process current packet, returning `True` when a valid message is reassembled.
+        The parsed message can retrieved via the `message` field (if it's not `None`).
+        In case of a checksum error or if the reassembly is not over, return `False`.
+
+        May raise `WireBufferError` if there is a concurrent payload reassembly in progress.
+        """
         if self.get_channel_state() == ChannelState.UNALLOCATED:
-            return None
+            return False
         try:
-            return self.reassembler.get_next_message(memoryview(packet))
+            return self.reassembler.handle_packet(memoryview(packet))
         except WireBufferError:
             self.reassembler.reset()
             raise

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

@@ -7,6 +7,7 @@
     BROADCAST_CHANNEL_ID,
     ChannelCache,
     iter_allocated_channels,
+    update_channel_last_used,
 )
 from trezor import io, loop, utils
 
@@ -60,7 +61,7 @@ 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) -> tuple[Channel, memoryview]:
+    async def get_next_message(self) -> Channel:
         packet = bytearray(self._iface.RX_PACKET_LEN)
         while True:
             packet_len = await self._read
@@ -84,11 +85,11 @@ async def get_next_message(self) -> tuple[Channel, memoryview]:
                 continue
 
             try:
-                message = channel.handle_packet(packet)
-                if message is not None:
-                    # `message` must be handled ASAP without blocking,
+                if channel.reassemble(packet):
+                    update_channel_last_used(channel.channel_id)
+                    # The reassembled message must be handled ASAP without blocking,
                     # since it may point to the global read buffer.
-                    return channel, message
+                    return channel
             except WireBufferError:
                 await channel.write_error(ThpErrorType.TRANSPORT_BUSY)
                 continue

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

@@ -12,7 +12,6 @@
     KEY_LENGTH,
     SESSION_ID_LENGTH,
     TAG_LENGTH,
-    update_channel_last_used,
     update_session_last_used,
 )
 from trezor import config, loop, protobuf, utils
@@ -39,7 +38,7 @@
     ThpUnallocatedSessionError,
 )
 from . import alternating_bit_protocol as ABP
-from . import checksum, control_byte, get_encoded_device_properties, session_manager
+from . import control_byte, get_encoded_device_properties, session_manager
 from .checksum import CHECKSUM_LENGTH
 from .crypto import PUBKEY_LENGTH, Handshake
 from .session_context import SeedlessSessionContext
@@ -82,8 +81,6 @@ async def handle_received_message(
     ctrl_byte, _, payload_length = ustruct.unpack(">BHH", message_buffer)
     message_length = payload_length + PacketHeader.INIT_LENGTH
 
-    _check_checksum(message_length, message_buffer)
-
     # Synchronization process
     seq_bit = control_byte.get_seq_bit(ctrl_byte)
     ack_bit = control_byte.get_ack_bit(ctrl_byte)
@@ -95,8 +92,6 @@ async def handle_received_message(
             ack_bit,
             iface=ctx.iface,
         )
-    # 0: Update "last-time used"
-    update_channel_last_used(ctx.channel_id)
 
     # 1: Handle ACKs
     if control_byte.is_ack(ctrl_byte):
@@ -161,18 +156,6 @@ def _send_ack(channel: Channel, ack_bit: int) -> Awaitable[None]:
     return channel.ctx.write_payload(header, b"")
 
 
-def _check_checksum(message_length: int, message_buffer: utils.BufferType) -> None:
-    if __debug__:
-        log.debug(__name__, "check_checksum")
-    if not checksum.is_valid(
-        checksum=message_buffer[message_length - CHECKSUM_LENGTH : message_length],
-        data=memoryview(message_buffer)[: message_length - CHECKSUM_LENGTH],
-    ):
-        if __debug__:
-            log.debug(__name__, "Invalid checksum, ignoring message.")
-        raise ThpError("Invalid checksum, ignoring message.")
-
-
 async def handle_ack(ctx: Channel, ack_bit: int) -> None:
     if not ABP.is_ack_valid(ctx.channel_cache, ack_bit):
         return