fix(ffi): make FFI struct fields private to close Drop soundness hole

arrow-array/src/ffi_stream.rs

@@ -97,17 +97,24 @@ const ENOSYS: i32 = 38;
 #[derive(Debug)]
 #[allow(non_camel_case_types)]
 pub struct FFI_ArrowArrayStream {
+    // Fields are private so safe code can't install a bogus callback that import
+    // or [`Drop`] would invoke. Write the release fields with the unsafe setters.
     /// C function to get schema from the stream
-    pub get_schema:
-        Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
+    get_schema: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowSchema) -> c_int>,
     /// C function to get next array from the stream
-    pub get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
+    get_next: Option<unsafe extern "C" fn(arg1: *mut Self, out: *mut FFI_ArrowArray) -> c_int>,
     /// C function to get the error from last operation on the stream
-    pub get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
+    get_last_error: Option<unsafe extern "C" fn(arg1: *mut Self) -> *const c_char>,
     /// C function to release the stream
-    pub release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
-    /// Private data used by the stream
-    pub private_data: *mut c_void,
+    ///
+    /// Private so safe code can't install a callback that [`Drop`] would invoke.
+    /// Use [`FFI_ArrowArrayStream::release`] and [`FFI_ArrowArrayStream::set_release`].
+    release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
+    /// Private data used by the stream, owned by the release callback.
+    ///
+    /// Private for the same reason as `release`. Use
+    /// [`FFI_ArrowArrayStream::private_data`] and [`FFI_ArrowArrayStream::set_private_data`].
+    private_data: *mut c_void,
 }
 
 unsafe impl Send for FFI_ArrowArrayStream {}
@@ -213,6 +220,48 @@ impl FFI_ArrowArrayStream {
             private_data: std::ptr::null_mut(),
         }
     }
+
+    /// Returns the producer-provided release callback, if any.
+    ///
+    /// Lets a consumer wrap release: save this callback, install its own with
+    /// [`FFI_ArrowArrayStream::set_release`], and chain back to it on drop. See
+    /// <https://github.com/apache/arrow-rs/issues/9771>.
+    pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
+        self.release
+    }
+
+    /// Returns the opaque producer-provided private data pointer.
+    ///
+    /// See [`FFI_ArrowArrayStream::release`] for the intended use.
+    pub fn private_data(&self) -> *mut c_void {
+        self.private_data
+    }
+
+    /// Replaces the release callback, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// [`Drop`] calls this callback with a pointer to `self`. The new callback
+    /// must correctly release this stream (usually by chaining to the returned
+    /// one) and must match the [`FFI_ArrowArrayStream::private_data`] it reads.
+    /// A wrong callback is undefined behavior on drop.
+    pub unsafe fn set_release(
+        &mut self,
+        release: Option<unsafe extern "C" fn(arg1: *mut Self)>,
+    ) -> Option<unsafe extern "C" fn(arg1: *mut Self)> {
+        std::mem::replace(&mut self.release, release)
+    }
+
+    /// Replaces the private data pointer, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// The old pointer is returned without being freed; the caller owns it from
+    /// here. The new pointer must match what the current
+    /// [`FFI_ArrowArrayStream::release`] callback expects.
+    pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void {
+        std::mem::replace(&mut self.private_data, private_data)
+    }
 }
 
 struct ExportedArrayStream {
@@ -566,4 +615,48 @@ mod tests {
 
         Ok(())
     }
+
+    // A consumer wraps the release callback with its own, then chains back to
+    // the original on drop. This is the same wrap-release pattern the
+    // release/private_data accessors exist for (#9771).
+    static STREAM_WRAPPER_RAN: std::sync::atomic::AtomicBool =
+        std::sync::atomic::AtomicBool::new(false);
+
+    struct StreamWrapperData {
+        original_release: Option<unsafe extern "C" fn(*mut FFI_ArrowArrayStream)>,
+        original_private_data: *mut c_void,
+    }
+
+    unsafe extern "C" fn wrapping_release(stream: *mut FFI_ArrowArrayStream) {
+        use std::sync::atomic::Ordering;
+        let stream = unsafe { &mut *stream };
+        let data = unsafe { Box::from_raw(stream.private_data() as *mut StreamWrapperData) };
+        STREAM_WRAPPER_RAN.store(true, Ordering::SeqCst);
+        unsafe { stream.set_release(data.original_release) };
+        unsafe { stream.set_private_data(data.original_private_data) };
+        if let Some(release) = stream.release() {
+            unsafe { release(stream) };
+        }
+    }
+
+    #[test]
+    fn test_wrap_release_callback() {
+        use std::sync::atomic::Ordering;
+
+        let batch_reader = TestRecordBatchReader::new(
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])),
+            Box::new(std::iter::empty()),
+        );
+        let mut stream = FFI_ArrowArrayStream::new(batch_reader);
+
+        let data = Box::new(StreamWrapperData {
+            original_release: stream.release(),
+            original_private_data: stream.private_data(),
+        });
+        unsafe { stream.set_release(Some(wrapping_release)) };
+        unsafe { stream.set_private_data(Box::into_raw(data) as *mut c_void) };
+
+        drop(stream); // runs wrapping_release, which chains to the original
+        assert!(STREAM_WRAPPER_RAN.load(Ordering::SeqCst));
+    }
 }

arrow-data/src/ffi.rs

@@ -37,31 +37,41 @@ use std::ffi::c_void;
 #[repr(C)]
 #[derive(Debug)]
 pub struct FFI_ArrowArray {
+    // Fields are private so safe code can't set them to values that break import
+    // or the release callback (e.g. a bad `buffers` pointer gets dereferenced on
+    // import). Read them with the getters; write the callback fields with the
+    // unsafe setters.
     /// Logical length of the array
-    pub length: i64,
+    length: i64,
     /// Number of null items in the array
-    pub null_count: i64,
+    null_count: i64,
     /// logical offset inside the array
-    pub offset: i64,
+    offset: i64,
     /// Number of physical buffers backing this array
-    pub n_buffers: i64,
+    n_buffers: i64,
     /// Number of children this array has
-    pub n_children: i64,
+    n_children: i64,
     /// C array of pointers to the start of each physical buffer backing this array
-    pub buffers: *mut *const c_void,
+    buffers: *mut *const c_void,
     /// C array of pointers to each child array of this array
-    pub children: *mut *mut FFI_ArrowArray,
+    children: *mut *mut FFI_ArrowArray,
     /// Pointer to the underlying array of dictionary values
-    pub dictionary: *mut FFI_ArrowArray,
-    /// Pointer to a producer-provided release callback
-    pub release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)>,
+    dictionary: *mut FFI_ArrowArray,
+    /// Producer-provided release callback.
+    ///
+    /// Private so safe code can't install a callback that [`Drop`] would invoke.
+    /// Use [`FFI_ArrowArray::release`] and [`FFI_ArrowArray::set_release`].
+    release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)>,
     /// Opaque pointer to producer-provided private data
     /// When exported, this MUST contain everything that is owned by this array.
     /// For example, any buffer pointed to in `buffers` must be here, as well
     /// as the `buffers` pointer itself.
     /// In other words, everything in [FFI_ArrowArray] must be owned by
     /// `private_data` and can assume that they do not outlive `private_data`.
-    pub private_data: *mut c_void,
+    ///
+    /// Private for the same reason as `release`. Use
+    /// [`FFI_ArrowArray::private_data`] and [`FFI_ArrowArray::set_private_data`].
+    private_data: *mut c_void,
 }
 
 impl Drop for FFI_ArrowArray {
@@ -251,6 +261,48 @@ impl FFI_ArrowArray {
         }
     }
 
+    /// Returns the producer-provided release callback, if any.
+    ///
+    /// Lets a consumer wrap release: save this callback, install its own with
+    /// [`FFI_ArrowArray::set_release`], and chain back to it on drop. See
+    /// <https://github.com/apache/arrow-rs/issues/9771>.
+    pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)> {
+        self.release
+    }
+
+    /// Returns the opaque producer-provided private data pointer.
+    ///
+    /// See [`FFI_ArrowArray::release`] for the intended use.
+    pub fn private_data(&self) -> *mut c_void {
+        self.private_data
+    }
+
+    /// Replaces the release callback, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// [`Drop`] calls this callback with a pointer to `self`. The new callback
+    /// must correctly release this array (usually by chaining to the returned
+    /// one) and must match the [`FFI_ArrowArray::private_data`] it reads. A
+    /// wrong callback is undefined behavior on drop.
+    pub unsafe fn set_release(
+        &mut self,
+        release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)>,
+    ) -> Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowArray)> {
+        std::mem::replace(&mut self.release, release)
+    }
+
+    /// Replaces the private data pointer, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// The old pointer is returned without being freed; the caller owns it from
+    /// here. The new pointer must match what the current
+    /// [`FFI_ArrowArray::release`] callback expects.
+    pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void {
+        std::mem::replace(&mut self.private_data, private_data)
+    }
+
     /// the length of the array
     #[inline]
     pub fn len(&self) -> usize {

arrow-schema/src/ffi.rs

@@ -75,25 +75,35 @@ bitflags! {
 #[derive(Debug)]
 #[allow(non_camel_case_types)]
 pub struct FFI_ArrowSchema {
+    // Fields are private so safe code can't set them to values that break the
+    // release callback (e.g. a `format`/`name` that isn't a `CString`, or a bad
+    // pointer). Read them with the getters; write the callback fields with the
+    // unsafe setters.
     /// Null-terminated, UTF8-encoded string describing the data type
-    pub format: *const c_char,
+    format: *const c_char,
     /// Null-terminated, UTF8-encoded string of the field or array name
-    pub name: *const c_char,
+    name: *const c_char,
     /// Binary string describing the type’s metadata
-    pub metadata: *const c_char,
+    metadata: *const c_char,
     /// A bitfield of flags enriching the type description
     /// Refer to [Arrow Flags](https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.flags)
-    pub flags: i64,
+    flags: i64,
     /// The number of children this type has
-    pub n_children: i64,
+    n_children: i64,
     /// C array of pointers to each child type of this type
-    pub children: *mut *mut FFI_ArrowSchema,
+    children: *mut *mut FFI_ArrowSchema,
     /// Pointer to the type of dictionary values
-    pub dictionary: *mut FFI_ArrowSchema,
-    /// Pointer to a producer-provided release callback
-    pub release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)>,
-    /// Opaque pointer to producer-provided private data
-    pub private_data: *mut c_void,
+    dictionary: *mut FFI_ArrowSchema,
+    /// Producer-provided release callback.
+    ///
+    /// Private so safe code can't install a callback that [`Drop`] would invoke.
+    /// Use [`FFI_ArrowSchema::release`] and [`FFI_ArrowSchema::set_release`].
+    release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)>,
+    /// Opaque producer-provided private data, owned by the release callback.
+    ///
+    /// Private for the same reason as `release`. Use
+    /// [`FFI_ArrowSchema::private_data`] and [`FFI_ArrowSchema::set_private_data`].
+    private_data: *mut c_void,
 }
 
 struct SchemaPrivateData {
@@ -276,6 +286,48 @@ impl FFI_ArrowSchema {
         }
     }
 
+    /// Returns the producer-provided release callback, if any.
+    ///
+    /// Lets a consumer wrap release: save this callback, install its own with
+    /// [`FFI_ArrowSchema::set_release`], and chain back to it on drop. See
+    /// <https://github.com/apache/arrow-rs/issues/9771>.
+    pub fn release(&self) -> Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)> {
+        self.release
+    }
+
+    /// Returns the opaque producer-provided private data pointer.
+    ///
+    /// See [`FFI_ArrowSchema::release`] for the intended use.
+    pub fn private_data(&self) -> *mut c_void {
+        self.private_data
+    }
+
+    /// Replaces the release callback, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// [`Drop`] calls this callback with a pointer to `self`. The new callback
+    /// must correctly release this schema (usually by chaining to the returned
+    /// one) and must match the [`FFI_ArrowSchema::private_data`] it reads. A
+    /// wrong callback is undefined behavior on drop.
+    pub unsafe fn set_release(
+        &mut self,
+        release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)>,
+    ) -> Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)> {
+        std::mem::replace(&mut self.release, release)
+    }
+
+    /// Replaces the private data pointer, returning the previous one.
+    ///
+    /// # Safety
+    ///
+    /// The old pointer is returned without being freed; the caller owns it from
+    /// here. The new pointer must match what the current
+    /// [`FFI_ArrowSchema::release`] callback expects.
+    pub unsafe fn set_private_data(&mut self, private_data: *mut c_void) -> *mut c_void {
+        std::mem::replace(&mut self.private_data, private_data)
+    }
+
     /// Returns the format of this schema.
     pub fn format(&self) -> &str {
         assert!(!self.format.is_null());
@@ -861,6 +913,7 @@ impl TryFrom<Schema> for FFI_ArrowSchema {
 mod tests {
     use super::*;
     use crate::Fields;
+    use std::sync::atomic::{AtomicBool, Ordering};
 
     fn round_trip_type(dtype: DataType) {
         let c_schema = FFI_ArrowSchema::try_from(&dtype).unwrap();
@@ -1029,4 +1082,41 @@ mod tests {
         let field = Field::try_from(&c_schema).unwrap();
         assert_eq!(field.name(), "");
     }
+
+    // A consumer wraps the release callback with its own, then chains back to
+    // the original on drop. This is the cross-thread use case from #9771 and is
+    // what the release/private_data accessors exist for.
+    static WRAPPER_RAN: AtomicBool = AtomicBool::new(false);
+
+    struct WrapperData {
+        original_release: Option<unsafe extern "C" fn(*mut FFI_ArrowSchema)>,
+        original_private_data: *mut c_void,
+    }
+
+    unsafe extern "C" fn wrapping_release(schema: *mut FFI_ArrowSchema) {
+        let schema = unsafe { &mut *schema };
+        let data = unsafe { Box::from_raw(schema.private_data() as *mut WrapperData) };
+        WRAPPER_RAN.store(true, Ordering::SeqCst);
+        // restore the originals, then let the original callback free everything
+        unsafe { schema.set_release(data.original_release) };
+        unsafe { schema.set_private_data(data.original_private_data) };
+        if let Some(release) = schema.release() {
+            unsafe { release(schema) };
+        }
+    }
+
+    #[test]
+    fn test_wrap_release_callback() {
+        let mut schema = FFI_ArrowSchema::try_from(&DataType::Int32).unwrap();
+
+        let data = Box::new(WrapperData {
+            original_release: schema.release(),
+            original_private_data: schema.private_data(),
+        });
+        unsafe { schema.set_release(Some(wrapping_release)) };
+        unsafe { schema.set_private_data(Box::into_raw(data) as *mut c_void) };
+
+        drop(schema); // runs wrapping_release, which chains to the original
+        assert!(WRAPPER_RAN.load(Ordering::SeqCst));
+    }
 }