added std feature to enforce panic/unwind safety; moved inline(never) comment to debugging section; make no-panic feature for no branch of psm macro? handle closure panic with match?

Author: nstilt1

zeroize_stack/Cargo.toml

@@ -20,6 +20,8 @@ psm = "0.1.26"
 zeroize = { version = "1.0" }
 
 [features]
+default = ["std"]
+std = []
 
 [package.metadata.docs.rs]
 all-features = true

zeroize_stack/src/lib.rs

@@ -1,3 +1,11 @@
+#![no_std]
+#![cfg_attr(docsrs, feature(doc_auto_cfg))]
+#![doc(
+    html_logo_url = "https://raw.githubusercontent.com/RustCrypto/media/6ee8e381/logo.svg",
+    html_favicon_url = "https://raw.githubusercontent.com/RustCrypto/media/6ee8e381/logo.svg"
+)]
+#![warn(missing_docs, unused_qualifications)]
+
 //! # zeroize_stack
 //!
 //! A crate for sanitizing stack memory after sensitive operations—sometimes referred to as _Stack Bleaching_.
@@ -31,8 +39,16 @@
 //! caller must ensure:
 //! - The stack size provided is large enough for the closure to run with.
 //! - The closure does not unwind or return control flow by any means other than
-//!   directly returning.
+//!   directly returning. `std` users do not need to worry about this due to 
+//!   the existence of `catch_unwind`.
 //!
+//! ## `nostd` Support
+//! 
+//! This crate is compatible with `nostd` environments, but it is less safe
+//! in the event that your stack-switched stack panics. Panicking on a separate 
+//! stack can cause undefined behavior (UB), but if it can be caught with 
+//! `std::panic::catch_unwind`, that aspect of the safety should be more safe.
+//! 
 //! ## Use Cases
 //!
 //! - Cryptographic routines
@@ -45,6 +61,27 @@ extern crate alloc;
 
 use alloc::{vec, vec::Vec};
 
+#[cfg(feature = "std")]
+extern crate std;
+#[cfg(feature = "std")]
+use core::any::Any;
+#[cfg(feature = "std")]
+use std::{
+    boxed::Box,
+    panic::catch_unwind,
+};
+#[cfg(feature = "std")]
+type StackSwitchResult<T> = Result<T, Box<dyn Any + Send>>;
+#[cfg(not(feature = "std"))]
+type StackSwitchResult<T> = T;
+
+use core::panic::{AssertUnwindSafe, UnwindSafe};
+
+#[derive(Debug)]
+enum Error {
+    StackPanicked
+}
+
 psm::psm_stack_manipulation! {
     yes {
         /// Executes a function/closure and clears the function's stack by using
@@ -60,34 +97,56 @@ psm::psm_stack_manipulation! {
         ///   some architectures might consume more memory in the stack, such as SPARC.
         /// * `crypto_fn` - the code to run while on the separate stack.
         ///
+        /// ## Panicking
+        /// 
+        /// This function panics when `psm` detects that `on_stack` is unavailable.
+        /// 
+        /// ## Errors
+        /// 
+        /// With the `std` feature enabled, this function will result in an error when 
+        /// the closure panics. You may want to log these errors securely, privately, 
+        /// as cryptography panics could be a little revealing if displayed to 
+        /// the end user.
+        /// 
+        /// ## Debugging
+        /// 
+        /// Using `#[inline(never)]` on the closure's function definition could 
+        /// make it easier to debug as the function should show up.
+        /// 
         /// # Safety
-        ///
-        /// * `crypto_fn` should be marked as `#[inline(never)]`, preventing register
-        ///   reuse and stack layout changes.
+        /// 
         /// * The stack needs to be large enough for `crypto_fn()` to execute without
         ///   overflow.
-        /// * `crypto_fn()` must not unwind or return control flow by any other means
+        /// * `nostd` only: `crypto_fn()` must not unwind or return control flow by any other means
         ///   than by directly returning.
-        pub unsafe fn exec_on_sanitized_stack<F, R>(stack_size_kb: isize, crypto_fn: F) -> R
+        pub unsafe fn exec_on_sanitized_stack<F, R>(stack_size_kb: isize, crypto_fn: F) -> StackSwitchResult<R>
         where
-            F: FnOnce() -> R,
+            F: FnOnce() -> R + UnwindSafe,
         {
             assert!(
                 stack_size_kb * 1024 > 0,
                 "Stack size must be greater than 0 kb and `* 1024` must not overflow `isize`"
             );
-            let mut stack = create_aligned_vec(stack_size_kb as usize, core::mem::align_of::<u128>());
+            let mut stack = create_aligned_vec(stack_size_kb as usize, align_of::<u128>());
+            
             let res = unsafe {
                 psm::on_stack(stack.as_mut_ptr(), stack.len(), || {
-                    crypto_fn()
+                    #[cfg(not(feature = "std"))]
+                    {
+                        crypto_fn()
+                    }
+                    #[cfg(feature = "std")]
+                    {
+                        catch_unwind(AssertUnwindSafe(crypto_fn))
+                    }
                 })
             };
             stack.zeroize();
             res
         }
     }
     no {
-        pub unsafe fn exec_on_sanitized_stack<F, R>(_stack_size_kb: isize, _crypto_fn: F) -> R
+        pub unsafe fn exec_on_sanitized_stack<F, R>(_stack_size_kb: isize, _crypto_fn: F) -> StackSwitchResult
         where
             F: FnOnce() -> R,
         {

zeroize_stack/tests/zeroize_stack.rs

@@ -13,7 +13,7 @@ mod stack_sanitization_tests {
     #[test]
     fn stack_sanitization_v2() {
         let result = unsafe { exec_on_sanitized_stack(4, || dummy_fn()) };
-        assert_eq!(result.1, 12345);
+        assert_eq!(result.unwrap().1, 12345);
         // results in segmentation fault, which is somewhat normal... just wanted
         // to try it
         // assert_eq!(unsafe {*result.0}, 42);