From c01a86dd9b1b1ae07a3f9819694e9dc2c9b8c47c Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Tue, 25 Jul 2023 10:47:55 +0200
Subject: [PATCH] refactor(core): remove debuglink support from main
 handle_session

It is being moved into a custom handler for the debug app.
---
 core/src/session.py              |  3 +-
 core/src/trezor/wire/__init__.py | 55 +++++++++-----------------------
 2 files changed, 16 insertions(+), 42 deletions(-)

diff --git a/core/src/session.py b/core/src/session.py
index 4f161c36c..1ecbc467b 100644
--- a/core/src/session.py
+++ b/core/src/session.py
@@ -22,9 +22,8 @@ workflow.start_default()
 
 # initialize the wire codec
 wire.setup(usb.iface_wire)
-if __debug__:
-    wire.setup(usb.iface_debug, is_debug_session=True)
 
+# start the event loop
 loop.run()
 
 if __debug__:
diff --git a/core/src/trezor/wire/__init__.py b/core/src/trezor/wire/__init__.py
index 09991914d..86f7a5569 100644
--- a/core/src/trezor/wire/__init__.py
+++ b/core/src/trezor/wire/__init__.py
@@ -53,9 +53,9 @@ if TYPE_CHECKING:
 EXPERIMENTAL_ENABLED = False
 
 
-def setup(iface: WireInterface, is_debug_session: bool = False) -> None:
+def setup(iface: WireInterface) -> None:
     """Initialize the wire stack on passed USB interface."""
-    loop.schedule(handle_session(iface, codec_v1.SESSION_ID, is_debug_session))
+    loop.schedule(handle_session(iface, codec_v1.SESSION_ID))
 
 
 def wrap_protobuf_load(
@@ -88,7 +88,7 @@ if __debug__:
 
 
 async def _handle_single_message(
-    ctx: context.Context, msg: codec_v1.Message, use_workflow: bool
+    ctx: context.Context, msg: codec_v1.Message
 ) -> codec_v1.Message | None:
     """Handle a message that was loaded from USB by the caller.
 
@@ -149,14 +149,7 @@ async def _handle_single_message(
         # communication inside, but it should eventually return a
         # response message, or raise an exception (a rather common
         # thing to do).  Exceptions are handled in the code below.
-        if use_workflow:
-            # Spawn a workflow around the task. This ensures that concurrent
-            # workflows are shut down.
-            res_msg = await workflow.spawn(context.with_context(ctx, task))
-        else:
-            # For debug messages, ignore workflow processing and just await
-            # results of the handler.
-            res_msg = await task
+        res_msg = await workflow.spawn(context.with_context(ctx, task))
 
     except context.UnexpectedMessage as exc:
         # Workflow was trying to read a message from the wire, and
@@ -192,22 +185,10 @@ async def _handle_single_message(
     return None
 
 
-async def handle_session(
-    iface: WireInterface, session_id: int, is_debug_session: bool = False
-) -> None:
-    if __debug__ and is_debug_session:
-        ctx_buffer = WIRE_BUFFER_DEBUG
-    else:
-        ctx_buffer = WIRE_BUFFER
-
-    ctx = context.Context(iface, session_id, ctx_buffer)
+async def handle_session(iface: WireInterface, session_id: int) -> None:
+    ctx = context.Context(iface, session_id, WIRE_BUFFER)
     next_msg: codec_v1.Message | None = None
 
-    if __debug__ and is_debug_session:
-        import apps.debug
-
-        apps.debug.DEBUG_CONTEXT = ctx
-
     # Take a mark of modules that are imported at this point, so we can
     # roll back and un-import any others.
     modules = utils.unimport_begin()
@@ -230,27 +211,21 @@ async def handle_session(
                 next_msg = None
 
             try:
-                next_msg = await _handle_single_message(
-                    ctx, msg, use_workflow=not is_debug_session
-                )
+                next_msg = await _handle_single_message(ctx, msg)
             except Exception as exc:
                 # Log and ignore. The session handler can only exit explicitly in the
                 # following finally block.
                 if __debug__:
                     log.exception(__name__, exc)
             finally:
-                if not __debug__ or not is_debug_session:
-                    # Unload modules imported by the workflow.  Should not raise.
-                    # This is not done for the debug session because the snapshot taken
-                    # in a debug session would clear modules which are in use by the
-                    # workflow running on wire.
-                    utils.unimport_end(modules)
-
-                    if next_msg is None and msg.type not in AVOID_RESTARTING_FOR:
-                        # Shut down the loop if there is no next message waiting.
-                        # Let the session be restarted from `main`.
-                        loop.clear()
-                        return  # pylint: disable=lost-exception
+                # Unload modules imported by the workflow.  Should not raise.
+                utils.unimport_end(modules)
+
+                if next_msg is None and msg.type not in AVOID_RESTARTING_FOR:
+                    # Shut down the loop if there is no next message waiting.
+                    # Let the session be restarted from `main`.
+                    loop.clear()
+                    return  # pylint: disable=lost-exception
 
         except Exception as exc:
             # Log and try again. The session handler can only exit explicitly via