Reimplement event loop restarts in THP

core/embed/upymod/qstrdefsport.h

@@ -421,7 +421,6 @@ Q(session_manager)
 Q(storage.cache_thp)
 Q(storage.cache_thp_keys)
 Q(thp)
-Q(transmission_loop)
 Q(trezor.enums.ThpMessageType)
 Q(trezor.enums.ThpPairingMethod)
 Q(trezor.wire.thp)
@@ -438,7 +437,6 @@ Q(trezor.wire.thp.pairing_context)
 Q(trezor.wire.thp.received_message_handler)
 Q(trezor.wire.thp.session_context)
 Q(trezor.wire.thp.session_manager)
-Q(trezor.wire.thp.transmission_loop)
 Q(trezor.wire.thp.ui)
 Q(trezor.wire.thp.writer)
 Q(ui)

core/src/apps/base.py

@@ -231,7 +231,7 @@ async def handle_ThpCreateNewSession(
 
         Returns an appropriate `Failure` message if session creation fails.
         """
-        from trezor import log, loop
+        from trezor import log
         from trezor.enums import FailureType
         from trezor.messages import Failure
         from trezor.wire import NotInitialized
@@ -281,9 +281,6 @@ async def handle_ThpCreateNewSession(
                 message.passphrase if message.passphrase is not None else "",
             )
 
-        channel.sessions[new_session.session_id] = new_session
-        loop.schedule(new_session.handle())
-
         return Success(message="New session created.")
 
     async def handle_ThpCredentialRequest(

core/src/apps/management/reboot_to_bootloader.py

@@ -94,7 +94,7 @@ async def reboot_to_bootloader(msg: RebootToBootloader) -> NoReturn:
         boot_args = None
 
     ctx = get_context()
-    await ctx.write_force(Success(message="Rebooting"))
+    await ctx.write(Success(message="Rebooting"))
     # make sure the outgoing USB buffer is flushed
     await loop.wait(ctx.iface.iface_num() | io.POLL_WRITE)
 

core/src/apps/management/wipe_device.py

@@ -13,7 +13,7 @@
 
 async def wipe_device(msg: WipeDevice) -> NoReturn:
     import storage
-    from trezor import TR, config, loop, translations
+    from trezor import TR, config, translations
     from trezor.enums import ButtonRequestType
     from trezor.messages import Success
     from trezor.pin import render_empty_loader
@@ -49,7 +49,7 @@ async def wipe_device(msg: WipeDevice) -> NoReturn:
     translations.deinit()
     translations.erase()
     try:
-        await get_context().write_force(Success(message="Device wiped"))
+        await get_context().write(Success(message="Device wiped"))
     except Exception:
         if __debug__:
             log.debug(__name__, "Failed to send Success message after wipe.")
@@ -58,6 +58,5 @@ async def wipe_device(msg: WipeDevice) -> NoReturn:
 
     # reload settings
     reload_settings_from_storage()
-    loop.clear()
     if __debug__:
         log.debug(__name__, "Device wipe - finished")

core/src/apps/thp/pairing.py

@@ -221,7 +221,7 @@ async def _handle_code_entry_is_selected(ctx: PairingContext) -> None:
     if ctx.code_entry_secret is None:
         await _handle_code_entry_is_selected_first_time(ctx)
     else:
-        await ctx.write_force(ThpPairingPreparationsFinished())
+        await ctx.write(ThpPairingPreparationsFinished())
 
 
 async def _handle_code_entry_is_selected_first_time(ctx: PairingContext) -> None:
@@ -251,15 +251,15 @@ async def _handle_code_entry_is_selected_first_time(ctx: PairingContext) -> None
     )
     assert ctx.code_code_entry is not None
     ctx.cpace.generate_keys(f"{ctx.code_code_entry:06}".encode("ascii"))
-    await ctx.write_force(
+    await ctx.write(
         ThpCodeEntryCpaceTrezor(cpace_trezor_public_key=ctx.cpace.trezor_public_key)
     )
 
 
 @check_state_and_log(ChannelState.TP1)
 async def _handle_nfc_is_selected(ctx: PairingContext) -> None:
     ctx.nfc_secret = random.bytes(16)
-    await ctx.write_force(ThpPairingPreparationsFinished())
+    await ctx.write(ThpPairingPreparationsFinished())
 
 
 @check_state_and_log(ChannelState.TP1)
@@ -271,7 +271,7 @@ async def _handle_qr_code_is_selected(ctx: PairingContext) -> None:
     sha_ctx.update(ctx.qr_code_secret)
 
     ctx.code_qr_code = sha_ctx.digest()[:16]
-    await ctx.write_force(ThpPairingPreparationsFinished())
+    await ctx.write(ThpPairingPreparationsFinished())
 
 
 @check_state_and_log(ChannelState.TP3)

core/src/storage/cache_thp.py

@@ -1,24 +1,16 @@
 import builtins
 from micropython import const
-from typing import TYPE_CHECKING
 
 from storage.cache_common import (
     CHANNEL_HOST_STATIC_PUBKEY,
     CHANNEL_ID,
-    CHANNEL_IFACE,
     CHANNEL_STATE,
     CHANNEL_SYNC,
     SESSION_ID,
     SESSION_STATE,
     DataCache,
 )
 
-if TYPE_CHECKING:
-    from typing import Iterable, Tuple
-
-    pass
-
-
 # THP specific constants
 _MAX_CHANNELS_COUNT = const(10)
 _MAX_SESSIONS_COUNT = const(20)
@@ -183,13 +175,14 @@ def update_session_last_used(channel_id: bytes, session_id: bytes) -> None:
             return
 
 
-def iter_allocated_channels(iface_num: int) -> Iterable[ChannelCache]:
+def find_allocated_channel(cid: int) -> ChannelCache | None:
     for channel in _CHANNELS:
         state = channel.get_int(CHANNEL_STATE, _UNALLOCATED_STATE)
         if state == _UNALLOCATED_STATE:
             continue
-        if channel.get_int(CHANNEL_IFACE) == iface_num:
-            yield channel
+        if channel.get_int(CHANNEL_ID) == cid:
+            return channel
+    return None
 
 
 def get_allocated_session(
@@ -393,7 +386,7 @@ def clear_all() -> None:
         channel.clear()
 
 
-def clear_all_except_one_session_keys(excluded: Tuple[bytes, bytes]) -> None:
+def clear_all_except_one_session_keys(excluded: tuple[bytes, bytes]) -> None:
     cid, sid = excluded
 
     for channel in _CHANNELS:

core/src/trezor/wire/__init__.py

@@ -82,9 +82,10 @@ def setup(iface: WireInterface) -> None:
 
 
 if utils.USE_THP:
-    # memory_manager is imported to create READ/WRITE buffers
-    # in more stable area of memory
-    from .thp import memory_manager  # noqa: F401
+    from .thp.memory_manager import ThpBuffer
+
+    # Allocate THP read/write buffers in more stable area of memory
+    THP_BUFFERS_PROVIDER = Provider((ThpBuffer(), ThpBuffer()))
 
     if __debug__:
         _THP_CHANNELS = []
@@ -99,19 +100,21 @@ def find_thp_channel(channel_id: bytes) -> Channel | None:
             return None
 
     async def handle_session(iface: WireInterface) -> None:
-        ctx = ThpContext.load_from_cache(iface)
+        ctx = ThpContext(iface)
         if __debug__:
             _THP_CHANNELS.append(ctx._channels)
-
-        while True:
-            try:
-                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
-                raise  # the traceback will be printed by `loop._step()`
+        try:
+            channel = await ctx.get_next_message()
+            while await received_message_handler.handle_received_message(channel):
+                pass
+        finally:
+            # Wait for all active workflows to finish.
+            await workflow.join_all()
+            if __debug__:
+                import apps.debug
+
+                await apps.debug.close_session()
+            loop.clear()
 
 else:
 

core/src/trezor/wire/errors.py

@@ -14,11 +14,6 @@ def __init__(self, message: str) -> None:
         self.message = message
 
 
-class WireBufferError(Error):
-    def __init__(self, message: str = "Buffer error") -> None:
-        super().__init__(FailureType.BufferError, message)
-
-
 class UnexpectedMessage(Error):
     def __init__(self, message: str) -> None:
         super().__init__(FailureType.UnexpectedMessage, message)

core/src/trezor/wire/protocol_common.py

@@ -74,13 +74,10 @@ async def read(
         """
         ...
 
-    async def write(self, msg: protobuf.MessageType) -> None:
+    def write(self, msg: protobuf.MessageType) -> Awaitable[None]:
         """Write a message to the wire."""
         ...
 
-    def write_force(self, msg: protobuf.MessageType) -> Awaitable[None]:
-        return self.write(msg)
-
     async def call(
         self,
         msg: protobuf.MessageType,

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

@@ -44,18 +44,10 @@ class ThpDecryptionError(ThpError):
     pass
 
 
-class ThpInvalidDataError(ThpError):
-    pass
-
-
 class ThpDeviceLockedError(ThpError):
     pass
 
 
-class ThpUnallocatedChannelError(ThpError):
-    pass
-
-
 class ThpUnallocatedSessionError(ThpError):
 
     def __init__(self, session_id: int) -> None: