From 34e2c4523862bb558646fbe72e8a79e60577a842 Mon Sep 17 00:00:00 2001 From: grdddj Date: Fri, 12 May 2023 15:13:14 +0200 Subject: [PATCH] WIP - improve TOIF handling and checks --- core/embed/rust/src/micropython/obj.rs | 17 ++++++ core/embed/rust/src/ui/display/toif.rs | 7 +++ core/embed/rust/src/ui/layout/util.rs | 61 ++++++---------------- core/embed/rust/src/ui/model_tr/layout.rs | 4 +- core/mocks/generated/trezorui2.pyi | 4 +- core/src/apps/management/apply_settings.py | 4 +- 6 files changed, 47 insertions(+), 50 deletions(-) diff --git a/core/embed/rust/src/micropython/obj.rs b/core/embed/rust/src/micropython/obj.rs index 0cbd02894..76039834e 100644 --- a/core/embed/rust/src/micropython/obj.rs +++ b/core/embed/rust/src/micropython/obj.rs @@ -303,6 +303,23 @@ impl TryFrom<(Obj, Obj)> for Obj { } } +impl TryFrom<(Obj, Obj, Obj)> for Obj { + type Error = Error; + + fn try_from(val: (Obj, Obj, Obj)) -> Result { + // SAFETY: + // - Should work with any micropython objects. + // EXCEPTION: Will raise if allocation fails. + let values = [val.0, val.1, val.2]; + let obj = catch_exception(|| unsafe { ffi::mp_obj_new_tuple(3, values.as_ptr()) })?; + if obj.is_null() { + Err(Error::AllocationFailed) + } else { + Ok(obj) + } + } +} + // // # Additional conversions based on the methods above. // diff --git a/core/embed/rust/src/ui/display/toif.rs b/core/embed/rust/src/ui/display/toif.rs index c1c5d77f0..8d5e694df 100644 --- a/core/embed/rust/src/ui/display/toif.rs +++ b/core/embed/rust/src/ui/display/toif.rs @@ -181,6 +181,13 @@ impl<'i> Toif<'i> { } } + pub const fn is_grayscale(&self) -> bool { + matches!( + self.format(), + ToifFormat::GrayScaleOH | ToifFormat::GrayScaleEH + ) + } + pub const fn width(&self) -> i16 { u16::from_le_bytes([self.data[4], self.data[5]]) as i16 } diff --git a/core/embed/rust/src/ui/layout/util.rs b/core/embed/rust/src/ui/layout/util.rs index 102d592fd..8c43bf5e9 100644 --- a/core/embed/rust/src/ui/layout/util.rs +++ b/core/embed/rust/src/ui/layout/util.rs @@ -21,13 +21,7 @@ use heapless::Vec; #[cfg(feature = "jpeg")] use crate::ui::display::tjpgd::{jpeg_info, jpeg_test}; -use crate::{ - micropython::{ - buffer::get_buffer, - ffi::{mp_obj_new_int, mp_obj_new_tuple}, - }, - ui::display::toif::Toif, -}; +use crate::{micropython::buffer::get_buffer, ui::display::toif::Toif}; pub fn iter_into_objs(iterable: Obj) -> Result<[Obj; N], Error> { let err = Error::ValueError(cstr!("Invalid iterable length")); @@ -224,27 +218,15 @@ pub extern "C" fn upy_disable_animation(disable: Obj) -> Obj { #[cfg(feature = "jpeg")] pub extern "C" fn upy_jpeg_info(data: Obj) -> Obj { let block = || { - let buffer = unsafe { get_buffer(data) }; - - if let Ok(buffer) = buffer { - let info = jpeg_info(buffer); - - if let Some(info) = info { - let obj = unsafe { - let values = [ - mp_obj_new_int(info.0.x as _), - mp_obj_new_int(info.0.y as _), - mp_obj_new_int(info.1 as _), - ]; - mp_obj_new_tuple(3, values.as_ptr()) - }; - - Ok(obj) - } else { - Err(Error::ValueError(cstr!("Invalid image format."))) - } + let buffer = unsafe { get_buffer(data) }?; + + if let Some(info) = jpeg_info(buffer) { + let w = info.0.x as u16; + let h = info.0.y as u16; + let mcu_h = info.1 as u16; + (w.into(), h.into(), mcu_h.into()).try_into() } else { - Err(Error::ValueError(cstr!("Buffer error."))) + Err(Error::ValueError(cstr!("Invalid image format."))) } }; @@ -253,26 +235,15 @@ pub extern "C" fn upy_jpeg_info(data: Obj) -> Obj { pub extern "C" fn upy_toif_info(data: Obj) -> Obj { let block = || { - let buffer = unsafe { get_buffer(data) }; - - if let Ok(buffer) = buffer { - let toif = Toif::new(buffer); - - if let Some(toif) = toif { - let obj = unsafe { - let values = [ - mp_obj_new_int(toif.width() as _), - mp_obj_new_int(toif.height() as _), - ]; - mp_obj_new_tuple(2, values.as_ptr()) - }; + let buffer = unsafe { get_buffer(data) }?; - Ok(obj) - } else { - Err(Error::ValueError(cstr!("Invalid image format."))) - } + if let Some(toif) = Toif::new(buffer) { + let w = toif.width() as u16; + let h = toif.height() as u16; + let is_grayscale = toif.is_grayscale(); + (w.into(), h.into(), is_grayscale.into()).try_into() } else { - Err(Error::ValueError(cstr!("Buffer error."))) + Err(Error::ValueError(cstr!("Invalid image format."))) } }; diff --git a/core/embed/rust/src/ui/model_tr/layout.rs b/core/embed/rust/src/ui/model_tr/layout.rs index 23eae6556..2640c670a 100644 --- a/core/embed/rust/src/ui/model_tr/layout.rs +++ b/core/embed/rust/src/ui/model_tr/layout.rs @@ -1241,8 +1241,8 @@ pub static mp_module_trezorui2: Module = obj_module! { /// """Disable animations, debug builds only.""" Qstr::MP_QSTR_disable_animation => obj_fn_1!(upy_disable_animation).as_obj(), - /// def toif_info(data: bytes) -> tuple[int, int]: - /// """Get TOIF image dimensions (width: int, height: int).""" + /// def toif_info(data: bytes) -> tuple[int, int, bool]: + /// """Get TOIF image dimensions and format (width: int, height: int, is_grayscale: bool).""" Qstr::MP_QSTR_toif_info => obj_fn_1!(upy_toif_info).as_obj(), /// def confirm_action( diff --git a/core/mocks/generated/trezorui2.pyi b/core/mocks/generated/trezorui2.pyi index ac3ee0ea5..f40adef88 100644 --- a/core/mocks/generated/trezorui2.pyi +++ b/core/mocks/generated/trezorui2.pyi @@ -10,8 +10,8 @@ def disable_animation(disable: bool) -> None: # rust/src/ui/model_tr/layout.rs -def toif_info(data: bytes) -> tuple[int, int]: - """Get TOIF image dimensions (width: int, height: int).""" +def toif_info(data: bytes) -> tuple[int, int, bool]: + """Get TOIF image dimensions and format (width: int, height: int, is_grayscale: bool).""" # rust/src/ui/model_tr/layout.rs diff --git a/core/src/apps/management/apply_settings.py b/core/src/apps/management/apply_settings.py index ef2e1d762..ea8749f7f 100644 --- a/core/src/apps/management/apply_settings.py +++ b/core/src/apps/management/apply_settings.py @@ -35,11 +35,13 @@ def _validate_homescreen(homescreen: bytes) -> None: def _validate_homescreen_tr(homescreen: bytes) -> None: try: - w, h = trezorui2.toif_info(homescreen) + w, h, is_grayscale = trezorui2.toif_info(homescreen) except ValueError: raise DataError("Invalid homescreen") if w != 128 or h != 64: raise DataError("Homescreen must be 128x64 pixel large") + if not is_grayscale: + raise DataError("Homescreen must be grayscale") def _validate_homescreen_tt(homescreen: bytes) -> None: