← Watch feed
refactor(core): extract `ButtonRequest` handler into a separate method
Public commit record
77/100 · Adequate
What the developer wrote
Authored by Roman Zeyde
refactor(core): extract `ButtonRequest` handler into a separate method
Also, add a `Context.read()` overload for unexpected messages.
[no changelog]
✓ Specific, descriptive subject✓ Names a concrete action or component✓ Uses a recognizable type or scope✓ Provides an explanatory body
What changed, and why it matters
This commit is a code cleanup that moves the handling of on-screen button prompts into a separate helper function and adds a way to tell the message-reading code to reject any unexpected message. The change itself does not appear to fix a specific security bug, but it touches the code that decides which messages the device accepts while a user is confirming an action. There is no vendor statement that this is a security fix, and no independent researcher is credited.
Recommended action
Treat as a routine refactor. If auditing, verify that `Context.read(None)` callers cannot accidentally pass `None` where a real expected-type set is required, and that the new `button_request_handler()` still waits for `ButtonAck` before allowing the layout result to be returned, preserving the existing synchronization behavior.
Security signals we found
01
Behavior change in wire protocol read path: `expected_types=None` now causes every message to be treated as unexpected
02
Refactor of button-request ACK synchronization, which the existing comment says prevents THP channel desync
03
No changelog entry and commit is titled as a refactor
Evidence from the diff
The patch refactors ButtonRequest handling in core/src/trezor/ui/__init__.py into a standalone button_request_handler() coroutine and adds a Context.read(None) overload that always raises UnexpectedMessageException. Implementations in codec_context.py, pairing_context.py, and session_context.py now treat expected_types=None as “no expected types,” causing any received message to be rejected as unexpected. The refactor also changes the button-request mailbox from carrying raw (code, name) tuples to carrying constructed ButtonRequest objects, and introduces put_button_request() to centralize that conversion. The commit message calls this a refactor and explicitly says “[no changelog].”
Changed components
core/src/trezor/ui/__init__.pycore/src/trezor/ui/layouts/common.pycore/src/trezor/wire/codec/codec_context.pycore/src/trezor/wire/protocol_common.pycore/src/trezor/wire/thp/pairing_context.pycore/src/trezor/wire/thp/session_context.pyInspect captured patch +62 / −50
diff --git a/core/src/trezor/ui/__init__.py b/core/src/trezor/ui/__init__.py
index a0342b24..59b195cf 100644
--- a/core/src/trezor/ui/__init__.py
+++ b/core/src/trezor/ui/__init__.py
@@ -157,7 +157,7 @@ class Layout(Generic[T]):
self.timers: dict[int, loop.Task] = {}
self.result_box: loop.mailbox[Any] = loop.mailbox()
self.button_request_ack_pending: bool = False
- self.button_request_box: loop.mailbox[ButtonRequestMsg] = loop.mailbox()
+ self.button_request_box: loop.mailbox[ButtonRequest | None] = loop.mailbox()
self.button_request_task: loop.Task | None = None
self.transition_out: AttachType | None = None
self.backlight_level = BacklightLevels.NORMAL
@@ -275,9 +275,20 @@ class Layout(Generic[T]):
# else we are (a) still running or (b) already finished
is_done = None
try:
- if self.context is not None and self.result_box.is_empty():
+ if (ctx := self.context) is not None and self.result_box.is_empty():
is_done = loop.mailbox() # (see below)
- self.button_request_task = self._handle_button_requests(is_done)
+
+ def _button_request_task() -> Generator[Any, Any, None]:
+ try:
+ yield from button_request_handler(
+ context=ctx,
+ button_requests=self.button_request_box,
+ ack_callback=self._button_request_acked,
+ )
+ finally:
+ is_done.put(None)
+
+ self.button_request_task = _button_request_task()
self._start_task(self.button_request_task)
elif __debug__ and not self.button_request_box.is_empty():
log.debug(
@@ -292,7 +303,7 @@ class Layout(Generic[T]):
if is_done is not None:
# Make sure ButtonRequest is ACKed, before the result is returned.
# Otherwise, THP channel may become desynced (due to two consecutive writes).
- self.button_request_box.put(None, replace=True)
+ self.put_button_request(None)
task = loop.spawn(_waiting_screen())
try:
await is_done
@@ -365,10 +376,16 @@ class Layout(Generic[T]):
"don't forget to yield your input flow from time to time ^_^"
)
- # in production, we don't want this to fail, hence replace=True
- self.button_request_box.put(res, replace=True)
+ self.put_button_request(res)
return True
+ def put_button_request(self, msg: ButtonRequestMsg | None) -> None:
+ br = msg and ButtonRequest(
+ code=msg[0], name=msg[1], pages=self.layout.page_count()
+ )
+ # in production, we don't want this to fail, hence replace=True
+ self.button_request_box.put(br, replace=True)
+
def _paint(self) -> None:
"""Paint the layout and ensure that homescreen cache is properly invalidated."""
import storage.cache as storage_cache
@@ -473,41 +490,12 @@ class Layout(Generic[T]):
finally:
touch.close()
- async def _handle_button_requests(self, is_done: loop.mailbox[None] | None) -> None:
- try:
- if self.context is None:
- return
- while True:
- # The following task will raise `UnexpectedMessageException` on any message.
- unexpected_read = self.context.read(())
- result = await loop.race(unexpected_read, self.button_request_box)
- if result is None:
- return # exit the loop when the layout is done.
- assert isinstance(result, tuple)
- br_code, br_name = result
-
- if __debug__:
- log.info(__name__, "ButtonRequest sent: %s", br_name)
- await self.context.call(
- ButtonRequest(
- code=br_code, pages=self.layout.page_count(), name=br_name
- ),
- ButtonAck,
- )
- if __debug__:
- log.info(__name__, "ButtonRequest acked: %s", br_name)
-
- if (
- self.button_request_ack_pending
- and self.state is LayoutState.TRANSITIONING
- ):
- self.button_request_ack_pending = False
- self.state = LayoutState.ATTACHED
- if __debug__:
- self.notify_debuglink(self)
- finally:
- if is_done is not None:
- is_done.put(None)
+ def _button_request_acked(self) -> None:
+ if self.button_request_ack_pending and self.state is LayoutState.TRANSITIONING:
+ self.button_request_ack_pending = False
+ self.state = LayoutState.ATTACHED
+ if __debug__:
+ self.notify_debuglink(self)
if utils.USE_BLE:
@@ -581,6 +569,26 @@ class Layout(Generic[T]):
self.layout.__del__()
+async def button_request_handler(
+ context: Context,
+ button_requests: loop.mailbox[ButtonRequest | None],
+ ack_callback: Callable[[], None],
+) -> None:
+ while True:
+ # The following task will raise `UnexpectedMessageException` on any message.
+ unexpected_read = context.read(None)
+ br = await loop.race(unexpected_read, button_requests)
+ if br is None:
+ return # exit the loop when the layout is done.
+
+ if __debug__:
+ log.info(__name__, "ButtonRequest sent: %s", br.name)
+ await context.call(br, ButtonAck)
+ if __debug__:
+ log.info(__name__, "ButtonRequest acked: %s", br.name)
+ ack_callback()
+
+
class ProgressLayout:
"""Progress layout.
diff --git a/core/src/trezor/ui/layouts/common.py b/core/src/trezor/ui/layouts/common.py
index 5ba2d2e2..ef988f05 100644
--- a/core/src/trezor/ui/layouts/common.py
+++ b/core/src/trezor/ui/layouts/common.py
@@ -65,7 +65,7 @@ async def interact(
layout.start()
if br_name is not None:
# store the first button request to be sent
- layout.button_request_box.put((br_code, br_name))
+ layout.put_button_request((br_code, br_name))
# wait for the layout result
result = await layout.get_result()
diff --git a/core/src/trezor/wire/codec/codec_context.py b/core/src/trezor/wire/codec/codec_context.py
index ed83ac13..d3a65fd0 100644
--- a/core/src/trezor/wire/codec/codec_context.py
+++ b/core/src/trezor/wire/codec/codec_context.py
@@ -43,7 +43,7 @@ class CodecContext(Context):
async def read(
self,
- expected_types: Container[int],
+ expected_types: Container[int] | None,
expected_type: type[protobuf.MessageType] | None = None,
) -> protobuf.MessageType:
if __debug__:
@@ -59,7 +59,7 @@ class CodecContext(Context):
# If we got a message with unexpected type, raise the message via
# `UnexpectedMessageError` and let the session handler deal with it.
- if msg.type not in expected_types:
+ if not expected_types or msg.type not in expected_types:
raise UnexpectedMessageException(msg)
if expected_type is None:
diff --git a/core/src/trezor/wire/protocol_common.py b/core/src/trezor/wire/protocol_common.py
index 822ad286..384781c2 100644
--- a/core/src/trezor/wire/protocol_common.py
+++ b/core/src/trezor/wire/protocol_common.py
@@ -5,7 +5,7 @@ from trezor import protobuf
if TYPE_CHECKING:
from buffer_types import AnyBytes
from trezorio import WireInterface
- from typing import Awaitable, Container, TypeVar, overload
+ from typing import Awaitable, Container, Literal, NoReturn, TypeVar, overload
from storage.cache_common import DataCache
@@ -52,6 +52,10 @@ class Context:
if TYPE_CHECKING:
+ # Will always raise `UnexpectedMessageException` after receiving a message.
+ @overload
+ async def read(self, expected_types: Literal[None]) -> NoReturn: ...
+
@overload
async def read(
self, expected_types: Container[int]
@@ -64,7 +68,7 @@ class Context:
async def read(
self,
- expected_types: Container[int],
+ expected_types: Container[int] | None,
expected_type: type[protobuf.MessageType] | None = None,
) -> protobuf.MessageType:
"""Read a message from the wire.
diff --git a/core/src/trezor/wire/thp/pairing_context.py b/core/src/trezor/wire/thp/pairing_context.py
index 0e4466fa..0d0a2d83 100644
--- a/core/src/trezor/wire/thp/pairing_context.py
+++ b/core/src/trezor/wire/thp/pairing_context.py
@@ -91,7 +91,7 @@ class PairingContext(Context):
async def read(
self,
- expected_types: Container[int],
+ expected_types: Container[int] | None,
expected_type: type[protobuf.MessageType] | None = None,
) -> protobuf.MessageType:
if __debug__:
@@ -107,7 +107,7 @@ class PairingContext(Context):
)
_, message = await self.channel_ctx.decrypt_message()
- if message.type not in expected_types:
+ if not expected_types or message.type not in expected_types:
from trezor.messages import Cancel
if message.type == Cancel.MESSAGE_WIRE_TYPE:
diff --git a/core/src/trezor/wire/thp/session_context.py b/core/src/trezor/wire/thp/session_context.py
index 7714feaf..dd76ffb0 100644
--- a/core/src/trezor/wire/thp/session_context.py
+++ b/core/src/trezor/wire/thp/session_context.py
@@ -77,7 +77,7 @@ class GenericSessionContext(Context):
async def read(
self,
- expected_types: Container[int],
+ expected_types: Container[int] | None,
expected_type: type[protobuf.MessageType] | None = None,
) -> protobuf.MessageType:
if __debug__:
@@ -93,7 +93,7 @@ class GenericSessionContext(Context):
)
message = await self._read_next_message()
- if message.type not in expected_types:
+ if not expected_types or message.type not in expected_types:
if __debug__:
log.debug(
__name__,
Risk score
Our methodology →Why this scored 19/100
Human-validated context
Community notes
Notes can correct, qualify, or add evidence to the AI analysis. Every note shown here has been validated by a human moderator.
No validated notes yet.
The AI analysis stands alone for now. Submit a note if you can add evidence or important context.