feat(core/rust): unwrap! macro to trigger nicer behavior than the built-in

core/embed/rust/src/lib.rs

@@ -5,6 +5,9 @@
 #![allow(dead_code)]
 
 mod error;
+// use trezorhal for its macros early
+#[macro_use]
+mod trezorhal;
 #[cfg(feature = "micropython")]
 #[macro_use]
 mod micropython;
@@ -13,7 +16,6 @@ mod protobuf;
 mod time;
 #[cfg(feature = "ui_debug")]
 mod trace;
-mod trezorhal;
 
 #[cfg(feature = "ui")]
 #[macro_use]
@@ -25,23 +27,12 @@ pub mod ui;
 /// default `panic` below is that this "debug" version
 /// takes around 10 kB more space in the flash region.
 fn panic_debug(panic_info: &core::panic::PanicInfo) -> ! {
-    use cstr_core::CStr;
-
-    // SAFETY: Safe because we are passing in \0-terminated literals.
-    let empty = unsafe { CStr::from_bytes_with_nul_unchecked(b"\0") };
-    let msg = unsafe { CStr::from_bytes_with_nul_unchecked(b"rs\0") };
-
     // Filling at least the file and line information, if available.
     // TODO: find out how to display message from panic_info.message()
     if let Some(location) = panic_info.location() {
-        let file = location.file();
+        trezorhal::common::__fatal_error("", "rs", location.file(), location.line(), "");
-        let line = location.line();
-        let mut file_str = heapless::String::<100>::from(file);
-        file_str.push('\0').unwrap();
-        let file_cstr = unsafe { CStr::from_bytes_with_nul_unchecked(file_str.as_bytes()) };
-        trezorhal::common::fatal_error(empty, msg, file_cstr, line as _, empty);
     } else {
-        trezorhal::common::fatal_error(empty, msg, empty, 0, empty);
+        trezorhal::common::__fatal_error("", "rs", "", 0, "");
     }
 }
 
@@ -51,17 +42,13 @@ fn panic_debug(panic_info: &core::panic::PanicInfo) -> ! {
 #[panic_handler]
 /// Default panic handling. Not showing any details - thus saving flash space.
 fn panic(_info: &core::panic::PanicInfo) -> ! {
-    use cstr_core::CStr;
+    // TODO: as of Rust 1.63 / nightly 2022-08, ignoring the `_info` parameter does
-
+    // not help with saving flash space -- the `fmt` machinery still gets
-    // Although it would be ideal to use the original error message, ignoring it
+    // compiled in. We can avoid that by using unstable Cargo arguments:
-    // lets us avoid the `fmt` machinery and its code size and is also important for
+    //   -Zbuild-std=core -Zbuild-std-features=panic_immediate_abort
-    // security reasons, as we do not always controls the message contents. We
+    // Doing that will compile every panic!() to a single udf instruction which
-    // should also avoid printing "panic" or "rust" on the user screen to avoid any
+    // raises a Hard Fault on hardware.
-    // confusion.
+    //
-
+    // Otherwise, use `unwrap!` macro from trezorhal.
-    // SAFETY: Safe because we are passing in \0-terminated literals.
+    trezorhal::common::__fatal_error("", "rs", "", 0, "");
-    let empty = unsafe { CStr::from_bytes_with_nul_unchecked(b"\0") };
-    let msg = unsafe { CStr::from_bytes_with_nul_unchecked(b"rs\0") };
-
-    trezorhal::common::fatal_error(empty, msg, empty, 0, empty);
 }

core/embed/rust/src/micropython/buffer.rs

@@ -175,7 +175,7 @@ impl AsRef<str> for StrBuffer {
         // Rust seems to be stricter in what it considers UTF-8 though.
         // In case there's a mismatch, this code will cleanly panic
         // before attempting to use the data.
-        str::from_utf8(self.0.as_ref()).unwrap()
+        unwrap!(str::from_utf8(self.0.as_ref()), "Invalid internal UTF-8.")
     }
 }
 

core/embed/rust/src/micropython/map.rs

@@ -127,15 +127,14 @@ impl Map {
         unsafe {
             let map = self as *mut Self;
             // EXCEPTION: Will raise if allocation fails.
-            let elem = catch_exception(|| {
+            let elem = unwrap!(catch_exception(|| {
                 ffi::mp_map_lookup(
                     map,
                     index,
                     ffi::_mp_map_lookup_kind_t_MP_MAP_LOOKUP_ADD_IF_NOT_FOUND,
                 )
             })?
-            .as_mut()
+            .as_mut()); // `MP_MAP_LOOKUP_ADD_IF_NOT_FOUND` should always return a non-null pointer.
-            .unwrap(); // `MP_MAP_LOOKUP_ADD_IF_NOT_FOUND` should always return a non-null pointer.
             elem.value = value;
         }
         Ok(())

core/embed/rust/src/micropython/obj.rs

@@ -310,14 +310,14 @@ impl TryFrom<(Obj, Obj)> for Obj {
 impl From<u8> for Obj {
     fn from(val: u8) -> Self {
         // `u8` will fit into smallint so no error should happen here.
-        u32::from(val).try_into().unwrap()
+        unwrap!(u32::from(val).try_into())
     }
 }
 
 impl From<u16> for Obj {
     fn from(val: u16) -> Self {
         // `u16` will fit into smallint so no error should happen here.
-        u32::from(val).try_into().unwrap()
+        unwrap!(u32::from(val).try_into())
     }
 }
 

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

@@ -218,7 +218,7 @@ unsafe extern "C" fn msg_def_obj_attr(self_in: Obj, attr: ffi::qstr, dest: *mut
         match attr {
             Qstr::MP_QSTR_MESSAGE_NAME => {
                 // Return the QSTR name of this message def.
-                let name = Qstr::from_u16(find_name_by_msg_offset(this.def.offset).unwrap());
+                let name = Qstr::from_u16(unwrap!(find_name_by_msg_offset(this.def.offset)));
                 unsafe {
                     dest.write(name.into());
                 };

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

@@ -1,24 +1,100 @@
-use cstr_core::CStr;
+mod ffi {
-
+    extern "C" {
-extern "C" {
+        // trezorhal/common.c
-    // trezorhal/common.c
+        pub fn __fatal_error(
-    fn __fatal_error(
+            expr: *const cty::c_char,
-        expr: *const cty::c_char,
+            msg: *const cty::c_char,
-        msg: *const cty::c_char,
+            file: *const cty::c_char,
-        file: *const cty::c_char,
+            line: i32,
-        line: i32,
+            func: *const cty::c_char,
-        func: *const cty::c_char,
+        ) -> !;
-    ) -> !;
+    }
 }
 
-pub fn fatal_error(expr: &CStr, msg: &CStr, file: &CStr, line: i32, func: &CStr) -> ! {
+pub fn __fatal_error(expr: &str, msg: &str, file: &str, line: u32, func: &str) -> ! {
+    const MAX_LEN: usize = 50 + 1; // Leave space for the null terminator.
+
+    fn as_cstr_buf(s: &str) -> [cty::c_char; MAX_LEN] {
+        let mut buf = [0 as cty::c_char; MAX_LEN];
+        for (i, c) in s.as_bytes().iter().enumerate() {
+            if i >= MAX_LEN {
+                break;
+            }
+            buf[i] = *c as cty::c_char;
+        }
+        buf[MAX_LEN - 1] = 0;
+        buf
+    }
+
+    let expr_buf = as_cstr_buf(expr);
+    let msg_buf = as_cstr_buf(msg);
+    let file_buf = as_cstr_buf(file);
+    let func_buf = as_cstr_buf(func);
+
     unsafe {
-        __fatal_error(
+        ffi::__fatal_error(
-            expr.as_ptr(),
+            expr_buf.as_ptr(),
-            msg.as_ptr(),
+            msg_buf.as_ptr(),
-            file.as_ptr(),
+            file_buf.as_ptr(),
-            line,
+            line as i32,
-            func.as_ptr(),
+            func_buf.as_ptr(),
         );
     }
 }
+
+pub trait UnwrapOrFatalError<T> {
+    fn unwrap_or_fatal_error(self, expr: &str, msg: &str, file: &str, line: u32, func: &str) -> T;
+}
+
+impl<T> UnwrapOrFatalError<T> for Option<T> {
+    fn unwrap_or_fatal_error(self, expr: &str, msg: &str, file: &str, line: u32, func: &str) -> T {
+        match self {
+            Some(x) => x,
+            None => __fatal_error(expr, msg, file, line, func),
+        }
+    }
+}
+
+impl<T, E> UnwrapOrFatalError<T> for Result<T, E> {
+    fn unwrap_or_fatal_error(self, expr: &str, msg: &str, file: &str, line: u32, func: &str) -> T {
+        match self {
+            Ok(x) => x,
+            Err(_) => __fatal_error(expr, msg, file, line, func),
+        }
+    }
+}
+
+macro_rules! function_name {
+    () => {{
+        fn f() {}
+        fn type_name_of<T>(_: T) -> &'static str {
+            core::any::type_name::<T>()
+        }
+        let name = type_name_of(f);
+        name.get(..name.len() - 3).unwrap_or("")
+    }};
+}
+
+macro_rules! unwrap {
+    ($e:expr, $msg:expr) => {{
+        use crate::trezorhal::common::UnwrapOrFatalError;
+        $e.unwrap_or_fatal_error("unwrap failed", $msg, file!(), line!(), function_name!())
+    }};
+    ($expr:expr) => {
+        unwrap!($expr, "")
+    };
+}
+
+macro_rules! ensure {
+    ($what:expr, $error:expr) => {
+        if !($what) {
+            fatal_error!(stringify!($what), $error);
+        }
+    };
+}
+
+macro_rules! fatal_error {
+    ($expr:expr, $msg:expr) => {{
+        crate::trezorhal::common::__fatal_error($expr, $msg, file!(), line!(), function_name!());
+    }};
+}

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

@@ -1,4 +1,6 @@
 pub mod bip39;
+#[macro_use]
+#[allow(unused_macros)]
 pub mod common;
 #[cfg(feature = "ui")]
 pub mod display;

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

@@ -62,7 +62,7 @@ pub fn rect_fill_rounded(r: Rect, fg_color: Color, bg_color: Color, radius: u8)
 /// NOTE: Cannot start at odd x-coordinate. In this case icon is shifted 1px
 /// left.
 pub fn icon_top_left(top_left: Point, data: &[u8], fg_color: Color, bg_color: Color) {
-    let toif_info = display::toif_info(data).unwrap();
+    let toif_info = unwrap!(display::toif_info(data), "Invalid TOIF data");
     assert!(toif_info.grayscale);
     display::icon(
         top_left.x,
@@ -76,7 +76,7 @@ pub fn icon_top_left(top_left: Point, data: &[u8], fg_color: Color, bg_color: Co
 }
 
 pub fn icon(center: Point, data: &[u8], fg_color: Color, bg_color: Color) {
-    let toif_info = display::toif_info(data).unwrap();
+    let toif_info = unwrap!(display::toif_info(data), "Invalid TOIF data");
     assert!(toif_info.grayscale);
 
     let r = Rect::from_center_and_size(
@@ -95,7 +95,7 @@ pub fn icon(center: Point, data: &[u8], fg_color: Color, bg_color: Color) {
 }
 
 pub fn icon_rust(center: Point, data: &[u8], fg_color: Color, bg_color: Color) {
-    let toif_info = display::toif_info(data).unwrap();
+    let toif_info = unwrap!(display::toif_info(data), "Invalid TOIF data");
     assert!(toif_info.grayscale);
 
     let r = Rect::from_center_and_size(
@@ -119,14 +119,14 @@ pub fn icon_rust(center: Point, data: &[u8], fg_color: Color, bg_color: Color) {
 
             if clamped.contains(p) {
                 if x % 2 == 0 {
-                    ctx.uncompress(&mut dest).unwrap();
+                    unwrap!(ctx.uncompress(&mut dest), "Decompression failed");
                     pixeldata(colortable[(dest[0] >> 4) as usize]);
                 } else {
                     pixeldata(colortable[(dest[0] & 0xF) as usize]);
                 }
             } else if x % 2 == 0 {
                 //continue unzipping but dont write to display
-                ctx.uncompress(&mut dest).unwrap();
+                unwrap!(ctx.uncompress(&mut dest), "Decompression failed");
             }
         }
     }
@@ -135,7 +135,7 @@ pub fn icon_rust(center: Point, data: &[u8], fg_color: Color, bg_color: Color) {
 }
 
 pub fn image(center: Point, data: &[u8]) {
-    let toif_info = display::toif_info(data).unwrap();
+    let toif_info = unwrap!(display::toif_info(data), "Invalid TOIF data");
     assert!(!toif_info.grayscale);
 
     let r = Rect::from_center_and_size(
@@ -307,7 +307,7 @@ pub fn rect_rounded2_partial(
     let mut icon_width = 0;
 
     if let Some((icon_bytes, icon_color)) = icon {
-        let toif_info = display::toif_info(icon_bytes).unwrap();
+        let toif_info = unwrap!(display::toif_info(icon_bytes), "Invalid TOIF data");
         assert!(toif_info.grayscale);
 
         if toif_info.width <= MAX_ICON_SIZE && toif_info.height <= MAX_ICON_SIZE {
@@ -318,7 +318,7 @@ pub fn rect_rounded2_partial(
             icon_area_clamped = icon_area.clamp(constant::screen());
 
             let mut ctx = uzlib::UzlibContext::new(&icon_bytes[12..], false);
-            ctx.uncompress(&mut icon_data).unwrap();
+            unwrap!(ctx.uncompress(&mut icon_data), "Decompression failed");
             icon_colortable = get_color_table(icon_color, bg_color);
             icon_width = toif_info.width.into();
             use_icon = true;

core/embed/rust/src/ui/model_tt/component/keyboard/common.rs

@@ -111,7 +111,7 @@ impl MultiTapKeyboard {
 
         assert!(!key_text.is_empty());
         // Now we can be sure that a looped iterator will return a value
-        let ch = key_text.chars().cycle().nth(press).unwrap();
+        let ch = unwrap!(key_text.chars().cycle().nth(press));
         if is_pending {
             TextEdit::ReplaceLast(ch)
         } else {

core/embed/rust/src/ui/model_tt/component/keyboard/slip39.rs

@@ -141,7 +141,7 @@ impl Component for Slip39Input {
             {
                 assert!(!Self::keys()[key].is_empty());
                 // Now we can be sure that the looped iterator will return a value.
-                let ch = Self::keys()[key].chars().cycle().nth(press).unwrap();
+                let ch = unwrap!(Self::keys()[key].chars().cycle().nth(press));
                 text.pop();
                 text.push(ch)
                     .assert_if_debugging_ui("Text buffer is too small");
@@ -196,7 +196,7 @@ impl Slip39Input {
     /// ```
     fn key_digit(key: usize) -> char {
         let index = key + 1;
-        char::from_digit(index as u32, 10).unwrap()
+        unwrap!(char::from_digit(index as u32, 10))
     }
 
     fn complete_word_from_dictionary(&mut self, ctx: &mut EventCtx) {