use crate::{ platform::{ScreenCaptureFrame, ScreenCaptureSource, ScreenCaptureStream}, px, size, Pixels, Size, }; use anyhow::{anyhow, Result}; use block::ConcreteBlock; use cocoa::{ base::{id, nil, YES}, foundation::NSArray, }; use core_foundation::base::TCFType; use ctor::ctor; use futures::channel::oneshot; use media::core_media::{CMSampleBuffer, CMSampleBufferRef}; use metal::NSInteger; use objc::{ class, declare::ClassDecl, msg_send, runtime::{Class, Object, Sel}, sel, sel_impl, }; use std::{cell::RefCell, ffi::c_void, mem, ptr, rc::Rc}; #[derive(Clone)] pub struct MacScreenCaptureSource { sc_display: id, } pub struct MacScreenCaptureStream { sc_stream: id, sc_stream_output: id, } #[link(name = "ScreenCaptureKit", kind = "framework")] extern "C" {} static mut DELEGATE_CLASS: *const Class = ptr::null(); static mut OUTPUT_CLASS: *const Class = ptr::null(); const FRAME_CALLBACK_IVAR: &str = "frame_callback"; #[allow(non_upper_case_globals)] const SCStreamOutputTypeScreen: NSInteger = 0; impl ScreenCaptureSource for MacScreenCaptureSource { fn resolution(&self) -> Result> { unsafe { let width: i64 = msg_send![self.sc_display, width]; let height: i64 = msg_send![self.sc_display, height]; Ok(size(px(width as f32), px(height as f32))) } } fn stream( &self, frame_callback: Box, ) -> oneshot::Receiver>> { unsafe { let stream: id = msg_send![class!(SCStream), alloc]; let filter: id = msg_send![class!(SCContentFilter), alloc]; let configuration: id = msg_send![class!(SCStreamConfiguration), alloc]; let delegate: id = msg_send![DELEGATE_CLASS, alloc]; let output: id = msg_send![OUTPUT_CLASS, alloc]; let excluded_windows = NSArray::array(nil); let filter: id = msg_send![filter, initWithDisplay:self.sc_display excludingWindows:excluded_windows]; let configuration: id = msg_send![configuration, init]; let delegate: id = msg_send![delegate, init]; let output: id = msg_send![output, init]; output.as_mut().unwrap().set_ivar( FRAME_CALLBACK_IVAR, Box::into_raw(Box::new(frame_callback)) as *mut c_void, ); let stream: id = msg_send![stream, initWithFilter:filter configuration:configuration delegate:delegate]; let (mut tx, rx) = oneshot::channel(); let mut error: id = nil; let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id]; if error != nil { let message: id = msg_send![error, localizedDescription]; tx.send(Err(anyhow!("failed to add stream output {message:?}"))) .ok(); return rx; } let tx = Rc::new(RefCell::new(Some(tx))); let handler = ConcreteBlock::new({ move |error: id| { let result = if error == nil { let stream = MacScreenCaptureStream { sc_stream: stream, sc_stream_output: output, }; Ok(Box::new(stream) as Box) } else { let message: id = msg_send![error, localizedDescription]; Err(anyhow!("failed to stop screen capture stream {message:?}")) }; if let Some(tx) = tx.borrow_mut().take() { tx.send(result).ok(); } } }); let handler = handler.copy(); let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler]; rx } } } impl Drop for MacScreenCaptureSource { fn drop(&mut self) { unsafe { let _: () = msg_send![self.sc_display, release]; } } } impl ScreenCaptureStream for MacScreenCaptureStream {} impl Drop for MacScreenCaptureStream { fn drop(&mut self) { unsafe { let mut error: id = nil; let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _]; if error != nil { let message: id = msg_send![error, localizedDescription]; log::error!("failed to add stream output {message:?}"); } let handler = ConcreteBlock::new(move |error: id| { if error != nil { let message: id = msg_send![error, localizedDescription]; log::error!("failed to stop screen capture stream {message:?}"); } }); let block = handler.copy(); let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block]; let _: () = msg_send![self.sc_stream, release]; let _: () = msg_send![self.sc_stream_output, release]; } } } pub(crate) fn get_sources() -> oneshot::Receiver>>> { unsafe { let (mut tx, rx) = oneshot::channel(); let tx = Rc::new(RefCell::new(Some(tx))); let block = ConcreteBlock::new(move |shareable_content: id, error: id| { let Some(mut tx) = tx.borrow_mut().take() else { return; }; let result = if error == nil { let displays: id = msg_send![shareable_content, displays]; let mut result = Vec::new(); for i in 0..displays.count() { let display = displays.objectAtIndex(i); let source = MacScreenCaptureSource { sc_display: msg_send![display, retain], }; result.push(Box::new(source) as Box); } Ok(result) } else { let msg: id = msg_send![error, localizedDescription]; Err(anyhow!("Failed to register: {:?}", msg)) }; tx.send(result).ok(); }); let block = block.copy(); let _: () = msg_send![ class!(SCShareableContent), getShareableContentExcludingDesktopWindows:YES onScreenWindowsOnly:YES completionHandler:block]; rx } } #[ctor] unsafe fn build_classes() { let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap(); decl.add_method( sel!(outputVideoEffectDidStartForStream:), output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id), ); decl.add_method( sel!(outputVideoEffectDidStopForStream:), output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id), ); decl.add_method( sel!(stream:didStopWithError:), stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id), ); DELEGATE_CLASS = decl.register(); let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap(); decl.add_method( sel!(stream:didOutputSampleBuffer:ofType:), stream_did_output_sample_buffer_of_type as extern "C" fn(&Object, Sel, id, id, NSInteger), ); decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR); OUTPUT_CLASS = decl.register(); } extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {} extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {} extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {} extern "C" fn stream_did_output_sample_buffer_of_type( this: &Object, _: Sel, _stream: id, sample_buffer: id, buffer_type: NSInteger, ) { if buffer_type != SCStreamOutputTypeScreen { return; } unsafe { let sample_buffer = sample_buffer as CMSampleBufferRef; let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer); if let Some(buffer) = sample_buffer.image_buffer() { let callback: Box> = Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _); callback(ScreenCaptureFrame(buffer)); mem::forget(callback); } } }