Fix scstream double free (#74)

src/stream/internal/cleanup.rs

@@ -2,35 +2,54 @@ use objc::{
     runtime::{self, Object},
     Encoding,
 };
+use std::sync::atomic::{self, Ordering};
 
-use super::{output_handler::OutputTraitWrapper, delegate::StreamDelegateTraitWrapper};
+use super::{delegate::StreamDelegateTraitWrapper, output_handler::OutputTraitWrapper};
 
 const MAX_HANDLERS: usize = 10;
 
 #[repr(C)]
-pub struct Cleanup([*mut Object; MAX_HANDLERS], *mut Object);
+pub struct Cleanup {
+    handlers: [*mut Object; MAX_HANDLERS],
+    internal: *mut Object,
+    rc: atomic::AtomicUsize,
+}
 
 impl Cleanup {
     pub const fn new(delegate: *mut Object) -> Self {
-        Self([std::ptr::null_mut(); MAX_HANDLERS], delegate)
+        Self {
+            handlers: [std::ptr::null_mut(); MAX_HANDLERS],
+            internal: delegate,
+            rc: atomic::AtomicUsize::new(1),
+        }
     }
     pub fn add_handler(&mut self, handler: *mut Object) {
-        let index = self.0.iter().position(|&x| x.is_null()).unwrap();
-        self.0[index] = handler;
+        let index = self.handlers.iter().position(|&x| x.is_null()).unwrap();
+        self.handlers[index] = handler;
     }
 
     fn iter(&self) -> impl Iterator<Item = &*mut Object> {
-        self.0.iter().take_while(|&&x| !x.is_null())
+        self.handlers.iter().take_while(|&&x| !x.is_null())
     }
 
     #[allow(clippy::needless_pass_by_ref_mut)]
     pub fn drop_handlers(&mut self) {
-        if !self.1.is_null() {
+        if self.rc.fetch_sub(1, Ordering::Release) != 1 {
+            return;
+        }
+
+        /*
+         * See https://github.com/rust-lang/rust/blob/e1884a8e3c3e813aada8254edfa120e85bf5ffca/library/alloc/src/sync.rs#L1440-L1467
+         * why the fence is needed.
+         */
+        atomic::fence(Ordering::Acquire);
+
+        if !self.internal.is_null() {
             unsafe {
-                (*self.1)
+                (*self.internal)
                     .get_mut_ivar::<StreamDelegateTraitWrapper>("stream_delegate_wrapper")
                     .drop_trait();
-                runtime::object_dispose(self.1);
+                runtime::object_dispose(self.internal);
             };
         }
         self.iter().for_each(|handler| {
@@ -42,6 +61,13 @@ impl Cleanup {
             };
         });
     }
+
+    pub fn retain(&self) {
+        let old_rc = self.rc.fetch_add(1, Ordering::Relaxed);
+        if old_rc >= isize::MAX as usize {
+            std::process::abort();
+        }
+    }
 }
 
 unsafe impl objc::Encode for Cleanup {

src/stream/internal/delegate.rs

@@ -13,7 +13,7 @@ use crate::{
     utils::objc::get_concrete_from_void,
 };
 
-use super::stream::SCStream;
+use super::stream::get_concrete_stream_from_void;
 
 declare_trait_wrapper!(StreamDelegateTraitWrapper, SCStreamDelegateTrait);
 
@@ -25,7 +25,7 @@ extern "C" fn did_stop_with_error(
     error: *const c_void,
 ) {
     let handler = unsafe { this.get_ivar::<StreamDelegateTraitWrapper>("stream_delegate_wrapper") };
-    let stream = unsafe { get_concrete_from_void::<SCStream>(stream_ref) };
+    let stream = unsafe { get_concrete_stream_from_void(stream_ref) };
     let error: CFError = unsafe { get_concrete_from_void(error) };
     handler.did_stop_with_error(stream, error);
 }
@@ -37,7 +37,7 @@ extern "C" fn output_video_effect_did_start_for_stream(
     stream_ref: *const c_void,
 ) {
     let handler = unsafe { this.get_ivar::<StreamDelegateTraitWrapper>("stream_delegate_wrapper") };
-    let stream = unsafe { get_concrete_from_void::<SCStream>(stream_ref) };
+    let stream = unsafe { get_concrete_stream_from_void(stream_ref) };
     handler.output_video_effect_did_start_for_stream(stream);
 }
 type OutputVideoEffectDidStopForStreamMethod = extern "C" fn(&Object, Sel, *const c_void);
@@ -47,7 +47,7 @@ extern "C" fn output_video_effect_did_stop_for_stream(
     stream_ref: *const c_void,
 ) {
     let handler = unsafe { this.get_ivar::<StreamDelegateTraitWrapper>("stream_delegate_wrapper") };
-    let stream = unsafe { get_concrete_from_void::<SCStream>(stream_ref) };
+    let stream = unsafe { get_concrete_stream_from_void(stream_ref) };
     handler.output_video_effect_did_stop_for_stream(stream);
 }
 

src/stream/internal/stream.rs

@@ -11,7 +11,7 @@ use crate::{
 };
 use core_foundation::{base, error::CFError};
 use core_foundation::{
-    base::{CFTypeID, TCFType},
+    base::{CFTypeID, TCFType, TCFTypeRef},
     impl_TCFType,
 };
 use dispatch::ffi::{dispatch_get_global_queue, DISPATCH_QUEUE_PRIORITY_DEFAULT};
@@ -157,6 +157,23 @@ impl SCStream {
                 .map_err(|_| create_sc_error("Could not receive from completion handler"))?
         }
     }
+
+    pub fn internal_clone(&self) -> Self {
+        unsafe {
+            (*self.as_concrete_TypeRef().cast::<Object>())
+                .get_mut_ivar::<Cleanup>("cleanup")
+                .retain();
+        }
+        Clone::clone(&self)
+    }
+}
+
+pub unsafe fn get_concrete_stream_from_void(void_ptr: *const c_void) -> SCStream {
+    let stream = SCStream::wrap_under_get_rule(SCStreamRef::from_void_ptr(void_ptr));
+    (*stream.as_concrete_TypeRef().cast::<Object>())
+        .get_mut_ivar::<Cleanup>("cleanup")
+        .retain();
+    stream
 }
 
 #[cfg(test)]

src/stream/mod.rs

@@ -63,6 +63,10 @@ impl SCStream {
     pub fn stop_capture(&self) -> Result<(), CFError> {
         self.internal_stop_capture()
     }
+
+    pub fn clone(&self) -> Self {
+        self.internal_clone()
+    }
 }
 
 #[cfg(test)]
@@ -141,4 +145,25 @@ mod stream_test {
         stream.stop_capture()?;
         Ok(())
     }
+
+    #[test]
+    fn clone_stream() -> Result<(), CFError> {
+        let (tx, _) = channel();
+        let stream = {
+            let config = SCStreamConfiguration::new()
+                .set_captures_audio(true)?
+                .set_width(100)?
+                .set_height(100)?;
+
+            let display = SCShareableContent::get().unwrap().displays().remove(0);
+            let filter = SCContentFilter::new().with_display_excluding_windows(&display, &[]);
+            let mut stream = SCStream::new(&filter, &config);
+            stream.add_output_handler(TestStreamOutput { sender: tx }, SCStreamOutputType::Audio);
+            stream
+        };
+
+        let _stream_clone = stream.clone();
+
+        Ok(())
+    }
 }