WIP feat(core/ui): sending button requests from rust

[no changelog]

core/embed/rust/librust_qstr.h

@@ -93,6 +93,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_bounds;
   MP_QSTR_button;
   MP_QSTR_button_event;
+  MP_QSTR_button_request;
   MP_QSTR_buttons__abort;
   MP_QSTR_buttons__access;
   MP_QSTR_buttons__again;

core/embed/rust/src/ui/button_request.rs

@@ -0,0 +1,44 @@
+// ButtonRequestType from messages-common.proto
+// Eventually this should be generated
+#[derive(Clone, Copy)]
+#[repr(u16)]
+pub enum ButtonRequestCode {
+    Other = 1,
+    FeeOverThreshold = 2,
+    ConfirmOutput = 3,
+    ResetDevice = 4,
+    ConfirmWord = 5,
+    WipeDevice = 6,
+    ProtectCall = 7,
+    SignTx = 8,
+    FirmwareCheck = 9,
+    Address = 10,
+    PublicKey = 11,
+    MnemonicWordCount = 12,
+    MnemonicInput = 13,
+    UnknownDerivationPath = 15,
+    RecoveryHomepage = 16,
+    Success = 17,
+    Warning = 18,
+    PassphraseEntry = 19,
+    PinEntry = 20,
+}
+
+impl ButtonRequestCode {
+    pub fn num(&self) -> u16 {
+        *self as u16
+    }
+
+    pub fn with_type(self, br_type: &'static str) -> Option<ButtonRequest> {
+        Some(ButtonRequest {
+            code: self,
+            br_type,
+        })
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct ButtonRequest {
+    pub code: ButtonRequestCode,
+    pub br_type: &'static str,
+}

core/embed/rust/src/ui/component/base.rs

@@ -6,6 +6,7 @@ use crate::{
     strutil::TString,
     time::Duration,
     ui::{
+        button_request::{ButtonRequest, ButtonRequestCode},
         component::{maybe::PaintOverlapping, MsgMap},
         display::{self, Color},
         geometry::{Offset, Rect},
@@ -487,6 +488,7 @@ pub struct EventCtx {
     paint_requested: bool,
     anim_frame_scheduled: bool,
     page_count: Option<usize>,
+    button_request: Option<ButtonRequest>,
     root_repaint_requested: bool,
 }
 
@@ -513,6 +515,7 @@ impl EventCtx {
                                     * `Child::marked_for_paint` being true. */
             anim_frame_scheduled: false,
             page_count: None,
+            button_request: None,
             root_repaint_requested: false,
         }
     }
@@ -570,6 +573,16 @@ impl EventCtx {
         self.page_count
     }
 
+    pub fn send_button_request(&mut self, code: ButtonRequestCode, br_type: &'static str) {
+        #[cfg(feature = "ui_debug")]
+        assert!(self.button_request.is_none());
+        self.button_request = Some(ButtonRequest { code, br_type });
+    }
+
+    pub fn button_request(&mut self) -> Option<ButtonRequest> {
+        self.button_request.take()
+    }
+
     pub fn pop_timer(&mut self) -> Option<(TimerToken, Duration)> {
         self.timers.pop()
     }
@@ -579,6 +592,9 @@ impl EventCtx {
         self.paint_requested = false;
         self.anim_frame_scheduled = false;
         self.page_count = None;
+        #[cfg(feature = "ui_debug")]
+        assert!(self.button_request.is_none());
+        self.button_request = None;
         self.root_repaint_requested = false;
     }
 

core/embed/rust/src/ui/component/button_request.rs

@@ -0,0 +1,84 @@
+use crate::ui::{
+    button_request::{ButtonRequest, ButtonRequestCode},
+    component::{Component, Event, EventCtx},
+    geometry::Rect,
+    shape::Renderer,
+};
+
+/// Component that sends a ButtonRequest after receiving Event::Attach. The
+/// request is only sent once.
+#[derive(Clone, Copy)]
+pub struct OneButtonRequest<T> {
+    button_request: Option<ButtonRequest>,
+    pub inner: T,
+}
+
+impl<T> OneButtonRequest<T> {
+    pub fn new(button_request: ButtonRequest, inner: T) -> Self {
+        Self {
+            button_request: Some(button_request),
+            inner,
+        }
+    }
+}
+
+impl<T: Component> Component for OneButtonRequest<T> {
+    type Msg = T::Msg;
+
+    fn place(&mut self, bounds: Rect) -> Rect {
+        self.inner.place(bounds)
+    }
+
+    fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<Self::Msg> {
+        if matches!(event, Event::Attach) {
+            if let Some(button_request) = self.button_request.take() {
+                ctx.send_button_request(button_request.code, button_request.br_type)
+            }
+        }
+        self.inner.event(ctx, event)
+    }
+
+    fn paint(&mut self) {
+        self.inner.paint()
+    }
+
+    fn render<'s>(&'s self, target: &mut impl Renderer<'s>) {
+        self.inner.render(target)
+    }
+}
+
+#[cfg(feature = "ui_debug")]
+impl<T: crate::trace::Trace> crate::trace::Trace for OneButtonRequest<T> {
+    fn trace(&self, t: &mut dyn crate::trace::Tracer) {
+        self.inner.trace(t)
+    }
+}
+
+pub trait ButtonRequestExt {
+    fn one_button_request(
+        self,
+        code: ButtonRequestCode,
+        br_type: &'static str,
+    ) -> OneButtonRequest<Self>
+    where
+        Self: Sized,
+    {
+        OneButtonRequest::new(ButtonRequest { code, br_type }, self)
+    }
+}
+
+impl<T: Component> ButtonRequestExt for T {}
+
+#[cfg(all(feature = "micropython", feature = "touch"))]
+impl<T> crate::ui::flow::Swipable for OneButtonRequest<T>
+where
+    T: Component + crate::ui::flow::Swipable,
+{
+    fn can_swipe(&self, direction: crate::ui::component::SwipeDirection) -> bool {
+        self.inner.can_swipe(direction)
+    }
+
+    fn swiped(&mut self, ctx: &mut EventCtx, direction: crate::ui::component::SwipeDirection) {
+        self.inner.swiped(ctx, direction)
+    }
+}

core/embed/rust/src/ui/component/mod.rs

@@ -3,6 +3,7 @@
 pub mod bar;
 pub mod base;
 pub mod border;
+pub mod button_request;
 pub mod connect;
 pub mod empty;
 pub mod image;
@@ -24,6 +25,7 @@ pub mod timeout;
 pub use bar::Bar;
 pub use base::{Child, Component, ComponentExt, Event, EventCtx, Never, Root, TimerToken};
 pub use border::Border;
+pub use button_request::{ButtonRequestExt, OneButtonRequest};
 pub use empty::Empty;
 #[cfg(all(feature = "jpeg", feature = "micropython"))]
 pub use jpeg::Jpeg;

core/embed/rust/src/ui/layout/obj.rs

@@ -17,6 +17,7 @@ use crate::{
     },
     time::Duration,
     ui::{
+        button_request::ButtonRequest,
         component::{Component, Event, EventCtx, Never, Root, TimerToken},
         constant,
         display::{sync, Color},
@@ -248,6 +249,17 @@ impl LayoutObj {
         self.inner.borrow().page_count.into()
     }
 
+    fn obj_button_request(&self) -> Result<Obj, Error> {
+        let inner = &mut *self.inner.borrow_mut();
+
+        match inner.event_ctx.button_request() {
+            None => Ok(Obj::const_none()),
+            Some(ButtonRequest { code, br_type }) => {
+                (code.num().into(), br_type.try_into()?).try_into()
+            }
+        }
+    }
+
     #[cfg(feature = "ui_debug")]
     fn obj_bounds(&self) {
         use crate::ui::display;
@@ -280,6 +292,7 @@ impl LayoutObj {
                 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_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(),
             }),
         };
         &TYPE
@@ -470,6 +483,14 @@ extern "C" fn ui_layout_page_count(this: Obj) -> Obj {
     unsafe { util::try_or_raise(block) }
 }
 
+extern "C" fn ui_layout_button_request(this: Obj) -> Obj {
+    let block = || {
+        let this: Gc<LayoutObj> = this.try_into()?;
+        Ok(this.obj_button_request()?)
+    };
+    unsafe { util::try_or_raise(block) }
+}
+
 #[cfg(feature = "ui_debug")]
 #[no_mangle]
 pub extern "C" fn ui_debug_layout_type() -> &'static Type {

core/embed/rust/src/ui/mod.rs

@@ -2,6 +2,7 @@
 pub mod macros;
 
 pub mod animation;
+pub mod button_request;
 pub mod component;
 pub mod constant;
 pub mod display;

core/embed/rust/src/ui/model_mercury/flow/get_address.rs

@@ -1,6 +1,7 @@
 use crate::{
     error,
     ui::{
+        button_request::ButtonRequestCode,
         component::{
             image::BlendedImage,
             text::paragraphs::{Paragraph, Paragraphs},
@@ -116,7 +117,8 @@ impl GetAddress {
                 )
                 .with_subtitle("address".into())
                 .with_menu_button()
-                .map(|msg| matches!(msg, FrameMsg::Button(_)).then_some(FlowMsg::Info)),
+                .map(|msg| matches!(msg, FrameMsg::Button(_)).then_some(FlowMsg::Info))
+                .one_button_request(ButtonRequestCode::Address, "show_address"),
             )?
             .add(
                 Frame::left_aligned(

core/src/apps/base.py

@@ -240,9 +240,10 @@ async def handle_EndSession(msg: EndSession) -> Success:
 async def handle_Ping(msg: Ping) -> Success:
     if msg.button_protection:
         from trezor.enums import ButtonRequestType as B
-        from trezor.ui.layouts import confirm_action
+        from trezor.ui.layouts import confirm_action, flow_demo
 
-        await confirm_action("ping", TR.words__confirm, "ping", br_code=B.ProtectCall)
+        await flow_demo()
+        # await confirm_action("ping", TR.words__confirm, "ping", br_code=B.ProtectCall)
     return Success(message=msg.message)
 
 

core/src/trezor/ui/layouts/common.py

@@ -39,4 +39,5 @@ async def interact(
         # We know for certain how many pages the layout will have
         pages = layout.page_count()  # type: ignore [Cannot access member "page_count" for type "LayoutType"]
     await button_request(br_type, br_code, pages)
-    return await context.wait(layout)
+    # return await context.wait(layout)
+    return await layout

core/src/trezor/ui/layouts/mercury/__init__.py

@@ -1,9 +1,10 @@
 from typing import TYPE_CHECKING
 
 import trezorui2
-from trezor import TR, io, loop, ui
+from trezor import TR, io, log, loop, ui, workflow
 from trezor.enums import ButtonRequestType
-from trezor.wire import ActionCancelled
+from trezor.messages import ButtonAck, ButtonRequest
+from trezor.wire import ActionCancelled, context
 from trezor.wire.context import wait as ctx_wait
 
 from ..common import button_request, interact
@@ -36,7 +37,10 @@ class RustLayout(ui.Layout):
     def __init__(self, layout: Any):
         self.layout = layout
         self.timer = loop.Timer()
+        self.br_chan = loop.chan()
+
         self.layout.attach_timer_fn(self.set_timer)
+        self._send_button_request()
 
     def set_timer(self, token: int, deadline: int) -> None:
         self.timer.schedule(deadline, token)
@@ -63,6 +67,9 @@ class RustLayout(ui.Layout):
                 self.handle_swipe(),
                 self.handle_click_signal(),
                 self.handle_result_signal(),
+                context.with_context(
+                    context.get_context(), self.handle_button_request()
+                ),
             )
 
         async def handle_result_signal(self) -> None:
@@ -110,6 +117,7 @@ class RustLayout(ui.Layout):
                     (io.TOUCH_END, orig_x + 2 * off_x, orig_y + 2 * off_y),
                 ):
                     msg = self.layout.touch_event(event, x, y)
+                    self._send_button_request()
                     self._paint()
                     if msg is not None:
                         raise ui.Result(msg)
@@ -129,10 +137,12 @@ class RustLayout(ui.Layout):
             from apps.debug import notify_layout_change
 
             self.layout.touch_event(io.TOUCH_START, x, y)
+            self._send_button_request()
             self._paint()
             if hold_ms is not None:
                 await loop.sleep(hold_ms)
             msg = self.layout.touch_event(io.TOUCH_END, x, y)
+            self._send_button_request()
 
             if msg is not None:
                 debug_storage.new_layout_event_id = event_id
@@ -159,7 +169,13 @@ class RustLayout(ui.Layout):
     else:
 
         def create_tasks(self) -> tuple[loop.AwaitableTask, ...]:
-            return self.handle_timers(), self.handle_input_and_rendering()
+            return (
+                self.handle_timers(),
+                self.handle_input_and_rendering(),
+                context.with_context(
+                    context.get_context(), self.handle_button_request()
+                ),
+            )
 
     def _first_paint(self) -> None:
         ui.backlight_fade(ui.style.BACKLIGHT_NONE)
@@ -199,6 +215,7 @@ class RustLayout(ui.Layout):
             msg = None
             if event in (io.TOUCH_START, io.TOUCH_MOVE, io.TOUCH_END):
                 msg = self.layout.touch_event(event, x, y)
+                self._send_button_request()
             if msg is not None:
                 raise ui.Result(msg)
             self._paint()
@@ -208,10 +225,25 @@ class RustLayout(ui.Layout):
             # Using `yield` instead of `await` to avoid allocations.
             token = yield self.timer
             msg = self.layout.timer(token)
+            self._send_button_request()
             if msg is not None:
                 raise ui.Result(msg)
             self._paint()
 
+    async def handle_button_request(self):
+        while True:
+            br_code, br_type, page_count = await self.br_chan.take()
+            log.debug(__name__, "ButtonRequest.type=%s", br_type)
+            await context.maybe_call(
+                ButtonRequest(code=br_code, pages=page_count), ButtonAck
+            )
+
+    def _send_button_request(self):
+        res = self.layout.button_request()
+        if res is not None:
+            br_code, br_type = res
+            self.br_chan.publish((br_code, br_type, self.layout.page_count()))
+
     def page_count(self) -> int:
         return self.layout.page_count()
 
@@ -279,14 +311,10 @@ async def confirm_action(
 
 
 async def flow_demo() -> None:
-    await raise_if_not_confirmed(
-        interact(
-            RustLayout(trezorui2.flow_get_address()),
-            "get_address",
-            BR_TYPE_OTHER,
-        ),
-        ActionCancelled,
-    )
+    from trezor import workflow
+
+    workflow.close_others()
+    await raise_if_not_confirmed(RustLayout(trezorui2.flow_get_address()))
 
 
 async def confirm_single(

core/src/trezor/ui/layouts/mercury/reset.py

@@ -83,15 +83,14 @@ async def select_word(
     while len(words) < 3:
         words.append(words[-1])
 
-    result = await ctx_wait(
-        RustLayout(
-            trezorui2.select_word(
-                title=TR.reset__select_word_x_of_y_template.format(
-                    checked_index + 1, count
-                ),
-                description=description,
-                words=(words[0], words[1], words[2]),
-            )
+    # result = await ctx_wait(
+    result = await RustLayout(
+        trezorui2.select_word(
+            title=TR.reset__select_word_x_of_y_template.format(
+                checked_index + 1, count
+            ),
+            description=description,
+            words=(words[0], words[1], words[2]),
         )
     )
     if __debug__ and isinstance(result, str):