Implement unwind safety

Author: Coder-256

src/lib.rs

@@ -44,10 +44,11 @@ assert!(nums[0] == "2");
 
 #![warn(missing_docs)]
 
+use std::convert::TryFrom;
 use std::error::Error;
 use std::fmt;
-use std::mem;
 use std::os::raw::c_void;
+use std::panic::AssertUnwindSafe;
 use std::time::Duration;
 
 use crate::ffi::*;
@@ -79,68 +80,83 @@ impl fmt::Display for WaitTimeout {
 impl Error for WaitTimeout {}
 
 fn time_after_delay(delay: Duration) -> dispatch_time_t {
-    delay
-        .as_secs()
-        .checked_mul(1_000_000_000)
-        .and_then(|i| i.checked_add(delay.subsec_nanos() as u64))
-        .and_then(|i| {
-            if i < (i64::max_value() as u64) {
-                Some(i as i64)
-            } else {
-                None
-            }
-        })
-        .map_or(DISPATCH_TIME_FOREVER, |i| unsafe {
-            dispatch_time(DISPATCH_TIME_NOW, i)
-        })
+    i64::try_from(delay.as_nanos()).map_or(DISPATCH_TIME_FOREVER, |i| unsafe {
+        dispatch_time(DISPATCH_TIME_NOW, i)
+    })
 }
 
 fn context_and_function<F>(closure: F) -> (*mut c_void, dispatch_function_t)
 where
-    F: FnOnce(),
+    F: 'static + FnOnce(),
 {
-    extern "C" fn work_execute_closure<F>(context: Box<F>)
+    extern "C" fn work_execute_closure<F>(context: *mut c_void)
     where
         F: FnOnce(),
     {
-        (*context)();
+        let closure: Box<F> = unsafe { Box::from_raw(context as *mut _) };
+        if std::panic::catch_unwind(AssertUnwindSafe(closure)).is_err() {
+            // Abort to prevent unwinding across FFI
+            std::process::abort();
+        }
     }
 
     let closure = Box::new(closure);
-    let func: extern "C" fn(Box<F>) = work_execute_closure::<F>;
-    unsafe { (mem::transmute(closure), mem::transmute(func)) }
+    let func: dispatch_function_t = work_execute_closure::<F>;
+    (Box::into_raw(closure) as *mut c_void, func)
 }
 
-fn context_and_sync_function<F>(closure: &mut Option<F>) -> (*mut c_void, dispatch_function_t)
+fn with_context_and_sync_function<F, T>(
+    closure: F,
+    wrapper: impl FnOnce(*mut c_void, dispatch_function_t),
+) -> Option<T>
 where
-    F: FnOnce(),
+    F: FnOnce() -> T,
 {
-    extern "C" fn work_read_closure<F>(context: &mut Option<F>)
+    #[derive(Debug)]
+    struct SyncContext<F, T> {
+        closure: *mut F,
+        result: Option<std::thread::Result<T>>,
+    }
+
+    extern "C" fn work_execute_closure<F, T>(context: *mut c_void)
     where
-        F: FnOnce(),
+        F: FnOnce() -> T,
     {
-        // This is always passed Some, so it's safe to unwrap
-        let closure = context.take().unwrap();
-        closure();
+        let sync_context: &mut SyncContext<F, T> = unsafe { &mut *(context as *mut _) };
+        let closure = unsafe { Box::from_raw(sync_context.closure) };
+        sync_context.result = Some(std::panic::catch_unwind(AssertUnwindSafe(closure)));
     }
 
-    let context: *mut Option<F> = closure;
-    let func: extern "C" fn(&mut Option<F>) = work_read_closure::<F>;
-    unsafe { (context as *mut c_void, mem::transmute(func)) }
+    let mut sync_context: SyncContext<F, T> = SyncContext {
+        closure: Box::into_raw(Box::new(closure)),
+        result: None,
+    };
+    let func: dispatch_function_t = work_execute_closure::<F, T>;
+    wrapper(&mut sync_context as *mut _ as *mut c_void, func);
+
+    // If the closure panicked, resume unwinding
+    match sync_context.result.transpose() {
+        Ok(res) => res,
+        Err(unwind_payload) => std::panic::resume_unwind(unwind_payload),
+    }
 }
 
 fn context_and_apply_function<F>(closure: &F) -> (*mut c_void, extern "C" fn(*mut c_void, usize))
 where
     F: Fn(usize),
 {
-    extern "C" fn work_apply_closure<F>(context: &F, iter: usize)
+    extern "C" fn work_apply_closure<F>(context: *mut c_void, iter: usize)
     where
         F: Fn(usize),
     {
-        context(iter);
+        let context: &F = unsafe { &*(context as *const _) };
+        if std::panic::catch_unwind(AssertUnwindSafe(|| context(iter))).is_err() {
+            // Abort to prevent unwinding across FFI
+            std::process::abort();
+        }
     }
 
     let context: *const F = closure;
-    let func: extern "C" fn(&F, usize) = work_apply_closure::<F>;
-    unsafe { (context as *mut c_void, mem::transmute(func)) }
+    let func: extern "C" fn(*mut c_void, usize) = work_apply_closure::<F>;
+    (context as *mut c_void, func)
 }

src/once.rs

@@ -1,7 +1,7 @@
 use std::cell::UnsafeCell;
 
-use crate::context_and_sync_function;
 use crate::ffi::*;
+use crate::with_context_and_sync_function;
 
 /// A predicate used to execute a closure only once for the lifetime of an
 /// application.
@@ -34,11 +34,9 @@ impl Once {
         where
             F: FnOnce(),
         {
-            let mut work = Some(work);
-            let (context, work) = context_and_sync_function(&mut work);
-            unsafe {
+            with_context_and_sync_function(work, |context, work| unsafe {
                 dispatch_once_f(predicate, context, work);
-            }
+            });
         }
 
         unsafe {

src/queue.rs

@@ -6,7 +6,8 @@ use std::time::Duration;
 
 use crate::ffi::*;
 use crate::{
-    context_and_apply_function, context_and_function, context_and_sync_function, time_after_delay,
+    context_and_apply_function, context_and_function, time_after_delay,
+    with_context_and_sync_function,
 };
 
 /// The type of a dispatch queue.
@@ -134,21 +135,10 @@ impl Queue {
         F: Send + FnOnce() -> T,
         T: Send,
     {
-        let mut result = None;
-        {
-            let result_ref = &mut result;
-            let work = move || {
-                *result_ref = Some(work());
-            };
-
-            let mut work = Some(work);
-            let (context, work) = context_and_sync_function(&mut work);
-            unsafe {
-                dispatch_sync_f(self.ptr, context, work);
-            }
-        }
-        // This was set so it's safe to unwrap
-        result.unwrap()
+        with_context_and_sync_function(work, |context, work| unsafe {
+            dispatch_sync_f(self.ptr, context, work);
+        })
+        .unwrap()
     }
 
     /// Submits a closure for asynchronous execution on self and returns
@@ -197,7 +187,7 @@ impl Queue {
     {
         let slice_ptr = slice.as_mut_ptr();
         let work = move |i| unsafe {
-            work(&mut *slice_ptr.offset(i as isize));
+            work(&mut *slice_ptr.add(i));
         };
         let (context, work) = context_and_apply_function(&work);
         unsafe {
@@ -221,8 +211,8 @@ impl Queue {
         let dest_ptr = dest.as_mut_ptr();
 
         let work = move |i| unsafe {
-            let result = work(ptr::read(src_ptr.offset(i as isize)));
-            ptr::write(dest_ptr.offset(i as isize), result);
+            let result = work(ptr::read(src_ptr.add(i)));
+            ptr::write(dest_ptr.add(i), result);
         };
         let (context, work) = context_and_apply_function(&work);
         unsafe {
@@ -251,21 +241,10 @@ impl Queue {
         F: Send + FnOnce() -> T,
         T: Send,
     {
-        let mut result = None;
-        {
-            let result_ref = &mut result;
-            let work = move || {
-                *result_ref = Some(work());
-            };
-
-            let mut work = Some(work);
-            let (context, work) = context_and_sync_function(&mut work);
-            unsafe {
-                dispatch_barrier_sync_f(self.ptr, context, work);
-            }
-        }
-        // This was set so it's safe to unwrap
-        result.unwrap()
+        with_context_and_sync_function(work, |context, work| unsafe {
+            dispatch_barrier_sync_f(self.ptr, context, work);
+        })
+        .unwrap()
     }
 
     /// Submits a closure to be executed on self as a barrier and returns
@@ -523,4 +502,34 @@ mod tests {
         q.exec_sync(|| ());
         assert_eq!(*num.lock().unwrap(), 1);
     }
+
+    #[test]
+    #[should_panic(expected = "panic from exec_sync")]
+    fn test_panic_serial() {
+        let q = Queue::create("", QueueAttribute::Serial);
+
+        q.exec_sync(|| panic!("panic from exec_sync"));
+
+        panic!("exec_sync did NOT panic");
+    }
+
+    #[test]
+    #[should_panic(expected = "panic from exec_sync")]
+    fn test_panic_concurrent() {
+        let q = Queue::create("", QueueAttribute::Concurrent);
+
+        q.exec_sync(|| panic!("panic from exec_sync"));
+
+        panic!("exec_sync did NOT panic");
+    }
+
+    #[test]
+    #[should_panic(expected = "panic from barrier_sync")]
+    fn test_panic_barrier() {
+        let q = Queue::create("", QueueAttribute::Serial);
+
+        q.barrier_sync(|| panic!("panic from barrier_sync"));
+
+        panic!("barrier_sync did NOT panic");
+    }
 }