From af031eb4b1c65e40d22bfe7f3335fc252833acd0 Mon Sep 17 00:00:00 2001
From: Gary Guo <gary@garyguo.net>
Date: Thu, 26 Aug 2021 10:51:55 +0100
Subject: [PATCH] Panic utilities

---
 Cargo.toml   |  1 +
 src/abi.rs   | 17 +++++++++--
 src/lib.rs   |  9 ++++--
 src/panic.rs | 86 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 108 insertions(+), 5 deletions(-)
 create mode 100644 src/panic.rs

diff --git a/Cargo.toml b/Cargo.toml
index b688e20..8d4a937 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,7 @@ hide-trace = []
 personality = []
 personality-dummy = []
 print = ["libc"]
+panic = ["alloc"]
 system-alloc = []
 default = ["dwarf-expr", "hide-trace", "fde-phdr", "fde-registry"]
 
diff --git a/src/abi.rs b/src/abi.rs
index 00ad302..c20e80e 100644
--- a/src/abi.rs
+++ b/src/abi.rs
@@ -10,7 +10,7 @@ use crate::util::*;
 
 #[repr(transparent)]
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct UnwindReasonCode(c_int);
+pub struct UnwindReasonCode(pub c_int);
 
 #[allow(unused)]
 impl UnwindReasonCode {
@@ -27,7 +27,7 @@ impl UnwindReasonCode {
 
 #[repr(transparent)]
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct UnwindAction(c_int);
+pub struct UnwindAction(pub c_int);
 
 impl UnwindAction {
     pub const SEARCH_PHASE: Self = Self(1);
@@ -78,6 +78,19 @@ pub struct UnwindException {
     private_unused: [usize; Arch::UNWIND_PRIVATE_DATA_SIZE - 2],
 }
 
+impl UnwindException {
+    #[inline]
+    pub fn new() -> UnwindException {
+        UnwindException {
+            exception_class: 0,
+            exception_cleanup: None,
+            private_1: None,
+            private_2: 0,
+            private_unused: [0; Arch::UNWIND_PRIVATE_DATA_SIZE - 2],
+        }
+    }
+}
+
 pub type UnwindTraceFn =
     unsafe extern "C" fn(ctx: &mut UnwindContext<'_>, arg: *mut c_void) -> UnwindReasonCode;
 
diff --git a/src/lib.rs b/src/lib.rs
index c569acf..74476dc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,7 @@
     any(feature = "personality", feature = "personality-dummy"),
     feature(lang_items)
 )]
+#![cfg_attr(feature = "panic", feature(core_intrinsics))]
 #![warn(rust_2018_idioms)]
 #![warn(unsafe_op_in_unsafe_fn)]
 #![no_std]
@@ -12,7 +13,8 @@
 #[cfg(feature = "alloc")]
 extern crate alloc;
 
-mod abi;
+pub mod abi;
+
 mod arch;
 mod find_fde;
 mod frame;
@@ -26,7 +28,8 @@ mod personality;
 #[cfg(feature = "personality-dummy")]
 mod personality_dummy;
 
+#[cfg(feature = "panic")]
+pub mod panic;
+
 #[cfg(feature = "system-alloc")]
 mod system_alloc;
-
-pub use abi::*;
diff --git a/src/panic.rs b/src/panic.rs
new file mode 100644
index 0000000..c441254
--- /dev/null
+++ b/src/panic.rs
@@ -0,0 +1,86 @@
+use alloc::boxed::Box;
+use core::any::Any;
+use core::mem::ManuallyDrop;
+
+use crate::abi::*;
+
+#[repr(C)]
+struct Exception {
+    exception: UnwindException,
+    payload: Box<dyn Any + Send>,
+}
+
+const RUST_EXCEPTION_CLASS: u64 = u64::from_be_bytes(*b"MOZ\0RUST");
+
+pub fn begin_panic(payload: Box<dyn Any + Send>) -> UnwindReasonCode {
+    unsafe extern "C" fn exception_cleanup(
+        _unwind_code: UnwindReasonCode,
+        exception: *mut UnwindException,
+    ) {
+        unsafe {
+            let _ = Box::from_raw(exception as *mut Exception);
+        }
+        core::intrinsics::abort();
+    }
+
+    let mut unwind_ex = UnwindException::new();
+    unwind_ex.exception_class = RUST_EXCEPTION_CLASS;
+    unwind_ex.exception_cleanup = Some(exception_cleanup);
+    let exception = Box::new(Exception {
+        exception: unwind_ex,
+        payload,
+    });
+    _Unwind_RaiseException(unsafe { &mut *(Box::into_raw(exception) as *mut UnwindException) })
+}
+
+#[cold]
+unsafe fn cleanup(payload: *mut u8) -> Box<dyn Any + Send + 'static> {
+    let exception = payload as *mut UnwindException;
+    if unsafe { (*exception).exception_class } != RUST_EXCEPTION_CLASS {
+        unsafe { _Unwind_DeleteException(exception) };
+        core::intrinsics::abort();
+    }
+    let unwind_ex = unsafe { Box::from_raw(exception as *mut Exception) };
+    unwind_ex.payload
+}
+
+pub fn catch_unwind<R, F: FnOnce() -> R>(f: F) -> Result<R, Box<dyn Any + Send>> {
+    union Data<F, R> {
+        f: ManuallyDrop<F>,
+        r: ManuallyDrop<R>,
+        p: ManuallyDrop<Box<dyn Any + Send>>,
+    }
+
+    let mut data = Data {
+        f: ManuallyDrop::new(f),
+    };
+
+    let data_ptr = &mut data as *mut _ as *mut u8;
+    unsafe {
+        return if core::intrinsics::r#try(do_call::<F, R>, data_ptr, do_catch::<F, R>) == 0 {
+            Ok(ManuallyDrop::into_inner(data.r))
+        } else {
+            Err(ManuallyDrop::into_inner(data.p))
+        };
+    }
+
+    #[inline]
+    fn do_call<F: FnOnce() -> R, R>(data: *mut u8) {
+        unsafe {
+            let data = data as *mut Data<F, R>;
+            let data = &mut (*data);
+            let f = ManuallyDrop::take(&mut data.f);
+            data.r = ManuallyDrop::new(f());
+        }
+    }
+
+    #[inline]
+    fn do_catch<F: FnOnce() -> R, R>(data: *mut u8, payload: *mut u8) {
+        unsafe {
+            let data = data as *mut Data<F, R>;
+            let data = &mut (*data);
+            let obj = cleanup(payload);
+            data.p = ManuallyDrop::new(obj);
+        }
+    }
+}