perf(core): share wire buffer between wire and debuglink

[no changelog]

core/src/trezor/wire/__init__.py

@@ -130,8 +130,6 @@ DUMMY_CONTEXT = DummyContext()
 PROTOBUF_BUFFER_SIZE = 8192
 
 WIRE_BUFFER = bytearray(PROTOBUF_BUFFER_SIZE)
-if __debug__:
-    WIRE_BUFFER_DEBUG = bytearray(PROTOBUF_BUFFER_SIZE)
 
 
 class Context:
@@ -378,11 +376,7 @@ async def _handle_single_message(
 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(iface, session_id, ctx_buffer)
+    ctx = Context(iface, session_id, WIRE_BUFFER)
     next_msg: codec_v1.Message | None = None
 
     if __debug__ and is_debug_session:

core/src/trezor/wire/codec_v1.py

@@ -4,6 +4,8 @@ from micropython import const
 from trezor import io, loop, utils
 
 if False:
+    from typing import Any
+
     from trezorio import WireInterface
 
 _REP_LEN = const(64)
@@ -17,6 +19,35 @@ _REP_CONT_DATA = const(1)  # offset of data in the continuation report
 SESSION_ID = const(0)
 INVALID_TYPE = const(-1)
 
+# The wire buffer is shared between the main wire interface and debuglink
+# (see __init__.py). There's no obvious guarantee that both interfaces won't
+# use it at the same time, thus we check this at runtime in debug builds.
+if __debug__:
+
+    class BufferLock(object):
+        def __init__(self) -> None:
+            self.in_use = False
+
+        def __enter__(self) -> None:
+            assert not self.in_use, "global buffer already used by another context"
+            self.in_use = True
+
+        def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
+            self.in_use = False
+
+
+else:
+
+    class BufferLock(object):  # type: ignore
+        def __enter__(self) -> None:
+            pass
+
+        def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
+            pass
+
+
+buffer_lock = BufferLock()
+
 
 class CodecError(Exception):
     pass
@@ -41,31 +72,32 @@ async def read_message(iface: WireInterface, buffer: utils.BufferType) -> Messag
 
     read_and_throw_away = False
 
-    if msize > len(buffer):
-        # allocate a new buffer to fit the message
-        try:
-            mdata: utils.BufferType = bytearray(msize)
-        except MemoryError:
-            mdata = bytearray(_REP_LEN)
-            read_and_throw_away = True
-    else:
-        # reuse a part of the supplied buffer
-        mdata = memoryview(buffer)[:msize]
-
-    # buffer the initial data
-    nread = utils.memcpy(mdata, 0, report, _REP_INIT_DATA)
-
-    while nread < msize:
-        # wait for continuation report
-        report = await read
-        if report[0] != _REP_MARKER:
-            raise CodecError("Invalid magic")
-
-        # buffer the continuation data
-        if read_and_throw_away:
-            nread += len(report) - 1
+    with buffer_lock:
+        if msize > len(buffer):
+            # allocate a new buffer to fit the message
+            try:
+                mdata: utils.BufferType = bytearray(msize)
+            except MemoryError:
+                mdata = bytearray(_REP_LEN)
+                read_and_throw_away = True
         else:
-            nread += utils.memcpy(mdata, nread, report, _REP_CONT_DATA)
+            # reuse a part of the supplied buffer
+            mdata = memoryview(buffer)[:msize]
+
+        # buffer the initial data
+        nread = utils.memcpy(mdata, 0, report, _REP_INIT_DATA)
+
+        while nread < msize:
+            # wait for continuation report
+            report = await read
+            if report[0] != _REP_MARKER:
+                raise CodecError("Invalid magic")
+
+            # buffer the continuation data
+            if read_and_throw_away:
+                nread += len(report) - 1
+            else:
+                nread += utils.memcpy(mdata, nread, report, _REP_CONT_DATA)
 
     if read_and_throw_away:
         raise CodecError("Message too large")
@@ -76,30 +108,31 @@ async def read_message(iface: WireInterface, buffer: utils.BufferType) -> Messag
 async def write_message(iface: WireInterface, mtype: int, mdata: bytes) -> None:
     write = loop.wait(iface.iface_num() | io.POLL_WRITE)
 
-    # gather data from msg
-    msize = len(mdata)
-
-    # prepare the report buffer with header data
-    report = bytearray(_REP_LEN)
-    repofs = _REP_INIT_DATA
-    ustruct.pack_into(
-        _REP_INIT, report, 0, _REP_MARKER, _REP_MAGIC, _REP_MAGIC, mtype, msize
-    )
+    with buffer_lock:
+        # gather data from msg
+        msize = len(mdata)
 
-    nwritten = 0
-    while True:
-        # copy as much as possible to the report buffer
-        nwritten += utils.memcpy(report, repofs, mdata, nwritten)
+        # prepare the report buffer with header data
+        report = bytearray(_REP_LEN)
+        repofs = _REP_INIT_DATA
+        ustruct.pack_into(
+            _REP_INIT, report, 0, _REP_MARKER, _REP_MAGIC, _REP_MAGIC, mtype, msize
+        )
 
-        # write the report
+        nwritten = 0
         while True:
-            await write
-            n = iface.write(report)
-            if n == len(report):
+            # copy as much as possible to the report buffer
+            nwritten += utils.memcpy(report, repofs, mdata, nwritten)
+
+            # write the report
+            while True:
+                await write
+                n = iface.write(report)
+                if n == len(report):
+                    break
+
+            # if we have more data to write, use continuation reports for it
+            if nwritten < msize:
+                repofs = _REP_CONT_DATA
+            else:
                 break
-
-        # if we have more data to write, use continuation reports for it
-        if nwritten < msize:
-            repofs = _REP_CONT_DATA
-        else:
-            break