From 7265e6079e9e9fda8657e42019c3d294d32092e7 Mon Sep 17 00:00:00 2001
From: beerosagos <luca@shiftcrypto.ch>
Date: Mon, 16 Dec 2024 12:22:32 +0100
Subject: [PATCH] rust: handle panics to show meaningful error

In b09c333 we added the `panic_immediate_abort` to the rust cargo flags.
This caused all the panics previously handled by the handler defined in
lib.rs to just show the `hard fault` error message on the screen of the
bitbox.

This adds the manual handling of some of those panics, to provide useful
debug messages.
---
 src/rust/bitbox02-rust-c/src/lib.rs              |  2 +-
 src/rust/bitbox02-rust/src/general.rs            |  8 ++++++++
 src/rust/bitbox02-rust/src/general/screen.rs     |  4 ++--
 src/rust/bitbox02-rust/src/hww/api/restore.rs    | 11 ++++++++---
 src/rust/bitbox02-rust/src/workflow/unlock.rs    |  8 +++++---
 src/rust/bitbox02/src/ui/ui.rs                   |  5 +++--
 src/rust/bitbox02/src/ui/ui_stub.rs              |  2 +-
 src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs |  2 +-
 8 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/rust/bitbox02-rust-c/src/lib.rs b/src/rust/bitbox02-rust-c/src/lib.rs
index 094a41a26..5c15e6ce2 100644
--- a/src/rust/bitbox02-rust-c/src/lib.rs
+++ b/src/rust/bitbox02-rust-c/src/lib.rs
@@ -47,7 +47,7 @@ mod workflow;
 fn panic(info: &core::panic::PanicInfo) -> ! {
     ::util::log::log!("{}", info);
     #[cfg(feature = "firmware")]
-    bitbox02_rust::print_debug!(0, "Error: {}", info);
+    bitbox02_rust::print_screen!(0, "Error: {}", info);
     loop {}
 }
 
diff --git a/src/rust/bitbox02-rust/src/general.rs b/src/rust/bitbox02-rust/src/general.rs
index bfe3bc2d0..c4d3388a5 100644
--- a/src/rust/bitbox02-rust/src/general.rs
+++ b/src/rust/bitbox02-rust/src/general.rs
@@ -14,3 +14,11 @@
 
 #[macro_use]
 pub mod screen;
+
+/// displays the input error message on the screen and enters
+/// an infinite loop.
+#[allow(clippy::empty_loop)]
+pub fn abort(err: &str) -> ! {
+    print_screen!(0, "Error: {}", err);
+    loop {}
+}
diff --git a/src/rust/bitbox02-rust/src/general/screen.rs b/src/rust/bitbox02-rust/src/general/screen.rs
index ddf93717a..1aedafb50 100644
--- a/src/rust/bitbox02-rust/src/general/screen.rs
+++ b/src/rust/bitbox02-rust/src/general/screen.rs
@@ -31,11 +31,11 @@ pub fn print_debug_internal(duration: Duration, msg: &str) {
 /// ```no_run
 /// # #[macro_use] extern crate bitbox02_rust; fn main() {
 /// let my_str = "abc";
-/// print_debug!(1000, "{}", &my_str);
+/// print_screen!(1000, "{}", &my_str);
 /// # }
 /// ```
 #[macro_export]
-macro_rules! print_debug {
+macro_rules! print_screen {
     ($duration:expr, $($arg:tt)*) => ({
         extern crate alloc;
         let duration = core::time::Duration::from_millis($duration);
diff --git a/src/rust/bitbox02-rust/src/hww/api/restore.rs b/src/rust/bitbox02-rust/src/hww/api/restore.rs
index c0a43f42f..6cc65c451 100644
--- a/src/rust/bitbox02-rust/src/hww/api/restore.rs
+++ b/src/rust/bitbox02-rust/src/hww/api/restore.rs
@@ -17,6 +17,7 @@ use crate::pb;
 
 use pb::response::Response;
 
+use crate::general::abort;
 use crate::workflow::{confirm, mnemonic, password, status, unlock};
 
 pub async fn from_file(request: &pb::RestoreBackupRequest) -> Result<Response, Error> {
@@ -75,7 +76,9 @@ pub async fn from_file(request: &pb::RestoreBackupRequest) -> Result<Response, E
     }
 
     bitbox02::memory::set_initialized().or(Err(Error::Memory))?;
-    bitbox02::keystore::unlock(&password).expect("restore_from_file: unlock failed");
+    if bitbox02::keystore::unlock(&password).is_err() {
+        abort("restore_from_file: unlock failed");
+    };
 
     // Ignore non-critical error.
     let _ = bitbox02::memory::set_device_name(&metadata.name);
@@ -144,9 +147,11 @@ pub async fn from_mnemonic(
     }
 
     bitbox02::memory::set_initialized().or(Err(Error::Memory))?;
-
     // This should never fail.
-    bitbox02::keystore::unlock(&password).expect("restore_from_mnemonic: unlock failed");
+    if bitbox02::keystore::unlock(&password).is_err() {
+        abort("restore_from_mnemonic: unlock failed");
+    };
+
     unlock::unlock_bip39().await;
     Ok(Response::Success(pb::Success {}))
 }
diff --git a/src/rust/bitbox02-rust/src/workflow/unlock.rs b/src/rust/bitbox02-rust/src/workflow/unlock.rs
index 3c629225e..fa6f0cf70 100644
--- a/src/rust/bitbox02-rust/src/workflow/unlock.rs
+++ b/src/rust/bitbox02-rust/src/workflow/unlock.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::general::abort;
 use crate::workflow::confirm;
 use crate::workflow::password;
 use crate::workflow::status::status;
@@ -116,9 +117,10 @@ pub async fn unlock_bip39() {
         }
     }
 
-    bitbox02::ui::with_lock_animation(|| {
-        keystore::unlock_bip39(&mnemonic_passphrase).expect("bip39 unlock failed");
-    });
+    let result = bitbox02::ui::with_lock_animation(|| keystore::unlock_bip39(&mnemonic_passphrase));
+    if result.is_err() {
+        abort("bip39 unlock failed");
+    }
 }
 
 /// Invokes the unlock workflow. This function does not finish until the keystore is unlocked, or
diff --git a/src/rust/bitbox02/src/ui/ui.rs b/src/rust/bitbox02/src/ui/ui.rs
index 1aae8acca..deb04b4b2 100644
--- a/src/rust/bitbox02/src/ui/ui.rs
+++ b/src/rust/bitbox02/src/ui/ui.rs
@@ -422,10 +422,11 @@ pub fn trinary_input_string_set_input(component: &mut Component, word: &str) {
     }
 }
 
-pub fn with_lock_animation<F: Fn()>(f: F) {
+pub fn with_lock_animation<F: Fn() -> R, R>(f: F) -> R {
     unsafe { bitbox02_sys::lock_animation_start() };
-    f();
+    let result = f();
     unsafe { bitbox02_sys::lock_animation_stop() };
+    result
 }
 
 pub fn screen_stack_pop_all() {
diff --git a/src/rust/bitbox02/src/ui/ui_stub.rs b/src/rust/bitbox02/src/ui/ui_stub.rs
index 2e1fba914..19279ec1f 100644
--- a/src/rust/bitbox02/src/ui/ui_stub.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub.rs
@@ -151,7 +151,7 @@ pub fn trinary_input_string_set_input(_component: &mut Component, _word: &str) {
     panic!("not implemented")
 }
 
-pub fn with_lock_animation<F: Fn()>(f: F) {
+pub fn with_lock_animation<F: Fn() -> R, R>(f: F) -> R {
     f()
 }
 
diff --git a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
index 56f0b0193..b51509309 100644
--- a/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
+++ b/src/rust/bitbox02/src/ui/ui_stub_c_unit_tests.rs
@@ -162,7 +162,7 @@ pub fn trinary_input_string_set_input(_component: &mut Component, _word: &str) {
     panic!("not implemented")
 }
 
-pub fn with_lock_animation<F: Fn()>(f: F) {
+pub fn with_lock_animation<F: Fn() -> R, R>(f: F) -> R {
     f()
 }