From b75ce4500663519cf128415aac2d6944af247159 Mon Sep 17 00:00:00 2001 From: Martin Milata Date: Fri, 12 Apr 2024 21:13:28 +0200 Subject: [PATCH] WIP: sending ButtonRequest from rust [skip_ci] --- core/embed/rust/librust_qstr.h | 1 + core/embed/rust/src/ui/button_request.rs | 44 +++++++++++++++++++ core/embed/rust/src/ui/component/base.rs | 16 +++++++ core/embed/rust/src/ui/flow/base.rs | 8 ++-- core/embed/rust/src/ui/flow/flow.rs | 40 +++++++++++++++-- core/embed/rust/src/ui/layout/obj.rs | 22 ++++++++++ core/embed/rust/src/ui/mod.rs | 1 + .../src/ui/model_mercury/flow/get_address.rs | 36 ++++++++------- .../src/trezor/ui/layouts/mercury/__init__.py | 39 +++++++++++----- 9 files changed, 172 insertions(+), 35 deletions(-) create mode 100644 core/embed/rust/src/ui/button_request.rs diff --git a/core/embed/rust/librust_qstr.h b/core/embed/rust/librust_qstr.h index c4637ce96..c39c68eba 100644 --- a/core/embed/rust/librust_qstr.h +++ b/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; diff --git a/core/embed/rust/src/ui/button_request.rs b/core/embed/rust/src/ui/button_request.rs new file mode 100644 index 000000000..9bf8065a3 --- /dev/null +++ b/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 { + Some(ButtonRequest { + code: self, + br_type, + }) + } +} + +#[derive(Clone, Copy)] +pub struct ButtonRequest { + pub code: ButtonRequestCode, + pub br_type: &'static str, +} diff --git a/core/embed/rust/src/ui/component/base.rs b/core/embed/rust/src/ui/component/base.rs index 10e27fd76..0d3612ec0 100644 --- a/core/embed/rust/src/ui/component/base.rs +++ b/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, + button_request: Option, 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 { + 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; } diff --git a/core/embed/rust/src/ui/flow/base.rs b/core/embed/rust/src/ui/flow/base.rs index a01f8f0bc..ff398ac61 100644 --- a/core/embed/rust/src/ui/flow/base.rs +++ b/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 { Nothing, /// Initiate transition to another state, end event processing. - /// NOTE: it might make sense to include Option here - Goto(Q, SwipeDirection), + /// Optionally send button request when transition finishes. + Goto(Q, SwipeDirection, Option), /// Yield a message to the caller of the flow (i.e. micropython), end event /// processing. @@ -76,7 +76,7 @@ impl Decision { /// 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 diff --git a/core/embed/rust/src/ui/flow/flow.rs b/core/embed/rust/src/ui/flow/flow.rs index 2bb622786..ab1ddb322 100644 --- a/core/embed/rust/src/ui/flow/flow.rs +++ b/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 { swipe: Swipe, /// Animation parameter. anim_offset: Offset, + /// Button request sent after showing first state. + initial_button_request: Option, } struct Transition { + /// State we are transitioning _from_. prev_state: Q, + /// Animation progress. animation: Animation, + /// Direction of the slide animation. direction: SwipeDirection, + /// Send ButtonRequest when finished. + button_request: Option, } impl SwipeFlow { @@ -45,10 +53,26 @@ impl SwipeFlow { 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, + ) { self.transition = Some(Transition { prev_state: self.state, animation: Animation::new( @@ -58,6 +82,7 @@ impl SwipeFlow { Instant::now(), ), direction, + button_request, }); self.state = state; ctx.request_anim_frame(); @@ -91,6 +116,9 @@ impl SwipeFlow { 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 SwipeFlow { // 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 Component for SwipeFlow { } fn event(&mut self, ctx: &mut EventCtx, event: Event) -> Option { + 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 Component for SwipeFlow { 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), diff --git a/core/embed/rust/src/ui/layout/obj.rs b/core/embed/rust/src/ui/layout/obj.rs index 6c00316ad..e7512fd31 100644 --- a/core/embed/rust/src/ui/layout/obj.rs +++ b/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 { + 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 = 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 { diff --git a/core/embed/rust/src/ui/mod.rs b/core/embed/rust/src/ui/mod.rs index 5bed9d88f..c0ba8d010 100644 --- a/core/embed/rust/src/ui/mod.rs +++ b/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; diff --git a/core/embed/rust/src/ui/model_mercury/flow/get_address.rs b/core/embed/rust/src/ui/model_mercury/flow/get_address.rs index c56037f4f..d557d3897 100644 --- a/core/embed/rust/src/ui/model_mercury/flow/get_address.rs +++ b/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 { 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 { 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()) } } diff --git a/core/src/trezor/ui/layouts/mercury/__init__.py b/core/src/trezor/ui/layouts/mercury/__init__.py index 2550c6b9c..cb00ae827 100644 --- a/core/src/trezor/ui/layouts/mercury/__init__.py +++ b/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(