Fix intermittent SIGSEGV in cross-platform FFI plugin tests (drasi-project#376)

ffi hardening

Author: danielgerlag

components/host-sdk/src/identity_bridge.rs

@@ -44,57 +44,66 @@ impl IdentityProviderVtableBuilder {
             context_json: *const u8,
             context_len: usize,
         ) -> FfiCredentialsResult {
-            let wrapper = unsafe { &*(state as *const HostIdentityProviderState) };
-            let provider = wrapper.provider.clone();
+            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                let wrapper = unsafe { &*(state as *const HostIdentityProviderState) };
+                let provider = wrapper.provider.clone();
 
-            // Deserialize context from JSON
-            let context = if context_json.is_null() || context_len == 0 {
-                CredentialContext::default()
-            } else {
-                let json_bytes = unsafe { std::slice::from_raw_parts(context_json, context_len) };
-                let json_str = std::str::from_utf8(json_bytes).unwrap_or("{}");
-                let properties: std::collections::HashMap<String, String> =
-                    match serde_json::from_str(json_str) {
-                        Ok(p) => p,
-                        Err(e) => {
-                            log::warn!("Failed to deserialize credential context JSON: {e}");
-                            std::collections::HashMap::new()
-                        }
-                    };
-                CredentialContext { properties }
-            };
+                // Deserialize context from JSON
+                let context = if context_json.is_null() || context_len == 0 {
+                    CredentialContext::default()
+                } else {
+                    let json_bytes =
+                        unsafe { std::slice::from_raw_parts(context_json, context_len) };
+                    let json_str = std::str::from_utf8(json_bytes).unwrap_or("{}");
+                    let properties: std::collections::HashMap<String, String> =
+                        match serde_json::from_str(json_str) {
+                            Ok(p) => p,
+                            Err(e) => {
+                                log::warn!("Failed to deserialize credential context JSON: {e}");
+                                std::collections::HashMap::new()
+                            }
+                        };
+                    CredentialContext { properties }
+                };
 
-            // This vtable is called from the plugin side (extern "C") which may
-            // already be inside a tokio runtime. We spawn a thread and create a
-            // lightweight current-thread runtime to avoid nesting runtimes.
-            let result = std::thread::spawn(move || {
-                let rt = tokio::runtime::Builder::new_current_thread()
-                    .enable_all()
-                    .build()
-                    .map_err(|e| anyhow::anyhow!("Failed to create runtime: {e}"))?;
-                rt.block_on(provider.get_credentials(&context))
-            })
-            .join();
+                // This vtable is called from the plugin side (extern "C") which may
+                // already be inside a tokio runtime. We spawn a thread and create a
+                // lightweight current-thread runtime to avoid nesting runtimes.
+                let result = std::thread::spawn(move || {
+                    let rt = tokio::runtime::Builder::new_current_thread()
+                        .enable_all()
+                        .build()
+                        .map_err(|e| anyhow::anyhow!("Failed to create runtime: {e}"))?;
+                    rt.block_on(provider.get_credentials(&context))
+                })
+                .join();
 
-            match result {
-                Ok(Ok(creds)) => FfiCredentialsResult::ok(credentials_to_ffi(creds)),
-                Ok(Err(e)) => FfiCredentialsResult::err(e.to_string()),
-                Err(_) => FfiCredentialsResult::err("get_credentials thread panicked".into()),
-            }
+                match result {
+                    Ok(Ok(creds)) => FfiCredentialsResult::ok(credentials_to_ffi(creds)),
+                    Ok(Err(e)) => FfiCredentialsResult::err(e.to_string()),
+                    Err(_) => FfiCredentialsResult::err("get_credentials thread panicked".into()),
+                }
+            }))
+            .unwrap_or_else(|_| FfiCredentialsResult::err("get_credentials_fn panicked".into()))
         }
 
         extern "C" fn clone_fn(state: *const c_void) -> *mut c_void {
-            let wrapper = unsafe { &*(state as *const HostIdentityProviderState) };
-            let cloned = Box::new(HostIdentityProviderState {
-                provider: wrapper.provider.clone(),
-            });
-            Box::into_raw(cloned) as *mut c_void
+            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                let wrapper = unsafe { &*(state as *const HostIdentityProviderState) };
+                let cloned = Box::new(HostIdentityProviderState {
+                    provider: wrapper.provider.clone(),
+                });
+                Box::into_raw(cloned) as *mut c_void
+            }))
+            .unwrap_or(std::ptr::null_mut())
         }
 
         extern "C" fn drop_fn(state: *mut c_void) {
-            if !state.is_null() {
-                unsafe { drop(Box::from_raw(state as *mut HostIdentityProviderState)) };
-            }
+            let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                if !state.is_null() {
+                    unsafe { drop(Box::from_raw(state as *mut HostIdentityProviderState)) };
+                }
+            }));
         }
 
         let wrapper = Box::new(HostIdentityProviderState { provider });

components/host-sdk/src/proxies/bootstrap_provider.rs

@@ -155,26 +155,30 @@ fn build_ffi_bootstrap_sender(event_tx: BootstrapEventSender) -> FfiBootstrapSen
     let state = Box::into_raw(Box::new(std_tx)) as *mut c_void;
 
     extern "C" fn send_fn(state: *mut c_void, event: *mut FfiBootstrapEvent) -> i32 {
-        let tx = unsafe { &*(state as *const std::sync::mpsc::Sender<BootstrapEvent>) };
-        if event.is_null() {
-            return -1;
-        }
-        let ffi_event = unsafe { &*event };
-        let bootstrap_event = unsafe { *Box::from_raw(ffi_event.opaque as *mut BootstrapEvent) };
-        // Free the FFI envelope but not the opaque (we took ownership)
-        unsafe { drop(Box::from_raw(event)) };
-        match tx.send(bootstrap_event) {
-            Ok(()) => 0,
-            Err(_) => -1,
-        }
+        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            let tx = unsafe { &*(state as *const std::sync::mpsc::Sender<BootstrapEvent>) };
+            if event.is_null() {
+                return -1;
+            }
+            let ffi_event = unsafe { &*event };
+            let bootstrap_event =
+                unsafe { *Box::from_raw(ffi_event.opaque as *mut BootstrapEvent) };
+            // Free the FFI envelope but not the opaque (we took ownership)
+            unsafe { drop(Box::from_raw(event)) };
+            match tx.send(bootstrap_event) {
+                Ok(()) => 0,
+                Err(_) => -1,
+            }
+        }))
+        .unwrap_or(-1)
     }
 
     extern "C" fn drop_fn(state: *mut c_void) {
-        unsafe {
+        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe {
             drop(Box::from_raw(
                 state as *mut std::sync::mpsc::Sender<BootstrapEvent>,
             ))
-        };
+        }));
     }
 
     FfiBootstrapSender {

components/host-sdk/src/proxies/reaction.rs

@@ -36,6 +36,16 @@ pub struct ReactionProxy {
     _library: Arc<Library>,
     cached_id: String,
     cached_type_name: String,
+    /// Per-instance callback context for plugin-emitted log/lifecycle callbacks.
+    ///
+    /// Stored as an `Arc` whose strong count was bumped by `Arc::into_raw` when
+    /// the raw pointer was handed to the plugin. The host's `Arc` is kept here
+    /// so the proxy holds at least two strong references; on Drop the host's
+    /// `Arc` is `mem::forget`-ed unconditionally so any **late** log/lifecycle
+    /// callback emitted by the plugin (after `stop()` returns) still finds a
+    /// valid pointer. The cdylib itself is intentionally leaked process-wide
+    /// (see `host-sdk/src/loader.rs`), so the small per-instance `Arc` leak is
+    /// acceptable in exchange for closing the late-callback UAF window.
     _callback_ctx: std::sync::Mutex<Option<Arc<crate::callbacks::InstanceCallbackContext>>>,
     /// Channel for push-based result delivery. Created on start, closed on stop/drop.
     result_tx:
@@ -195,7 +205,13 @@ impl Reaction for ReactionProxy {
             update_tx: context.update_tx.clone(),
         });
 
-        let ctx_ptr = Arc::as_ptr(&per_instance_ctx) as *mut c_void;
+        // Bug C fix: hand the plugin a strong reference (Arc::into_raw bumps
+        // the refcount) so log/lifecycle callbacks emitted late by the plugin
+        // (e.g. from inside stop_fn or from internal tasks shutting down) do
+        // not deref freed memory. The matching `mem::forget` happens in Drop
+        // and intentionally leaks one strong ref per instance.
+        let ctx_for_plugin = per_instance_ctx.clone();
+        let ctx_ptr = Arc::into_raw(ctx_for_plugin) as *mut c_void;
 
         if let Ok(mut guard) = self._callback_ctx.lock() {
             *guard = Some(per_instance_ctx);
@@ -257,19 +273,18 @@ impl Reaction for ReactionProxy {
     }
 
     async fn stop(&self) -> anyhow::Result<()> {
-        // Close the sender so the forwarder's callback returns null
+        // Bug B fix: close ONLY the sender. Dropping the sender is sufficient
+        // to unblock the forwarder's `rx.recv()` (it returns Err, the callback
+        // returns null, and the forwarder breaks). Do NOT also drop the
+        // receiver here — the callback may still be holding `context.rx.lock()`
+        // for a recv() that is racing this stop, and removing the receiver
+        // mid-flight creates a race against the forwarder's in-flight
+        // `enqueue_query_result(qr).await` against the reaction's
+        // shutting-down priority queue.
         {
             let mut guard = self.result_tx.lock().expect("result_tx lock poisoned");
             *guard = None;
         }
-        // Also drop the receiver to unblock the callback if it's blocked in recv()
-        if let Ok(guard) = self._push_ctx.lock() {
-            if let Some(ref ctx) = *guard {
-                if let Ok(mut rx_guard) = ctx.rx.lock() {
-                    *rx_guard = None;
-                }
-            }
-        }
 
         let state = drasi_plugin_sdk::ffi::SendMutPtr(self.vtable.state);
         let stop_fn = self.vtable.stop_fn;
@@ -327,15 +342,10 @@ impl Drop for ReactionProxy {
         if let Ok(mut guard) = self.result_tx.lock() {
             *guard = None;
         }
-        // Also drop the receiver inside the push context to unblock the callback
-        // if it's blocked in recv() (belt-and-suspenders with the sender drop).
-        if let Ok(guard) = self._push_ctx.lock() {
-            if let Some(ref ctx) = *guard {
-                if let Ok(mut rx_guard) = ctx.rx.lock() {
-                    *rx_guard = None;
-                }
-            }
-        }
+        // Bug B fix: do NOT drop the receiver here. Sender drop alone unblocks
+        // recv(); leaving the receiver in place avoids racing a callback that
+        // is currently holding `context.rx.lock()`. The receiver lives until
+        // the leaked push-ctx Arc is collected (see the `mem::forget` below).
 
         // Wait for the forwarder task to fully exit its processing loop.
         //
@@ -381,6 +391,20 @@ impl Drop for ReactionProxy {
                 std::mem::forget(ctx);
             }
         }
+
+        // Bug C fix: leak the per-instance callback context Arc unconditionally.
+        // The strong reference handed to the plugin via `Arc::into_raw` in
+        // initialize() is never reclaimed — late log/lifecycle callbacks
+        // emitted by the plugin (during stop_fn or from internal tasks) must
+        // still find a valid pointer. The cdylib itself is intentionally
+        // leaked process-wide (see host-sdk/src/loader.rs), so this small
+        // per-instance Arc leak is the price of closing the late-callback
+        // UAF window.
+        if let Ok(mut guard) = self._callback_ctx.lock() {
+            if let Some(ctx) = guard.take() {
+                std::mem::forget(ctx);
+            }
+        }
     }
 }
 

components/host-sdk/src/proxies/source.rs

@@ -40,23 +40,30 @@ use crate::state_store_bridge::StateStoreVtableBuilder;
 /// Runs the pinned future on a new OS thread with a current-thread tokio runtime.
 /// This avoids nesting issues with the host's multi-thread runtime.
 extern "C" fn host_executor(future_ptr: *mut c_void) -> *mut c_void {
-    // Wrap the raw pointer to make it Send-safe for std::thread::spawn
-    let send_ptr = drasi_plugin_sdk::ffi::SendMutPtr(future_ptr);
-    let result = std::thread::spawn(move || {
-        let boxed_future = unsafe {
-            Box::from_raw(send_ptr.as_ptr()
-                as *mut std::pin::Pin<Box<dyn std::future::Future<Output = *mut c_void>>>)
-        };
-        let rt = tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build()
-            .expect("failed to build tokio runtime for source proxy");
-        // Wrap the result in SendMutPtr to satisfy Send bound
-        drasi_plugin_sdk::ffi::SendMutPtr(rt.block_on(*boxed_future))
-    })
-    .join()
-    .expect("host executor thread panicked");
-    result.as_ptr()
+    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        // Wrap the raw pointer to make it Send-safe for std::thread::spawn
+        let send_ptr = drasi_plugin_sdk::ffi::SendMutPtr(future_ptr);
+        let result = std::thread::spawn(move || {
+            let boxed_future = unsafe {
+                Box::from_raw(send_ptr.as_ptr()
+                    as *mut std::pin::Pin<Box<dyn std::future::Future<Output = *mut c_void>>>)
+            };
+            let rt = match tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+            {
+                Ok(rt) => rt,
+                Err(_) => return drasi_plugin_sdk::ffi::SendMutPtr(std::ptr::null_mut()),
+            };
+            // Wrap the result in SendMutPtr to satisfy Send bound
+            drasi_plugin_sdk::ffi::SendMutPtr(rt.block_on(*boxed_future))
+        })
+        .join()
+        .map(|p| p.as_ptr())
+        .unwrap_or(std::ptr::null_mut());
+        result
+    }))
+    .unwrap_or(std::ptr::null_mut())
 }
 
 /// Wraps a `SourceVtable` into a DrasiLib `Source` trait implementation.
@@ -265,10 +272,15 @@ impl Source for SourceProxy {
             update_tx: context.update_tx.clone(),
         });
 
-        // Use as_ptr (no refcount increment) — the Arc in _callback_ctx keeps
-        // the context alive for the proxy's lifetime. This avoids the memory
-        // leak that Arc::into_raw would cause (no matching from_raw).
-        let ctx_ptr = Arc::as_ptr(&per_instance_ctx) as *mut c_void;
+        // Bug C fix: hand the plugin a strong reference (Arc::into_raw bumps
+        // the refcount) so log/lifecycle callbacks emitted late by the plugin
+        // (e.g. from inside stop_fn or from internal tasks shutting down) do
+        // not deref freed memory. The matching `mem::forget` happens in Drop
+        // and intentionally leaks one strong ref per instance — acceptable
+        // because the cdylib itself is intentionally process-leaked (see
+        // host-sdk/src/loader.rs).
+        let ctx_for_plugin = per_instance_ctx.clone();
+        let ctx_ptr = Arc::into_raw(ctx_for_plugin) as *mut c_void;
 
         // Store the Arc so it stays alive as long as this proxy
         if let Ok(mut guard) = self._callback_ctx.lock() {
@@ -316,6 +328,17 @@ impl Drop for SourceProxy {
         let drop_fn = self.vtable.drop_fn;
         let state = drasi_plugin_sdk::ffi::SendMutPtr(self.vtable.state);
         let _ = std::thread::spawn(move || (drop_fn)(state.as_ptr())).join();
+
+        // Bug C fix: leak the per-instance callback context Arc unconditionally.
+        // The strong reference handed to the plugin via `Arc::into_raw` in
+        // initialize() is never reclaimed — late log/lifecycle callbacks
+        // emitted by the plugin (during stop_fn or from internal tasks) must
+        // still find a valid pointer. Matches the pattern in ReactionProxy.
+        if let Ok(mut guard) = self._callback_ctx.lock() {
+            if let Some(ctx) = guard.take() {
+                std::mem::forget(ctx);
+            }
+        }
     }
 }