improve T2B1 homescreen checks

core/embed/rust/Cargo.toml

@@ -90,8 +90,9 @@ default_features = false
 version = "0.3.3"
 
 [dependencies.cstr_core]
-version = "0.2.4"
-default_features = false
+version = "0.2.6"
+default-features = false
+features = ["nightly"]
 
 # Build dependencies
 

core/embed/rust/librust_qstr.h

@@ -36,6 +36,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_button_event;
   MP_QSTR_cancel_arrow;
   MP_QSTR_case_sensitive;
+  MP_QSTR_check_homescreen_format;
   MP_QSTR_chunkify;
   MP_QSTR_coinjoin_authorized;
   MP_QSTR_confirm_action;
@@ -84,8 +85,6 @@ static void _librust_qstrs(void) {
   MP_QSTR_info_button;
   MP_QSTR_is_type_of;
   MP_QSTR_items;
-  MP_QSTR_jpeg_info;
-  MP_QSTR_jpeg_test;
   MP_QSTR_label;
   MP_QSTR_lines;
   MP_QSTR_max_count;
@@ -140,7 +139,6 @@ static void _librust_qstrs(void) {
   MP_QSTR_time_ms;
   MP_QSTR_timer;
   MP_QSTR_title;
-  MP_QSTR_toif_info;
   MP_QSTR_total_amount;
   MP_QSTR_total_fee_new;
   MP_QSTR_total_label;

core/embed/rust/src/error.rs

@@ -29,7 +29,7 @@ pub enum Error {
 #[macro_export]
 macro_rules! value_error {
     ($msg:expr) => {
-        Error::ValueError(cstr!($msg))
+        Error::ValueError(cstr_core::cstr!($msg))
     };
 }
 

core/embed/rust/src/protobuf/error.rs

@@ -2,13 +2,11 @@ use cstr_core::cstr;
 
 use crate::{error::Error, micropython::qstr::Qstr};
 
-// XXX const version of `from_bytes_with_nul_unchecked` is nightly-only.
-
-pub fn experimental_not_enabled() -> Error {
+pub const fn experimental_not_enabled() -> Error {
     value_error!("Experimental features are disabled.")
 }
 
-pub fn unknown_field_type() -> Error {
+pub const fn unknown_field_type() -> Error {
     value_error!("Unknown field type.")
 }
 
@@ -20,6 +18,6 @@ pub fn invalid_value(field: Qstr) -> Error {
     Error::ValueErrorParam(cstr!("Invalid value for field."), field.into())
 }
 
-pub fn end_of_buffer() -> Error {
+pub const fn end_of_buffer() -> Error {
     value_error!("End of buffer.")
 }

core/embed/rust/src/trezorhal/storage.rs

@@ -3,7 +3,7 @@
 use super::ffi;
 use crate::error::Error;
 use core::ptr;
-use cstr_core::{cstr, CStr};
+use cstr_core::CStr;
 
 /// Result of PIN delay callback.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]

core/embed/rust/src/ui/display/toif.rs

@@ -1,4 +1,5 @@
 use crate::{
+    error::Error,
     trezorhal::{
         display::ToifFormat,
         uzlib::{UzlibContext, UZLIB_WINDOW_SIZE},
@@ -198,16 +199,16 @@ pub struct Toif<'i> {
 }
 
 impl<'i> Toif<'i> {
-    pub const fn new(data: &'i [u8]) -> Option<Self> {
+    pub const fn new(data: &'i [u8]) -> Result<Self, Error> {
         if data.len() < TOIF_HEADER_LENGTH || data[0] != b'T' || data[1] != b'O' || data[2] != b'I'
         {
-            return None;
+            return Err(value_error!("Invalid TOIF header."));
         }
         let zdatalen = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
         if zdatalen + TOIF_HEADER_LENGTH != data.len() {
-            return None;
+            return Err(value_error!("Invalid TOIF length."));
         }
-        Some(Self {
+        Ok(Self {
             data,
             empty_right_column: false,
         })
@@ -305,8 +306,8 @@ pub struct Icon {
 impl Icon {
     pub const fn new(data: &'static [u8]) -> Self {
         let toif = match Toif::new(data) {
-            Some(t) => t,
-            None => panic!("Invalid image."),
+            Ok(t) => t,
+            _ => panic!("Invalid image."),
         };
         assert!(matches!(toif.format(), ToifFormat::GrayScaleEH));
         Self {

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

@@ -18,13 +18,8 @@ use crate::{
         util::set_animation_disabled,
     },
 };
-use cstr_core::cstr;
 use heapless::Vec;
 
-#[cfg(feature = "jpeg")]
-use crate::ui::display::tjpgd::{jpeg_info, jpeg_test};
-use crate::{micropython::buffer::get_buffer, ui::display::toif::Toif};
-
 pub fn iter_into_array<T, E, const N: usize>(iterable: Obj) -> Result<[T; N], Error>
 where
     T: TryFrom<Obj, Error = E>,
@@ -201,52 +196,6 @@ pub extern "C" fn upy_disable_animation(disable: Obj) -> Obj {
     unsafe { try_or_raise(block) }
 }
 
-#[cfg(feature = "jpeg")]
-pub extern "C" fn upy_jpeg_info(data: Obj) -> Obj {
-    let block = || {
-        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(value_error!("Invalid image format."))
-        }
-    };
-
-    unsafe { try_or_raise(block) }
-}
-
-pub extern "C" fn upy_toif_info(data: Obj) -> Obj {
-    let block = || {
-        let buffer = unsafe { get_buffer(data) }?;
-
-        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(value_error!("Invalid image format."))
-        }
-    };
-
-    unsafe { try_or_raise(block) }
-}
-
-#[cfg(feature = "jpeg")]
-pub extern "C" fn upy_jpeg_test(data: Obj) -> Obj {
-    let block = || {
-        let buffer = unsafe { get_buffer(data) }?;
-        let result = jpeg_test(buffer);
-        Ok(result.into())
-    };
-
-    unsafe { try_or_raise(block) }
-}
-
 pub fn get_user_custom_image() -> Result<Gc<[u8]>, Error> {
     let len = get_avatar_len()?;
     let mut data = Gc::<[u8]>::new_slice(len)?;

core/embed/rust/src/ui/model_tr/component/homescreen.rs

@@ -1,8 +1,9 @@
 use crate::{
     strutil::StringType,
-    trezorhal::usb::usb_configured,
+    trezorhal::{display::ToifFormat, usb::usb_configured},
     ui::{
         component::{Child, Component, Event, EventCtx, Label},
+        constant::{HEIGHT, WIDTH},
         display::{rect_fill, toif::Toif, Font, Icon},
         event::USBEvent,
         geometry::{Alignment2D, Insets, Offset, Point, Rect},
@@ -64,9 +65,13 @@ where
     }
 
     fn paint_homescreen_image(&self) {
-        if let Ok(user_custom_image) = get_user_custom_image() {
-            let toif_data = unwrap!(Toif::new(user_custom_image.as_ref()));
-            toif_data.draw(TOP_CENTER, Alignment2D::TOP_CENTER, theme::FG, theme::BG);
+        let homescreen_bytes = get_user_custom_image().ok();
+        let homescreen = homescreen_bytes
+            .as_ref()
+            .and_then(|data| Toif::new(data.as_ref()).ok())
+            .filter(check_homescreen_format);
+        if let Some(toif) = homescreen {
+            toif.draw(TOP_CENTER, Alignment2D::TOP_CENTER, theme::FG, theme::BG);
         } else {
             paint_default_image();
         }
@@ -314,6 +319,10 @@ where
     }
 }
 
+pub fn check_homescreen_format(toif: &Toif) -> bool {
+    toif.format() == ToifFormat::GrayScaleEH && toif.width() == WIDTH && toif.height() == HEIGHT
+}
+
 // DEBUG-ONLY SECTION BELOW
 
 #[cfg(feature = "ui_debug")]

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

@@ -51,7 +51,7 @@ pub use flow::Flow;
 pub use flow_pages::{FlowPages, Page};
 pub use frame::{Frame, ScrollableContent, ScrollableFrame};
 #[cfg(feature = "micropython")]
-pub use homescreen::{ConfirmHomescreen, Homescreen, Lockscreen};
+pub use homescreen::{check_homescreen_format, ConfirmHomescreen, Homescreen, Lockscreen};
 pub use input_methods::{
     number_input::NumberInput,
     passphrase::PassphraseEntry,

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

@@ -36,10 +36,9 @@ use crate::{
         layout::{
             obj::{ComponentMsgObj, LayoutObj},
             result::{CANCELLED, CONFIRMED, INFO},
-            util::{
-                iter_into_array, iter_into_vec, upy_disable_animation, upy_toif_info, ConfirmBlob,
-            },
+            util::{iter_into_array, iter_into_vec, upy_disable_animation, ConfirmBlob},
         },
+        model_tr::component::check_homescreen_format,
     },
 };
 
@@ -1649,6 +1648,20 @@ extern "C" fn new_confirm_firmware_update(
     unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
 }
 
+pub extern "C" fn upy_check_homescreen_format(data: Obj) -> Obj {
+    let block = || {
+        // SAFETY: buffer does not outlive this function
+        let buffer = unsafe { get_buffer(data) }?;
+
+        Ok(display::toif::Toif::new(buffer)
+            .map(|toif| check_homescreen_format(&toif))
+            .unwrap_or(false)
+            .into())
+    };
+
+    unsafe { util::try_or_raise(block) }
+}
+
 #[no_mangle]
 pub static mp_module_trezorui2: Module = obj_module! {
     Qstr::MP_QSTR___name__ => Qstr::MP_QSTR_trezorui2.to_obj(),
@@ -1666,9 +1679,9 @@ 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, 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 check_homescreen_format(data: bytes) -> bool:
+    ///     """Check homescreen format and dimensions."""
+    Qstr::MP_QSTR_check_homescreen_format => obj_fn_1!(upy_check_homescreen_format).as_obj(),
 
     /// def confirm_action(
     ///     *,

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

@@ -17,7 +17,10 @@ use crate::{
     trezorhal::{buffers::BufferJpegWork, display::ToifFormat, uzlib::UZLIB_WINDOW_SIZE},
     ui::{
         constant::HEIGHT,
-        display::{tjpgd::BufferInput, toif::Toif},
+        display::{
+            tjpgd::{jpeg_test, BufferInput},
+            toif::Toif,
+        },
         model_tt::component::homescreen::render::{
             HomescreenJpeg, HomescreenToif, HOMESCREEN_TOIF_SIZE,
         },
@@ -367,6 +370,10 @@ where
     }
 }
 
+pub fn check_homescreen_format(buffer: &[u8]) -> bool {
+    is_image_jpeg(buffer) && jpeg_test(buffer)
+}
+
 fn is_image_jpeg(buffer: &[u8]) -> bool {
     let jpeg = jpeg_info(buffer);
     if let Some((size, mcu_height)) = jpeg {
@@ -380,7 +387,7 @@ fn is_image_jpeg(buffer: &[u8]) -> bool {
 
 fn is_image_toif(buffer: &[u8]) -> bool {
     let toif = Toif::new(buffer);
-    if let Some(toif) = toif {
+    if let Ok(toif) = toif {
         if toif.size().x == HOMESCREEN_TOIF_SIZE
             && toif.size().y == HOMESCREEN_TOIF_SIZE
             && toif.format() == ToifFormat::FullColorBE

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

@@ -32,7 +32,7 @@ pub use error::ErrorScreen;
 pub use fido::{FidoConfirm, FidoMsg};
 pub use frame::{Frame, FrameMsg};
 #[cfg(feature = "micropython")]
-pub use homescreen::{Homescreen, HomescreenMsg, Lockscreen};
+pub use homescreen::{check_homescreen_format, Homescreen, HomescreenMsg, Lockscreen};
 pub use keyboard::{
     bip39::Bip39Input,
     mnemonic::{MnemonicInput, MnemonicKeyboard, MnemonicKeyboardMsg},

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

@@ -1,5 +1,4 @@
 use core::{cmp::Ordering, convert::TryInto};
-use cstr_core::cstr;
 
 use crate::{
     error::Error,
@@ -38,11 +37,9 @@ use crate::{
         layout::{
             obj::{ComponentMsgObj, LayoutObj},
             result::{CANCELLED, CONFIRMED, INFO},
-            util::{
-                iter_into_array, upy_disable_animation, upy_jpeg_info, upy_jpeg_test, ConfirmBlob,
-                PropsList,
-            },
+            util::{iter_into_array, upy_disable_animation, ConfirmBlob, PropsList},
         },
+        model_tt::component::check_homescreen_format,
     },
 };
 
@@ -1598,6 +1595,16 @@ extern "C" fn draw_welcome_screen() -> Obj {
     Obj::const_none()
 }
 
+pub extern "C" fn upy_check_homescreen_format(data: Obj) -> Obj {
+    let block = || {
+        let buffer = unsafe { get_buffer(data) }?;
+
+        Ok(check_homescreen_format(buffer).into())
+    };
+
+    unsafe { util::try_or_raise(block) }
+}
+
 #[no_mangle]
 extern "C" fn new_confirm_firmware_update(
     n_args: usize,
@@ -1645,13 +1652,9 @@ 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 jpeg_info(data: bytes) -> tuple[int, int, int]:
-    ///     """Get JPEG image dimensions (width: int, height: int, mcu_height: int)."""
-    Qstr::MP_QSTR_jpeg_info => obj_fn_1!(upy_jpeg_info).as_obj(),
-
-    /// def jpeg_test(data: bytes) -> bool:
-    ///     """Test JPEG image."""
-    Qstr::MP_QSTR_jpeg_test => obj_fn_1!(upy_jpeg_test).as_obj(),
+    /// def check_homescreen_format(data: bytes) -> bool:
+    ///     """Check homescreen format and dimensions."""
+    Qstr::MP_QSTR_check_homescreen_format => obj_fn_1!(upy_check_homescreen_format).as_obj(),
 
     /// def confirm_action(
     ///     *,