fix(core): improve T2B1 homescreen checks

[no changelog]

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/ui/layout/util.rs

@@ -20,10 +20,6 @@ use crate::{
 };
 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>,
@@ -200,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 {

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

@@ -37,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,
     },
 };
 
@@ -1592,6 +1590,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,
@@ -1639,13 +1647,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(
     ///     *,

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, bool]:
-    """Get TOIF image dimensions and format (width: int, height: int, is_grayscale: bool)."""
+def check_homescreen_format(data: bytes) -> bool:
+    """Check homescreen format and dimensions."""
 
 
 # rust/src/ui/model_tr/layout.rs
@@ -462,13 +462,8 @@ def disable_animation(disable: bool) -> None:
 
 
 # rust/src/ui/model_tt/layout.rs
-def jpeg_info(data: bytes) -> tuple[int, int, int]:
-    """Get JPEG image dimensions (width: int, height: int, mcu_height: int)."""
-
-
-# rust/src/ui/model_tt/layout.rs
-def jpeg_test(data: bytes) -> bool:
-    """Test JPEG image."""
+def check_homescreen_format(data: bytes) -> bool:
+    """Check homescreen format and dimensions."""
 
 
 # rust/src/ui/model_tt/layout.rs

core/src/apps/management/apply_settings.py

@@ -1,7 +1,6 @@
 from typing import TYPE_CHECKING
 
 import trezorui2
-from trezor import utils
 from trezor.enums import ButtonRequestType
 from trezor.ui.layouts import confirm_action
 from trezor.wire import DataError
@@ -13,38 +12,6 @@
 
 BRT_PROTECT_CALL = ButtonRequestType.ProtectCall  # CACHE
 
-if utils.INTERNAL_MODEL in ("T1B1", "T2B1"):
-
-    def _validate_homescreen_model_specific(homescreen: bytes) -> None:
-        from trezor.ui import HEIGHT, WIDTH
-
-        try:
-            w, h, is_grayscale = trezorui2.toif_info(homescreen)
-        except ValueError:
-            raise DataError("Invalid homescreen")
-        if w != WIDTH or h != HEIGHT:
-            raise DataError(f"Homescreen must be {WIDTH}x{HEIGHT} pixel large")
-        if not is_grayscale:
-            raise DataError("Homescreen must be grayscale")
-
-else:
-
-    def _validate_homescreen_model_specific(homescreen: bytes) -> None:
-        from trezor.ui import HEIGHT, WIDTH
-
-        try:
-            w, h, mcu_height = trezorui2.jpeg_info(homescreen)
-        except ValueError:
-            raise DataError("Invalid homescreen")
-        if w != WIDTH or h != HEIGHT:
-            raise DataError(f"Homescreen must be {WIDTH}x{HEIGHT} pixel large")
-        if mcu_height > 16:
-            raise DataError("Unsupported jpeg type")
-        try:
-            trezorui2.jpeg_test(homescreen)
-        except ValueError:
-            raise DataError("Invalid homescreen")
-
 
 def _validate_homescreen(homescreen: bytes) -> None:
     import storage.device as storage_device
@@ -56,8 +23,8 @@ def _validate_homescreen(homescreen: bytes) -> None:
         raise DataError(
             f"Homescreen is too large, maximum size is {storage_device.HOMESCREEN_MAXSIZE} bytes"
         )
-
-    _validate_homescreen_model_specific(homescreen)
+    if not trezorui2.check_homescreen_format(homescreen):
+        raise DataError("Wrong homescreen format")
 
 
 async def apply_settings(msg: ApplySettings) -> Success: