From e03d404dca28caae901f9a9e350680c8e4143c3e Mon Sep 17 00:00:00 2001 From: tychovrahe Date: Tue, 28 May 2024 19:36:51 +0200 Subject: [PATCH] fix(core): ensure drop is called on layout objects [no changelog] --- core/embed/rust/librust_qstr.h | 1 + core/embed/rust/src/ui/component/base.rs | 38 +++++++++++----- core/embed/rust/src/ui/layout/obj.rs | 43 +++++++++++++++---- .../ui/model_mercury/component/homescreen.rs | 5 --- .../embed/rust/src/ui/model_mercury/layout.rs | 3 ++ core/embed/rust/src/ui/model_tt/layout.rs | 3 ++ core/mocks/generated/trezorui2.pyi | 4 ++ core/src/apps/homescreen/__init__.py | 19 ++++++-- core/src/boot.py | 1 + .../src/trezor/ui/layouts/mercury/__init__.py | 3 ++ core/src/trezor/ui/layouts/tr/__init__.py | 3 ++ core/src/trezor/ui/layouts/tt/__init__.py | 3 ++ 12 files changed, 98 insertions(+), 28 deletions(-) diff --git a/core/embed/rust/librust_qstr.h b/core/embed/rust/librust_qstr.h index 427d6ad64..257254cc5 100644 --- a/core/embed/rust/librust_qstr.h +++ b/core/embed/rust/librust_qstr.h @@ -22,6 +22,7 @@ static void _librust_qstrs(void) { MP_QSTR_NORMAL; MP_QSTR_TR; MP_QSTR_TranslationsHeader; + MP_QSTR___del__; MP_QSTR___dict__; MP_QSTR___name__; MP_QSTR_account; diff --git a/core/embed/rust/src/ui/component/base.rs b/core/embed/rust/src/ui/component/base.rs index 36f5963c1..ce77c93bf 100644 --- a/core/embed/rust/src/ui/component/base.rs +++ b/core/embed/rust/src/ui/component/base.rs @@ -208,29 +208,45 @@ where /// Same as `Child` but also handles screen clearing when layout is first /// painted. pub struct Root { - inner: Child, + inner: Option>, marked_for_clear: bool, } impl Root { pub fn new(component: T) -> Self { Self { - inner: Child::new(component), + inner: Some(Child::new(component)), marked_for_clear: true, } } + pub fn inner_mut(&mut self) -> &mut Child { + if let Some(ref mut c) = self.inner { + c + } else { + fatal_error!("deallocated", "Root object is deallocated") + } + } + pub fn inner(&self) -> &Child { - &self.inner + if let Some(ref c) = self.inner { + c + } else { + fatal_error!("deallocated", "Root object is deallocated") + } } pub fn skip_paint(&mut self) { - self.inner.skip_paint() + self.inner_mut().skip_paint(); } pub fn clear_screen(&mut self) { self.marked_for_clear = true; } + + pub fn delete(&mut self) { + self.inner = None; + } } impl Component for Root @@ -240,15 +256,15 @@ where type Msg = T::Msg; fn place(&mut self, bounds: Rect) -> Rect { - self.inner.place(bounds) + self.inner_mut().place(bounds) } fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option { - let msg = self.inner.event(ctx, event); + let msg = self.inner_mut().event(ctx, event); if ctx.needs_repaint_root() { self.marked_for_clear = true; let mut dummy_ctx = EventCtx::new(); - let paint_msg = self.inner.event(&mut dummy_ctx, Event::RequestPaint); + let paint_msg = self.inner_mut().event(&mut dummy_ctx, Event::RequestPaint); assert!(paint_msg.is_none()); assert!(dummy_ctx.timers.is_empty()); } @@ -256,7 +272,7 @@ where } fn paint(&mut self) { - if self.marked_for_clear && self.inner.will_paint() { + if self.marked_for_clear && self.inner().will_paint() { self.marked_for_clear = false; display::clear() } @@ -264,12 +280,12 @@ where } fn render<'s>(&self, target: &mut impl Renderer<'s>) { - self.inner.render(target); + self.inner().render(target); } #[cfg(feature = "ui_bounds")] fn bounds(&self, sink: &mut dyn FnMut(Rect)) { - self.inner.bounds(sink) + self.inner().bounds(sink) } } @@ -279,7 +295,7 @@ where T: crate::trace::Trace, { fn trace(&self, t: &mut dyn crate::trace::Tracer) { - self.inner.trace(t) + self.inner().trace(t); } } diff --git a/core/embed/rust/src/ui/layout/obj.rs b/core/embed/rust/src/ui/layout/obj.rs index 01130a3cc..154b7f507 100644 --- a/core/embed/rust/src/ui/layout/obj.rs +++ b/core/embed/rust/src/ui/layout/obj.rs @@ -54,6 +54,7 @@ pub trait ObjComponent: MaybeTrace { fn obj_bounds(&self, _sink: &mut dyn FnMut(Rect)) {} fn obj_skip_paint(&mut self) {} fn obj_request_clear(&mut self) {} + fn obj_delete(&mut self) {} } impl ObjComponent for Root @@ -105,6 +106,10 @@ where fn obj_request_clear(&mut self) { self.clear_screen() } + + fn obj_delete(&mut self) { + self.delete() + } } /// `LayoutObj` is a GC-allocated object exported to MicroPython, with type @@ -133,15 +138,25 @@ impl LayoutObj { let root = unsafe { Gc::from_raw(Gc::into_raw(Gc::new(wrapped_root)?) as *mut dyn ObjComponent) }; - Gc::new(Self { - base: Self::obj_type().as_base(), - inner: RefCell::new(LayoutObjInner { - root, - event_ctx: EventCtx::new(), - timer_fn: Obj::const_none(), - page_count: 1, - }), - }) + // SAFETY: This is a Python object and hase a base as first element + unsafe { + Gc::new_with_custom_finaliser(Self { + base: Self::obj_type().as_base(), + inner: RefCell::new(LayoutObjInner { + root, + event_ctx: EventCtx::new(), + timer_fn: Obj::const_none(), + page_count: 1, + }), + }) + } + } + + pub fn obj_delete(&self) { + let mut inner = self.inner.borrow_mut(); + + // SAFETY: `inner.root` is unique because of the `inner.borrow_mut()`. + unsafe { Gc::as_mut(&mut inner.root) }.obj_delete(); } pub fn skip_first_paint(&self) { @@ -291,6 +306,7 @@ impl LayoutObj { Qstr::MP_QSTR_request_complete_repaint => obj_fn_1!(ui_layout_request_complete_repaint).as_obj(), Qstr::MP_QSTR_trace => obj_fn_2!(ui_layout_trace).as_obj(), Qstr::MP_QSTR_bounds => obj_fn_1!(ui_layout_bounds).as_obj(), + Qstr::MP_QSTR___del__ => obj_fn_1!(ui_layout_delete).as_obj(), Qstr::MP_QSTR_page_count => obj_fn_1!(ui_layout_page_count).as_obj(), Qstr::MP_QSTR_button_request => obj_fn_1!(ui_layout_button_request).as_obj(), }), @@ -526,3 +542,12 @@ extern "C" fn ui_layout_bounds(this: Obj) -> Obj { extern "C" fn ui_layout_bounds(_this: Obj) -> Obj { Obj::const_none() } + +extern "C" fn ui_layout_delete(this: Obj) -> Obj { + let block = || { + let this: Gc = this.try_into()?; + this.obj_delete(); + Ok(Obj::const_none()) + }; + unsafe { util::try_or_raise(block) } +} diff --git a/core/embed/rust/src/ui/model_mercury/component/homescreen.rs b/core/embed/rust/src/ui/model_mercury/component/homescreen.rs index 7ddcb3db9..4cad5d68e 100644 --- a/core/embed/rust/src/ui/model_mercury/component/homescreen.rs +++ b/core/embed/rust/src/ui/model_mercury/component/homescreen.rs @@ -14,7 +14,6 @@ use crate::{ shape::{self, Renderer}, }, }; -use core::mem; use crate::ui::{ component::Label, @@ -344,10 +343,6 @@ impl Component for Lockscreen<'_> { } if let Event::Touch(TouchEvent::TouchEnd(_)) = event { - let bg_img = mem::replace(&mut self.bg_image, None); - if let Some(bg_img) = bg_img { - drop(bg_img); - } return Some(HomescreenMsg::Dismissed); } diff --git a/core/embed/rust/src/ui/model_mercury/layout.rs b/core/embed/rust/src/ui/model_mercury/layout.rs index d1306c9ef..a7c7f559c 100644 --- a/core/embed/rust/src/ui/model_mercury/layout.rs +++ b/core/embed/rust/src/ui/model_mercury/layout.rs @@ -1532,6 +1532,9 @@ pub static mp_module_trezorui2: Module = obj_module! { /// def page_count(self) -> int: /// """Return the number of pages in the layout object.""" /// + /// def __del__(self) -> None: + /// """Calls drop on contents of the root component.""" + /// /// class UiResult: /// """Result of a UI operation.""" /// pass diff --git a/core/embed/rust/src/ui/model_tt/layout.rs b/core/embed/rust/src/ui/model_tt/layout.rs index aeba4a33d..bcce7c6ed 100644 --- a/core/embed/rust/src/ui/model_tt/layout.rs +++ b/core/embed/rust/src/ui/model_tt/layout.rs @@ -1679,6 +1679,9 @@ pub static mp_module_trezorui2: Module = obj_module! { /// def button_request(self) -> tuple[int, str] | None: /// """Return (code, type) of button request made during the last event or timer pass.""" /// + /// def __del__(self) -> None: + /// """Calls drop on contents of the root component.""" + /// /// class UiResult: /// """Result of a UI operation.""" /// pass diff --git a/core/mocks/generated/trezorui2.pyi b/core/mocks/generated/trezorui2.pyi index 67604acc0..0bd64c86a 100644 --- a/core/mocks/generated/trezorui2.pyi +++ b/core/mocks/generated/trezorui2.pyi @@ -49,6 +49,8 @@ class LayoutObj(Generic[T]): """Paint bounds of individual components on screen.""" def page_count(self) -> int: """Return the number of pages in the layout object.""" + def __del__(self) -> None: + """Calls drop on contents of the root component.""" # rust/src/ui/model_mercury/layout.rs @@ -1122,6 +1124,8 @@ class LayoutObj(Generic[T]): """Return the number of pages in the layout object.""" def button_request(self) -> tuple[int, str] | None: """Return (code, type) of button request made during the last event or timer pass.""" + def __del__(self) -> None: + """Calls drop on contents of the root component.""" # rust/src/ui/model_tt/layout.rs diff --git a/core/src/apps/homescreen/__init__.py b/core/src/apps/homescreen/__init__.py index f07a7af62..3f29f1b42 100644 --- a/core/src/apps/homescreen/__init__.py +++ b/core/src/apps/homescreen/__init__.py @@ -12,7 +12,11 @@ from apps.common.authorization import is_set_any_session async def busyscreen() -> None: - await Busyscreen(busy_expiry_ms()) + obj = Busyscreen(busy_expiry_ms()) + try: + await obj + finally: + obj.__del__() async def homescreen() -> None: @@ -42,12 +46,17 @@ async def homescreen() -> None: elif storage.device.get_experimental_features(): notification = TR.homescreen__title_experimental_mode - await Homescreen( + obj = Homescreen( label=label, notification=notification, notification_is_error=notification_is_error, hold_to_lock=config.has_pin(), ) + try: + await obj + finally: + obj.__del__() + lock_device() @@ -58,10 +67,14 @@ async def _lockscreen(screensaver: bool = False) -> None: # Only show the lockscreen UI if the device can in fact be locked, or if it is # and OLED device (in which case the lockscreen is a screensaver). if can_lock_device() or screensaver: - await Lockscreen( + obj = Lockscreen( label=storage.device.get_label(), coinjoin_authorized=is_set_any_session(MessageType.AuthorizeCoinJoin), ) + try: + await obj + finally: + obj.__del__() # Otherwise proceed directly to unlock() call. If the device is already unlocked, # it should be a no-op storage-wise, but it resets the internal configuration # to an unlocked state. diff --git a/core/src/boot.py b/core/src/boot.py index f606da3dc..8b94b379c 100644 --- a/core/src/boot.py +++ b/core/src/boot.py @@ -62,6 +62,7 @@ async def bootscreen() -> None: label=storage.device.get_label(), bootscreen=True ) await lockscreen + lockscreen.__del__() await verify_user_pin() storage.init_unlocked() allow_all_loader_messages() diff --git a/core/src/trezor/ui/layouts/mercury/__init__.py b/core/src/trezor/ui/layouts/mercury/__init__.py index 4a3197875..3ce312c81 100644 --- a/core/src/trezor/ui/layouts/mercury/__init__.py +++ b/core/src/trezor/ui/layouts/mercury/__init__.py @@ -40,6 +40,9 @@ class RustLayout(ui.Layout): self._send_button_request() self.backlight_level = ui.BacklightLevels.NORMAL + def __del__(self): + self.layout.__del__() + def set_timer(self, token: int, deadline: int) -> None: self.timer.schedule(deadline, token) diff --git a/core/src/trezor/ui/layouts/tr/__init__.py b/core/src/trezor/ui/layouts/tr/__init__.py index aaaa19823..8e72314d9 100644 --- a/core/src/trezor/ui/layouts/tr/__init__.py +++ b/core/src/trezor/ui/layouts/tr/__init__.py @@ -44,6 +44,9 @@ class RustLayout(LayoutParentType[T]): self.layout.attach_timer_fn(self.set_timer) self._send_button_request() + def __del__(self): + self.layout.__del__() + def set_timer(self, token: int, deadline: int) -> None: self.timer.schedule(deadline, token) diff --git a/core/src/trezor/ui/layouts/tt/__init__.py b/core/src/trezor/ui/layouts/tt/__init__.py index 25aefe2c5..6d15501d5 100644 --- a/core/src/trezor/ui/layouts/tt/__init__.py +++ b/core/src/trezor/ui/layouts/tt/__init__.py @@ -46,6 +46,9 @@ class RustLayout(LayoutParentType[T]): self._send_button_request() self.backlight_level = ui.BacklightLevels.NORMAL + def __del__(self): + self.layout.__del__() + def set_timer(self, token: int, deadline: int) -> None: self.timer.schedule(deadline, token)