Add segfault catcher. (#1053)

* Add `SafeRunner` for catching segfaults.

* Add API to manually abort.

* Fix MacOS build.

* Hide everything behind a feature. Fix various bugs triggered on early returns.

* Add some documentation.

Cargo.toml

@@ -54,6 +54,7 @@ scarb = ["build-cli", "dep:scarb-ui", "dep:scarb-metadata"]
 with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []
+with-segfault-catcher = []
 
 # the aquamarine dep is only used in docs and cannot be detected as used by cargo udeps
 [package.metadata.cargo-udeps.ignore]

src/error.rs

@@ -90,6 +90,10 @@ pub enum Error {
 
     #[error("Failed to parse a Cairo/Sierra program: {0}")]
     ProgramParser(String),
+
+    #[cfg(feature = "with-segfault-catcher")]
+    #[error(transparent)]
+    SafeRunner(crate::utils::safe_runner::SafeRunnerError),
 }
 
 impl Error {

src/executor.rs

@@ -9,6 +9,7 @@ use crate::{
     error::{panic::ToNativeAssertError, Error},
     execution_result::{BuiltinStats, ExecutionResult},
     native_panic,
+    runtime::BUILTIN_COSTS,
     starknet::{handler::StarknetSyscallHandlerCallbacks, StarknetSyscallHandler},
     types::TypeBuilder,
     utils::{libc_free, BuiltinCosts, RangeExt},
@@ -140,23 +141,13 @@ fn invoke_dynamic(
         .map(|syscall_handler| StarknetSyscallHandlerCallbacks::new(syscall_handler));
     // We only care for the previous syscall handler if we actually modify it
     #[cfg(feature = "with-cheatcode")]
-    let previous_syscall_handler = syscall_handler.as_mut().map(|syscall_handler| {
-        let previous_syscall_handler = crate::starknet::SYSCALL_HANDLER_VTABLE.get();
-        let syscall_handler_ptr = std::ptr::addr_of!(*syscall_handler) as *mut ();
-        crate::starknet::SYSCALL_HANDLER_VTABLE.set(syscall_handler_ptr);
-
-        previous_syscall_handler
-    });
+    let syscall_handler_guard = syscall_handler
+        .as_mut()
+        .map(|syscall_handler| SyscallHandlerGuard::install(syscall_handler as *mut _));
 
-    // Order matters, for the libfunc impl
-    let builtin_costs_stack: [u64; 7] = BuiltinCosts::default().into();
-    // Note: the ptr into a slice is valid, it can be used with cast()
-    // Care should be taken if you dereference it and take the .as_ptr() of the slice, since when you
-    // deref it, it will be a copy on the stack, so you will get the ptr of the value in the stack.
-    let builtin_costs: *mut [u64; 7] = Box::into_raw(Box::new(builtin_costs_stack));
     // We may be inside a recursive contract, save the possible saved builtin costs to restore it after our call.
-    let old_builtincosts_ptr =
-        crate::runtime::cairo_native__set_costs_builtin(builtin_costs.cast());
+    let builtin_costs = BuiltinCosts::default();
+    let builtin_costs_guard = BuiltinCostsGuard::install(builtin_costs);
 
     // Generate argument list.
     let mut iter = args.iter();
@@ -223,21 +214,24 @@ fn invoke_dynamic(
     #[cfg(target_arch = "aarch64")]
     let mut ret_registers = [0; 4];
 
-    unsafe {
+    #[allow(unused_mut)]
+    let mut run_trampoline = || unsafe {
         invoke_trampoline(
             function_ptr,
             invoke_data.as_ptr().cast(),
             invoke_data.len() >> 3,
             ret_registers.as_mut_ptr(),
         );
-    }
+    };
+    #[cfg(feature = "with-segfault-catcher")]
+    crate::utils::safe_runner::run_safely(run_trampoline).map_err(Error::SafeRunner)?;
+    #[cfg(not(feature = "with-segfault-catcher"))]
+    run_trampoline();
 
-    // If the syscall handler was changed, then reset the previous one.
-    // It's only necessary to restore the pointer if it's been modified i.e. if previous_syscall_handler is Some(...)
+    // Restore the previous syscall handler and builtin costs.
     #[cfg(feature = "with-cheatcode")]
-    if let Some(previous_syscall_handler) = previous_syscall_handler {
-        crate::starknet::SYSCALL_HANDLER_VTABLE.set(previous_syscall_handler);
-    }
+    drop(syscall_handler_guard);
+    drop(builtin_costs_guard);
 
     // Parse final gas.
     unsafe fn read_value<T>(ptr: &mut NonNull<()>) -> &T {
@@ -337,16 +331,6 @@ fn invoke_dynamic(
             debug_name: None,
         });
 
-    // Restore the old ptr and get back our builtincost box and free it.
-    let our_builtincosts_ptr =
-        crate::runtime::cairo_native__set_costs_builtin(old_builtincosts_ptr);
-
-    if !our_builtincosts_ptr.is_null() && old_builtincosts_ptr.is_aligned() {
-        unsafe {
-            let _ = Box::<[u64; 7]>::from_raw(our_builtincosts_ptr.cast_mut().cast());
-        };
-    }
-
     #[cfg(feature = "with-mem-tracing")]
     crate::utils::mem_tracing::report_stats();
 
@@ -357,6 +341,45 @@ fn invoke_dynamic(
     })
 }
 
+#[cfg(feature = "with-cheatcode")]
+#[derive(Debug)]
+struct SyscallHandlerGuard(*mut ());
+
+#[cfg(feature = "with-cheatcode")]
+impl SyscallHandlerGuard {
+    // NOTE: It is the caller's responsibility to ensure that the syscall handler is alive until the
+    //   guard is dropped.
+    pub fn install<T>(value: *mut T) -> Self {
+        let previous_value = crate::starknet::SYSCALL_HANDLER_VTABLE.get();
+        let syscall_handler_ptr = value as *mut ();
+        crate::starknet::SYSCALL_HANDLER_VTABLE.set(syscall_handler_ptr);
+
+        Self(previous_value)
+    }
+}
+
+#[cfg(feature = "with-cheatcode")]
+impl Drop for SyscallHandlerGuard {
+    fn drop(&mut self) {
+        crate::starknet::SYSCALL_HANDLER_VTABLE.set(self.0);
+    }
+}
+
+#[derive(Debug)]
+struct BuiltinCostsGuard(BuiltinCosts);
+
+impl BuiltinCostsGuard {
+    pub fn install(value: BuiltinCosts) -> Self {
+        Self(BUILTIN_COSTS.replace(value))
+    }
+}
+
+impl Drop for BuiltinCostsGuard {
+    fn drop(&mut self) {
+        BUILTIN_COSTS.set(self.0);
+    }
+}
+
 /// Parses the result by reading from the return ptr the given type.
 fn parse_result(
     type_id: &ConcreteTypeId,

src/executor/contract.rs

@@ -36,7 +36,7 @@ use crate::{
     context::NativeContext,
     error::{panic::ToNativeAssertError, Error, Result},
     execution_result::{BuiltinStats, ContractExecutionResult},
-    executor::invoke_trampoline,
+    executor::{invoke_trampoline, BuiltinCostsGuard},
     metadata::{gas::MetadataComputationConfig, runtime_bindings::setup_runtime},
     module::NativeModule,
     starknet::{handler::StarknetSyscallHandlerCallbacks, StarknetSyscallHandler},
@@ -333,11 +333,11 @@ impl AotContractExecutor {
         };
         let function_ptr = self.find_function_ptr(&function_id, true)?;
 
-        let builtin_costs: [u64; 7] = builtin_costs.unwrap_or_default().into();
-
+        // Initialize syscall handler and builtin costs.
         // We may be inside a recursive contract, save the possible saved builtin costs to restore it after our call.
-        let old_builtincosts_ptr =
-            crate::runtime::cairo_native__set_costs_builtin(builtin_costs.as_ptr());
+        let mut syscall_handler = StarknetSyscallHandlerCallbacks::new(&mut syscall_handler);
+        let builtin_costs = builtin_costs.unwrap_or_default();
+        let builtin_costs_guard = BuiltinCostsGuard::install(builtin_costs);
 
         //  it can vary from contract to contract thats why we need to store/ load it.
         let builtins_size: usize = self.contract_info.entry_points[&selector]
@@ -356,18 +356,13 @@ impl AotContractExecutor {
             .as_ptr()
             .to_bytes(&mut invoke_data, |_| unreachable!())?;
 
-        let mut syscall_handler = StarknetSyscallHandlerCallbacks::new(&mut syscall_handler);
-
         for b in &self.contract_info.entry_points[&selector].builtins {
             match b {
                 BuiltinType::Gas => {
                     gas.to_bytes(&mut invoke_data, |_| unreachable!())?;
                 }
                 BuiltinType::BuiltinCosts => {
-                    // todo: check if valid
-                    builtin_costs
-                        .as_ptr()
-                        .to_bytes(&mut invoke_data, |_| unreachable!())?;
+                    builtin_costs.to_bytes(&mut invoke_data, |_| unreachable!())?;
                 }
                 BuiltinType::System => {
                     (&mut syscall_handler as *mut StarknetSyscallHandlerCallbacks<_>)
@@ -444,14 +439,19 @@ impl AotContractExecutor {
         #[cfg(target_arch = "aarch64")]
         let mut ret_registers = [0; 4];
 
-        unsafe {
+        #[allow(unused_mut)]
+        let mut run_trampoline = || unsafe {
             invoke_trampoline(
                 function_ptr,
                 invoke_data.as_ptr().cast(),
                 invoke_data.len() >> 3,
                 ret_registers.as_mut_ptr(),
             );
-        }
+        };
+        #[cfg(feature = "with-segfault-catcher")]
+        crate::utils::safe_runner::run_safely(run_trampoline).map_err(Error::SafeRunner)?;
+        #[cfg(not(feature = "with-segfault-catcher"))]
+        run_trampoline();
 
         // Parse final gas.
         unsafe fn read_value<T>(ptr: &mut NonNull<()>) -> &T {
@@ -579,7 +579,7 @@ impl AotContractExecutor {
         };
 
         // Restore the original builtin costs pointer.
-        crate::runtime::cairo_native__set_costs_builtin(old_builtincosts_ptr);
+        drop(builtin_costs_guard);
 
         #[cfg(feature = "with-mem-tracing")]
         crate::utils::mem_tracing::report_stats();

src/runtime.rs

@@ -1,5 +1,6 @@
 #![allow(non_snake_case)]
 
+use crate::utils::BuiltinCosts;
 use cairo_lang_sierra_gas::core_libfunc_cost::{
     DICT_SQUASH_REPEATED_ACCESS_COST, DICT_SQUASH_UNIQUE_KEY_COST,
 };
@@ -22,7 +23,7 @@ use std::{
     io::Write,
     mem::{forget, ManuallyDrop},
     os::fd::FromRawFd,
-    ptr::{self, null, null_mut},
+    ptr,
     rc::Rc,
 };
 use std::{ops::Mul, vec::IntoIter};
@@ -162,7 +163,7 @@ impl Clone for FeltDict {
 
             layout: self.layout,
             elements: if self.mappings.is_empty() {
-                null_mut()
+                ptr::null_mut()
             } else {
                 unsafe {
                     alloc(Layout::from_size_align_unchecked(
@@ -255,7 +256,7 @@ pub unsafe extern "C" fn cairo_native__dict_new(
         mappings: HashMap::default(),
 
         layout: Layout::from_size_align_unchecked(size as usize, align as usize),
-        elements: null_mut(),
+        elements: ptr::null_mut(),
 
         dup_fn,
         drop_fn,
@@ -578,27 +579,24 @@ pub unsafe extern "C" fn cairo_native__libfunc__ec__ec_state_try_finalize_nz(
 }
 
 thread_local! {
-    // We can use cell because a ptr is copy.
-    static BUILTIN_COSTS: Cell<*const u64> = const {
-        Cell::new(null())
+    pub(crate) static BUILTIN_COSTS: Cell<BuiltinCosts> = const {
+        // These default values shouldn't be accessible, they will be overriden before entering
+        // compiled code.
+        Cell::new(BuiltinCosts {
+            r#const: 0,
+            pedersen: 0,
+            bitwise: 0,
+            ecop: 0,
+            poseidon: 0,
+            add_mod: 0,
+            mul_mod: 0,
+        })
     };
 }
 
-/// Store the gas builtin in the internal thread local. Returns the old pointer, to restore it after execution.
-/// Not a runtime metadata method, it should be called before the program is executed.
-pub extern "C" fn cairo_native__set_costs_builtin(ptr: *const u64) -> *const u64 {
-    let old = BUILTIN_COSTS.get();
-    BUILTIN_COSTS.set(ptr);
-    old
-}
-
 /// Get the gas builtin from the internal thread local.
-pub extern "C" fn cairo_native__get_costs_builtin() -> *const u64 {
-    if BUILTIN_COSTS.get().is_null() {
-        // We shouldn't panic here, but we can print a big message.
-        eprintln!("BUILTIN_COSTS POINTER IS NULL!");
-    }
-    BUILTIN_COSTS.get()
+pub extern "C" fn cairo_native__get_costs_builtin() -> *const [u64; 7] {
+    BUILTIN_COSTS.with(|x| x.as_ptr()) as *const [u64; 7]
 }
 
 // Utility methods for the print runtime function
@@ -878,7 +876,7 @@ mod tests {
         };
 
         let key = Felt::ONE.to_bytes_le();
-        let mut ptr = null_mut::<u64>();
+        let mut ptr = ptr::null_mut::<u64>();
 
         assert_eq!(
             unsafe { cairo_native__dict_get(dict, &key, (&raw mut ptr).cast()) },

src/utils.rs

@@ -34,6 +34,8 @@ mod block_ext;
 pub mod mem_tracing;
 mod program_registry_ext;
 mod range_ext;
+#[cfg(feature = "with-segfault-catcher")]
+pub mod safe_runner;
 
 #[cfg(target_os = "macos")]
 pub const SHARED_LIBRARY_EXT: &str = "dylib";
@@ -52,7 +54,10 @@ pub static HALF_PRIME: LazyLock<BigUint> = LazyLock::new(|| {
         .expect("hardcoded half prime constant should be valid")
 });
 
+// Order matters, for the libfunc impl
+// https://github.com/starkware-libs/sequencer/blob/1b7252f8a30244d39614d7666aa113b81291808e/crates/blockifier/src/execution/entry_point_execution.rs#L208
 #[derive(Debug, Clone, Copy, Deserialize, Serialize)]
+#[repr(C)]
 pub struct BuiltinCosts {
     pub r#const: u64,
     pub pedersen: u64,
@@ -63,22 +68,6 @@ pub struct BuiltinCosts {
     pub mul_mod: u64,
 }
 
-impl From<BuiltinCosts> for [u64; 7] {
-    // Order matters, for the libfunc impl
-    // https://github.com/starkware-libs/sequencer/blob/1b7252f8a30244d39614d7666aa113b81291808e/crates/blockifier/src/execution/entry_point_execution.rs#L208
-    fn from(value: BuiltinCosts) -> Self {
-        [
-            value.r#const,
-            value.pedersen,
-            value.bitwise,
-            value.ecop,
-            value.poseidon,
-            value.add_mod,
-            value.mul_mod,
-        ]
-    }
-}
-
 impl Default for BuiltinCosts {
     fn default() -> Self {
         Self {
@@ -93,6 +82,30 @@ impl Default for BuiltinCosts {
     }
 }
 
+impl crate::arch::AbiArgument for BuiltinCosts {
+    fn to_bytes(
+        &self,
+        buffer: &mut Vec<u8>,
+        find_dict_overrides: impl Copy
+            + Fn(
+                &cairo_lang_sierra::ids::ConcreteTypeId,
+            ) -> (
+                Option<extern "C" fn(*mut std::ffi::c_void, *mut std::ffi::c_void)>,
+                Option<extern "C" fn(*mut std::ffi::c_void)>,
+            ),
+    ) -> crate::error::Result<()> {
+        self.r#const.to_bytes(buffer, find_dict_overrides)?;
+        self.pedersen.to_bytes(buffer, find_dict_overrides)?;
+        self.bitwise.to_bytes(buffer, find_dict_overrides)?;
+        self.ecop.to_bytes(buffer, find_dict_overrides)?;
+        self.poseidon.to_bytes(buffer, find_dict_overrides)?;
+        self.add_mod.to_bytes(buffer, find_dict_overrides)?;
+        self.mul_mod.to_bytes(buffer, find_dict_overrides)?;
+
+        Ok(())
+    }
+}
+
 #[cfg(feature = "with-mem-tracing")]
 #[allow(unused_imports)]
 pub(crate) use self::mem_tracing::{