feat(core): allow using LayoutObj as a context manager
What the developer wrote
Authored by Roman Zeyde
feat(core): allow using LayoutObj as a context manager
[no changelog]
What changed, and why it matters
This commit adds a Python 'with' statement feature to Trezor's Rust-based UI layout objects. It lets code explicitly clean up a UI layout when leaving a block, and raises an error if someone tries to use an already-destroyed layout. There is no direct security bug visible in the change; it is a feature addition for explicit resource management.
No immediate action required. Review callers that adopt the context manager to ensure obj_delete() semantics are safe when invoked via __exit__, including during exception unwinding, and confirm no use-after-free can occur if __exit__ is called multiple times.
Security signals we found
Adds explicit context-manager lifecycle control to LayoutObj
__enter__ raises RuntimeError if root layout already dropped
__exit__ unconditionally calls obj_delete() on the LayoutObj
No input validation changes beyond argument count check
No changelog entry provided
Evidence from the diff
The patch exposes enter and exit methods on LayoutObj in the Rust/MicroPython binding. enter checks that the root component is still present and returns self. exit calls obj_delete() to drop the root component. It also registers the corresponding qstrs and updates the Python stub. The change is additive and does not alter existing deallocation paths.
Changed components
core/embed/rust/src/ui/layout/obj.rscore/embed/rust/src/ui/api/firmware_micropython.rscore/embed/rust/librust_qstr.hcore/mocks/generated/trezorui_api.pyiInspect captured patch +41 / −0
diff --git a/core/embed/rust/librust_qstr.h b/core/embed/rust/librust_qstr.h
index 812984ea..384ff3ea 100644
--- a/core/embed/rust/librust_qstr.h
+++ b/core/embed/rust/librust_qstr.h
@@ -78,6 +78,8 @@ static void _librust_qstrs(void) {
MP_QSTR_WipeDevice;
MP_QSTR___del__;
MP_QSTR___dict__;
+ MP_QSTR___enter__;
+ MP_QSTR___exit__;
MP_QSTR___name__;
MP_QSTR_about_items;
MP_QSTR_account;
diff --git a/core/embed/rust/src/ui/api/firmware_micropython.rs b/core/embed/rust/src/ui/api/firmware_micropython.rs
index 27a95641..572b240b 100644
--- a/core/embed/rust/src/ui/api/firmware_micropython.rs
+++ b/core/embed/rust/src/ui/api/firmware_micropython.rs
@@ -1441,6 +1441,12 @@ pub static mp_module_trezorui_api: Module = obj_module! {
/// def __del__(self) -> None:
/// """Calls drop on contents of the root component."""
///
+ /// def __enter__(self) -> LayoutObj[T]:
+ /// """Enters a context manager (checking the root component is not dropped)."""
+ ///
+ /// def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ /// """Exits a context manager (dropping the root component)."""
+ ///
/// class UiResult:
/// """Result of a UI operation."""
/// pass
diff --git a/core/embed/rust/src/ui/layout/obj.rs b/core/embed/rust/src/ui/layout/obj.rs
index 121c9bc3..68b7819c 100644
--- a/core/embed/rust/src/ui/layout/obj.rs
+++ b/core/embed/rust/src/ui/layout/obj.rs
@@ -423,6 +423,10 @@ impl LayoutObj {
Qstr::MP_QSTR_button_request => obj_fn_1!(ui_layout_button_request).as_obj(),
Qstr::MP_QSTR_get_transition_out => obj_fn_1!(ui_layout_get_transition_out).as_obj(),
Qstr::MP_QSTR_return_value => obj_fn_1!(ui_layout_return_value).as_obj(),
+
+ // Allow using LayoutObj as context manager for explicit deallocation.
+ Qstr::MP_QSTR___enter__ => obj_fn_1!(ui_layout_enter).as_obj(),
+ Qstr::MP_QSTR___exit__ => obj_fn_var!(4, 4, ui_layout_exit).as_obj(),
}),
};
&TYPE
@@ -718,3 +722,28 @@ extern "C" fn ui_layout_delete(this: Obj) -> Obj {
};
unsafe { util::try_or_raise(block) }
}
+
+extern "C" fn ui_layout_enter(this: Obj) -> Obj {
+ let block = || {
+ let obj: Gc<LayoutObj> = this.try_into()?;
+ // Raise an exception if the root layout has been already dropped.
+ obj.inner_mut()
+ .root
+ .as_ref()
+ .ok_or(Error::RuntimeError(c"No root layout on enter"))?;
+ Ok(this)
+ };
+ unsafe { util::try_or_raise(block) }
+}
+
+extern "C" fn ui_layout_exit(n_args: usize, args: *const Obj) -> Obj {
+ let block = |args: &[Obj], _kwargs: &Map| {
+ if args.len() != 4 {
+ return Err(Error::TypeError);
+ }
+ let this: Gc<LayoutObj> = args[0].try_into()?;
+ this.inner_mut().obj_delete();
+ Ok(Obj::const_none())
+ };
+ unsafe { util::try_with_args_and_kwargs(n_args, args, &Map::EMPTY, block) }
+}
diff --git a/core/mocks/generated/trezorui_api.pyi b/core/mocks/generated/trezorui_api.pyi
index 70fb04bf..1c794394 100644
--- a/core/mocks/generated/trezorui_api.pyi
+++ b/core/mocks/generated/trezorui_api.pyi
@@ -80,6 +80,10 @@ class LayoutObj(Generic[T]):
"""Retrieve the return value of the layout object."""
def __del__(self) -> None:
"""Calls drop on contents of the root component."""
+ def __enter__(self) -> LayoutObj[T]:
+ """Enters a context manager (checking the root component is not dropped)."""
+ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ """Exits a context manager (dropping the root component)."""
# rust/src/ui/api/firmware_micropython.rs
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.