WIP: sending ButtonRequest from rust

[skip_ci]

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/flow/base.rs

@@ -1,4 +1,4 @@
-use crate::ui::{component::EventCtx, geometry::Offset};
+use crate::ui::{button_request::ButtonRequest, component::EventCtx, geometry::Offset};
 use num_traits::ToPrimitive;
 
 #[derive(Copy, Clone)]
@@ -55,8 +55,8 @@ pub enum Decision<Q> {
     Nothing,
 
     /// Initiate transition to another state, end event processing.
-    /// NOTE: it might make sense to include Option<ButtonRequest> here
-    Goto(Q, SwipeDirection),
+    /// Optionally send button request when transition finishes.
+    Goto(Q, SwipeDirection, Option<ButtonRequest>),
 
     /// Yield a message to the caller of the flow (i.e. micropython), end event
     /// processing.
@@ -76,7 +76,7 @@ impl<Q> Decision<Q> {
 /// triggered by events and swipes.
 pub trait FlowState
 where
-    Self: Sized + Copy + PartialEq + Eq + ToPrimitive,
+    Self: Sized + Copy + Eq + ToPrimitive,
 {
     /// There needs to be a mapping from states to indices of the FlowStore
     /// array. Default implementation works for states that are enums, the

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

@@ -3,6 +3,7 @@ use crate::{
     time::{Duration, Instant},
     ui::{
         animation::Animation,
+        button_request::{ButtonRequest, ButtonRequestCode},
         component::{Component, Event, EventCtx},
         flow::{base::Decision, swipe::Swipe, FlowMsg, FlowState, FlowStore, SwipeDirection},
         geometry::{Offset, Rect},
@@ -29,12 +30,19 @@ pub struct SwipeFlow<Q, S> {
     swipe: Swipe,
     /// Animation parameter.
     anim_offset: Offset,
+    /// Button request sent after showing first state.
+    initial_button_request: Option<ButtonRequest>,
 }
 
 struct Transition<Q> {
+    /// State we are transitioning _from_.
     prev_state: Q,
+    /// Animation progress.
     animation: Animation<Offset>,
+    /// Direction of the slide animation.
     direction: SwipeDirection,
+    /// Send ButtonRequest when finished.
+    button_request: Option<ButtonRequest>,
 }
 
 impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
@@ -45,10 +53,26 @@ impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
             transition: None,
             swipe: Swipe::new().down().up().left().right(),
             anim_offset: Offset::zero(),
+            initial_button_request: None,
         })
     }
 
-    fn goto(&mut self, ctx: &mut EventCtx, direction: SwipeDirection, state: Q) {
+    pub fn with_initial_button_request(
+        mut self,
+        code: ButtonRequestCode,
+        br_type: &'static str,
+    ) -> Self {
+        self.initial_button_request = code.with_type(br_type);
+        self
+    }
+
+    fn goto(
+        &mut self,
+        ctx: &mut EventCtx,
+        direction: SwipeDirection,
+        state: Q,
+        button_request: Option<ButtonRequest>,
+    ) {
         self.transition = Some(Transition {
             prev_state: self.state,
             animation: Animation::new(
@@ -58,6 +82,7 @@ impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
                 Instant::now(),
             ),
             direction,
+            button_request,
         });
         self.state = state;
         ctx.request_anim_frame();
@@ -91,6 +116,9 @@ impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
     fn handle_transition(&mut self, ctx: &mut EventCtx) {
         if let Some(transition) = &self.transition {
             if transition.animation.finished(Instant::now()) {
+                if let Some(ButtonRequest { code, br_type }) = transition.button_request {
+                    ctx.send_button_request(code, br_type);
+                }
                 self.transition = None;
                 unwrap!(self.store.clone(None)); // Free the clone.
 
@@ -110,7 +138,7 @@ impl<Q: FlowState, S: FlowStore> SwipeFlow<Q, S> {
             // can render both states in the transition animation.
             unwrap!(self.store.clone(Some(i)));
             self.store.map_swipable(i, |s| s.swiped(ctx, direction));
-            Decision::Goto(self.state, direction)
+            Decision::Goto(self.state, direction, None)
         } else {
             Decision::Nothing
         }
@@ -139,9 +167,13 @@ impl<Q: FlowState, S: FlowStore> Component for SwipeFlow<Q, S> {
     }
 
     fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option<Self::Msg> {
+        if let (Event::Attach, Some(br)) = (event, self.initial_button_request) {
+            ctx.send_button_request(br.code, br.br_type)
+        }
         // TODO: are there any events we want to send to all? timers perhaps?
         if let Event::Timer(EventCtx::ANIM_FRAME_TIMER) = event {
             self.handle_transition(ctx);
+            return None;
         }
         // Ignore events while transition is running.
         if self.transition.is_some() {
@@ -158,8 +190,8 @@ impl<Q: FlowState, S: FlowStore> Component for SwipeFlow<Q, S> {
 
         match decision {
             Decision::Nothing => None,
-            Decision::Goto(next_state, direction) => {
-                self.goto(ctx, direction, next_state);
+            Decision::Goto(next_state, direction, button_request) => {
+                self.goto(ctx, direction, next_state, button_request);
                 None
             }
             Decision::Return(msg) => Some(msg),

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},
@@ -169,6 +170,7 @@ impl LayoutObj {
             unsafe { Gc::as_mut(&mut inner.root) }.obj_place(constant::screen());
         }
 
+        assert!(inner.event_ctx.button_request().is_none());
         // Clear the leftover flags from the previous event pass.
         inner.event_ctx.clear();
 
@@ -248,6 +250,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 +293,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 +484,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::{ButtonRequest, ButtonRequestCode},
         component::{
             image::BlendedImage,
             text::paragraphs::{Paragraph, Paragraphs},
@@ -34,19 +35,19 @@ impl FlowState for GetAddress {
     fn handle_swipe(&self, direction: SwipeDirection) -> Decision<Self> {
         match (self, direction) {
             (GetAddress::Address, SwipeDirection::Left) => {
-                Decision::Goto(GetAddress::Menu, direction)
+                Decision::Goto(GetAddress::Menu, direction, None)
             }
             (GetAddress::Address, SwipeDirection::Up) => {
-                Decision::Goto(GetAddress::Success, direction)
+                Decision::Goto(GetAddress::Success, direction, None)
             }
             (GetAddress::Menu, SwipeDirection::Right) => {
-                Decision::Goto(GetAddress::Address, direction)
+                Decision::Goto(GetAddress::Address, direction, None)
             }
             (GetAddress::QrCode, SwipeDirection::Right) => {
-                Decision::Goto(GetAddress::Menu, direction)
+                Decision::Goto(GetAddress::Menu, direction, None)
             }
             (GetAddress::AccountInfo, SwipeDirection::Right) => {
-                Decision::Goto(GetAddress::Menu, direction)
+                Decision::Goto(GetAddress::Menu, direction, None)
             }
             (GetAddress::Cancel, SwipeDirection::Up) => Decision::Return(FlowMsg::Cancelled),
             _ => Decision::Nothing,
@@ -56,35 +57,37 @@ impl FlowState for GetAddress {
     fn handle_event(&self, msg: FlowMsg) -> Decision<Self> {
         match (self, msg) {
             (GetAddress::Address, FlowMsg::Info) => {
-                Decision::Goto(GetAddress::Menu, SwipeDirection::Left)
+                Decision::Goto(GetAddress::Menu, SwipeDirection::Left, None)
             }
 
-            (GetAddress::Menu, FlowMsg::Choice(0)) => {
-                Decision::Goto(GetAddress::QrCode, SwipeDirection::Left)
-            }
+            (GetAddress::Menu, FlowMsg::Choice(0)) => Decision::Goto(
+                GetAddress::QrCode,
+                SwipeDirection::Left,
+                ButtonRequestCode::ProtectCall.with_type("test_br_qr"),
+            ),
 
             (GetAddress::Menu, FlowMsg::Choice(1)) => {
-                Decision::Goto(GetAddress::AccountInfo, SwipeDirection::Left)
+                Decision::Goto(GetAddress::AccountInfo, SwipeDirection::Left, None)
             }
 
             (GetAddress::Menu, FlowMsg::Choice(2)) => {
-                Decision::Goto(GetAddress::Cancel, SwipeDirection::Left)
+                Decision::Goto(GetAddress::Cancel, SwipeDirection::Left, None)
             }
 
             (GetAddress::Menu, FlowMsg::Cancelled) => {
-                Decision::Goto(GetAddress::Address, SwipeDirection::Right)
+                Decision::Goto(GetAddress::Address, SwipeDirection::Right, None)
             }
 
             (GetAddress::QrCode, FlowMsg::Cancelled) => {
-                Decision::Goto(GetAddress::Menu, SwipeDirection::Right)
+                Decision::Goto(GetAddress::Menu, SwipeDirection::Right, None)
             }
 
             (GetAddress::AccountInfo, FlowMsg::Cancelled) => {
-                Decision::Goto(GetAddress::Menu, SwipeDirection::Right)
+                Decision::Goto(GetAddress::Menu, SwipeDirection::Right, None)
             }
 
             (GetAddress::Cancel, FlowMsg::Cancelled) => {
-                Decision::Goto(GetAddress::Menu, SwipeDirection::Right)
+                Decision::Goto(GetAddress::Menu, SwipeDirection::Right, None)
             }
 
             (GetAddress::Success, _) => Decision::Return(FlowMsg::Confirmed),
@@ -181,7 +184,8 @@ impl GetAddress {
                 ),
                 |_| Some(FlowMsg::Confirmed),
             )?;
-        let res = SwipeFlow::new(GetAddress::Address, store)?;
+        let res = SwipeFlow::new(GetAddress::Address, store)?
+            .with_initial_button_request(ButtonRequestCode::Address, "show_address");
         Ok(LayoutObj::new(res)?.into())
     }
 }

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
 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,7 @@ class RustLayout(ui.Layout):
                 self.handle_swipe(),
                 self.handle_click_signal(),
                 self.handle_result_signal(),
+                self.handle_button_request(),
             )
 
         async def handle_result_signal(self) -> None:
@@ -110,6 +115,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 +135,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 +167,7 @@ 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(), self.handle_button_request()
 
     def _first_paint(self) -> None:
         ui.backlight_fade(ui.style.BACKLIGHT_NONE)
@@ -199,6 +207,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 +217,24 @@ 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()
 
+    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)
+            # FIXME: does nothing because CURRENT_CONTEXT is None
+            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 +302,8 @@ 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,
-    )
+    # workflow.close_others()?
+    await raise_if_not_confirmed(RustLayout(trezorui2.flow_get_address()))
 
 
 async def confirm_single(