From 97108665242b852d2a8b04c1cb8bdb5b1c91e2be Mon Sep 17 00:00:00 2001
From: Berkus Decker <berkus+github@metta.systems>
Date: Fri, 28 Jul 2023 22:19:24 +0300
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20Update=20panics,=20exit=20Q?=
 =?UTF-8?q?EMU=20on=20exceptions?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 machine/src/arch/aarch64/traps.rs | 13 ++++++--
 machine/src/lib.rs                |  1 +
 machine/src/panic.rs              | 51 +++++++++++++++++++++++++++++--
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/machine/src/arch/aarch64/traps.rs b/machine/src/arch/aarch64/traps.rs
index 5d16068..839eb28 100644
--- a/machine/src/arch/aarch64/traps.rs
+++ b/machine/src/arch/aarch64/traps.rs
@@ -117,7 +117,10 @@ pub struct ExceptionContext {
 unsafe extern "C" fn default_exception_handler() -> ! {
     println!("Unexpected exception. Halting CPU.");
 
-    endless_sleep()
+    #[cfg(not(qemu))]
+    endless_sleep();
+    #[cfg(qemu)]
+    qemu::semihosting::exit_failure()
 }
 
 // To implement an exception handler, override it by defining the respective
@@ -141,7 +144,7 @@ unsafe extern "C" fn current_el0_synchronous(e: &mut ExceptionContext) {
 #[no_mangle]
 unsafe extern "C" fn current_elx_synchronous(e: &mut ExceptionContext) {
     println!("[!] KERNEL synchronous exception happened.");
-    synchronous_common(e);
+    synchronous_common(e)
 }
 
 // unsafe extern "C" fn current_elx_irq(e: &mut ExceptionContext);
@@ -152,7 +155,11 @@ unsafe extern "C" fn current_elx_synchronous(e: &mut ExceptionContext) {
 unsafe extern "C" fn current_elx_serror(e: &mut ExceptionContext) {
     println!("[!] KERNEL serror exception happened.");
     synchronous_common(e);
-    endless_sleep()
+
+    #[cfg(not(qemu))]
+    endless_sleep();
+    #[cfg(qemu)]
+    qemu::semihosting::exit_failure()
 }
 
 fn cause_to_string(cause: u64) -> &'static str {
diff --git a/machine/src/lib.rs b/machine/src/lib.rs
index 0e9a8fd..1d1460b 100644
--- a/machine/src/lib.rs
+++ b/machine/src/lib.rs
@@ -7,6 +7,7 @@
 #![feature(strict_provenance)]
 #![feature(stmt_expr_attributes)]
 #![feature(slice_ptr_get)]
+#![feature(panic_info_message)]
 #![feature(nonnull_slice_from_raw_parts)]
 #![feature(custom_test_frameworks)]
 #![test_runner(crate::tests::test_runner)]
diff --git a/machine/src/panic.rs b/machine/src/panic.rs
index 3ca9075..912f8d5 100644
--- a/machine/src/panic.rs
+++ b/machine/src/panic.rs
@@ -1,12 +1,57 @@
-pub fn handler(info: &core::panic::PanicInfo) -> ! {
+//! A panic handler for hardware and for QEMU.
+use core::panic::PanicInfo;
+
+pub fn handler(info: &PanicInfo) -> ! {
+    // Protect against panic infinite loops if any of the following code panics itself.
+    panic_prevent_reenter();
     // @todo This may fail to print if the panic message is too long for local print buffer.
-    crate::println!("{}", info);
+    crate::println!("\nPanic: {}\n", info);
     crate::endless_sleep()
 }
 
+/// We have two separate handlers because other crates may use machine crate as a dependency for
+/// running their tests, and this means machine could be compiled with different features.
 pub fn handler_for_tests(info: &core::panic::PanicInfo) -> ! {
     crate::println!("\n[failed]\n");
+    // Protect against panic infinite loops if any of the following code panics itself.
+    panic_prevent_reenter();
     // @todo This may fail to print if the panic message is too long for local print buffer.
-    crate::println!("\nError: {}\n", info);
+    crate::println!("\nPanic: {}\n", info);
     crate::qemu::semihosting::exit_failure()
 }
+
+//--------------------------------------------------------------------------------------------------
+// Private Code
+//--------------------------------------------------------------------------------------------------
+
+/// Stop immediately if called a second time.
+///
+/// # Note
+///
+/// Using atomics here relieves us from needing to use `unsafe` for the static variable.
+///
+/// On `AArch64`, which is the only implemented architecture at the time of writing this,
+/// [`AtomicBool::load`] and [`AtomicBool::store`] are lowered to ordinary load and store
+/// instructions. They are therefore safe to use even with MMU + caching deactivated.
+///
+/// [`AtomicBool::load`]: core::sync::atomic::AtomicBool::load
+/// [`AtomicBool::store`]: core::sync::atomic::AtomicBool::store
+fn panic_prevent_reenter() {
+    use core::sync::atomic::{AtomicBool, Ordering};
+
+    #[cfg(not(target_arch = "aarch64"))]
+    compile_error!("Add the target_arch to above check if the following code is safe to use");
+
+    static PANIC_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
+
+    if !PANIC_IN_PROGRESS.load(Ordering::Relaxed) {
+        PANIC_IN_PROGRESS.store(true, Ordering::Relaxed);
+
+        return;
+    }
+
+    #[cfg(qemu)]
+    crate::qemu::semihosting::exit_failure();
+    #[cfg(not(qemu))]
+    crate::endless_sleep()
+}