WIP - improve TOIF handling and checks

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<Self, Self::Error> {
+        // 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.
 //

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
     }

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<const N: usize>(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) };
+        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.")))
-            }
+        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) };
+        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())
-                };
-
-                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.")))
         }
     };
 

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(

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

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: