refactor(python): move invalidation & refresh into workflow() decorator
What the developer wrote
Authored by Roman Zeyde
refactor(python): move invalidation & refresh into workflow() decorator
Otherwise, THP ACK piggybacking may result in unnecessary re-transmissions.
[no changelog]
What changed, and why it matters
This is a small internal cleanup in Trezor's Python library. It moves the 'refresh device info' and 'invalidate client' steps into a shared decorator so they happen after a workflow finishes, rather than inside each individual function. The stated reason is to avoid unnecessary re-transmissions in the new THP (Trezor Host Protocol) when an ACK is piggybacked. There is no direct security bug being fixed; it is a reliability/refactoring change.
No immediate security action required. Treat as a normal reliability/refactoring patch. Reviewers may want to confirm that moving `refresh_features()` outside the session context does not change error-handling behavior for failed workflows.
Security signals we found
Refactor of session lifecycle / feature refresh timing
Comment references THP ACK piggybacking and retransmission avoidance
No explicit vulnerability, CVE, or security fix described in commit message
Evidence from the diff
The commit refactors workflow() in tools.py to accept refresh_features=True and invalidate_client=True parameters. It then removes explicit session.refresh_features() and session.client._invalidate() calls from many device operations in device.py, and moves one refresh_features() call in client.py outside the with session: block. The decorator now performs these actions after the wrapped function returns and the session context has exited. A code comment explains that this separation is needed because the MicroPython event loop may restart after the workflow, and sending GetFeatures in a separate interaction avoids THP retransmissions.
Changed components
python/src/trezorlib/tools.pypython/src/trezorlib/device.pypython/src/trezorlib/client.pyInspect captured patch +29 / −26
diff --git a/python/src/trezorlib/client.py b/python/src/trezorlib/client.py
index 7a7481d6..6d805adb 100644
--- a/python/src/trezorlib/client.py
+++ b/python/src/trezorlib/client.py
@@ -562,7 +562,7 @@ class TrezorClient(t.Generic[SessionType], metaclass=ABCMeta):
session = _use_session or self._get_any_session()
with session:
session.call_raw(messages.LockDevice())
- self.refresh_features()
+ self.refresh_features()
def ensure_unlocked(self) -> None:
"""Ensure the device is unlocked."""
diff --git a/python/src/trezorlib/device.py b/python/src/trezorlib/device.py
index 35641b2b..f0a8cf2e 100644
--- a/python/src/trezorlib/device.py
+++ b/python/src/trezorlib/device.py
@@ -41,7 +41,7 @@ ENTROPY_CHECK_MIN_VERSION = (2, 8, 7)
HOMESCREEN_STREAMING_MIN_VERSION = (2, 8, 11)
-@workflow()
+@workflow(refresh_features=True)
def apply_settings(
session: "Session",
label: Optional[str] = None,
@@ -82,7 +82,6 @@ def apply_settings(
else:
settings.homescreen = homescreen
session.call(settings, expect=messages.Success)
- session.refresh_features()
def _send_chunked_data(
@@ -99,7 +98,7 @@ def _send_chunked_data(
response = session.call(messages.DataChunkAck(data_chunk=chunk))
-@workflow()
+@workflow(refresh_features=True)
def change_language(
session: "Session",
language_data: bytes,
@@ -114,40 +113,34 @@ def change_language(
_send_chunked_data(session, response, language_data)
else:
messages.Success.ensure_isinstance(response)
- session.refresh_features() # changing the language in features
-@workflow()
+@workflow(refresh_features=True)
def apply_flags(session: "Session", flags: int) -> None:
session.call(messages.ApplyFlags(flags=flags), expect=messages.Success)
- session.refresh_features()
-@workflow()
+@workflow(refresh_features=True)
def change_pin(session: "Session", remove: bool = False) -> None:
session.call(messages.ChangePin(remove=remove), expect=messages.Success)
- session.refresh_features()
-@workflow()
+@workflow(refresh_features=True)
def change_wipe_code(session: "Session", remove: bool = False) -> None:
session.call(messages.ChangeWipeCode(remove=remove), expect=messages.Success)
- session.refresh_features()
-@workflow()
+@workflow(refresh_features=True)
def sd_protect(session: "Session", operation: messages.SdProtectOperationType) -> None:
session.call(messages.SdProtect(operation=operation), expect=messages.Success)
- session.refresh_features()
-@workflow()
+@workflow(invalidate_client=True)
def wipe(session: "Session") -> None:
session.call(messages.WipeDevice(), expect=messages.Success)
- session.client._invalidate()
-@workflow()
+@workflow(refresh_features=True)
def recover(
session: "Session",
word_count: int = 24,
@@ -224,9 +217,7 @@ def recover(
res = session.call(messages.Cancel())
# check that the result is a Success
- res = messages.Success.ensure_isinstance(res)
- # reinitialize the device
- session.refresh_features()
+ messages.Success.ensure_isinstance(res)
def is_slip39_backup_type(backup_type: messages.BackupType) -> bool:
@@ -319,7 +310,7 @@ def _get_external_entropy() -> bytes:
return secrets.token_bytes(32)
-@workflow()
+@workflow(refresh_features=True)
def setup(
session: "Session",
*,
@@ -425,7 +416,6 @@ def setup(
_reset_no_entropycheck(session, msg, _get_entropy)
xpubs = []
- session.refresh_features()
return xpubs
@@ -558,7 +548,7 @@ def _reset_with_entropycheck(
return xpubs
-@workflow()
+@workflow(refresh_features=True)
def backup(
session: "Session",
group_threshold: Optional[int] = None,
@@ -574,7 +564,6 @@ def backup(
),
expect=messages.Success,
)
- session.refresh_features()
@workflow()
@@ -622,7 +611,7 @@ def unlock_bootloader(session: "Session") -> None:
session.call(messages.UnlockBootloader(), expect=messages.Success)
-@workflow()
+@workflow(refresh_features=True)
def set_busy(session: "Session", expiry_ms: Optional[int]) -> None:
"""Sets or clears the busy state of the device.
@@ -630,7 +619,6 @@ def set_busy(session: "Session", expiry_ms: Optional[int]) -> None:
Setting `expiry_ms=None` clears the busy state.
"""
session.call(messages.SetBusy(expiry_ms=expiry_ms), expect=messages.Success)
- session.refresh_features()
@workflow()
diff --git a/python/src/trezorlib/tools.py b/python/src/trezorlib/tools.py
index d94d87ae..135d950d 100644
--- a/python/src/trezorlib/tools.py
+++ b/python/src/trezorlib/tools.py
@@ -381,6 +381,8 @@ class workflow(t.Generic[P, R]):
from_version: tuple[int, int, int] | None = None,
capability: messages.Capability | None = None,
capabilities: set[messages.Capability] | None = None,
+ refresh_features: bool = False,
+ invalidate_client: bool = False,
) -> None:
self.from_version = from_version
if capability is not None and capabilities is not None:
@@ -390,6 +392,8 @@ class workflow(t.Generic[P, R]):
elif capabilities is None:
capabilities = set()
self.capabilities = capabilities
+ self.refresh_features = refresh_features
+ self.invalidate_client = invalidate_client
self.func: SessionFunc[P, R] | None = None
def __call__(self, func: SessionFunc[P, R]) -> SessionFunc[P, R]:
@@ -400,5 +404,16 @@ class workflow(t.Generic[P, R]):
__tracebackhide__ = True # for pytest # pylint: disable=W0612
if self.func is None:
raise RuntimeError("workflow decorator must be used with a function")
+
with session:
- return self.func(session, *args, **kwargs)
+ result = self.func(session, *args, **kwargs)
+
+ if self.invalidate_client:
+ session.client._invalidate()
+
+ if self.refresh_features:
+ # MicroPython event loop may get restarted after running the workflow above,
+ # so `GetFeatures` will be sent in a separate interaction (to avoid THP retransmissions).
+ session.refresh_features()
+
+ return result
Why this scored 18/100
Community notes
Notes can correct, qualify, or add evidence to the AI analysis. Every note shown here has been validated by a human moderator.
The AI analysis stands alone for now. Submit a note if you can add evidence or important context.